Merge branch 'main' into num_func_v2

Author: mergify[bot]

src/query/service/src/pipelines/executor/pipeline_executor.rs

@@ -178,11 +178,11 @@ impl PipelineExecutor {
             }
         }
 
-        if let Err(error_code) = self.graph.check_finished() {
-            let may_error = Some(error_code);
-            (self.on_finished_callback)(&may_error)?;
-            return Err(may_error.unwrap());
-        }
+        // if let Err(error_code) = self.graph.check_finished() {
+        //     let may_error = Some(error_code);
+        //     (self.on_finished_callback)(&may_error)?;
+        //     return Err(may_error.unwrap());
+        // }
 
         (self.on_finished_callback)(&None)?;
         Ok(())

src/query/service/src/pipelines/executor/pipeline_pulling_executor.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
 use std::sync::mpsc::Receiver;
 use std::sync::mpsc::RecvTimeoutError;
 use std::sync::mpsc::SyncSender;
@@ -23,53 +24,93 @@ use common_base::base::Thread;
 use common_datablocks::DataBlock;
 use common_exception::ErrorCode;
 use common_exception::Result;
+use common_pipeline_core::SinkPipeBuilder;
+use parking_lot::Condvar;
+use parking_lot::Mutex;
 
 use crate::pipelines::executor::executor_settings::ExecutorSettings;
 use crate::pipelines::executor::PipelineExecutor;
 use crate::pipelines::processors::port::InputPort;
 use crate::pipelines::processors::processor::ProcessorPtr;
 use crate::pipelines::processors::Sink;
 use crate::pipelines::processors::Sinker;
