fix(query): fix thread leak and remove all usage (#17081)

* fix(query): fix thread leak and remove all usage * fix(query): fix thread leak and remove all usage * fix(query): fix thread leak and remove all usage
databendlabs · Dec 19, 2024 · 327cd75 · 327cd75
1 parent e131edf
commit 327cd75
Show file tree

Hide file tree

Showing 6 changed files with 27 additions and 29 deletions.
diff --git a/src/query/catalog/src/table_context.rs b/src/query/catalog/src/table_context.rs
@@ -24,7 +24,6 @@ use std::time::SystemTime;
 use dashmap::DashMap;
 use databend_common_base::base::Progress;
 use databend_common_base::base::ProgressValues;
-use databend_common_base::runtime::Runtime;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_exception::ResultExt;
@@ -374,8 +373,6 @@ pub trait TableContext: Send + Sync {
     fn is_temp_table(&self, catalog_name: &str, database_name: &str, table_name: &str) -> bool;
     fn get_shared_settings(&self) -> Arc<Settings>;
 
-    fn get_runtime(&self) -> Result<Arc<Runtime>>;
-
     fn add_m_cte_temp_table(&self, database_name: &str, table_name: &str);
 
     async fn drop_m_cte_temp_table(&self) -> Result<()>;

diff --git a/src/query/service/src/sessions/query_ctx.rs b/src/query/service/src/sessions/query_ctx.rs
@@ -35,7 +35,6 @@ use databend_common_base::base::Progress;
 use databend_common_base::base::ProgressValues;
 use databend_common_base::runtime::profile::Profile;
 use databend_common_base::runtime::profile::ProfileStatisticsName;
-use databend_common_base::runtime::Runtime;
 use databend_common_base::runtime::TrySpawn;
 use databend_common_base::JoinHandle;
 use databend_common_catalog::catalog::CATALOG_DEFAULT;
@@ -1476,10 +1475,6 @@ impl TableContext for QueryContext {
                 .is_temp_table(database_name, table_name)
     }
 
-    fn get_runtime(&self) -> Result<Arc<Runtime>> {
-        self.shared.try_get_runtime()
-    }
-
     fn add_m_cte_temp_table(&self, database_name: &str, table_name: &str) {
         self.m_cte_temp_table
             .write()

diff --git a/src/query/service/tests/it/sql/exec/get_table_bind_test.rs b/src/query/service/tests/it/sql/exec/get_table_bind_test.rs
@@ -22,7 +22,6 @@ use dashmap::DashMap;
 use databend_common_base::base::tokio;
 use databend_common_base::base::Progress;
 use databend_common_base::base::ProgressValues;
-use databend_common_base::runtime::Runtime;
 use databend_common_catalog::catalog::Catalog;
 use databend_common_catalog::cluster_info::Cluster;
 use databend_common_catalog::database::Database;
@@ -997,10 +996,6 @@ impl TableContext for CtxDelegation {
         false
     }
 
-    fn get_runtime(&self) -> Result<Arc<Runtime>> {
-        todo!()
-    }
-
     fn add_m_cte_temp_table(&self, _database_name: &str, _table_name: &str) {
         todo!()
     }

diff --git a/src/query/service/tests/it/storages/fuse/operations/commit.rs b/src/query/service/tests/it/storages/fuse/operations/commit.rs
@@ -21,7 +21,6 @@ use dashmap::DashMap;
 use databend_common_base::base::tokio;
 use databend_common_base::base::Progress;
 use databend_common_base::base::ProgressValues;
-use databend_common_base::runtime::Runtime;
 use databend_common_catalog::catalog::Catalog;
 use databend_common_catalog::cluster_info::Cluster;
 use databend_common_catalog::database::Database;
@@ -876,11 +875,6 @@ impl TableContext for CtxDelegation {
     fn is_temp_table(&self, _catalog_name: &str, _database_name: &str, _table_name: &str) -> bool {
         false
     }
-
-    fn get_runtime(&self) -> Result<Arc<Runtime>> {
-        todo!()
-    }
-
     fn add_m_cte_temp_table(&self, _database_name: &str, _table_name: &str) {
         todo!()
     }

diff --git a/src/query/service/tests/it/storages/fuse/pruning_pipeline.rs b/src/query/service/tests/it/storages/fuse/pruning_pipeline.rs
@@ -16,6 +16,8 @@ use std::sync::Arc;
 
 use databend_common_ast::ast::Engine;
 use databend_common_base::base::tokio;
+use databend_common_base::runtime::GlobalIORuntime;
+use databend_common_base::runtime::Runtime;
 use databend_common_base::runtime::TrySpawn;
 use databend_common_catalog::plan::PartInfoPtr;
 use databend_common_catalog::plan::PartStatistics;
@@ -96,17 +98,25 @@ async fn apply_block_pruning(
     )?;
     prune_pipeline.set_max_threads(1);
     prune_pipeline.set_on_init(move || {
-        ctx.get_runtime()?.try_spawn(
-            async move {
+        // We cannot use the runtime associated with the query to avoid increasing its lifetime.
+        GlobalIORuntime::instance().spawn(async move {
+            // avoid block global io runtime
+            let runtime = Runtime::with_worker_threads(2, None)?;
+
+            let join_handler = runtime.spawn(async move {
                 let segment_pruned_result =
                     fuse_pruner.clone().segment_pruning(segment_locs).await?;
                 for segment in segment_pruned_result {
                     let _ = segment_tx.send(Ok(segment)).await;
                 }
                 Ok::<_, ErrorCode>(())
-            },
-            None,
-        )?;
+            });
+            join_handler
+                .await
+                .unwrap()
+                .expect("Join error while in prune pipeline");
+            Ok::<_, ErrorCode>(())
+        });
         Ok(())
     });
 

diff --git a/src/query/storages/fuse/src/operations/read_partitions.rs b/src/query/storages/fuse/src/operations/read_partitions.rs
@@ -238,8 +238,11 @@ impl FuseTable {
             derterministic_cache_key.clone(),
         )?;
         prune_pipeline.set_on_init(move || {
-            ctx.get_runtime()?.try_spawn(
-                async move {
+            // We cannot use the runtime associated with the query to avoid increasing its lifetime.
+            GlobalIORuntime::instance().spawn(async move {
+                // avoid block global io runtime
+                let runtime = Runtime::with_worker_threads(2, None)?;
+                let join_handler = runtime.spawn(async move {
                     let segment_pruned_result =
                         pruner.clone().segment_pruning(lazy_init_segments).await?;
                     for segment in segment_pruned_result {
@@ -249,9 +252,13 @@ impl FuseTable {
                         }
                     }
                     Ok::<_, ErrorCode>(())
-                },
-                None,
-            )?;
+                });
+
+                if let Err(cause) = join_handler.await {
+                    log::warn!("Join error while in prune pipeline, cause: {:?}", cause);
+                }
+                Ok::<_, ErrorCode>(())
+            });
             Ok(())
         });