refactor(query): enable little vacuum after query finish (#15434)

* refactor(query): enable little vacuum after query finish

* refactor(query): enable little vacuum after query finish

* refactor(query): enable little vacuum after query finish

* refactor(query): enable little vacuum after query finish

* refactor(query): enable little vacuum after query finish

* refactor(query): enable little vacuum after query finish

* refactor(query): enable little vacuum after query finish

* refactor(query): enable little vacuum after query finish

* refactor(query): enable little vacuum after query finish

* refactor(query): enable little vacuum after query finish

* refactor(query): enable little vacuum after query finish

* refactor(query): enable little vacuum after query finish

* refactor(query): enable little vacuum after query finish

Author: zhang2014

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

@@ -57,7 +57,7 @@ impl VacuumHandler for RealVacuumHandler {
         temporary_dir: String,
         retain: Option<Duration>,
         vacuum_limit: Option<usize>,
-    ) -> Result<Vec<String>> {
+    ) -> Result<usize> {
         do_vacuum_temporary_files(temporary_dir, retain, vacuum_limit).await
     }
 }

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

@@ -13,50 +13,187 @@
 // limitations under the License.
 
 use std::time::Duration;
+use std::time::Instant;
 use std::time::SystemTime;
 use std::time::UNIX_EPOCH;
 
 use databend_common_exception::Result;
 use databend_common_storage::DataOperator;
+use futures_util::stream;
 use futures_util::TryStreamExt;
