refactor(processor): save pipeline executor into query context

Author: zhang2014

src/common/base/src/base/catch_unwind.rs

@@ -0,0 +1,29 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use common_exception::ErrorCode;
+use common_exception::Result;
+
+pub fn catch_unwind<F: FnOnce() -> R, R>(f: F) -> Result<R> {
+    match std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)) {
+        Ok(res) => Ok(res),
+        Err(cause) => match cause.downcast_ref::<&'static str>() {
+            None => match cause.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())),
+        },
+    }
+}

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

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+mod catch_unwind;
 mod global_runtime;
 mod net;
 mod profiling;
@@ -27,6 +28,7 @@ mod string_func;
 mod thread;
 mod uniq_id;
 
+pub use catch_unwind::catch_unwind;
 pub use global_runtime::GlobalIORuntime;
 pub use net::get_free_tcp_port;
 pub use net::get_free_udp_port;
@@ -53,6 +55,7 @@ pub use string_func::mask_string;
 pub use string_func::replace_nth_char;
 pub use string_func::unescape_for_key;
 pub use thread::Thread;
+pub use thread::ThreadJoinHandle;
 pub use tokio;
 pub use uniq_id::GlobalSequence;
 pub use uniq_id::GlobalUniqName;

src/common/base/src/base/thread.rs

@@ -15,12 +15,38 @@
 use std::thread::Builder;
 use std::thread::JoinHandle;
 
+use common_exception::ErrorCode;
+use common_exception::Result;
+
 use super::runtime_tracker::ThreadTracker;
 
 pub struct Thread;
 
+pub struct ThreadJoinHandle<T> {
+    inner: JoinHandle<T>,
+}
+
+impl<T> ThreadJoinHandle<T> {
+    pub fn create(inner: JoinHandle<T>) -> ThreadJoinHandle<T> {
+        ThreadJoinHandle { inner }
+    }
+
+    pub fn join(self) -> Result<T> {
+        match self.inner.join() {
+            Ok(res) => Ok(res),
+            Err(cause) => match cause.downcast_ref::<&'static str>() {
+                None => match cause.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())),
+            },
+        }
+    }
+}
+
 impl Thread {
-    pub fn named_spawn<F, T>(mut name: Option<String>, f: F) -> JoinHandle<T>
+    pub fn named_spawn<F, T>(mut name: Option<String>, f: F) -> ThreadJoinHandle<T>
     where
         F: FnOnce() -> T,
         F: Send + 'static,
@@ -42,18 +68,18 @@ impl Thread {
             thread_builder = thread_builder.name(named);
         }
 
-        match ThreadTracker::current_runtime_tracker() {
+        ThreadJoinHandle::create(match ThreadTracker::current_runtime_tracker() {
             None => thread_builder.spawn(f).unwrap(),
             Some(runtime_tracker) => thread_builder
                 .spawn(move || {
                     ThreadTracker::create(runtime_tracker);
                     f()
                 })
                 .unwrap(),
-        }
+        })
     }
 
-    pub fn spawn<F, T>(f: F) -> JoinHandle<T>
+    pub fn spawn<F, T>(f: F) -> ThreadJoinHandle<T>
     where
         F: FnOnce() -> T,
         F: Send + 'static,

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

@@ -506,9 +506,7 @@ impl QueryCoordinator {
     pub fn shutdown_query(&mut self) {
         if let Some(query_info) = &self.info {
             if let Some(query_executor) = &query_info.query_executor {
-                if let Err(cause) = query_executor.finish() {
-                    tracing::error!("Cannot shutdown query, because {:?}", cause);
-                }
+                query_executor.finish(None);
             }
         }
     }
@@ -553,14 +551,9 @@ impl QueryCoordinator {
             }
         }
 
-        let query_need_abort = info.query_ctx.query_need_abort();
         let executor_settings = ExecutorSettings::try_create(&info.query_ctx.get_settings())?;
 
-        let executor = PipelineCompleteExecutor::from_pipelines(
-            query_need_abort,
-            pipelines,
-            executor_settings,
-        )?;
+        let executor = PipelineCompleteExecutor::from_pipelines(pipelines, executor_settings)?;
 
         self.fragment_exchanges.clear();
         let info_mut = self.info.as_mut().expect("Query info is None");

src/query/service/src/interpreters/async_insert_queue.rs

