diff --git a/src/flow/src/batching_mode/batching_execution.rs b/src/flow/src/batching_mode/batching_execution.rs index 944101935a..0a699a227b 100644 --- a/src/flow/src/batching_mode/batching_execution.rs +++ b/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, + guard: BatchingExecutionGuard, task: &BatchingTask, engine: &QueryEngineRef, frontend: &Arc, max_window_cnt: Option, ) -> 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] diff --git a/src/flow/src/batching_mode/engine.rs b/src/flow/src/batching_mode/engine.rs index 458df50a4a..5bb69c3a2c 100644 --- a/src/flow/src/batching_mode/engine.rs +++ b/src/flow/src/batching_mode/engine.rs @@ -1043,6 +1043,8 @@ fn abort_flow_task(flow_id: FlowId, task: Option, 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, + _guard: crate::BatchingExecutionGuard, task: &BatchingTask, _engine: &QueryEngineRef, _frontend: &Arc, @@ -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>>; @@ -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 diff --git a/src/flow/src/batching_mode/task.rs b/src/flow/src/batching_mode/task.rs index a040848c8e..a7fa07bcbd 100644 --- a/src/flow/src/batching_mode/task.rs +++ b/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>, QueryContextRef)>, +} + +impl BatchingExecutionGuard { + fn new(lock: OwnedMutexGuard<()>) -> Self { + Self { + _lock: lock, + restore: None, + } + } + + fn restore_query_context(&mut self, state: Arc>, 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, @@ -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 { self.state.read().unwrap().last_execution_time_millis() } @@ -1332,20 +1368,22 @@ impl BatchingTask { frontend_client: &Arc, max_window_cnt: Option, ) -> 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, max_window_cnt: Option, ) -> 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, scheduled_time_secs: i64, ) -> ExecuteOnceOutcome { - let _execution_guard = self.execution_lock.lock().await; - - struct QueryContextRestoreGuard { - state: Arc>, - old_ctx: Option, - } - - 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 diff --git a/src/flow/src/batching_mode/task/test.rs b/src/flow/src/batching_mode/task/test.rs index e9f8e3e89d..d19538a3c9 100644 --- a/src/flow/src/batching_mode/task/test.rs +++ b/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, + _guard: BatchingExecutionGuard, _task: &BatchingTask, _engine: &QueryEngineRef, _frontend: &Arc, @@ -90,6 +92,56 @@ impl crate::BatchingExecution for CountingExecution { } } +struct RetainingExecution { + calls: std::sync::atomic::AtomicUsize, + started: Arc, + release: Arc, + finished: Arc, +} + +#[async_trait::async_trait] +impl crate::BatchingExecution for RetainingExecution { + async fn execute_once( + self: Arc, + guard: BatchingExecutionGuard, + _task: &BatchingTask, + _engine: &QueryEngineRef, + _frontend: &Arc, + _max_window_cnt: Option, + ) -> 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>>, + dropped: std::sync::Mutex>>, +} + +struct DropAck(Option>); + +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 { + 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 = + 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", diff --git a/src/flow/src/lib.rs b/src/flow/src/lib.rs index 5bfaf8114f..c063515606 100644 --- a/src/flow/src/lib.rs +++ b/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;