Merge pull request #10501 from dantengsky/fix-thread-leakings

fix: thread leaking

Author: BohuTANG

src/common/base/src/runtime/global_runtime.rs

@@ -21,6 +21,15 @@ use crate::runtime::Runtime;
 
 pub struct GlobalIORuntime;
 
+pub struct GlobalQueryRuntime(pub Runtime);
+
+impl GlobalQueryRuntime {
+    #[inline(always)]
+    pub fn runtime<'a>(self: &'a Arc<Self>) -> &'a Runtime {
+        &self.0
+    }
+}
+
 impl GlobalIORuntime {
     pub fn init(num_cpus: usize) -> Result<()> {
         let thread_num = std::cmp::max(num_cpus, num_cpus::get() / 2);
@@ -37,3 +46,18 @@ impl GlobalIORuntime {
         GlobalInstance::get()
     }
 }
+
+impl GlobalQueryRuntime {
+    pub fn init(num_cpus: usize) -> Result<()> {
+        let thread_num = std::cmp::max(num_cpus, num_cpus::get() / 2);
+        let thread_num = std::cmp::max(2, thread_num);
+
+        let rt = Runtime::with_worker_threads(thread_num, Some("g-query-worker".to_owned()))?;
+        GlobalInstance::set(Arc::new(GlobalQueryRuntime(rt)));
+        Ok(())
+    }
+
+    pub fn instance() -> Arc<GlobalQueryRuntime> {
+        GlobalInstance::get()
+    }
+}

src/common/base/src/runtime/mod.rs

@@ -23,6 +23,7 @@ mod thread_pool;
 pub use catch_unwind::catch_unwind;
 pub use catch_unwind::CatchUnwindFuture;
 pub use global_runtime::GlobalIORuntime;
+pub use global_runtime::GlobalQueryRuntime;
 pub use runtime::execute_futures_in_parallel;
 pub use runtime::match_join_handle;
 pub use runtime::Dropper;

src/query/service/src/api/rpc/exchange/statistics_sender.rs

@@ -18,6 +18,7 @@ use std::sync::Arc;
 
 use async_channel::Receiver;
 use async_channel::Sender;
+use common_base::runtime::GlobalQueryRuntime;
 use common_base::runtime::TrySpawn;
 use common_catalog::table_context::TableContext;
 use common_exception::ErrorCode;
@@ -63,8 +64,8 @@ impl StatisticsSender {
         let shutdown_flag = self.shutdown_flag.clone();
         let shutdown_flag_receiver = self.shutdown_flag_receiver.clone();
 
-        let spawner = ctx.clone();
-        spawner.spawn(async move {
+        let spawner = GlobalQueryRuntime::instance();
+        spawner.runtime().spawn(async move {
             let mut recv = Box::pin(flight_exchange.recv());
             let mut notified = Box::pin(shutdown_flag_receiver.recv());
 

src/query/service/src/global_services.rs

@@ -14,6 +14,7 @@
 
 use common_base::base::GlobalInstance;
 use common_base::runtime::GlobalIORuntime;
+use common_base::runtime::GlobalQueryRuntime;
 use common_catalog::catalog::CatalogManager;
 use common_config::GlobalConfig;
 use common_config::InnerConfig;
@@ -47,6 +48,7 @@ impl GlobalServices {
 
         QueryLogger::init(app_name_shuffle, &config.log)?;
         GlobalIORuntime::init(config.storage.num_cpus as usize)?;
+        GlobalQueryRuntime::init(config.storage.num_cpus as usize)?;
 
         // Cluster discovery.
         ClusterDiscovery::init(config.clone()).await?;

src/query/service/src/servers/http/v1/query/http_query.rs

@@ -20,6 +20,7 @@ use std::time::Instant;
 use common_base::base::tokio;
 use common_base::base::tokio::sync::Mutex as TokioMutex;
 use common_base::base::tokio::sync::RwLock;
+use common_base::runtime::GlobalQueryRuntime;
 use common_base::runtime::TrySpawn;
 use common_catalog::table_context::StageAttachment;
 use common_exception::ErrorCode;
@@ -278,28 +279,37 @@ impl HttpQuery {
         let query_id_clone = id.clone();
 
         let schema = ExecuteState::get_schema(&sql, ctx.clone()).await?;
-        ctx.try_spawn(async move {
-            let state = state_clone.clone();
-            if let Err(e) =
-                ExecuteState::try_start_query(state, &sql, session, ctx_clone.clone(), block_sender)
-                    .await
-            {
-                InterpreterQueryLog::fail_to_start(ctx_clone.clone(), e.clone());
-                let state = ExecuteStopped {
-                    stats: Progresses::default(),
-                    reason: Err(e.clone()),
-                    stop_time: Instant::now(),
-                    affect: ctx_clone.get_affect(),
-                };
-                tracing::info!(
-                    "http query {}, change state to Stopped, fail to start {:?}",
-                    &query_id,
-                    e
-                );
-                Executor::start_to_stop(&state_clone, ExecuteState::Stopped(Box::new(state))).await;
-                block_sender_closer.close();
-            }
-        })?;
+        let http_query_runtime_instance = GlobalQueryRuntime::instance();
+        http_query_runtime_instance
+            .runtime()
+            .try_spawn(async move {
+                let state = state_clone.clone();
+                if let Err(e) = ExecuteState::try_start_query(
+                    state,
+                    &sql,
+                    session,
+                    ctx_clone.clone(),
+                    block_sender,
+                )
+                .await
+                {
+                    InterpreterQueryLog::fail_to_start(ctx_clone.clone(), e.clone());
+                    let state = ExecuteStopped {
+                        stats: Progresses::default(),
+                        reason: Err(e.clone()),
+                        stop_time: Instant::now(),
+                        affect: ctx_clone.get_affect(),
+                    };
+                    tracing::info!(
+                        "http query {}, change state to Stopped, fail to start {:?}",
+                        &query_id,
+                        e
+                    );
+                    Executor::start_to_stop(&state_clone, ExecuteState::Stopped(Box::new(state)))
+                        .await;
+                    block_sender_closer.close();
+                }
+            })?;
 
         let format_settings = ctx.get_format_settings()?;
         let data = Arc::new(TokioMutex::new(PageManager::new(
@@ -319,6 +329,7 @@ impl HttpQuery {
             config,
             expire_state: Arc::new(TokioMutex::new(ExpireState::Working)),
         };
+
         Ok(Arc::new(query))
     }
 

src/query/service/src/sessions/query_ctx_shared.rs

@@ -270,6 +270,11 @@ impl QueryContextShared {
         }
     }
 
+    pub fn get_runtime(&self) -> Option<Arc<Runtime>> {
+        let query_runtime = self.runtime.read();
+        (*query_runtime).clone()
+    }
+
     pub fn attach_query_str(&self, kind: String, query: String) {
         {
             let mut running_query = self.running_query.write();

src/query/service/src/sessions/session_info.rs

@@ -34,7 +34,7 @@ impl Session {
 
         let shared_query_context = &session_ctx.get_query_context_shared();
         if let Some(shared) = shared_query_context {
-            if let Ok(runtime) = shared.try_get_runtime() {
+            if let Some(runtime) = shared.get_runtime() {
                 let mem_stat = runtime.get_tracker();
                 memory_usage = mem_stat.get_memory_usage();
             }