refactor(flow): transfer serialized round ownership to execution

Signed-off-by: discord9 <discord9@163.com>

src/flow/src/batching_mode/batching_execution.rs

@@ -21,17 +21,25 @@ use table::TableRef;
 
 use crate::Result;
 use crate::batching_mode::frontend_client::FrontendClient;
-use crate::batching_mode::task::{BatchingTask, ExecuteOnceOutcome};
+use crate::batching_mode::task::{BatchingExecutionGuard, BatchingTask, ExecuteOnceOutcome};
 
 #[async_trait::async_trait]
 pub trait BatchingExecution: Send + Sync + 'static {
+    /// Execute one round while retaining the guard through all task-state updates.
+    /// An implementation that continues after caller cancellation must retain the
+    /// guard with that work and make it stoppable through [`Self::stop`].
     async fn execute_once(
-        &self,
+        self: Arc<Self>,
+        guard: BatchingExecutionGuard,
         task: &BatchingTask,
         engine: &QueryEngineRef,
         frontend: &Arc<FrontendClient>,
         max_window_cnt: Option<usize>,
     ) -> ExecuteOnceOutcome;
+
+    /// Retire this execution instance, rejecting new work and requesting that any
+    /// retained local work stop. This is not an acknowledgement of remote quiescence.
+    fn stop(&self) {}
 }
 
 #[async_trait::async_trait]

src/flow/src/batching_mode/engine.rs

@@ -1043,6 +1043,8 @@ fn abort_flow_task(flow_id: FlowId, task: Option<BatchingTask>, action: &str) ->
         return false;
     };
 
+    task.stop_execution();
+
     if let Some(handle) = task.state.write().unwrap().task_handle.take() {
         handle.abort();
         debug!("Aborted {action} flow task {flow_id}");
@@ -1256,12 +1258,14 @@ mod tests {
 
     struct TestExecution {
         manual_calls: std::sync::atomic::AtomicUsize,
+        stops: std::sync::atomic::AtomicUsize,
     }
 
     #[async_trait::async_trait]
     impl crate::BatchingExecution for TestExecution {
         async fn execute_once(
-            &self,
+            self: Arc<Self>,
+            _guard: crate::BatchingExecutionGuard,
             task: &BatchingTask,
             _engine: &QueryEngineRef,
             _frontend: &Arc<FrontendClient>,
@@ -1283,6 +1287,10 @@ mod tests {
                 result: Ok(None),
             }
         }
+
+        fn stop(&self) {
+            self.stops.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+        }
     }
 
     type TestExecutionResult = crate::Result<Option<Arc<dyn crate::BatchingExecution>>>;
@@ -1322,6 +1330,7 @@ mod tests {
 
         let execution = Arc::new(TestExecution {
             manual_calls: Default::default(),
+            stops: Default::default(),
         });
         let entered = Arc::new(Notify::new());
         let release = Arc::new(Notify::new());
@@ -1412,6 +1421,7 @@ mod tests {
                 release: None,
                 result: std::sync::Mutex::new(Some(Ok(Some(Arc::new(TestExecution {
                     manual_calls: Default::default(),
+                    stops: Default::default(),
                 }))))),
             })))
             .await;
@@ -1941,11 +1951,30 @@ GROUP BY l.number, time_window
     }
 
     #[tokio::test]
