refactor(query): isolate http timeout tasks (#19838)

Author: youngsofun

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

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

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

@@ -39,6 +39,7 @@ pub use defer::defer;
 pub use executor_stats::ExecutorStats;
 pub use executor_stats::ExecutorStatsSlot;
 pub use executor_stats::ExecutorStatsSnapshot;
+pub use global_runtime::GlobalControlRuntime;
 pub use global_runtime::GlobalIORuntime;
 pub use global_runtime::GlobalQueryRuntime;
 pub use memory::GLOBAL_MEM_STAT;

src/query/service/src/global_services.rs

@@ -18,6 +18,7 @@ use std::sync::Arc;
 use databend_common_base::base::BuildInfoRef;
 use databend_common_base::base::GlobalInstance;
 use databend_common_base::runtime::GLOBAL_QUERIES_MANAGER;
+use databend_common_base::runtime::GlobalControlRuntime;
 use databend_common_base::runtime::GlobalIORuntime;
 use databend_common_base::runtime::GlobalQueryRuntime;
 use databend_common_catalog::catalog::CatalogCreator;
@@ -110,6 +111,7 @@ impl GlobalServices {
 
         // 3. runtime init.
         GlobalIORuntime::init(config.storage.num_cpus as usize)?;
+        GlobalControlRuntime::init()?;
         GlobalQueryRuntime::init(config.storage.num_cpus as usize)?;
 
         // 4. cluster discovery init.

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

@@ -24,7 +24,7 @@ use std::time::SystemTime;
 use chrono::SecondsFormat;
 use databend_common_base::JoinHandle;
 use databend_common_base::base::GlobalInstance;
-use databend_common_base::runtime::GlobalIORuntime;
+use databend_common_base::runtime::GlobalControlRuntime;
 use databend_common_config::InnerConfig;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
@@ -161,34 +161,36 @@ impl HttpQueryManager {
         // it may cannot destroy with final or kill when we hold ref of Arc<HttpQuery>
         let http_query_weak = Arc::downgrade(&query);
 
-        GlobalIORuntime::instance().spawn(async move {
-            loop {
-                let expire_res = match http_query_weak.upgrade() {
-                    None => {
-                        break;
-                    }
-                    Some(query) => query.check_timeout().await,
-                };
+        GlobalControlRuntime::instance()
+            .runtime()
+            .spawn(async move {
+                loop {
+                    let expire_res = match http_query_weak.upgrade() {
+                        None => {
+                            break;
+                        }
+                        Some(query) => query.check_timeout().await,
+                    };
 
-                match expire_res {
-                    TimeoutResult::TimedOut => {
-                        _ = self_clone
-                            .close_query(&query_id_clone, CloseReason::TimedOut, &None, false)
-                            .await
-                            .ok();
-                    }
-                    TimeoutResult::Sleep(t) => {
-                        sleep(t).await;
-                    }
-                    TimeoutResult::Remove => {
-                        let mut queries = self_clone.queries.write();
-                        queries.remove(&query_id_clone);
-                        log::info!("Query {query_id_clone} removed");
-                        break;
+                    match expire_res {
+                        TimeoutResult::TimedOut => {
+                            _ = self_clone
+                                .close_query(&query_id_clone, CloseReason::TimedOut, &None, false)
+                                .await
+                                .ok();
+                        }
+                        TimeoutResult::Sleep(t) => {
+                            sleep(t).await;
+                        }
+                        TimeoutResult::Remove => {
+                            let mut queries = self_clone.queries.write();
+                            queries.remove(&query_id_clone);
+                            log::info!("Query {query_id_clone} removed");
+                            break;
+                        }
                     }
                 }
-            }
-        });
+            });
         Ok(query)
     }
 
@@ -230,16 +232,18 @@ impl HttpQueryManager {
         let deleter = {
             let self_clone = self.clone();
             let last_query_id_clone = last_query_id.clone();
-            GlobalIORuntime::instance().spawn(async move {
-                sleep(Duration::from_secs(timeout_secs)).await;
-                if self_clone.get_txn(&last_query_id_clone).is_some() {
-                    log::info!(
-                        "Transaction timed out after {} seconds, last_query_id = {}",
-                        timeout_secs,
-                        last_query_id_clone
-                    );
-                }
-            })
+            GlobalControlRuntime::instance()
+                .runtime()
+                .spawn(async move {
+                    sleep(Duration::from_secs(timeout_secs)).await;
+                    if self_clone.get_txn(&last_query_id_clone).is_some() {
+                        log::info!(
+                            "Transaction timed out after {} seconds, last_query_id = {}",
+                            timeout_secs,
+                            last_query_id_clone
+                        );
+                    }
+                })
         };
         txn_managers.insert(last_query_id, (txn_mgr, deleter));
     }

src/query/service/src/servers/http/v1/session/client_session_manager.rs

@@ -19,7 +19,7 @@ use std::time::Duration;
 use std::time::SystemTime;
 
 use databend_common_base::base::GlobalInstance;
-use databend_common_base::runtime::GlobalIORuntime;
+use databend_common_base::runtime::GlobalControlRuntime;
 use databend_common_cache::Cache;
 use databend_common_cache::LruCache;
 use databend_common_config::InnerConfig;
@@ -121,7 +121,9 @@ impl ClientSessionManager {
         });
         GlobalInstance::set(mgr.clone());
 
-        GlobalIORuntime::instance().spawn(async move { Self::check_timeout(mgr).await });
+        GlobalControlRuntime::instance()
+            .runtime()
+            .spawn(async move { Self::check_timeout(mgr).await });
         Ok(())
     }