feat: optimize vacuum drop table, execute in parallel (#15478)

lichuang · dantengsky · web-flow · commit 6981624de89a · 2024-05-11T19:21:57.000+08:00
* feat: optimize vacuum drop table, execute in parallel

* feat: optimize vacuum drop table, execute in parallel

* feat: optimize vacuum drop table, execute in parallel

* feat: optimize vacuum drop table, execute in parallel

* feat: optimize vacuum drop table, execute in parallel

* Update src/query/ee/src/storages/fuse/operations/vacuum_drop_tables.rs

---------

Co-authored-by: dantengsky &lt;dantengsky@gmail.com&gt;
diff --git a/src/query/ee/src/storages/fuse/operations/handler.rs b/src/query/ee/src/storages/fuse/operations/handler.rs
@@ -46,10 +46,11 @@ impl VacuumHandler for RealVacuumHandler {
 
     async fn do_vacuum_drop_tables(
         &self,
+        threads_nums: usize,
         tables: Vec<Arc<dyn Table>>,
         dry_run_limit: Option<usize>,
     ) -> Result<Option<Vec<VacuumDropFileInfo>>> {
-        do_vacuum_drop_tables(tables, dry_run_limit).await
+        do_vacuum_drop_tables(threads_nums, tables, dry_run_limit).await
     }
 
     async fn do_vacuum_temporary_files(
diff --git a/src/query/ee/src/storages/fuse/operations/vacuum_drop_tables.rs b/src/query/ee/src/storages/fuse/operations/vacuum_drop_tables.rs
@@ -15,6 +15,7 @@
 use std::sync::Arc;
 use std::time::Instant;
 
+use databend_common_base::runtime::execute_futures_in_parallel;
 use databend_common_catalog::table::Table;
 use databend_common_exception::Result;
 use databend_common_meta_app::schema::TableInfo;
@@ -28,108 +29,131 @@ use opendal::Operator;
 
 #[async_backtrace::framed]
 pub async fn do_vacuum_drop_table(
-    table_info: &TableInfo,
-    operator: &Operator,
+    tables: Vec<(TableInfo, Operator)>,
     dry_run_limit: Option<usize>,
 ) -> Result<Option<Vec<VacuumDropFileInfo>>> {
-    let dir = format!("{}/", FuseTable::parse_storage_prefix(table_info)?);
+    let mut list_files = vec![];
+    for (table_info, operator) in tables {
+        let dir = format!("{}/", FuseTable::parse_storage_prefix(&table_info)?);
 
-    info!(
-        "vacuum drop table {:?} dir {:?}, is_external_table:{:?}",
-        table_info.name,
-        dir,
-        table_info.meta.storage_params.is_some()
-    );
+        info!(
+            "vacuum drop table {:?} dir {:?}, is_external_table:{:?}",
+            table_info.name,
+            dir,
+            table_info.meta.storage_params.is_some()
+        );
 
-    let start = Instant::now();
+        let start = Instant::now();
 
-    let ret = match dry_run_limit {
-        None => {
-            operator.remove_all(&dir).await?;
-            Ok(None)
-        }
-        Some(dry_run_limit) => {
-            let mut ds = operator
-                .lister_with(&dir)
-                .recursive(true)
-                .metakey(Metakey::Mode)
-                .metakey(Metakey::ContentLength)
-                .await?;
-            let mut list_files = Vec::new();
-            while let Some(de) = ds.try_next().await? {
-                let meta = de.metadata();
-                if EntryMode::FILE == meta.mode() {
-                    list_files.push((
-                        table_info.name.clone(),
-                        de.name().to_string(),
-                        meta.content_length(),
-                    ));
-                    if list_files.len() >= dry_run_limit {
-                        break;
+        match dry_run_limit {
+            None => {
+                operator.remove_all(&dir).await?;
+            }
+            Some(dry_run_limit) => {
+                let mut ds = operator
+                    .lister_with(&dir)
+                    .recursive(true)
+                    .metakey(Metakey::Mode)
+                    .metakey(Metakey::ContentLength)
+                    .await?;
+
+                while let Some(de) = ds.try_next().await? {
+                    let meta = de.metadata();
+                    if EntryMode::FILE == meta.mode() {
+                        list_files.push((
+                            table_info.name.clone(),
+                            de.name().to_string(),
+                            meta.content_length(),
+                        ));
+                        if list_files.len() >= dry_run_limit {
+                            break;
+                        }
                     }
                 }
             }
+        };
 
-            Ok(Some(list_files))
-        }
-    };
-
-    info!(
-        "vacuum drop table {:?} dir {:?}, cost:{} sec",
-        table_info.name,
-        dir,
-        start.elapsed().as_secs()
-    );
-    ret
+        info!(
+            "vacuum drop table {:?} dir {:?}, cost:{} sec",
+            table_info.name,
+            dir,
+            start.elapsed().as_secs()
+        );
+    }
+    Ok(if dry_run_limit.is_some() {
+        Some(list_files)
+    } else {
+        None
+    })
 }
 
 #[async_backtrace::framed]
 pub async fn do_vacuum_drop_tables(
+    threads_nums: usize,
     tables: Vec<Arc<dyn Table>>,
     dry_run_limit: Option<usize>,
 ) -> Result<Option<Vec<VacuumDropFileInfo>>> {
     let start = Instant::now();
     let tables_len = tables.len();
     info!("do_vacuum_drop_tables {} tables", tables_len);
-    let mut list_files = Vec::new();
-    let mut left_limit = dry_run_limit;
+
+    let batch_size = (tables_len / threads_nums).min(50).max(1);
+
+    let mut table_vecs = Vec::with_capacity(tables.len());
     for table in tables {
-        // only operate fuse table
-        let ret = if let Ok(fuse_table) = FuseTable::try_from_table(table.as_ref()) {
-            let table_info = table.get_table_info();
-            let operator = fuse_table.get_operator_ref();
-            do_vacuum_drop_table(table_info, operator, left_limit).await?
-        } else {
-            info!(
-                "ignore table {}, which is not of FUSE engine. Table engine {}",
-                table.get_table_info().name,
-                table.engine()
-            );
-            continue;
-        };
-        if let Some(ret) = ret {
-            list_files.extend(ret);
-            if list_files.len() >= dry_run_limit.unwrap() {
+        let (table_info, operator) =
+            if let Ok(fuse_table) = FuseTable::try_from_table(table.as_ref()) {
+                (fuse_table.get_table_info(), fuse_table.get_operator())
+            } else {
                 info!(
-                    "do_vacuum_drop_tables {} tables, cost:{} sec",
-                    tables_len,
-                    start.elapsed().as_secs()
+                    "ignore table {}, which is not of FUSE engine. Table engine {}",
+                    table.get_table_info().name,
+                    table.engine()
                 );
-                return Ok(Some(list_files));
-            } else {
-                left_limit = Some(dry_run_limit.unwrap() - list_files.len());
+                continue;
+            };
+
+        table_vecs.push((table_info.clone(), operator));
+    }
+
+    let result = if batch_size >= table_vecs.len() {
+        do_vacuum_drop_table(table_vecs, dry_run_limit).await?
+    } else {
+        let mut chunks = table_vecs.chunks(batch_size);
+        let dry_run_limit = dry_run_limit
+            .map(|dry_run_limit| (dry_run_limit / threads_nums).min(dry_run_limit).max(1));
+        let tasks = std::iter::from_fn(move || {
+            chunks
+                .next()
+                .map(|tables| do_vacuum_drop_table(tables.to_vec(), dry_run_limit))
+        });
+
+        let result = execute_futures_in_parallel(
+            tasks,
+            threads_nums,
+            threads_nums * 2,
+            "batch-vacuum-drop-tables-worker".to_owned(),
+        )
+        .await?;
+        if dry_run_limit.is_some() {
+            let mut ret_files = vec![];
+            for file in result {
+                // return error if any errors happens during `do_vacuum_drop_table`
+                if let Some(files) = file? {
+                    ret_files.extend(files);
+                }
             }
+            Some(ret_files)
+        } else {
+            None
         }
-    }
+    };
+
     info!(
         "do_vacuum_drop_tables {} tables, cost:{} sec",
         tables_len,
         start.elapsed().as_secs()
     );
 
-    Ok(if dry_run_limit.is_some() {
-        Some(list_files)
-    } else {
-        None
-    })
+    Ok(result)
 }
diff --git a/src/query/ee/tests/it/storages/fuse/operations/vacuum.rs b/src/query/ee/tests/it/storages/fuse/operations/vacuum.rs
@@ -21,6 +21,7 @@ use databend_common_meta_app::schema::TableInfo;
 use databend_common_meta_app::schema::TableMeta;
 use databend_common_meta_app::storage::StorageParams;
 use databend_common_storage::DataOperator;
+use databend_common_storages_fuse::TableContext;
 use databend_enterprise_query::storages::fuse::do_vacuum_drop_tables;
 use databend_enterprise_query::storages::fuse::operations::vacuum_drop_tables::do_vacuum_drop_table;
 use databend_enterprise_query::storages::fuse::operations::vacuum_temporary_files::do_vacuum_temporary_files;
@@ -69,10 +70,12 @@ async fn test_fuse_do_vacuum_drop_tables() -> Result<()> {
     let tbl = fixture.default_table_name();
     let qry = format!("drop table {}.{}", db, tbl);
     fixture.execute_command(&qry).await?;
+    let ctx = fixture.new_query_ctx().await?;
+    let threads_nums = ctx.get_settings().get_max_threads()? as usize;
 
     // verify dry run never delete files
     {
-        do_vacuum_drop_tables(vec![table.clone()], Some(100)).await?;
+        do_vacuum_drop_tables(threads_nums, vec![table.clone()], Some(100)).await?;
         check_data_dir(
             &fixture,
             "test_fuse_do_vacuum_drop_table: verify generate files",
@@ -88,7 +91,7 @@ async fn test_fuse_do_vacuum_drop_tables() -> Result<()> {
     }
 
     {
-        do_vacuum_drop_tables(vec![table], None).await?;
+        do_vacuum_drop_tables(threads_nums, vec![table], None).await?;
 
         // after vacuum drop tables, verify the files number
         check_data_dir(
@@ -243,7 +246,8 @@ async fn test_fuse_do_vacuum_drop_table_deletion_error() -> Result<()> {
     let faulty_accessor = std::sync::Arc::new(AccessorFaultyDeletion::new());
     let operator = OperatorBuilder::new(faulty_accessor.clone()).finish();
 
-    let result = do_vacuum_drop_table(&table_info, &operator, None).await;
+    let tables = vec![(table_info, operator)];
+    let result = do_vacuum_drop_table(tables, None).await;
     assert!(result.is_err());
 
     // verify that accessor.delete() was called
@@ -271,7 +275,8 @@ async fn test_fuse_do_vacuum_drop_table_external_storage() -> Result<()> {
     let accessor = std::sync::Arc::new(AccessorFaultyDeletion::new());
     let operator = OperatorBuilder::new(accessor.clone()).finish();
 
-    let result = do_vacuum_drop_table(&table_info, &operator, None).await;
+    let tables = vec![(table_info, operator)];
+    let result = do_vacuum_drop_table(tables, None).await;
     assert!(result.is_err());
 
     // verify that accessor.delete() was called
diff --git a/src/query/ee_features/vacuum_handler/src/vacuum_handler.rs b/src/query/ee_features/vacuum_handler/src/vacuum_handler.rs
@@ -38,6 +38,7 @@ pub trait VacuumHandler: Sync + Send {
 
     async fn do_vacuum_drop_tables(
         &self,
+        threads_nums: usize,
         tables: Vec<Arc<dyn Table>>,
         dry_run_limit: Option<usize>,
     ) -> Result<Option<Vec<VacuumDropFileInfo>>>;
@@ -75,11 +76,12 @@ impl VacuumHandlerWrapper {
     #[async_backtrace::framed]
     pub async fn do_vacuum_drop_tables(
         &self,
+        threads_nums: usize,
         tables: Vec<Arc<dyn Table>>,
         dry_run_limit: Option<usize>,
     ) -> Result<Option<Vec<VacuumDropFileInfo>>> {
         self.handler
-            .do_vacuum_drop_tables(tables, dry_run_limit)
+            .do_vacuum_drop_tables(threads_nums, tables, dry_run_limit)
             .await
     }
 
diff --git a/src/query/service/src/interpreters/interpreter_vacuum_drop_tables.rs b/src/query/service/src/interpreters/interpreter_vacuum_drop_tables.rs
@@ -158,8 +158,10 @@ impl Interpreter for VacuumDropTablesInterpreter {
             .collect::<Vec<_>>();
 
         let handler = get_vacuum_handler();
+        let threads_nums = self.ctx.get_settings().get_max_threads()? as usize;
         let files_opt = handler
             .do_vacuum_drop_tables(
+                threads_nums,
                 tables,
                 if self.plan.option.dry_run.is_some() {
                     Some(DRY_RUN_LIMIT)
