refactor(executor): add graph information to prepare adding points (#14681)

* chore: add graph information to prepare adding points

Signed-off-by: Liuqing Yue <dqhl76@gmail.com>

* fix: unit test

Signed-off-by: Liuqing Yue <dqhl76@gmail.com>

* chore: clean

Signed-off-by: Liuqing Yue <dqhl76@gmail.com>

---------

Signed-off-by: Liuqing Yue <dqhl76@gmail.com>

Author: dqhl76

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

@@ -228,10 +228,11 @@ impl ExecutingGraph {
     pub unsafe fn init_schedule_queue(
         locker: &StateLockGuard,
         capacity: usize,
+        graph: &Arc<RunningGraph>,
     ) -> Result<ScheduleQueue> {
         let mut schedule_queue = ScheduleQueue::with_capacity(capacity);
         for sink_index in locker.graph.externals(Direction::Outgoing) {
-            ExecutingGraph::schedule_queue(locker, sink_index, &mut schedule_queue)?;
+            ExecutingGraph::schedule_queue(locker, sink_index, &mut schedule_queue, graph)?;
         }
 
         Ok(schedule_queue)
@@ -244,6 +245,7 @@ impl ExecutingGraph {
         locker: &StateLockGuard,
         index: NodeIndex,
         schedule_queue: &mut ScheduleQueue,
+        graph: &Arc<RunningGraph>,
     ) -> Result<()> {
         let mut need_schedule_nodes = VecDeque::new();
         let mut need_schedule_edges = VecDeque::new();
@@ -304,11 +306,17 @@ impl ExecutingGraph {
                     }
                     Event::NeedData | Event::NeedConsume => State::Idle,
                     Event::Sync => {
-                        schedule_queue.push_sync(node.processor.clone());
+                        schedule_queue.push_sync(ProcessorWrapper {
+                            processor: node.processor.clone(),
+                            graph: graph.clone(),
+                        });
                         State::Processing
                     }
                     Event::Async => {
-                        schedule_queue.push_async(node.processor.clone());
+                        schedule_queue.push_async(ProcessorWrapper {
+                            processor: node.processor.clone(),
+                            graph: graph.clone(),
+                        });
                         State::Processing
                     }
                 };
@@ -322,9 +330,15 @@ impl ExecutingGraph {
     }
 }
 
+#[derive(Clone)]
+pub struct ProcessorWrapper {
+    pub processor: ProcessorPtr,
+    pub graph: Arc<RunningGraph>,
+}
+
 pub struct ScheduleQueue {
-    pub sync_queue: VecDeque<ProcessorPtr>,
-    pub async_queue: VecDeque<ProcessorPtr>,
+    pub sync_queue: VecDeque<ProcessorWrapper>,
+    pub async_queue: VecDeque<ProcessorWrapper>,
 }
 
 impl ScheduleQueue {
@@ -336,25 +350,15 @@ impl ScheduleQueue {
     }
 
     #[inline]
-    pub fn push_sync(&mut self, processor: ProcessorPtr) {
+    pub fn push_sync(&mut self, processor: ProcessorWrapper) {
         self.sync_queue.push_back(processor);
     }
 
     #[inline]
-    pub fn push_async(&mut self, processor: ProcessorPtr) {
+    pub fn push_async(&mut self, processor: ProcessorWrapper) {
         self.async_queue.push_back(processor);
     }
 
-    pub fn schedule_tail(mut self, global: &ExecutorTasksQueue, ctx: &mut ExecutorWorkerContext) {
-        let mut tasks = VecDeque::with_capacity(self.sync_queue.len());
-
-        while let Some(processor) = self.sync_queue.pop_front() {
-            tasks.push_back(ExecutorTask::Sync(processor));
-        }
-
-        global.push_tasks(ctx, tasks)
-    }
-
     pub fn schedule(
         mut self,
         global: &Arc<ExecutorTasksQueue>,
@@ -384,7 +388,7 @@ impl ScheduleQueue {
     }
 
     pub fn schedule_async_task(
-        proc: ProcessorPtr,
+        proc: ProcessorWrapper,
         query_id: Arc<String>,
         executor: &Arc<PipelineExecutor>,
         wakeup_worker_id: usize,
@@ -394,18 +398,20 @@ impl ScheduleQueue {
         unsafe {
             workers_condvar.inc_active_async_worker();
             let weak_executor = Arc::downgrade(executor);
-            let node_profile = executor.graph.get_node_profile(proc.id()).clone();
-            let process_future = proc.async_process();
+            let graph = proc.graph;
+            let node_profile = executor.graph.get_node_profile(proc.processor.id()).clone();
+            let process_future = proc.processor.async_process();
             executor.async_runtime.spawn(
                 query_id.as_ref().clone(),
                 TrackedFuture::create(ProcessorAsyncTask::create(
                     query_id,
                     wakeup_worker_id,
-                    proc.clone(),
+                    proc.processor.clone(),
                     global_queue,
                     workers_condvar,
                     weak_executor,
                     node_profile,
+                    graph,
                     process_future,
                 ))
                 .in_span(Span::enter_with_local_parent(std::any::type_name::<
@@ -420,36 +426,46 @@ impl ScheduleQueue {
             ctx.set_task(ExecutorTask::Sync(processor));
         }
     }
+
+    pub fn schedule_tail(mut self, global: &ExecutorTasksQueue, ctx: &mut ExecutorWorkerContext) {
+        let mut tasks = VecDeque::with_capacity(self.sync_queue.len());
+
+        while let Some(processor) = self.sync_queue.pop_front() {
+            tasks.push_back(ExecutorTask::Sync(processor));
+        }
+
+        global.push_tasks(ctx, tasks)
+    }
 }
 
 pub struct RunningGraph(ExecutingGraph);
 
 impl RunningGraph {
-    pub fn create(pipeline: Pipeline) -> Result<RunningGraph> {
+    pub fn create(pipeline: Pipeline) -> Result<Arc<RunningGraph>> {
         let graph_state = ExecutingGraph::create(pipeline)?;
         debug!("Create running graph:{:?}", graph_state);
-        Ok(RunningGraph(graph_state))
+        Ok(Arc::new(RunningGraph(graph_state)))
     }
 
-    pub fn from_pipelines(pipelines: Vec<Pipeline>) -> Result<RunningGraph> {
+    pub fn from_pipelines(pipelines: Vec<Pipeline>) -> Result<Arc<RunningGraph>> {
         let graph_state = ExecutingGraph::from_pipelines(pipelines)?;
         debug!("Create running graph:{:?}", graph_state);
-        Ok(RunningGraph(graph_state))
+        Ok(Arc::new(RunningGraph(graph_state)))
     }
 
     /// # Safety
     ///
     /// Method is thread unsafe and require thread safe call
-    pub unsafe fn init_schedule_queue(&self, capacity: usize) -> Result<ScheduleQueue> {
-        ExecutingGraph::init_schedule_queue(&self.0, capacity)
+    pub unsafe fn init_schedule_queue(self: Arc<Self>, capacity: usize) -> Result<ScheduleQueue> {
+        ExecutingGraph::init_schedule_queue(&self.0, capacity, &self)
     }
 
     /// # Safety
     ///
     /// Method is thread unsafe and require thread safe call
-    pub unsafe fn schedule_queue(&self, node_index: NodeIndex) -> Result<ScheduleQueue> {
+    pub unsafe fn schedule_queue(self: Arc<Self>, node_index: NodeIndex) -> Result<ScheduleQueue> {
         let mut schedule_queue = ScheduleQueue::with_capacity(0);
-        ExecutingGraph::schedule_queue(&self.0, node_index, &mut schedule_queue)?;
+        ExecutingGraph::schedule_queue(&self.0, node_index, &mut schedule_queue, &self)?;
         Ok(schedule_queue)
     }
 
@@ -627,15 +643,15 @@ impl Debug for ScheduleQueue {
 
             for item in &self.sync_queue {
                 sync_queue.push(QueueItem {
-                    id: item.id().index(),
-                    name: item.name().to_string(),
+                    id: item.processor.id().index(),
+                    name: item.processor.name().to_string(),
                 })
             }
 
             for item in &self.async_queue {
                 async_queue.push(QueueItem {
-                    id: item.id().index(),
-                    name: item.name().to_string(),
+                    id: item.processor.id().index(),
+                    name: item.processor.name().to_string(),
                 })
             }
 

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

@@ -21,12 +21,13 @@ use databend_common_exception::Result;
 use parking_lot::Mutex;
 use petgraph::prelude::NodeIndex;
 
+use crate::pipelines::executor::executor_graph::ProcessorWrapper;
 use crate::pipelines::executor::ExecutorTask;
 use crate::pipelines::executor::ExecutorWorkerContext;
+use crate::pipelines::executor::RunningGraph;
 use crate::pipelines::executor::WatchNotify;
 use crate::pipelines::executor::WorkersCondvar;
 use crate::pipelines::executor::WorkersWaitingStatus;
-use crate::pipelines::processors::ProcessorPtr;
 
 pub struct ExecutorTasksQueue {
     finished: Arc<AtomicBool>,
@@ -117,7 +118,7 @@ impl ExecutorTasksQueue {
         workers_condvar.wait(worker_id, self.finished.clone());
     }
 
-    pub fn init_sync_tasks(&self, tasks: VecDeque<ProcessorPtr>) {
+    pub fn init_sync_tasks(&self, tasks: VecDeque<ProcessorWrapper>) {
         let mut workers_tasks = self.workers_tasks.lock();
 
         let mut worker_id = 0;
@@ -197,18 +198,29 @@ pub struct CompletedAsyncTask {
     pub id: NodeIndex,
     pub worker_id: usize,
     pub res: Result<()>,
+    pub graph: Arc<RunningGraph>,
 }
 
 impl CompletedAsyncTask {
-    pub fn create(id: NodeIndex, worker_id: usize, res: Result<()>) -> Self {
-        CompletedAsyncTask { id, worker_id, res }
+    pub fn create(
+        id: NodeIndex,
+        worker_id: usize,
+        res: Result<()>,
+        graph: Arc<RunningGraph>,
+    ) -> Self {
+        CompletedAsyncTask {
+            id,
+            worker_id,
+            res,
+            graph,
+        }
     }
 }
 
 struct ExecutorTasks {
     tasks_size: usize,
     workers_waiting_status: WorkersWaitingStatus,
-    workers_sync_tasks: Vec<VecDeque<ProcessorPtr>>,
+    workers_sync_tasks: Vec<VecDeque<ProcessorWrapper>>,
     workers_completed_async_tasks: Vec<VecDeque<CompletedAsyncTask>>,
 }
 

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

@@ -24,14 +24,15 @@ use databend_common_pipeline_core::processors::Profile;
 use databend_common_pipeline_core::processors::ProfileStatisticsName;
 use petgraph::prelude::NodeIndex;
 
+use crate::pipelines::executor::executor_graph::ProcessorWrapper;
 use crate::pipelines::executor::CompletedAsyncTask;
+use crate::pipelines::executor::PipelineExecutor;
 use crate::pipelines::executor::RunningGraph;
 use crate::pipelines::executor::WorkersCondvar;
-use crate::pipelines::processors::ProcessorPtr;
 
 pub enum ExecutorTask {
     None,
-    Sync(ProcessorPtr),
+    Sync(ProcessorWrapper),
     AsyncCompleted(CompletedAsyncTask),
 }
 
@@ -73,12 +74,15 @@ impl ExecutorWorkerContext {
     }
 
     /// # Safety
-    pub unsafe fn execute_task(&mut self, graph: &RunningGraph) -> Result<NodeIndex> {
+    pub unsafe fn execute_task(
+        &mut self,
+        _: &Arc<PipelineExecutor>,
+    ) -> Result<Option<(NodeIndex, Arc<RunningGraph>)>> {
         match std::mem::replace(&mut self.task, ExecutorTask::None) {
             ExecutorTask::None => Err(ErrorCode::Internal("Execute none task.")),
-            ExecutorTask::Sync(processor) => self.execute_sync_task(processor, graph),
+            ExecutorTask::Sync(processor) => self.execute_sync_task(processor),
             ExecutorTask::AsyncCompleted(task) => match task.res {
-                Ok(_) => Ok(task.id),
+                Ok(_) => Ok(Some((task.id, task.graph))),
                 Err(cause) => Err(cause),
             },
         }
@@ -87,19 +91,18 @@ impl ExecutorWorkerContext {
     /// # Safety
     unsafe fn execute_sync_task(
         &mut self,
-        proc: ProcessorPtr,
-        graph: &RunningGraph,
-    ) -> Result<NodeIndex> {
-        Profile::track_profile(graph.get_node_profile(proc.id()));
+        proc: ProcessorWrapper,
+    ) -> Result<Option<(NodeIndex, Arc<RunningGraph>)>> {
+        Profile::track_profile(proc.graph.get_node_profile(proc.processor.id()));
 
         let instant = Instant::now();
 
-        proc.process()?;
+        proc.processor.process()?;
 
         let nanos = instant.elapsed().as_nanos();
         assume(nanos < 18446744073709551615_u128);
         Profile::record_usize_profile(ProfileStatisticsName::CpuTime, nanos as usize);
-        Ok(proc.id())
+        Ok(Some((proc.processor.id(), proc.graph)))
     }
 
     pub fn get_workers_condvar(&self) -> &Arc<WorkersCondvar> {
@@ -115,8 +118,8 @@ impl Debug for ExecutorTask {
                 ExecutorTask::Sync(p) => write!(
                     f,
                     "ExecutorTask::Sync {{ id: {}, name: {}}}",
-                    p.id().index(),
-                    p.name()
+                    p.processor.id().index(),
+                    p.processor.name()
                 ),
                 ExecutorTask::AsyncCompleted(_) => write!(f, "ExecutorTask::CompletedAsync"),
             }

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

@@ -53,7 +53,7 @@ pub type FinishedCallback =
 
 pub struct PipelineExecutor {
     threads_num: usize,
-    pub(crate) graph: RunningGraph,
+    pub(crate) graph: Arc<RunningGraph>,
     workers_condvar: Arc<WorkersCondvar>,
     pub async_runtime: Arc<Runtime>,
     pub global_tasks_queue: Arc<ExecutorTasksQueue>,
@@ -173,7 +173,7 @@ impl PipelineExecutor {
     }
 
     fn try_create(
-        graph: RunningGraph,
+        graph: Arc<RunningGraph>,
         threads_num: usize,
         on_init_callback: Mutex<Option<InitCallback>>,
         on_finished_callback: Mutex<Option<FinishedCallback>>,
@@ -228,7 +228,7 @@ impl PipelineExecutor {
 
     #[minitrace::trace]
     pub fn execute(self: &Arc<Self>) -> Result<()> {
-        self.init()?;
+        self.init(self.graph.clone())?;
 
         self.start_executor_daemon()?;
 
@@ -265,7 +265,7 @@ impl PipelineExecutor {
         Ok(())
     }
 
-    fn init(self: &Arc<Self>) -> Result<()> {
+    fn init(self: &Arc<Self>, graph: Arc<RunningGraph>) -> Result<()> {
         unsafe {
             // TODO: the on init callback cannot be killed.
             {
@@ -285,7 +285,7 @@ impl PipelineExecutor {
                 );
             }
 
-            let mut init_schedule_queue = self.graph.init_schedule_queue(self.threads_num)?;
+            let mut init_schedule_queue = graph.init_schedule_queue(self.threads_num)?;
 
             let mut wakeup_worker_id = 0;
             while let Some(proc) = init_schedule_queue.async_queue.pop_front() {
@@ -400,13 +400,13 @@ impl PipelineExecutor {
             }
 
             while !self.global_tasks_queue.is_finished() && context.has_task() {
-                let executed_pid = context.execute_task(&self.graph)?;
-
-                // Not scheduled graph if pipeline is finished.
-                if !self.global_tasks_queue.is_finished() {
-                    // We immediately schedule the processor again.
-                    let schedule_queue = self.graph.schedule_queue(executed_pid)?;
-                    schedule_queue.schedule(&self.global_tasks_queue, &mut context, self);
+                if let Some((executed_pid, graph)) = context.execute_task(self)? {
+                    // Not scheduled graph if pipeline is finished.
+                    if !self.global_tasks_queue.is_finished() {
+                        // We immediately schedule the processor again.
+                        let schedule_queue = graph.schedule_queue(executed_pid)?;
+                        schedule_queue.schedule(&self.global_tasks_queue, &mut context, self);
+                    }
                 }
             }
         }