-use crate::pipelines::Pipe;
 use crate::pipelines::Pipeline;
 use crate::pipelines::PipelineBuildResult;
 
 struct State {
-    sender: SyncSender<Result<Option<DataBlock>>>,
+    is_catch_error: AtomicBool,
+    finish_mutex: Mutex<bool>,
+    finish_condvar: Condvar,
+
+    catch_error: Mutex<Option<ErrorCode>>,
 }
 
 impl State {
-    pub fn create(sender: SyncSender<Result<Option<DataBlock>>>) -> Arc<State> {
-        Arc::new(State { sender })
+    pub fn create() -> Arc<State> {
+        Arc::new(State {
+            catch_error: Mutex::new(None),
+            is_catch_error: AtomicBool::new(false),
+            finish_mutex: Mutex::new(false),
+            finish_condvar: Condvar::new(),
+        })
+    }
+
+    pub fn finished(&self, message: Result<()>) {
+        if let Err(error) = message {
+            self.is_catch_error.store(true, Ordering::Release);
+            *self.catch_error.lock() = Some(error);
+        }
+
+        let mut mutex = self.finish_mutex.lock();
+        *mutex = true;
+        self.finish_condvar.notify_one();
+    }
+
+    pub fn wait_finish(&self) {
+        let mut mutex = self.finish_mutex.lock();
+
+        while !*mutex {
+            self.finish_condvar.wait(&mut mutex);
+        }
+    }
+
+    pub fn is_catch_error(&self) -> bool {
+        self.is_catch_error.load(Ordering::Relaxed)
+    }
+
+    pub fn get_catch_error(&self) -> ErrorCode {
+        let catch_error = self.catch_error.lock();
+
+        match catch_error.as_ref() {
+            None => ErrorCode::LogicalError("It's a bug."),
+            Some(catch_error) => catch_error.clone(),
+        }
     }
 }
 
 // Use this executor when the pipeline is pulling pipeline (exists source but not exists sink)
 pub struct PipelinePullingExecutor {
     state: Arc<State>,
     executor: Arc<PipelineExecutor>,
-    receiver: Receiver<Result<Option<DataBlock>>>,
+    receiver: Receiver<DataBlock>,
 }
 
 impl PipelinePullingExecutor {
-    fn wrap_pipeline(
-        pipeline: &mut Pipeline,
-        tx: SyncSender<Result<Option<DataBlock>>>,
-    ) -> Result<()> {
+    fn wrap_pipeline(pipeline: &mut Pipeline, tx: SyncSender<DataBlock>) -> Result<()> {
         if pipeline.is_pushing_pipeline()? || !pipeline.is_pulling_pipeline()? {
             return Err(ErrorCode::LogicalError(
                 "Logical error, PipelinePullingExecutor can only work on pulling pipeline.",
             ));
         }
 
-        pipeline.resize(1)?;
-        let input = InputPort::create();
+        let mut sink_pipe_builder = SinkPipeBuilder::create();
+
+        for _index in 0..pipeline.output_len() {
+            let input = InputPort::create();
+            sink_pipe_builder.add_sink(input.clone(), PullingSink::create(tx.clone(), input));
+        }
 
-        pipeline.add_pipe(Pipe::SimplePipe {
-            outputs_port: vec![],
-            inputs_port: vec![input.clone()],
-            processors: vec![PullingSink::create(tx, input)],
-        });
+        pipeline.add_pipe(sink_pipe_builder.finalize());
         Ok(())
     }
 
@@ -79,14 +120,13 @@ impl PipelinePullingExecutor {
         settings: ExecutorSettings,
     ) -> Result<PipelinePullingExecutor> {
         let (sender, receiver) = std::sync::mpsc::sync_channel(pipeline.output_len());
-        let state = State::create(sender.clone());
 
         Self::wrap_pipeline(&mut pipeline, sender)?;
         let executor = PipelineExecutor::create(query_need_abort, pipeline, settings)?;
         Ok(PipelinePullingExecutor {
             receiver,
-            state,
             executor,
+            state: State::create(),
         })
     }
 
@@ -97,15 +137,14 @@ impl PipelinePullingExecutor {
     ) -> Result<PipelinePullingExecutor> {
         let mut main_pipeline = build_res.main_pipeline;
         let (sender, receiver) = std::sync::mpsc::sync_channel(main_pipeline.output_len());
-        let state = State::create(sender.clone());
         Self::wrap_pipeline(&mut main_pipeline, sender)?;
 
         let mut pipelines = build_res.sources_pipelines;
         pipelines.push(main_pipeline);
 
         Ok(PipelinePullingExecutor {
             receiver,
-            state,
+            state: State::create(),
             executor: PipelineExecutor::from_pipelines(query_need_abort, pipelines, settings)?,
         })
     }
@@ -136,17 +175,7 @@ impl PipelinePullingExecutor {
 
     fn thread_function(state: Arc<State>, executor: Arc<PipelineExecutor>) -> impl Fn() {
         move || {
-            if let Err(cause) = executor.execute() {
-                if let Err(send_err) = state.sender.send(Err(cause)) {
-                    tracing::warn!("Send error {:?}", send_err);
-                }
-
-                return;
-            }
-
-            if let Err(send_err) = state.sender.send(Ok(None)) {
-                tracing::warn!("Send finish event error {:?}", send_err);
-            }
+            state.finished(executor.execute());
         }
     }
 
@@ -155,36 +184,30 @@ impl PipelinePullingExecutor {
     }
 
     pub fn pull_data(&mut self) -> Result<Option<DataBlock>> {
-        match self.receiver.recv() {
-            Ok(data_block) => data_block,
-            Err(_recv_err) => Err(ErrorCode::LogicalError("Logical error, receiver error.")),
-        }
-    }
+        loop {
+            return match self.receiver.recv_timeout(Duration::from_millis(100)) {
+                Ok(data_block) => Ok(Some(data_block)),
+                Err(RecvTimeoutError::Timeout) => {
+                    if self.state.is_catch_error() {
+                        return Err(self.state.get_catch_error());
+                    }
 
-    pub fn try_pull_data<F>(&mut self, f: F) -> Result<Option<DataBlock>>
-    where F: Fn() -> bool {
-        if !self.executor.is_finished() {
-            while !f() {
-                return match self.receiver.recv_timeout(Duration::from_millis(100)) {
-                    Ok(data_block) => data_block,
-                    Err(RecvTimeoutError::Timeout) => {
-                        continue;
+                    continue;
+                }
+                Err(_disconnected) => {
+                    if !self.executor.is_finished() {
+                        self.executor.finish()?;
                     }
-                    Err(RecvTimeoutError::Disconnected) => {
-                        Err(ErrorCode::LogicalError("Logical error, receiver error."))
+
+                    self.state.wait_finish();
+
+                    if self.state.is_catch_error() {
+                        return Err(self.state.get_catch_error());
                     }
-                };
-            }
-            Ok(None)
-        } else {
-            match self.receiver.try_recv() {
-                Ok(data_block) => data_block,
-                // puller will not pull again once it received a None
-                Err(err) => Err(ErrorCode::LogicalError(format!(
-                    "Logical error, try receiver error. after executor finish {}",
-                    err
-                ))),
-            }
+
+                    Ok(None)
+                }
+            };
         }
     }
 }
@@ -198,31 +221,31 @@ impl Drop for PipelinePullingExecutor {
 }
 
 struct PullingSink {
-    sender: SyncSender<Result<Option<DataBlock>>>,
+    sender: Option<SyncSender<DataBlock>>,
 }
 
 impl PullingSink {
-    pub fn create(
-        tx: SyncSender<Result<Option<DataBlock>>>,
-        input: Arc<InputPort>,
-    ) -> ProcessorPtr {
-        Sinker::create(input, PullingSink { sender: tx })
+    pub fn create(tx: SyncSender<DataBlock>, input: Arc<InputPort>) -> ProcessorPtr {
+        Sinker::create(input, PullingSink { sender: Some(tx) })
     }
 }
 
 impl Sink for PullingSink {
     const NAME: &'static str = "PullingExecutorSink";
 
     fn on_finish(&mut self) -> Result<()> {
+        drop(self.sender.take());
         Ok(())
     }
 
     fn consume(&mut self, data_block: DataBlock) -> Result<()> {
-        if let Err(cause) = self.sender.send(Ok(Some(data_block))) {
-            return Err(ErrorCode::LogicalError(format!(
-                "Logical error, cannot push data into SyncSender, cause {:?}",
-                cause
-            )));
+        if let Some(sender) = &self.sender {
+            if let Err(cause) = sender.send(data_block) {
+                return Err(ErrorCode::LogicalError(format!(
+                    "Logical error, cannot push data into SyncSender, cause {:?}",
+                    cause
+                )));
+            }
         }
 
         Ok(())