fix(query): fix unexpected panic message (#16221)

Author: zhang2014

src/common/base/src/lib.rs

@@ -37,7 +37,7 @@ pub mod version;
 
 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::JoinHandle;
 pub use runtime::GLOBAL_TASK;

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

@@ -39,13 +39,13 @@ pub use memory::MemStat;
 pub use memory::GLOBAL_MEM_STAT;
 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::JoinHandle;
 pub use runtime::Runtime;
 pub use runtime::TrySpawn;
 pub use runtime::GLOBAL_TASK;

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

@@ -15,28 +15,72 @@
 use std::backtrace::Backtrace;
 use std::future::Future;
 use std::panic::Location;
+use std::pin::Pin;
 use std::sync::Arc;
+use std::task::Context;
+use std::task::Poll;
 use std::time::Duration;
 use std::time::Instant;
 
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use futures::future;
+use futures::FutureExt;
 use log::warn;
 use tokio::runtime::Builder;
 use tokio::runtime::Handle;
 use tokio::sync::oneshot;
 use tokio::sync::OwnedSemaphorePermit;
 use tokio::sync::Semaphore;
-use tokio::task::JoinHandle;
 
+// use tokio::task::JoinHandle;
 use crate::runtime::catch_unwind::CatchUnwindFuture;
 use crate::runtime::drop_guard;
 use crate::runtime::memory::MemStat;
 use crate::runtime::Thread;
 use crate::runtime::ThreadJoinHandle;
 use crate::runtime::ThreadTracker;
 
+pub struct JoinHandle<Output> {
+    inner: tokio::task::JoinHandle<Output>,
+}
+
+impl<Output> JoinHandle<Output> {
+    pub fn create(inner: tokio::task::JoinHandle<Output>) -> Self {
+        Self { inner }
+    }
+
+    pub fn abort(&self) {
+        self.inner.abort();
+    }
+}
+
+impl<Output> Future for JoinHandle<Output> {
+    type Output = Result<Output>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match self.inner.poll_unpin(cx) {
+            Poll::Pending => Poll::Pending,
+            Poll::Ready(res) => match res {
+                Ok(res) => Poll::Ready(Ok(res)),
+                Err(error) => match error.is_panic() {
+                    true => {
+                        let cause = error.into_panic();
+                        Poll::Ready(Err(match cause.downcast_ref::<&'static str>() {
+                            None => match cause.downcast_ref::<String>() {
+                                None => ErrorCode::PanicError("Sorry, unknown panic message"),
+                                Some(message) => ErrorCode::PanicError(message.to_string()),
+                            },
+                            Some(message) => ErrorCode::PanicError(message.to_string()),
+                        }))
+                    }
+                    false => Poll::Ready(Err(ErrorCode::TokioError("Tokio task is cancelled"))),
+                },
+            },
+        }
+    }
+}
+
 /// Methods to spawn tasks.
 pub trait TrySpawn {
     /// Tries to spawn a new asynchronous task, returning a tokio::JoinHandle for it.
@@ -277,7 +321,7 @@ impl Runtime {
                     fut(permit).await
                 }),
             ));
-            handlers.push(handler)
+            handlers.push(JoinHandle::create(handler))
         }
 
         Ok(handlers)
@@ -292,11 +336,11 @@ impl Runtime {
         R: Send + 'static,
     {
         #[allow(clippy::disallowed_methods)]
-        match_join_handle(
+        let handle = JoinHandle::create(
             self.handle
                 .spawn_blocking(ThreadTracker::tracking_function(f)),
-        )
-        .await
+        );
+        handle.await.flatten()
     }
 }
 
@@ -317,7 +361,7 @@ impl TrySpawn for Runtime {
         };
 
         #[expect(clippy::disallowed_methods)]
-        Ok(self.handle.spawn(task))
+        Ok(JoinHandle::create(self.handle.spawn(task)))
     }
 }
 
@@ -353,26 +397,6 @@ impl Drop for Dropper {
     }
 }
 
-pub async fn match_join_handle<T>(handle: JoinHandle<Result<T>>) -> Result<T> {
-    match handle.await {
-        Ok(Ok(res)) => Ok(res),
-        Ok(Err(cause)) => Err(cause),
-        Err(join_error) => match join_error.is_cancelled() {
-            true => Err(ErrorCode::TokioError("Tokio error is cancelled.")),
-            false => {
-                let panic_error = join_error.into_panic();
-                match panic_error.downcast_ref::<&'static str>() {
-                    None => match panic_error.downcast_ref::<String>() {
-                        None => Err(ErrorCode::PanicError("Sorry, unknown panic message")),
-                        Some(message) => Err(ErrorCode::PanicError(message.to_string())),
-                    },
-                    Some(message) => Err(ErrorCode::PanicError(message.to_string())),
-                }
-            }
-        },
-    }
-}
-
 /// Run multiple futures parallel
 /// using a semaphore to limit the parallelism number, and a specified thread pool to run the futures.
 /// It waits for all futures to complete and returns their results.
@@ -397,9 +421,7 @@ where
     let join_handlers = runtime.try_spawn_batch(semaphore, futures).await?;
 
     // 3. get all the result.
-    future::try_join_all(join_handlers)
-        .await
-        .map_err(|e| ErrorCode::Internal(format!("try join all futures failure, {}", e)))
+    future::try_join_all(join_handlers).await
 }
 
 pub const GLOBAL_TASK: &str = "Zxv39PlwG1ahbF0APRUf03";

