feat(query): support query level async backtrace

Cargo.toml

@@ -196,7 +196,7 @@ reqwest = { version = "0.11.19", default-features = false, features = [
 tokio = { version = "1.26.0", features = ["full"] }
 
 # backtrace
-async-backtrace = "0.2.2"
+async-backtrace = { git = "https://github.com/zhang2014/async-backtrace.git", rev = "e7e1b5f" }
 
 # observability
 logcall = "0.1.5"

src/common/base/Cargo.toml

@@ -27,7 +27,7 @@ common-exception = { path = "../exception" }
 # GitHub dependencies
 
 # Crates.io dependencies
-async-backtrace = "0.2.2"
+async-backtrace = { git = "https://github.com/zhang2014/async-backtrace.git", rev = "e7e1b5f" }
 async-channel = "1.7.1"
 async-trait = "0.1.57"
 bytesize = "1.1.0"

src/common/base/src/lib.rs

@@ -33,5 +33,8 @@ pub mod rangemap;
 pub mod runtime;
 
 pub use runtime::dump_backtrace;
+pub use runtime::get_all_tasks;
 pub use runtime::match_join_handle;
 pub use runtime::set_alloc_error_hook;
+pub use runtime::AsyncTaskItem;
+pub use runtime::GLOBAL_TASK;

src/common/base/src/runtime/backtrace.rs

@@ -15,11 +15,30 @@
 use std::fmt::Write;
 
 #[derive(Debug)]
-struct AsyncTaskItem {
-    stack_frames: Vec<String>,
+pub struct AsyncTaskItem {
+    pub stack_frames: Vec<String>,
 }
 
 pub fn dump_backtrace(wait_for_running_tasks: bool) -> String {
+    let (tasks, polling_tasks) = get_all_tasks(wait_for_running_tasks);
+
+    let mut output = String::new();
+    for mut tasks in [tasks, polling_tasks] {
+        tasks.sort_by(|l, r| Ord::cmp(&l.stack_frames.len(), &r.stack_frames.len()));
+
+        for item in tasks.into_iter().rev() {
+            for frame in item.stack_frames {
+                writeln!(output, "{}", frame).unwrap();
+            }
+
+            writeln!(output).unwrap();
+        }
+    }
+
+    output
+}
+
+pub fn get_all_tasks(wait_for_running_tasks: bool) -> (Vec<AsyncTaskItem>, Vec<AsyncTaskItem>) {
     let tree = async_backtrace::taskdump_tree(wait_for_running_tasks);
 
     let mut tasks = vec![];
@@ -61,19 +80,5 @@ pub fn dump_backtrace(wait_for_running_tasks: bool) -> String {
             stack_frames: std::mem::take(&mut current_stack_frames),
         }),
     };
-
-    let mut output = String::new();
-    for mut tasks in [tasks, polling_tasks] {
-        tasks.sort_by(|l, r| Ord::cmp(&l.stack_frames.len(), &r.stack_frames.len()));
-
-        for item in tasks.into_iter().rev() {
-            for frame in item.stack_frames {
-                writeln!(output, "{}", frame).unwrap();
-            }
-
-            writeln!(output).unwrap();
-        }
-    }
-
-    output
+    (tasks, polling_tasks)
 }

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

@@ -22,6 +22,8 @@ mod thread;
 mod thread_pool;
 
 pub use backtrace::dump_backtrace;
+pub use backtrace::get_all_tasks;
+pub use backtrace::AsyncTaskItem;
 pub use catch_unwind::catch_unwind;
 pub use catch_unwind::CatchUnwindFuture;
 pub use global_runtime::GlobalIORuntime;
@@ -31,6 +33,7 @@ pub use runtime::match_join_handle;
 pub use runtime::Dropper;
 pub use runtime::Runtime;
 pub use runtime::TrySpawn;
+pub use runtime::GLOBAL_TASK;
 pub use runtime_tracker::set_alloc_error_hook;
 pub use runtime_tracker::LimitMemGuard;
 pub use runtime_tracker::MemStat;

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

@@ -39,7 +39,7 @@ pub trait TrySpawn {
     ///
     /// It allows to return an error before spawning the task.
     #[track_caller]
-    fn try_spawn<T>(&self, task: T) -> Result<JoinHandle<T::Output>>
+    fn try_spawn<T>(&self, id: impl Into<String>, task: T) -> Result<JoinHandle<T::Output>>
     where
         T: Future + Send + 'static,
         T::Output: Send + 'static;
@@ -48,32 +48,32 @@ pub trait TrySpawn {
     ///
     /// A default impl of this method just calls `try_spawn` and just panics if there is an error.
     #[track_caller]
-    fn spawn<T>(&self, task: T) -> JoinHandle<T::Output>
+    fn spawn<T>(&self, id: impl Into<String>, task: T) -> JoinHandle<T::Output>
     where
         T: Future + Send + 'static,
         T::Output: Send + 'static,
     {
-        self.try_spawn(task).unwrap()
+        self.try_spawn(id, task).unwrap()
     }
 }
 
 impl<S: TrySpawn> TrySpawn for Arc<S> {
     #[track_caller]
-    fn try_spawn<T>(&self, task: T) -> Result<JoinHandle<T::Output>>
+    fn try_spawn<T>(&self, id: impl Into<String>, task: T) -> Result<JoinHandle<T::Output>>
     where
         T: Future + Send + 'static,
         T::Output: Send + 'static,
     {
-        self.as_ref().try_spawn(task)
+        self.as_ref().try_spawn(id, task)
     }
 
     #[track_caller]
-    fn spawn<T>(&self, task: T) -> JoinHandle<T::Output>
+    fn spawn<T>(&self, id: impl Into<String>, task: T) -> JoinHandle<T::Output>
     where
         T: Future + Send + 'static,
         T::Output: Send + 'static,
     {
-        self.as_ref().spawn(task)
+        self.as_ref().spawn(id, task)
     }
 }
 
@@ -286,12 +286,19 @@ impl Runtime {
 
 impl TrySpawn for Runtime {
     #[track_caller]
-    fn try_spawn<T>(&self, task: T) -> Result<JoinHandle<T::Output>>
+    fn try_spawn<T>(&self, id: impl Into<String>, task: T) -> Result<JoinHandle<T::Output>>
     where
         T: Future + Send + 'static,
         T::Output: Send + 'static,
     {
-        Ok(self.handle.spawn(async_backtrace::location!().frame(task)))
+        let id = id.into();
+        let task = match id == GLOBAL_TASK {
+            true => async_backtrace::location!(String::from(GLOBAL_TASK_DESC)).frame(task),
+            false => {
+                async_backtrace::location!(format!("Running query {} spawn task", id)).frame(task)
+            }
+        };
+        Ok(self.handle.spawn(task))
     }
 }
 
@@ -373,3 +380,6 @@ where
         .await
         .map_err(|e| ErrorCode::Internal(format!("try join all futures failure, {}", e)))
 }
+
+pub const GLOBAL_TASK: &str = "Zxv39PlwG1ahbF0APRUf03";
+pub const GLOBAL_TASK_DESC: &str = "Global spawn task";

src/common/base/tests/it/runtime.rs

@@ -19,6 +19,7 @@ use std::time::Instant;
 
 use common_base::runtime::Runtime;
 use common_base::runtime::TrySpawn;
+use common_base::GLOBAL_TASK;
 use common_exception::Result;
 use once_cell::sync::Lazy;
 use rand::distributions::Distribution;
@@ -32,16 +33,16 @@ async fn test_runtime() -> Result<()> {
 
     let runtime = Runtime::with_default_worker_threads()?;
     let runtime_counter = Arc::clone(&counter);
-    let runtime_header = runtime.spawn(async move {
+    let runtime_header = runtime.spawn(GLOBAL_TASK, async move {
         let rt1 = Runtime::with_default_worker_threads().unwrap();
         let rt1_counter = Arc::clone(&runtime_counter);
-        let rt1_header = rt1.spawn(async move {
+        let rt1_header = rt1.spawn(GLOBAL_TASK, async move {
             let rt2 = Runtime::with_worker_threads(1, None).unwrap();
             let rt2_counter = Arc::clone(&rt1_counter);
-            let rt2_header = rt2.spawn(async move {
+            let rt2_header = rt2.spawn(GLOBAL_TASK, async move {
                 let rt3 = Runtime::with_default_worker_threads().unwrap();
                 let rt3_counter = Arc::clone(&rt2_counter);
-                let rt3_header = rt3.spawn(async move {
+                let rt3_header = rt3.spawn(GLOBAL_TASK, async move {
                     let mut num = rt3_counter.lock().unwrap();
                     *num += 1;
                 });
@@ -72,7 +73,7 @@ async fn test_runtime() -> Result<()> {
 async fn test_shutdown_long_run_runtime() -> Result<()> {
     let runtime = Runtime::with_default_worker_threads()?;
 
-    runtime.spawn(async move {
+    runtime.spawn(GLOBAL_TASK, async move {
         std::thread::sleep(Duration::from_secs(6));
     });
 

src/common/base/tests/it/runtime_tracker.rs

@@ -18,6 +18,7 @@ use common_base::runtime::MemStat;
 use common_base::runtime::Runtime;
 use common_base::runtime::TrackedFuture;
 use common_base::runtime::TrySpawn;
+use common_base::GLOBAL_TASK;
 use common_exception::Result;
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 8)]
@@ -29,9 +30,9 @@ async fn test_async_thread_tracker() -> Result<()> {
     let inner_runtime = Runtime::with_worker_threads(2, Some(String::from("Inner")))?;
 
     let memory_tracker = MemStat::create("test_async_thread_tracker".to_string());
-    let inner_join_handler = inner_runtime.spawn(TrackedFuture::create_with_mem_stat(
-        Some(memory_tracker.clone()),
-        async move {
+    let inner_join_handler = inner_runtime.spawn(
+        GLOBAL_TASK,
+        TrackedFuture::create_with_mem_stat(Some(memory_tracker.clone()), async move {
             let memory = vec![0_u8; 3 * 1024 * 1024];
             tokio::time::sleep(Duration::from_millis(100)).await;
             out_tx.send(()).await.unwrap();
@@ -57,10 +58,10 @@ async fn test_async_thread_tracker() -> Result<()> {
             tokio::time::sleep(Duration::from_millis(100)).await;
             out_tx.send(()).await.unwrap();
             inner_rx.recv().await.unwrap();
-        },
-    ));
+        }),
+    );
 
-    let outer_join_handler = outer_runtime.spawn(async move {
+    let outer_join_handler = outer_runtime.spawn(GLOBAL_TASK, async move {
         for (min_memory_usage, max_memory_usage) in [
             (2 * 1024 * 1024, 4 * 1024 * 1024),
             (2 * 1024 * 1024, 4 * 1024 * 1024),

src/common/storage/src/operator.rs

@@ -22,6 +22,7 @@ use anyhow::anyhow;
 use common_base::base::GlobalInstance;
 use common_base::runtime::GlobalIORuntime;
 use common_base::runtime::TrySpawn;
+use common_base::GLOBAL_TASK;
 use common_exception::ErrorCode;
 use common_meta_app::storage::StorageAzblobConfig;
 use common_meta_app::storage::StorageCosConfig;
@@ -433,7 +434,7 @@ impl DataOperator {
         // IO hang on reuse connection.
         let op = operator.clone();
         if let Err(cause) = GlobalIORuntime::instance()
-            .spawn(async move { op.check().await })
+            .spawn(GLOBAL_TASK, async move { op.check().await })
             .await
             .expect("join must succeed")
         {

src/meta/client/src/grpc_client.rs

@@ -37,6 +37,7 @@ use common_base::future::TimingFutureExt;
 use common_base::runtime::Runtime;
 use common_base::runtime::TrySpawn;
 use common_base::runtime::UnlimitedFuture;
+use common_base::GLOBAL_TASK;
 use common_grpc::ConnectionFactory;
 use common_grpc::GrpcConnectionError;
 use common_grpc::RpcClientConf;
@@ -368,13 +369,14 @@ impl MetaGrpcClient {
             rt: rt.clone(),
         });
 
-        rt.spawn(UnlimitedFuture::create(Self::worker_loop(
-            worker.clone(),
-            rx,
-        )));
-        rt.spawn(UnlimitedFuture::create(Self::auto_sync_endpoints(
-            worker, one_tx,
-        )));
+        rt.spawn(
+            GLOBAL_TASK,
+            UnlimitedFuture::create(Self::worker_loop(worker.clone(), rx)),
+        );
+        rt.spawn(
+            GLOBAL_TASK,
+            UnlimitedFuture::create(Self::auto_sync_endpoints(worker, one_tx)),
+        );
 
         Ok(handle)
     }

src/query/ee_features/table_lock/src/table_lock_heartbeat.rs

@@ -48,7 +48,7 @@ impl TableLockHeartbeat {
         let expire_secs = ctx.get_settings().get_table_lock_expire_secs()?;
         let sleep_range = (expire_secs * 1000 / 3)..=((expire_secs * 1000 / 3) * 2);
 
-        self.shutdown_handler = Some(GlobalIORuntime::instance().spawn({
+        self.shutdown_handler = Some(GlobalIORuntime::instance().spawn(ctx.get_id(), {
             let shutdown_flag = self.shutdown_flag.clone();
             let shutdown_notify = self.shutdown_notify.clone();
 

src/query/pipeline/sinks/Cargo.toml

@@ -13,6 +13,7 @@ test = false
 
 [dependencies]
 common-base = { path = "../../../common/base" }
+common-catalog = { path = "../../catalog" }
 common-exception = { path = "../../../common/exception" }
 common-expression = { path = "../../expression" }
 common-pipeline-core = { path = "../core" }

src/query/pipeline/sinks/src/async_sink.rs

@@ -19,6 +19,7 @@ use async_trait::async_trait;
 use async_trait::unboxed_simple;
 use common_base::runtime::GlobalIORuntime;
 use common_base::runtime::TrySpawn;
+use common_catalog::table_context::TableContext;
 use common_exception::Result;
 use common_expression::DataBlock;
 use common_pipeline_core::processors::port::InputPort;
@@ -47,15 +48,21 @@ pub struct AsyncSinker<T: AsyncSink + 'static> {
     inner: Option<T>,
     finished: bool,
     input: Arc<InputPort>,
+    query_id: String,
     input_data: Option<DataBlock>,
     called_on_start: bool,
     called_on_finish: bool,
 }
 
 impl<T: AsyncSink + 'static> AsyncSinker<T> {
-    pub fn create(input: Arc<InputPort>, inner: T) -> Box<dyn Processor> {
+    pub fn create(
+        input: Arc<InputPort>,
+        ctx: Arc<dyn TableContext>,
+        inner: T,
+    ) -> Box<dyn Processor> {
         Box::new(AsyncSinker {
             input,
+            query_id: ctx.get_id(),
             finished: false,
             input_data: None,
             inner: Some(inner),
@@ -69,7 +76,7 @@ impl<T: AsyncSink + 'static> Drop for AsyncSinker<T> {
     fn drop(&mut self) {
         if !self.called_on_start || !self.called_on_finish {
             if let Some(mut inner) = self.inner.take() {
-                GlobalIORuntime::instance().spawn({
+                GlobalIORuntime::instance().spawn(self.query_id.clone(), {
                     let called_on_start = self.called_on_start;
                     let called_on_finish = self.called_on_finish;
                     async move {