chore(query): improve hook vacuum temp after query (#17125)

* chore(query): improve ctx lock

* chore(query): improve vacuum

* send spill nums via progress in cluster

* send spill nums via progress in cluster

Author: sundy-li

src/query/ee/src/storages/fuse/operations/vacuum_temporary_files.rs

@@ -45,8 +45,15 @@ pub async fn do_vacuum_temporary_files(
     }
 
     match options {
-        VacuumTempOptions::QueryHook(nodes_num, query_id) => {
-            vacuum_query_hook(abort_checker, &temporary_dir, *nodes_num, query_id, limit).await
+        VacuumTempOptions::QueryHook(nodes, query_id) => {
+            vacuum_query_hook(
+                abort_checker,
+                &temporary_dir,
+                nodes.as_slice(),
+                query_id,
+                limit,
+            )
+            .await
         }
         VacuumTempOptions::VacuumCommand(duration) => {
             vacuum_by_duration(abort_checker, &temporary_dir, limit, duration).await
@@ -172,41 +179,53 @@ async fn vacuum_by_duration(
 async fn vacuum_query_hook(
     abort_checker: AbortChecker,
     temporary_dir: &str,
-    nodes_num: usize,
+    nodes: &[usize],
     query_id: &str,
     mut limit: usize,
 ) -> Result<usize> {
     let mut removed_total = 0;
+    let metas_f = nodes
+        .iter()
+        .map(|i| async move {
+            let operator = DataOperator::instance().operator();
+            let meta_file_path =
+                format!("{}/{}_{}{}", temporary_dir, query_id, i, SPILL_META_SUFFIX);
+            let buffer = operator.read(&meta_file_path).await?;
+            std::result::Result::<(String, Buffer), opendal::Error>::Ok((meta_file_path, buffer))
+        })
+        .collect::<Vec<_>>();
 
-    for i in 0..nodes_num {
-        if limit == 0 {
-            break;
-        }
+    let metas = futures_util::future::join_all(metas_f)
+        .await
+        .into_iter()
+        .filter_map(|x| x.is_ok().then(|| x.unwrap()));
+
+    for (meta_file_path, buffer) in metas {
         abort_checker.try_check_aborting()?;
-        let meta_file_path = format!("{}/{}_{}{}", temporary_dir, query_id, i, SPILL_META_SUFFIX);
-        let removed =
-            vacuum_by_meta(temporary_dir, &meta_file_path, limit, &mut removed_total).await?;
+        let removed = vacuum_by_meta_buffer(
+            &meta_file_path,
+            temporary_dir,
+            buffer,
+            limit,
+            &mut removed_total,
+        )
+        .await?;
         limit = limit.saturating_sub(removed);
     }
+
     Ok(removed_total)
 }
 
-async fn vacuum_by_meta(
-    temporary_dir: &str,
+async fn vacuum_by_meta_buffer(
     meta_file_path: &str,
+    temporary_dir: &str,
+    meta: Buffer,
     limit: usize,
     removed_total: &mut usize,
 ) -> Result<usize> {
     let operator = DataOperator::instance().operator();
-    let meta: Buffer;
-    let r = operator.read(meta_file_path).await;
-    match r {
-        Ok(r) => meta = r,
-        Err(e) if e.kind() == ErrorKind::NotFound => return Ok(0),
-        Err(e) => return Err(e.into()),
-    }
-    let meta = meta.to_bytes();
     let start_time = Instant::now();
+    let meta = meta.to_bytes();
     let files: Vec<String> = meta.lines().map(|x| Ok(x?)).collect::<Result<Vec<_>>>()?;
 
     let (to_be_removed, remain) = files.split_at(limit.min(files.len()));
@@ -240,6 +259,23 @@ async fn vacuum_by_meta(
     Ok(cur_removed)
 }
 
+async fn vacuum_by_meta(
+    temporary_dir: &str,
+    meta_file_path: &str,
+    limit: usize,
+    removed_total: &mut usize,
+) -> Result<usize> {
+    let operator = DataOperator::instance().operator();
+    let meta: Buffer;
+    let r = operator.read(meta_file_path).await;
+    match r {
+        Ok(r) => meta = r,
+        Err(e) if e.kind() == ErrorKind::NotFound => return Ok(0),
+        Err(e) => return Err(e.into()),
+    }
+    vacuum_by_meta_buffer(meta_file_path, temporary_dir, meta, limit, removed_total).await
+}
+
 async fn vacuum_by_list_dir(
     dir_path: &str,
     limit: usize,

src/query/ee_features/vacuum_handler/src/vacuum_handler.rs

@@ -58,8 +58,8 @@ pub trait VacuumHandler: Sync + Send {
 
 #[derive(Debug, Clone)]
 pub enum VacuumTempOptions {
-    // nodes_num, query_id
-    QueryHook(usize, String),
+    // nodes, query_id
+    QueryHook(Vec<usize>, String),
     VacuumCommand(Option<Duration>),
 }
 

src/query/service/src/clusters/cluster.rs

@@ -80,6 +80,7 @@ pub trait ClusterHelper {
     fn is_local(&self, node: &NodeInfo) -> bool;
     fn local_id(&self) -> String;
     fn ordered_index(&self) -> usize;
+    fn index_of_nodeid(&self, node_id: &str) -> Option<usize>;
 
     fn get_nodes(&self) -> Vec<Arc<NodeInfo>>;
 
@@ -125,6 +126,12 @@ impl ClusterHelper for Cluster {
             .unwrap_or(0)
     }
 
+    fn index_of_nodeid(&self, node_id: &str) -> Option<usize> {
+        let mut nodes = self.get_nodes();
+        nodes.sort_by(|a, b| a.id.cmp(&b.id));
+        nodes.iter().position(|x| x.id == node_id)
+    }
+
     fn get_nodes(&self) -> Vec<Arc<NodeInfo>> {
         self.nodes.to_vec()
     }

src/query/service/src/interpreters/hook/vacuum_hook.rs

@@ -12,19 +12,19 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use databend_common_base::runtime::GlobalIORuntime;
 use databend_common_catalog::table_context::TableContext;
+use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_license::license::Feature::Vacuum;
 use databend_common_license::license_manager::LicenseManagerSwitch;
-use databend_common_storage::DataOperator;
 use databend_enterprise_vacuum_handler::get_vacuum_handler;
 use databend_enterprise_vacuum_handler::vacuum_handler::VacuumTempOptions;
 use databend_storages_common_cache::TempDirManager;
 use log::warn;
-use opendal::Buffer;
 use rand::Rng;
 
 use crate::clusters::ClusterHelper;
@@ -43,16 +43,35 @@ pub fn hook_vacuum_temp_files(query_ctx: &Arc<QueryContext>) -> Result<()> {
     {
         let handler = get_vacuum_handler();
 
-        let cluster_nodes = query_ctx.get_cluster().get_nodes().len();
+        let cluster = query_ctx.get_cluster();
         let query_id = query_ctx.get_id();
 
+        let mut node_files = HashMap::new();
+        for node in cluster.nodes.iter() {
+            let num = query_ctx.get_spill_file_nums(Some(node.id.clone()));
+            if num != 0 {
+                if let Some(index) = cluster.index_of_nodeid(&node.id) {
+                    node_files.insert(index, num);
+                }
+            }
+        }
+        if node_files.is_empty() {
+            return Ok(());
+        }
+
+        log::info!(
+            "Vacuum temporary files by hook, node files: {:?}",
+            node_files
+        );
+
+        let nodes = node_files.keys().cloned().collect::<Vec<usize>>();
         let abort_checker = query_ctx.clone().get_abort_checker();
-        let _ = GlobalIORuntime::instance().block_on(async move {
+        let _ = GlobalIORuntime::instance().block_on::<(), ErrorCode, _>(async move {
             let removed_files = handler
                 .do_vacuum_temporary_files(
                     abort_checker,
                     spill_prefix.clone(),
-                    &VacuumTempOptions::QueryHook(cluster_nodes, query_id),
+                    &VacuumTempOptions::QueryHook(nodes, query_id),
                     vacuum_limit as usize,
                 )
                 .await;
@@ -61,15 +80,6 @@ pub fn hook_vacuum_temp_files(query_ctx: &Arc<QueryContext>) -> Result<()> {
                 log::warn!("Vacuum temporary files has error: {:?}", cause);
             }
 
-            if vacuum_limit != 0 && matches!(removed_files, Ok(res) if res == vacuum_limit as usize)
-            {
-                // Have not been removed files
-                let op = DataOperator::instance().operator();
-                op.create_dir(&format!("{}/", spill_prefix)).await?;
-                op.write(&format!("{}/finished", spill_prefix), Buffer::new())
-                    .await?;
-            }
-
             Ok(())
         });
     }

src/query/service/src/servers/flight/v1/exchange/statistics_receiver.rs

@@ -44,7 +44,7 @@ impl StatisticsReceiver {
         let mut exchange_handler = Vec::with_capacity(statistics_exchanges.len());
         let runtime = Runtime::with_worker_threads(2, Some(String::from("StatisticsReceiver")))?;
 
-        for (_source, exchange) in statistics_exchanges.into_iter() {
+        for (source_target, exchange) in statistics_exchanges.into_iter() {
             let rx = exchange.convert_to_receiver();
             exchange_handler.push(runtime.spawn({
                 let ctx = ctx.clone();
@@ -62,7 +62,11 @@ impl StatisticsReceiver {
                                 return Ok(());
                             }
                             Either::Left((Ok(false), recv)) => {
-                                match StatisticsReceiver::recv_data(&ctx, recv.await) {
+                                match StatisticsReceiver::recv_data(
+                                    &ctx,
+                                    &source_target,
+                                    recv.await,
+                                ) {
                                     Ok(true) => {
                                         return Ok(());
                                     }
@@ -71,7 +75,11 @@ impl StatisticsReceiver {
                                         return Err(cause);
                                     }
                                     _ => loop {
-                                        match StatisticsReceiver::recv_data(&ctx, rx.recv().await) {
+                                        match StatisticsReceiver::recv_data(
+                                            &ctx,
+                                            &source_target,
+                                            rx.recv().await,
+                                        ) {
                                             Ok(true) => {
                                                 return Ok(());
                                             }
@@ -86,7 +94,7 @@ impl StatisticsReceiver {
                                 }
                             }
                             Either::Right((res, left)) => {
-                                match StatisticsReceiver::recv_data(&ctx, res) {
+                                match StatisticsReceiver::recv_data(&ctx, &source_target, res) {
                                     Ok(true) => {
                                         return Ok(());
                                     }
@@ -113,7 +121,11 @@ impl StatisticsReceiver {
         })
     }
 
-    fn recv_data(ctx: &Arc<QueryContext>, recv_data: Result<Option<DataPacket>>) -> Result<bool> {
+    fn recv_data(
+        ctx: &Arc<QueryContext>,
+        source_target: &str,
+        recv_data: Result<Option<DataPacket>>,
+    ) -> Result<bool> {
         match recv_data {
             Ok(None) => Ok(true),
             Err(transport_error) => Err(transport_error),
@@ -122,7 +134,7 @@ impl StatisticsReceiver {
             Ok(Some(DataPacket::FragmentData(_))) => unreachable!(),
             Ok(Some(DataPacket::SerializeProgress(progress))) => {
                 for progress_info in progress {
-                    progress_info.inc(ctx);
+                    progress_info.inc(source_target, ctx);
                 }
 
                 Ok(false)

src/query/service/src/servers/flight/v1/exchange/statistics_sender.rs

@@ -230,6 +230,10 @@ impl StatisticsSender {
             progress_info.push(ProgressInfo::ResultProgress(result_progress_values));
         }
 
+        let spill_file_nums = ctx.get_spill_file_nums(None);
+        if spill_file_nums != 0 {
+            progress_info.push(ProgressInfo::SpillTotalFileNums(spill_file_nums));
+        }
         progress_info
     }
 

src/query/service/src/servers/flight/v1/packets/packet_data_progressinfo.rs

@@ -33,14 +33,18 @@ pub enum ProgressInfo {
     ScanProgress(ProgressValues),
     WriteProgress(ProgressValues),
     ResultProgress(ProgressValues),
+    SpillTotalFileNums(usize),
 }
 
 impl ProgressInfo {
-    pub fn inc(&self, ctx: &Arc<QueryContext>) {
+    pub fn inc(&self, source_target: &str, ctx: &Arc<QueryContext>) {
         match self {
             ProgressInfo::ScanProgress(values) => ctx.get_scan_progress().incr(values),
             ProgressInfo::WriteProgress(values) => ctx.get_write_progress().incr(values),
             ProgressInfo::ResultProgress(values) => ctx.get_result_progress().incr(values),
+            ProgressInfo::SpillTotalFileNums(values) => {
+                ctx.set_cluster_spill_file_nums(source_target, *values)
+            }
         };
     }
 
@@ -49,6 +53,11 @@ impl ProgressInfo {
             ProgressInfo::ScanProgress(values) => (1_u8, values),
             ProgressInfo::WriteProgress(values) => (2_u8, values),
             ProgressInfo::ResultProgress(values) => (3_u8, values),
+            ProgressInfo::SpillTotalFileNums(values) => {
+                bytes.write_u8(4)?;
+                bytes.write_u64::<BigEndian>(values as u64)?;
+                return Ok(());
+            }
         };
 
         bytes.write_u8(info_type)?;
@@ -59,6 +68,12 @@ impl ProgressInfo {
 
     pub fn read<T: Read>(bytes: &mut T) -> Result<ProgressInfo> {
         let info_type = bytes.read_u8()?;
+
+        if info_type == 4 {
+            let values = bytes.read_u64::<BigEndian>()? as usize;
+            return Ok(ProgressInfo::SpillTotalFileNums(values));
+        }
+
         let rows = bytes.read_u64::<BigEndian>()? as usize;
         let bytes = bytes.read_u64::<BigEndian>()? as usize;