-    async fn test_abort_flow_task_aborts_handle() {
+    async fn test_abort_flow_task_stops_execution_without_loop_handle() {
         let (task, _shutdown_tx) = new_test_task(42).await;
+        let execution = Arc::new(TestExecution {
+            manual_calls: Default::default(),
+            stops: Default::default(),
+        });
+        let task = task.with_execution(Some(execution.clone()));
+
+        assert!(!abort_flow_task(42, Some(task), "test"));
+        assert_eq!(execution.stops.load(std::sync::atomic::Ordering::SeqCst), 1);
+    }
+
+    #[tokio::test]
+    async fn test_abort_flow_task_stops_execution_before_aborting_handle() {
+        let (task, _shutdown_tx) = new_test_task(42).await;
+        let execution = Arc::new(TestExecution {
+            manual_calls: Default::default(),
+            stops: Default::default(),
+        });
+        let task = task.with_execution(Some(execution.clone()));
         let drop_rx = install_abort_observed_handle(&task).await;
 
         assert!(abort_flow_task(42, Some(task), "test"));
+        assert_eq!(execution.stops.load(std::sync::atomic::Ordering::SeqCst), 1);
 
         tokio::time::timeout(Duration::from_secs(1), drop_rx)
             .await

src/flow/src/batching_mode/task.rs

@@ -42,7 +42,7 @@ use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
 use table::TableRef;
 use table::table::adapter::DfTableProviderAdapter;
 use tokio::sync::oneshot::error::TryRecvError;
-use tokio::sync::{Mutex, oneshot};
+use tokio::sync::{Mutex, OwnedMutexGuard, oneshot};
 use tokio::time::Instant;
 
 use crate::batching_mode::BatchingModeOptions;
@@ -201,6 +201,36 @@ fn format_insert_target_columns(plan: &LogicalPlan) -> String {
         .join(", ")
 }
 
+/// Owns a whole serialized execution round. It may be moved to an execution
+/// collaborator, so cancellation of the caller cannot release the round early.
+pub struct BatchingExecutionGuard {
+    _lock: OwnedMutexGuard<()>,
+    restore: Option<(Arc<RwLock<TaskState>>, QueryContextRef)>,
+}
+
+impl BatchingExecutionGuard {
+    fn new(lock: OwnedMutexGuard<()>) -> Self {
+        Self {
+            _lock: lock,
+            restore: None,
+        }
+    }
+
+    fn restore_query_context(&mut self, state: Arc<RwLock<TaskState>>, old_ctx: QueryContextRef) {
+        self.restore = Some((state, old_ctx));
+    }
+}
+
+impl Drop for BatchingExecutionGuard {
+    fn drop(&mut self) {
+        if let Some((state, old_ctx)) = self.restore.take() {
+            if let Ok(mut state) = state.write() {
+                state.query_ctx = old_ctx;
+            }
+        }
+    }
+}
+
 #[derive(Clone)]
 pub struct BatchingTask {
     pub config: Arc<TaskConfig>,
@@ -361,6 +391,12 @@ impl BatchingTask {
         self
     }
 
+    pub(crate) fn stop_execution(&self) {
+        if let Some(execution) = &self.execution {
+            execution.stop();
+        }
+    }
+
     pub fn last_execution_time_millis(&self) -> Option<i64> {
         self.state.read().unwrap().last_execution_time_millis()
     }
@@ -1332,20 +1368,22 @@ impl BatchingTask {
         frontend_client: &Arc<FrontendClient>,
         max_window_cnt: Option<usize>,
     ) -> ExecuteOnceOutcome {
-        let _execution_guard = self.execution_lock.lock().await;
-        self.execute_once_unlocked(engine, frontend_client, max_window_cnt)
+        let guard = BatchingExecutionGuard::new(self.execution_lock.clone().lock_owned().await);
+        self.execute_once_with_guard(guard, engine, frontend_client, max_window_cnt)
             .await
     }
 
-    async fn execute_once_unlocked(
+    async fn execute_once_with_guard(
         &self,
+        guard: BatchingExecutionGuard,
         engine: &QueryEngineRef,
         frontend_client: &Arc<FrontendClient>,
         max_window_cnt: Option<usize>,
     ) -> ExecuteOnceOutcome {
         if let Some(execution) = &self.execution {
             return execution
-                .execute_once(self, engine, frontend_client, max_window_cnt)
+                .clone()
+                .execute_once(guard, self, engine, frontend_client, max_window_cnt)
                 .await;
         }
         self.execute_once_default_unlocked(engine, frontend_client, max_window_cnt)
@@ -1418,23 +1456,7 @@ impl BatchingTask {
         frontend_client: &Arc<FrontendClient>,
         scheduled_time_secs: i64,
     ) -> ExecuteOnceOutcome {
-        let _execution_guard = self.execution_lock.lock().await;
-
-        struct QueryContextRestoreGuard {
-            state: Arc<RwLock<TaskState>>,
-            old_ctx: Option<QueryContextRef>,
-        }
-
-        impl Drop for QueryContextRestoreGuard {
-            fn drop(&mut self) {
-                let Some(old_ctx) = self.old_ctx.take() else {
-                    return;
-                };
-                if let Ok(mut state) = self.state.write() {
-                    state.query_ctx = old_ctx;
-                }
-            }
-        }
+        let mut guard = BatchingExecutionGuard::new(self.execution_lock.clone().lock_owned().await);
 
         // Convert to milliseconds before touching the task state so an
         // unrepresentable scheduled time fails as an explicit error without
@@ -1462,21 +1484,12 @@ impl BatchingTask {
             state.query_ctx = Arc::new(new_ctx);
             old
         };
-        let restore_guard = QueryContextRestoreGuard {
-            state: self.state.clone(),
-            old_ctx: Some(old_ctx),
-        };
+        guard.restore_query_context(self.state.clone(), old_ctx);
 
-        let outcome = self
-            .execute_once_unlocked(engine, frontend_client, None)
-            .await;
-
-        // Restore while still holding `execution_lock` so no future manual
-        // flush can observe the temporary scheduled time. The guard also
-        // restores during unwind/cancellation.
-        drop(restore_guard);
-
-        outcome
+        // A collaborator may retain the guard in an owned child task. Its Drop
+        // restores the scheduled context before releasing serialization.
+        self.execute_once_with_guard(guard, engine, frontend_client, None)
+            .await
     }
 
     /// Generate the create table SQL

src/flow/src/batching_mode/task/test.rs

@@ -14,6 +14,7 @@
 
 use std::collections::{BTreeMap, BTreeSet, HashMap};
 use std::sync::Arc;
+use std::task::Poll;
 
 use catalog::memory::MemoryCatalogManager;
 use catalog::{DeregisterTableRequest, RegisterTableRequest};
@@ -67,7 +68,8 @@ struct CountingExecution {
 #[async_trait::async_trait]
 impl crate::BatchingExecution for CountingExecution {
     async fn execute_once(
-        &self,
+        self: Arc<Self>,
+        _guard: BatchingExecutionGuard,
         _task: &BatchingTask,
         _engine: &QueryEngineRef,
         _frontend: &Arc<FrontendClient>,
@@ -90,6 +92,56 @@ impl crate::BatchingExecution for CountingExecution {
     }
 }
 
+struct RetainingExecution {
+    calls: std::sync::atomic::AtomicUsize,
+    started: Arc<tokio::sync::Notify>,
+    release: Arc<tokio::sync::Notify>,
+    finished: Arc<tokio::sync::Notify>,
+}
+
+#[async_trait::async_trait]
+impl crate::BatchingExecution for RetainingExecution {
+    async fn execute_once(
+        self: Arc<Self>,
+        guard: BatchingExecutionGuard,
+        _task: &BatchingTask,
+        _engine: &QueryEngineRef,
+        _frontend: &Arc<FrontendClient>,
+        _max_window_cnt: Option<usize>,
+    ) -> ExecuteOnceOutcome {
+        if self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst) != 0 {
+            return ExecuteOnceOutcome {
+                new_query: None,
+                result: Ok(None),
+            };
+        }
+
+        let started = self.started.clone();
+        let release = self.release.clone();
+        let finished = self.finished.clone();
+        let child = tokio::spawn(async move {
+            started.notify_one();
+            release.notified().await;
+            drop(guard);
+            finished.notify_one();
+            ExecuteOnceOutcome {
+                new_query: None,
+                result: Ok(None),
+            }
+        });
+        match child.await {
+            Ok(outcome) => outcome,
+            Err(err) => ExecuteOnceOutcome {
+                new_query: None,
+                result: Err(Error::Unexpected {
+                    reason: format!("retaining test child failed: {err}"),
+                    location: snafu::location!(),
+                }),
+            },
+        }
+    }
+}
+
 #[tokio::test]
 async fn test_execution_delegate_dispatch_is_serialized() {
     let TestTaskParts {
@@ -120,6 +172,216 @@ async fn test_execution_delegate_dispatch_is_serialized() {
     );
 }
 
+#[tokio::test]
+async fn test_delegate_guard_survives_caller_cancellation_until_child_finishes() {
+    let TestTaskParts {
+        task, query_engine, ..
+    } = new_test_task_engine_and_plan_with_query("SELECT number, ts FROM numbers_with_ts", "sink")
+        .await;
+    let execution = Arc::new(RetainingExecution {
+        calls: Default::default(),
+        started: Arc::new(tokio::sync::Notify::new()),
+        release: Arc::new(tokio::sync::Notify::new()),
+        finished: Arc::new(tokio::sync::Notify::new()),
+    });
+    let task = task.with_execution(Some(execution.clone()));
+    let (frontend, _handler) = FrontendClient::from_empty_grpc_handler(QueryOptions::default());
+    let frontend = Arc::new(frontend);
+
+    let first_task = task.clone();
+    let first_engine = query_engine.clone();
+    let first_frontend = frontend.clone();
+    let first = tokio::spawn(async move {
+        first_task
+            .execute_once_serialized(&first_engine, &first_frontend, None)
+            .await
+    });
+    tokio::time::timeout(Duration::from_secs(1), execution.started.notified())
+        .await
+        .expect("delegate child did not retain the guard");
+    first.abort();
+    assert!(
+        first
+            .await
+            .expect_err("caller cancellation should abort")
+            .is_cancelled()
+    );
+
+    let second = task.execute_once_serialized(&query_engine, &frontend, None);
+    futures::pin_mut!(second);
+    assert!(
+        matches!(futures::poll!(second.as_mut()), Poll::Pending),
+        "the next round must remain pending while the retained guard is held"
+    );
+    assert_eq!(
+        execution.calls.load(std::sync::atomic::Ordering::SeqCst),
+        1,
+        "the pending waiter must not enter the collaborator"
+    );
+    execution.release.notify_one();
+    tokio::time::timeout(Duration::from_secs(1), execution.finished.notified())
+        .await
+        .expect("delegate child did not release its guard");
+    assert_eq!(
+        tokio::time::timeout(Duration::from_secs(1), second)
+            .await
+            .expect("next round should proceed after child release")
+            .unwrap(),
+        None
+    );
+    assert_eq!(execution.calls.load(std::sync::atomic::Ordering::SeqCst), 2);
+}
+
+#[tokio::test]
+async fn test_scheduled_context_is_retained_until_delegate_child_releases_guard() {
+    let TestTaskParts {
+        task, query_engine, ..
+    } = new_test_task_engine_and_plan_with_query("SELECT number, ts FROM numbers_with_ts", "sink")
+        .await;
+    let execution = Arc::new(RetainingExecution {
+        calls: Default::default(),
+        started: Arc::new(tokio::sync::Notify::new()),
+        release: Arc::new(tokio::sync::Notify::new()),
+        finished: Arc::new(tokio::sync::Notify::new()),
+    });
+    let task = task.with_execution(Some(execution.clone()));
+    let (frontend, _handler) = FrontendClient::from_empty_grpc_handler(QueryOptions::default());
+    let frontend = Arc::new(frontend);
+    let scheduled = 1_700_000_000;
+
+    let task_to_run = task.clone();
+    let engine_to_run = query_engine.clone();
+    let frontend_to_run = frontend.clone();
+    let execution_call = tokio::spawn(async move {
+        task_to_run
+            .execute_once_serialized_at_scheduled_time(&engine_to_run, &frontend_to_run, scheduled)
+            .await
+    });
+    tokio::time::timeout(Duration::from_secs(1), execution.started.notified())
+        .await
+        .expect("scheduled delegate child did not start");
+    assert_eq!(
+        task.state
+            .read()
+            .unwrap()
+            .query_ctx
+            .extension(FLOW_SCHEDULED_TIME_MILLIS),
+        Some("1700000000000"),
+        "scheduled context restored before the delegate child released the guard"
+    );
+    execution_call.abort();
+    match execution_call.await {
+        Err(error) => assert!(error.is_cancelled()),
+        Ok(_) => panic!("scheduled caller cancellation should abort"),
+    }
+    assert_eq!(
+        task.state
+            .read()
+            .unwrap()
+            .query_ctx
+            .extension(FLOW_SCHEDULED_TIME_MILLIS),
+        Some("1700000000000"),
+        "scheduled context restored after caller cancellation but before child release"
+    );
+
+    execution.release.notify_one();
+    tokio::time::timeout(Duration::from_secs(1), execution.finished.notified())
+        .await
+        .expect("scheduled delegate child did not release its guard");
+    assert_eq!(
+        task.state
+            .read()
+            .unwrap()
+            .query_ctx
+            .extension(FLOW_SCHEDULED_TIME_MILLIS),
+        None,
+        "scheduled context was not restored when child released guard"
+    );
+}
+
+struct BlockingDefaultExecutionHandler {
+    entered: std::sync::Mutex<Option<tokio::sync::oneshot::Sender<()>>>,
+    dropped: std::sync::Mutex<Option<tokio::sync::oneshot::Sender<()>>>,
+}
+
+struct DropAck(Option<tokio::sync::oneshot::Sender<()>>);
+
+impl Drop for DropAck {
+    fn drop(&mut self) {
+        if let Some(dropped) = self.0.take() {
+            let _ = dropped.send(());
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError
+    for BlockingDefaultExecutionHandler
+{
+    async fn do_query(
+        &self,
+        _query: api::v1::greptime_request::Request,
+        _ctx: QueryContextRef,
+    ) -> std::result::Result<Output, BoxedError> {
+        let _ack = DropAck(self.dropped.lock().unwrap().take());
+        if let Some(entered) = self.entered.lock().unwrap().take() {
+            let _ = entered.send(());
+        }
+        std::future::pending().await
+    }
+}
+
+#[tokio::test]
+async fn test_default_execution_remains_inline_and_cancellable() {
+    let query = "SELECT number, date_bin(INTERVAL '5 second', ts) AS time_window \
+                 FROM numbers_with_ts GROUP BY time_window, number";
+    let TestTaskParts {
+        task, query_engine, ..
+    } = new_time_window_test_task_with_query(query).await;
+    register_twe_sink(&query_engine, "missing_sink", 9200);
+    task.mark_all_windows_as_dirty().unwrap();
+
+    let (entered_tx, entered_rx) = tokio::sync::oneshot::channel();
+    let (dropped_tx, dropped_rx) = tokio::sync::oneshot::channel();
+    let handler: Arc<dyn crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError> =
+        Arc::new(BlockingDefaultExecutionHandler {
+            entered: std::sync::Mutex::new(Some(entered_tx)),
+            dropped: std::sync::Mutex::new(Some(dropped_tx)),
+        });
+    let frontend = Arc::new(FrontendClient::from_grpc_handler(
+        Arc::downgrade(&handler),
+        QueryOptions::default(),
+    ));
+
+    let task_to_cancel = task.clone();
+    let engine_to_cancel = query_engine.clone();
+    let frontend_to_cancel = frontend.clone();
+    let caller = tokio::spawn(async move {
+        task_to_cancel
+            .execute_once_serialized(&engine_to_cancel, &frontend_to_cancel, None)
+            .await
+    });
+    tokio::time::timeout(Duration::from_secs(1), entered_rx)
+        .await
+        .expect("default execution did not dispatch a frontend query")
+        .expect("default execution handler entry notification dropped");
+    caller.abort();
+    match caller.await {
+        Err(error) => assert!(error.is_cancelled()),
+        Ok(_) => panic!("default caller cancellation should abort"),
+    }
+    tokio::time::timeout(Duration::from_secs(1), dropped_rx)
+        .await
+        .expect("cancelling default execution did not drop the active frontend future")
+        .expect("default execution drop acknowledgement was not sent");
+    tokio::time::timeout(
+        Duration::from_secs(1),
+        task.execution_lock.clone().lock_owned(),
+    )
+    .await
+    .expect("default execution must not leave an owned child holding the lock");
+}
+
 async fn new_test_task_and_plan_with_missing_sink() -> (BatchingTask, LogicalPlan) {
     new_test_task_engine_and_plan_with_query(
         "SELECT number, ts FROM numbers_with_ts",

src/flow/src/lib.rs

@@ -49,7 +49,8 @@ pub use batching_mode::frontend_client::{
     FrontendClient, GrpcQueryHandlerWithBoxedError, PeerDesc,
 };
 pub use batching_mode::task::{
-    BatchingTask, DirtyRestore, ExecuteOnceOutcome, PlanInfo, QueryCoverage, TaskArgs,
+    BatchingExecutionGuard, BatchingTask, DirtyRestore, ExecuteOnceOutcome, PlanInfo,
+    QueryCoverage, TaskArgs,
 };
 pub use batching_mode::time_window::{TimeWindowExpr, find_time_window_expr};
 pub use batching_mode::utils::sql_to_df_plan;