+use log::info;
+use opendal::Entry;
+use opendal::EntryMode;
 use opendal::Metakey;
 
 #[async_backtrace::framed]
 pub async fn do_vacuum_temporary_files(
     temporary_dir: String,
     retain: Option<Duration>,
     limit: Option<usize>,
-) -> Result<Vec<String>> {
+) -> Result<usize> {
+    let limit = limit.unwrap_or(usize::MAX);
+    let expire_time = retain
+        .map(|x| x.as_millis())
+        .unwrap_or(1000 * 60 * 60 * 24 * 3) as i64;
+    let timestamp = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_millis() as i64;
+
     let operator = DataOperator::instance().operator();
 
+    let temporary_dir = format!("{}/", temporary_dir);
+
     let mut ds = operator
         .lister_with(&temporary_dir)
-        .recursive(true)
-        .metakey(Metakey::LastModified)
+        .metakey(Metakey::Mode | Metakey::LastModified)
         .await?;
 
-    let limit = limit.unwrap_or(usize::MAX);
-    let expire_time = retain.map(|x| x.as_millis()).unwrap_or(60 * 60 * 24 * 3) as i64;
-    let timestamp = SystemTime::now()
-        .duration_since(UNIX_EPOCH)
-        .unwrap()
-        .as_millis() as i64;
+    let mut removed_temp_files = 0;
 
-    let mut remove_temp_files_name = Vec::new();
-    while let Some(de) = ds.try_next().await? {
-        let meta = de.metadata();
+    while removed_temp_files < limit {
+        let instant = Instant::now();
+        let mut end_of_stream = true;
+        let mut remove_temp_files_path = Vec::with_capacity(1000);
 
-        if let Some(modified) = meta.last_modified() {
-            if timestamp - modified.timestamp_millis() >= expire_time {
-                operator.delete(de.path()).await?;
-                remove_temp_files_name.push(de.name().to_string());
+        while let Some(de) = ds.try_next().await? {
+            let meta = de.metadata();
+
+            match meta.mode() {
+                EntryMode::DIR => {
+                    let life_mills =
+                        match operator.is_exist(&format!("{}finished", de.path())).await? {
+                            true => 0,
+                            false => expire_time,
+                        };
+
+                    vacuum_finished_query(
+                        &mut removed_temp_files,
+                        &de,
+                        limit,
+                        timestamp,
+                        life_mills,
+                    )
+                    .await?;
+
+                    if removed_temp_files >= limit {
+                        end_of_stream = false;
+                        break;
+                    }
+                }
+                EntryMode::FILE => {
+                    if let Some(modified) = meta.last_modified() {
+                        if timestamp - modified.timestamp_millis() >= expire_time {
+                            removed_temp_files += 1;
+                            remove_temp_files_path.push(de.path().to_string());
+
+                            if removed_temp_files >= limit || remove_temp_files_path.len() >= 1000 {
+                                end_of_stream = false;
+                                break;
+                            }
+                        }
+                    }
+                }
+                EntryMode::Unknown => unreachable!(),
             }
+        }
+
+        if !remove_temp_files_path.is_empty() {
+            let cur_removed = remove_temp_files_path.len();
+            operator
+                .remove_via(stream::iter(remove_temp_files_path))
+                .await?;
+
+            info!(
+                "vacuum removed {} temp files in {:?}(elapsed: {:?})",
+                cur_removed,
+                temporary_dir,
+                instant.elapsed()
+            );
+        }
+
+        if end_of_stream {
+            break;
+        }
+    }
 
-            if remove_temp_files_name.len() >= limit {
-                break;
+    Ok(removed_temp_files)
+}
+
+async fn vacuum_finished_query(
+    removed_temp_files: &mut usize,
+    de: &Entry,
+    limit: usize,
+    timestamp: i64,
+    life_mills: i64,
+) -> Result<()> {
+    let operator = DataOperator::instance().operator();
+
+    let mut all_files_removed = true;
+    let mut ds = operator
+        .lister_with(de.path())
+        .metakey(Metakey::Mode | Metakey::LastModified)
+        .await?;
+
+    while *removed_temp_files < limit {
+        let instant = Instant::now();
+
+        let mut end_of_stream = true;
+        let mut all_each_files_removed = true;
+        let mut remove_temp_files_path = Vec::with_capacity(1001);
+
+        while let Some(de) = ds.try_next().await? {
+            let meta = de.metadata();
+            if meta.is_file() {
+                if de.name() == "finished" {
+                    continue;
+                }
+
+                if let Some(modified) = meta.last_modified() {
+                    if timestamp - modified.timestamp_millis() >= life_mills {
+                        *removed_temp_files += 1;
+                        remove_temp_files_path.push(de.path().to_string());
+
+                        if *removed_temp_files >= limit || remove_temp_files_path.len() >= 1000 {
+                            end_of_stream = false;
+                            break;
+                        }
+
+                        continue;
+                    }
+                }
             }
+
+            all_each_files_removed = false;
         }
+
+        all_files_removed &= all_each_files_removed;
+
+        if !remove_temp_files_path.is_empty() {
+            let cur_removed = remove_temp_files_path.len();
+
+            operator
+                .remove_via(stream::iter(remove_temp_files_path))
+                .await?;
+
+            info!(
+                "vacuum removed {} temp files in {:?}(elapsed: {:?})",
+                cur_removed,
+                de.path(),
+                instant.elapsed()
+            );
+        }
+
+        if end_of_stream {
+            break;
+        }
+    }
+
+    if all_files_removed {
+        operator.delete(&format!("{}finished", de.path())).await?;
+        operator.delete(de.path()).await?;
     }
 
-    Ok(remove_temp_files_name)
+    Ok(())
 }

src/query/ee/tests/it/storages/fuse/operations/vacuum.rs

@@ -117,7 +117,10 @@ async fn test_do_vacuum_temporary_files() -> Result<()> {
     operator.write("test_dir/test2", vec![1, 2]).await?;
     operator.write("test_dir/test3", vec![1, 2]).await?;
 
-    assert_eq!(3, operator.list("test_dir/").await?.len());
+    assert_eq!(
+        3,
+        operator.list_with("test_dir/").recursive(true).await?.len()
+    );
 
     tokio::time::sleep(Duration::from_secs(2)).await;
     do_vacuum_temporary_files(
@@ -129,8 +132,23 @@ async fn test_do_vacuum_temporary_files() -> Result<()> {
 
     assert_eq!(2, operator.list("test_dir/").await?.len());
 
-    do_vacuum_temporary_files("test_dir/".to_string(), Some(Duration::from_secs(2)), None).await?;
-    assert_eq!(0, operator.list("test_dir/").await?.len());
+    operator.write("test_dir/test4/test4", vec![1, 2]).await?;
+    operator.write("test_dir/test5/test5", vec![1, 2]).await?;
+    operator
+        .write("test_dir/test5/finished", vec![1, 2])
+        .await?;
+
+    do_vacuum_temporary_files(
+        "test_dir/".to_string(),
+        Some(Duration::from_secs(2)),
+        Some(2),
+    )
+    .await?;
+    assert_eq!(operator.list("test_dir/").await?.len(), 2);
+
+    tokio::time::sleep(Duration::from_secs(3)).await;
+    do_vacuum_temporary_files("test_dir/".to_string(), Some(Duration::from_secs(3)), None).await?;
+    assert!(operator.list_with("test_dir/").await?.is_empty());
 
     Ok(())
 }

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

@@ -47,7 +47,7 @@ pub trait VacuumHandler: Sync + Send {
         temporary_dir: String,
         retain: Option<Duration>,
         vacuum_limit: Option<usize>,
-    ) -> Result<Vec<String>>;
+    ) -> Result<usize>;
 }
 
 pub struct VacuumHandlerWrapper {
@@ -89,7 +89,7 @@ impl VacuumHandlerWrapper {
         temporary_dir: String,
         retain: Option<Duration>,
         vacuum_limit: Option<usize>,
-    ) -> Result<Vec<String>> {
+    ) -> Result<usize> {
         self.handler
             .do_vacuum_temporary_files(temporary_dir, retain, vacuum_limit)
             .await

src/query/pipeline/core/src/pipeline.rs

@@ -533,6 +533,9 @@ impl Drop for Pipeline {
     }
 }
 
-pub fn query_spill_prefix(tenant: &str) -> String {
-    format!("_query_spill/{}", tenant)
+pub fn query_spill_prefix(tenant: &str, query_id: &str) -> String {
+    match query_id.is_empty() {
+        true => format!("_query_spill/{}", tenant),
+        false => format!("_query_spill/{}/{}", tenant, query_id),
+    }
 }

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

@@ -14,6 +14,7 @@
 
 pub(crate) mod compact_hook;
 pub(crate) mod refresh_hook;
+pub(crate) mod vacuum_hook;
 
 #[allow(clippy::module_inception)]
 mod hook;

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

@@ -0,0 +1,69 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use databend_common_base::runtime::GlobalIORuntime;
+use databend_common_catalog::table_context::TableContext;
+use databend_common_exception::Result;
+use databend_common_license::license::Feature::Vacuum;
+use databend_common_license::license_manager::get_license_manager;
+use databend_common_pipeline_core::query_spill_prefix;
+use databend_common_storage::DataOperator;
+use databend_enterprise_vacuum_handler::get_vacuum_handler;
+
+use crate::sessions::QueryContext;
+
+pub fn hook_vacuum_temp_files(query_ctx: &Arc<QueryContext>) -> Result<()> {
+    let tenant = query_ctx.get_tenant();
+    let settings = query_ctx.get_settings();
+    let spill_prefix = query_spill_prefix(tenant.tenant_name(), &query_ctx.get_id());
+    let license_manager = get_license_manager();
+
+    if license_manager
+        .manager
+        .check_enterprise_enabled(query_ctx.get_license_key(), Vacuum)
+        .is_ok()
+    {
+        let handler = get_vacuum_handler();
+
+        let _ = GlobalIORuntime::instance().block_on(async move {
+            let vacuum_limit = match settings.get_max_vacuum_temp_files_after_query()? {
+                0 => None,
+                v => Some(v as usize),
+            };
+            let removed_files = handler
+                .do_vacuum_temporary_files(
+                    spill_prefix.clone(),
+                    Some(Duration::from_secs(0)),
+                    vacuum_limit,
+                )
+                .await;
+
+            if (removed_files.is_ok() && vacuum_limit.is_none())
+                && !matches!(removed_files, Ok(res) if Some(res) != vacuum_limit)
+            {
+                let op = DataOperator::instance().operator();
+                op.create_dir(&format!("{}/", spill_prefix)).await?;
+                op.write(&format!("{}/finished", spill_prefix), vec![])
+                    .await?;
+            }
+
+            Ok(())
+        });
+    }
+
+    Ok(())
+}

src/query/service/src/interpreters/interpreter.rs

@@ -32,6 +32,7 @@ use databend_common_sql::Planner;
 use log::error;
 use log::info;
 
+use crate::interpreters::hook::vacuum_hook::hook_vacuum_temp_files;
 use crate::interpreters::interpreter_txn_commit::CommitInterpreter;
 use crate::interpreters::InterpreterMetrics;
 use crate::interpreters::InterpreterQueryLog;
@@ -132,6 +133,8 @@ pub trait Interpreter: Sync + Send {
                     }
                 }
 
+                hook_vacuum_temp_files(&query_ctx)?;
+
                 let err_opt = match may_error {
                     Ok(_) => None,
                     Err(e) => Some(e.clone()),