src/query/service/src/servers/flight/v1/actions/init_query_fragments.rs

@@ -12,7 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use databend_common_base::match_join_handle;
 use databend_common_base::runtime::ThreadTracker;
 use databend_common_base::runtime::TrySpawn;
 use databend_common_exception::Result;
@@ -37,7 +36,7 @@ pub async fn init_query_fragments(fragments: QueryFragments) -> Result<()> {
         DataExchangeManager::instance().init_query_fragments_plan(&fragments)
     }));
 
-    if let Err(cause) = match_join_handle(join_handler).await {
+    if let Err(cause) = join_handler.await.flatten() {
         DataExchangeManager::instance().on_finished_query(&query_id);
         return Err(cause);
     }

src/query/service/src/servers/flight/v1/exchange/exchange_manager.rs

@@ -28,6 +28,7 @@ use databend_common_base::base::GlobalInstance;
 use databend_common_base::runtime::GlobalIORuntime;
 use databend_common_base::runtime::Thread;
 use databend_common_base::runtime::TrySpawn;
+use databend_common_base::JoinHandle;
 use databend_common_config::GlobalConfig;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
@@ -40,7 +41,6 @@ use parking_lot::Mutex;
 use parking_lot::ReentrantMutex;
 use petgraph::prelude::EdgeRef;
 use petgraph::Direction;
-use tokio::task::JoinHandle;
 use tonic::Status;
 
 use super::exchange_params::ExchangeParams;

src/query/service/src/servers/flight/v1/exchange/statistics_receiver.rs

@@ -17,10 +17,9 @@ use std::sync::Arc;
 
 use databend_common_base::base::tokio::sync::broadcast::channel;
 use databend_common_base::base::tokio::sync::broadcast::Sender;
-use databend_common_base::base::tokio::task::JoinHandle;
-use databend_common_base::match_join_handle;
 use databend_common_base::runtime::Runtime;
 use databend_common_base::runtime::TrySpawn;
+use databend_common_base::JoinHandle;
 use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::Result;
 use futures_util::future::select;
@@ -161,7 +160,7 @@ impl StatisticsReceiver {
         let mut exchanges_handler = std::mem::take(&mut self.exchange_handler);
         futures::executor::block_on(async move {
             while let Some(exchange_handler) = exchanges_handler.pop() {
-                match_join_handle(exchange_handler).await?;
+                exchange_handler.await.flatten()?;
             }
 
             Ok(())

src/query/service/src/servers/flight/v1/exchange/statistics_sender.rs

@@ -16,9 +16,9 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use async_channel::Sender;
-use databend_common_base::base::tokio::task::JoinHandle;
 use databend_common_base::base::tokio::time::sleep;
 use databend_common_base::runtime::TrySpawn;
+use databend_common_base::JoinHandle;
 use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;

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

@@ -26,12 +26,12 @@ use databend_common_base::base::tokio::time::sleep;
 use databend_common_base::base::GlobalInstance;
 use databend_common_base::runtime::GlobalIORuntime;
 use databend_common_base::runtime::TrySpawn;
+use databend_common_base::JoinHandle;
 use databend_common_config::InnerConfig;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_storages_common_txn::TxnManagerRef;
 use parking_lot::Mutex;
-use tokio::task;
 
 use super::expiring_map::ExpiringMap;
 use super::HttpQueryContext;
@@ -84,7 +84,7 @@ pub struct HttpQueryManager {
     pub(crate) queries: Arc<DashMap<String, Arc<HttpQuery>>>,
     pub(crate) removed_queries: Arc<parking_lot::Mutex<LimitedQueue<String>>>,
     #[allow(clippy::type_complexity)]
-    pub(crate) txn_managers: Arc<Mutex<HashMap<String, (TxnManagerRef, task::JoinHandle<()>)>>>,
+    pub(crate) txn_managers: Arc<Mutex<HashMap<String, (TxnManagerRef, JoinHandle<()>)>>>,
     pub(crate) sessions: Mutex<ExpiringMap<String, Arc<Session>>>,
 }
 

src/query/service/src/sessions/query_ctx.rs

@@ -32,12 +32,12 @@ use chrono::Utc;
 use chrono_tz::Tz;
 use dashmap::mapref::multiple::RefMulti;
 use dashmap::DashMap;
-use databend_common_base::base::tokio::task::JoinHandle;
 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::TrySpawn;
+use databend_common_base::JoinHandle;
 use databend_common_catalog::lock::LockTableOption;
 use databend_common_catalog::merge_into_join::MergeIntoJoin;
 use databend_common_catalog::plan::DataSourceInfo;

src/query/storages/fuse/src/operations/recluster.rs

@@ -275,19 +275,13 @@ impl FuseTable {
             }));
         }
 
-        match futures::future::try_join_all(works).await {
-            Err(e) => Err(ErrorCode::StorageOther(format!(
-                "segment pruning failure, {}",
-                e
-            ))),
-            Ok(workers) => {
-                let mut metas = vec![];
-                for worker in workers {
-                    let res = worker?;
-                    metas.extend(res);
-                }
-                Ok(metas)
-            }
+        let mut metas = vec![];
+        let workers = futures::future::try_join_all(works).await?;
+        for worker in workers {
+            let res = worker?;
+            metas.extend(res);
         }
+
+        Ok(metas)
     }
 }