@@ -409,12 +409,10 @@ impl AsyncInsertManager {
                 let mut pipelines = build_res.sources_pipelines;
                 pipelines.push(build_res.main_pipeline);
                 let executor_settings = ExecutorSettings::try_create(settings)?;
-                let executor = PipelineCompleteExecutor::from_pipelines(
-                    ctx.query_need_abort(),
-                    pipelines,
-                    executor_settings,
-                )?;
+                let executor =
+                    PipelineCompleteExecutor::from_pipelines(pipelines, executor_settings)?;
 
+                ctx.set_executor(Arc::downgrade(&executor.get_inner()));
                 executor.execute()?;
                 drop(executor);
                 let blocks = ctx.consume_precommit_blocks();

src/query/service/src/interpreters/interpreter.rs

@@ -56,20 +56,17 @@ pub trait Interpreter: Sync + Send {
         }
 
         let settings = ctx.get_settings();
-        let query_need_abort = ctx.query_need_abort();
         let executor_settings = ExecutorSettings::try_create(&settings)?;
         build_res.set_max_threads(settings.get_max_threads()? as usize);
 
         if build_res.main_pipeline.is_complete_pipeline()? {
             let mut pipelines = build_res.sources_pipelines;
             pipelines.push(build_res.main_pipeline);
 
-            let complete_executor = PipelineCompleteExecutor::from_pipelines(
-                query_need_abort,
-                pipelines,
-                executor_settings,
-            )?;
+            let complete_executor =
+                PipelineCompleteExecutor::from_pipelines(pipelines, executor_settings)?;
 
+            ctx.set_executor(Arc::downgrade(&complete_executor.get_inner()));
             complete_executor.execute()?;
             return Ok(Box::pin(DataBlockStream::create(
                 Arc::new(DataSchema::new(vec![])),
@@ -83,12 +80,12 @@ pub trait Interpreter: Sync + Send {
             return handle.execute(ctx.clone(), build_res, self.schema()).await;
         }
 
+        let pulling_executor =
+            PipelinePullingExecutor::from_pipelines(build_res, executor_settings)?;
+
+        ctx.set_executor(Arc::downgrade(&pulling_executor.get_inner()));
         Ok(Box::pin(Box::pin(ProcessorExecutorStream::create(
-            PipelinePullingExecutor::from_pipelines(
-                ctx.query_need_abort(),
-                build_res,
-                executor_settings,
-            )?,
+            pulling_executor,
         )?)))
     }
 

src/query/service/src/interpreters/interpreter_common.rs

@@ -113,13 +113,12 @@ pub fn append2table(
 }
 
 pub fn execute_pipeline(ctx: Arc<QueryContext>, mut res: PipelineBuildResult) -> Result<()> {
-    let query_need_abort = ctx.query_need_abort();
     let executor_settings = ExecutorSettings::try_create(&ctx.get_settings())?;
     res.set_max_threads(ctx.get_settings().get_max_threads()? as usize);
     let mut pipelines = res.sources_pipelines;
     pipelines.push(res.main_pipeline);
-    let executor =
-        PipelineCompleteExecutor::from_pipelines(query_need_abort, pipelines, executor_settings)?;
+    let executor = PipelineCompleteExecutor::from_pipelines(pipelines, executor_settings)?;
+    ctx.set_executor(Arc::downgrade(&executor.get_inner()));
     executor.execute()
 }
 

src/query/service/src/interpreters/interpreter_table_optimize.rs

@@ -70,13 +70,10 @@ impl Interpreter for OptimizeTableInterpreter {
             if let Some(mutator) = mutator {
                 let settings = ctx.get_settings();
                 pipeline.set_max_threads(settings.get_max_threads()? as usize);
-                let query_need_abort = ctx.query_need_abort();
                 let executor_settings = ExecutorSettings::try_create(&settings)?;
-                let executor = PipelineCompleteExecutor::try_create(
-                    query_need_abort,
-                    pipeline,
-                    executor_settings,
-                )?;
+                let executor = PipelineCompleteExecutor::try_create(pipeline, executor_settings)?;
+
+                ctx.set_executor(Arc::downgrade(&executor.get_inner()));
                 executor.execute()?;
                 drop(executor);
 

src/query/service/src/interpreters/interpreter_table_recluster.rs

@@ -74,13 +74,10 @@ impl Interpreter for ReclusterTableInterpreter {
 
             pipeline.set_max_threads(settings.get_max_threads()? as usize);
 
-            let query_need_abort = ctx.query_need_abort();
             let executor_settings = ExecutorSettings::try_create(&settings)?;
-            let executor = PipelineCompleteExecutor::try_create(
-                query_need_abort,
-                pipeline,
-                executor_settings,
-            )?;
+            let executor = PipelineCompleteExecutor::try_create(pipeline, executor_settings)?;
+
+            ctx.set_executor(Arc::downgrade(&executor.get_inner()));
             executor.execute()?;
             drop(executor);
 

src/query/service/src/lib.rs

@@ -24,6 +24,7 @@
 #![feature(option_get_or_insert_default)]
 #![feature(result_option_inspect)]
 #![feature(is_some_with)]
+#![feature(result_flattening)]
 
 extern crate core;