refactor(query): disallowed tokio spawn and block_on (databendlabs#14639)

* refactor(query): disallowed thread::spawn

* refactor(query): disallowed tokio::Handle::spawn

* refactor(query): disallowed tokio::Handle::spawn_blocking

* refactor(query): disallowed tokio block_on

* refactor(query): disallowed tokio block_on

* refactor(query): disallowed tokio block_on

* refactor(query): location futrue

* refactor(query): location futrue

---------

Co-authored-by: Bohu <[email protected]>

Author: zhang2014

Cargo.lock

@@ -6,10 +6,28 @@ disallowed-methods = [
     { path = "num_traits::sign::Signed::signum", reason = "This returns 1.0 for 0.0 but 0 for 0." },
     { path = "tokio::spawn", reason = "Please use `databend_common_base::runtime::spawn` instead." },
     { path = "tokio::task::spawn", reason = "Please use `databend_common_base::runtime::spawn` instead." },
+    { path = "tokio::runtime::Handle::spawn", reason = "Please use `databend_common_base::runtime::spawn` instead." },
+    { path = "tokio::runtime::Runtime::spawn", reason = "Please use `databend_common_base::runtime::Runtime::spawn` instead." },
+
+    { path = "tokio::task::spawn_local", reason = "Please use `databend_common_base::runtime::spawn_local` instead." },
+
+    { path = "tokio::runtime::Handle::block_on", reason = "Please use `databend_common_base::runtime::block_on` instead." },
+#    { path = "tokio::runtime::Runtime::block_on", reason = "Please use `databend_common_base::runtime::Runtime::block_on` instead." },
+
+    { path = "tokio::task::spawn_blocking", reason = "Please use `databend_common_base::runtime::spawn_blocking` instead." },
+    { path = "tokio::runtime::Handle::spawn_blocking", reason = "Please use `databend_common_base::runtime::spawn_blocking` instead." },
+    { path = "tokio::runtime::Runtime::spawn_blocking", reason = "Please use `databend_common_base::runtime::Runtime::spawn_blocking` instead." },
+
+    { path = "std::thread::spawn", reason = "Please use `databend_common_base::runtime::Thread::spawn` instead." },
 ]
 
 disallowed-types = [
     { path = "once_cell::sync::Lazy", reason = "Please use `std::sync::LazyLock` instead." },
+
+    { path = "tokio::task::Builder", reason = "Disallowed types, please remove it." },
+    { path = "tokio::task::LocalSet", reason = "Disallowed types, please remove it." },
+    { path = "tokio::task::join_set::JoinSet", reason = "Disallowed types, please remove it." },
+    { path = "tokio::task::join_set::Builder", reason = "Disallowed types, please remove it." },
 ]
 
 disallowed-macros = [

clippy.toml

@@ -58,7 +58,6 @@ databend-storages-common-table-meta = { path = "../query/storages/common/table_m
 # Crates.io dependencies
 anyerror = { workspace = true }
 anyhow = { workspace = true }
-async-backtrace = { workspace = true }
 chrono = { workspace = true }
 clap = { workspace = true }
 futures = { workspace = true }

src/binaries/Cargo.toml

@@ -37,7 +37,7 @@ fn main() {
             std::process::exit(cause.code() as i32);
         }
         Ok(rt) => {
-            if let Err(cause) = rt.block_on(async_backtrace::location!().frame(main_entrypoint())) {
+            if let Err(cause) = rt.block_on(main_entrypoint()) {
                 eprintln!("Databend Query start failure, cause: {:?}", cause);
                 std::process::exit(cause.code() as i32);
             }

src/binaries/query/ee_main.rs

@@ -38,7 +38,7 @@ fn main() {
             std::process::exit(cause.code() as i32);
         }
         Ok(rt) => {
-            if let Err(cause) = rt.block_on(async_backtrace::location!().frame(main_entrypoint())) {
+            if let Err(cause) = rt.block_on(main_entrypoint()) {
                 eprintln!("Databend Query start failure, cause: {:?}", cause);
                 std::process::exit(cause.code() as i32);
             }

src/binaries/query/oss_main.rs

@@ -19,7 +19,6 @@ mod global_runtime;
 mod runtime;
 mod runtime_tracker;
 mod thread;
-mod thread_pool;
 
 pub use backtrace::dump_backtrace;
 pub use backtrace::get_all_tasks;
@@ -28,9 +27,14 @@ pub use catch_unwind::catch_unwind;
 pub use catch_unwind::CatchUnwindFuture;
 pub use global_runtime::GlobalIORuntime;
 pub use global_runtime::GlobalQueryRuntime;
+pub use runtime::block_on;
 pub use runtime::execute_futures_in_parallel;
 pub use runtime::match_join_handle;
 pub use runtime::spawn;
+pub use runtime::spawn_blocking;
+pub use runtime::spawn_local;
+pub use runtime::try_block_on;
+pub use runtime::try_spawn_blocking;
 pub use runtime::Dropper;
 pub use runtime::Runtime;
 pub use runtime::TrySpawn;
@@ -44,5 +48,3 @@ pub use runtime_tracker::UnlimitedFuture;
 pub use runtime_tracker::GLOBAL_MEM_STAT;
 pub use thread::Thread;
 pub use thread::ThreadJoinHandle;
-pub use thread_pool::TaskJoinHandler;
-pub use thread_pool::ThreadPool;

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

@@ -14,8 +14,8 @@
 
 use std::backtrace::Backtrace;
 use std::future::Future;
+use std::panic::Location;
 use std::sync::Arc;
-use std::thread;
 use std::time::Duration;
 use std::time::Instant;
 
@@ -32,6 +32,8 @@ use tokio::task::JoinHandle;
 
 use crate::runtime::catch_unwind::CatchUnwindFuture;
 use crate::runtime::MemStat;
+use crate::runtime::Thread;
+use crate::runtime::ThreadJoinHandle;
 
 /// Methods to spawn tasks.
 pub trait TrySpawn {
@@ -101,7 +103,7 @@ impl Runtime {
         let handle = runtime.handle().clone();
 
         // Block the runtime to shutdown.
-        let join_handler = thread::spawn(move || {
+        let join_handler = Thread::spawn(move || {
             // We ignore channel is closed.
             let _ = runtime.block_on(recv_stop);
 
@@ -201,10 +203,16 @@ impl Runtime {
         self.handle.clone()
     }
 
+    #[track_caller]
     pub fn block_on<T, F>(&self, future: F) -> F::Output
     where F: Future<Output = Result<T>> + Send + 'static {
         let future = CatchUnwindFuture::create(future);
-        self.handle.block_on(future).flatten()
+        #[allow(clippy::disallowed_methods)]
+        tokio::task::block_in_place(|| {
+            self.handle
+                .block_on(location_future(future, std::panic::Location::caller()))
+                .flatten()
+        })
     }
 
     // For each future of `futures`, before being executed
@@ -260,6 +268,7 @@ impl Runtime {
             let permit = semaphore.acquire_owned().await.map_err(|e| {
                 ErrorCode::Internal(format!("semaphore closed, acquire permit failure. {}", e))
             })?;
+            #[expect(clippy::disallowed_methods)]
             let handler = self
                 .handle
                 .spawn(async_backtrace::location!().frame(async move {
@@ -280,6 +289,7 @@ impl Runtime {
         F: FnOnce() -> Result<R> + Send + 'static,
         R: Send + 'static,
     {
+        #[allow(clippy::disallowed_methods)]
         match_join_handle(self.handle.spawn_blocking(f)).await
     }
 }
@@ -298,6 +308,7 @@ impl TrySpawn for Runtime {
                 async_backtrace::location!(format!("Running query {} spawn task", id)).frame(task)
             }
         };
+        #[expect(clippy::disallowed_methods)]
         Ok(self.handle.spawn(task))
     }
 }
@@ -306,7 +317,7 @@ impl TrySpawn for Runtime {
 pub struct Dropper {
     name: Option<String>,
     close: Option<oneshot::Sender<()>>,
-    join_handler: Option<thread::JoinHandle<bool>>,
+    join_handler: Option<ThreadJoinHandle<bool>>,
 }
 
 impl Drop for Dropper {
@@ -389,6 +400,70 @@ pub fn spawn<F>(future: F) -> tokio::task::JoinHandle<F::Output>
 where
     F: Future + Send + 'static,
     F::Output: Send + 'static,
+{
+    #[expect(clippy::disallowed_methods)]
+    tokio::spawn(location_future(future, std::panic::Location::caller()))
+}
+
+#[track_caller]
+pub fn spawn_local<F>(future: F) -> tokio::task::JoinHandle<F::Output>
+where
+    F: Future + Send + 'static,
+    F::Output: Send + 'static,
+{
+    #[expect(clippy::disallowed_methods)]
+    tokio::task::spawn_local(location_future(future, std::panic::Location::caller()))
+}
+
+#[track_caller]
+pub fn spawn_blocking<F, R>(f: F) -> tokio::task::JoinHandle<R>
+where
+    F: FnOnce() -> R + Send + 'static,
+    R: Send + 'static,
+{
+    #[expect(clippy::disallowed_methods)]
+    tokio::runtime::Handle::current().spawn_blocking(f)
+}
+
+#[track_caller]
+pub fn try_spawn_blocking<F, R>(f: F) -> Result<tokio::task::JoinHandle<R>, F>
+where
+    F: FnOnce() -> R + Send + 'static,
+    R: Send + 'static,
+{
+    match tokio::runtime::Handle::try_current() {
+        Err(_) => Err(f),
+        #[expect(clippy::disallowed_methods)]
+        Ok(handler) => Ok(handler.spawn_blocking(f)),
+    }
+}
+
+#[track_caller]
+pub fn block_on<F: Future>(future: F) -> F::Output {
+    #[expect(clippy::disallowed_methods)]
+    tokio::task::block_in_place(|| {
+        tokio::runtime::Handle::current()
+            .block_on(location_future(future, std::panic::Location::caller()))
+    })
+}
+
+#[track_caller]
+pub fn try_block_on<F: Future>(future: F) -> Result<F::Output, F> {
+    match tokio::runtime::Handle::try_current() {
+        Err(_) => Err(future),
+        #[expect(clippy::disallowed_methods)]
+        Ok(handler) => Ok(tokio::task::block_in_place(|| {
+            handler.block_on(location_future(future, std::panic::Location::caller()))
+        })),
+    }
+}
+
+fn location_future<F: Future>(
+    future: F,
+    frame_location: &'static Location,
+) -> impl Future<Output = F::Output>
+where
+    F: Future,
 {
     // NOTE:
     // Frame name: https://play.rust-lang.org/?version=nightly&mode=debug&edition=2021&gist=689fbc84ab4be894c0cdd285bea24845
@@ -397,16 +472,12 @@ where
     let frame_name = std::any::type_name::<F>()
         .trim_end_matches("::{{closure}}")
         .to_string();
-    let frame_location = std::panic::Location::caller();
 
-    #[expect(clippy::disallowed_methods)]
-    tokio::spawn(
-        async_backtrace::location!(
-            frame_name,
-            frame_location.file(),
-            frame_location.line(),
-            frame_location.column()
-        )
-        .frame(future),
+    async_backtrace::location!(
+        frame_name,
+        frame_location.file(),
+        frame_location.line(),
+        frame_location.column()
     )
+    .frame(future)
 }

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

@@ -25,7 +25,6 @@ mod runtime;
 mod runtime_tracker;
 mod stoppable;
 mod string;
-mod thread_pool;
 
 // runtime tests depends on the memory stat collector.
 #[global_allocator]

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

@@ -109,7 +109,7 @@ pub fn build_operator<B: Builder>(builder: B) -> Result<Operator> {
         // Magic happens here. We will add a layer upon original
         // storage operator so that all underlying storage operations
         // will send to storage runtime.
-        .layer(RuntimeLayer::new(GlobalIORuntime::instance().inner()))
+        .layer(RuntimeLayer::new(GlobalIORuntime::instance()))
         .layer({
             let retry_timeout = env::var("_DATABEND_INTERNAL_RETRY_TIMEOUT")
                 .ok()