From d27dce55a131d01abcdcce75402d82dc94fbdc4c Mon Sep 17 00:00:00 2001 From: WenyXu Date: Tue, 13 Oct 2026 04:34:33 +0000 Subject: [PATCH] fix(operator): stop detached write retries when callers cancel Signed-off-by: WenyXu --- src/operator/src/bulk_insert.rs | 8 +- src/operator/src/delete.rs | 17 ++-- src/operator/src/insert.rs | 17 ++-- src/operator/src/write_dispatcher.rs | 32 ++++++- src/operator/src/write_dispatcher/tests.rs | 101 ++++++++++++++++++++- 5 files changed, 154 insertions(+), 21 deletions(-) diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index 29537959df5..d4a1ed0497c 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -35,6 +35,7 @@ use snafu::{OptionExt, ResultExt, ensure}; use store_api::storage::RegionId; use table::TableRef; use table::metadata::TableInfoRef; +use tokio_util::sync::CancellationToken; use crate::error::Result; use crate::insert::Inserter; @@ -72,7 +73,7 @@ impl Inserter { let writes = self.write_dispatcher.prepare(writes).await?; Ok(self .write_dispatcher - .execute(writes, None) + .execute(writes, None, None) .await .context(error::RequestInsertsSnafu)? .affected_rows) @@ -191,6 +192,7 @@ impl Inserter { Some( metrics::HANDLE_BULK_INSERT_ELAPSED.with_label_values(&["datanode_handle"]), ), + None, ) .await .context(error::RequestRegionSnafu) @@ -217,6 +219,9 @@ impl Inserter { .with_label_values(&["wait_all_datanode"]) .start_timer(); + // Dropping the caller stops retries without abandoning in-flight writes. + let cancellation = CancellationToken::new(); + let _guard = cancellation.clone().drop_guard(); let mut handles = Vec::with_capacity(region_masks.len()); for (region_number, mask) in region_masks { @@ -328,6 +333,7 @@ impl Inserter { &metrics::HANDLE_BULK_INSERT_ELAPSED .with_label_values(&["datanode_handle"]), ), + Some(&cancellation), ) .await .context(error::RequestRegionSnafu) diff --git a/src/operator/src/delete.rs b/src/operator/src/delete.rs index c389374b7ac..9bf14471696 100644 --- a/src/operator/src/delete.rs +++ b/src/operator/src/delete.rs @@ -27,6 +27,7 @@ use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt, ensure}; use table::TableRef; use table::requests::DeleteRequest as TableDeleteRequest; +use tokio_util::sync::CancellationToken; use crate::error::{ CatalogSnafu, InvalidDeleteRequestSnafu, JoinTaskSnafu, MissingTimeIndexColumnSnafu, @@ -141,12 +142,16 @@ impl Deleter { let dispatcher = self.write_dispatcher.clone(); let request = request_factory.build_delete(requests); let writes = dispatcher.prepare(vec![request]).await?; - let affected_rows = - common_runtime::spawn_global(async move { dispatcher.execute(writes, None).await }) - .await - .context(JoinTaskSnafu)? - .context(RequestDeletesSnafu)? - .affected_rows; + // Dropping the caller stops retries without abandoning in-flight writes. + let cancellation = CancellationToken::new(); + let _guard = cancellation.clone().drop_guard(); + let affected_rows = common_runtime::spawn_global(async move { + dispatcher.execute(writes, None, Some(&cancellation)).await + }) + .await + .context(JoinTaskSnafu)? + .context(RequestDeletesSnafu)? + .affected_rows; crate::metrics::DIST_DELETE_ROW_COUNT .with_label_values(&[ctx.get_db_string().as_str()]) .inc_by(affected_rows as u64); diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 2c01a35b8d9..4c0c1424846 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -80,6 +80,7 @@ use table::requests::{ validate_semantic_option, }; use table::table_reference::TableReference; +use tokio_util::sync::CancellationToken; use crate::batcher::PendingRowsBatcher; use crate::error::{ @@ -841,12 +842,16 @@ impl Inserter { let dispatcher = self.write_dispatcher.clone(); let request = request_factory.build_insert(normal_requests); let writes = dispatcher.prepare(vec![request]).await?; - let affected_rows = - common_runtime::spawn_global(async move { dispatcher.execute(writes, None).await }) - .await - .context(JoinTaskSnafu)? - .context(RequestInsertsSnafu)? - .affected_rows; + // Dropping the caller stops retries without abandoning in-flight writes. + let cancellation = CancellationToken::new(); + let _guard = cancellation.clone().drop_guard(); + let affected_rows = common_runtime::spawn_global(async move { + dispatcher.execute(writes, None, Some(&cancellation)).await + }) + .await + .context(JoinTaskSnafu)? + .context(RequestInsertsSnafu)? + .affected_rows; crate::metrics::DIST_INGEST_ROW_COUNT .with_label_values(&[ctx.get_db_string().as_str()]) .inc_by(affected_rows as u64); diff --git a/src/operator/src/write_dispatcher.rs b/src/operator/src/write_dispatcher.rs index 71c8d583afe..6e03eed0b2e 100644 --- a/src/operator/src/write_dispatcher.rs +++ b/src/operator/src/write_dispatcher.rs @@ -36,6 +36,7 @@ use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef}; use prometheus::Histogram; use snafu::{OptionExt, ResultExt}; use store_api::storage::RegionId; +use tokio_util::sync::CancellationToken; pub use self::options::WriteRetryOptions; use crate::error::{self, FindRegionLeaderSnafu, RequestRegionSnafu}; @@ -170,10 +171,12 @@ impl WriteDispatcher { } /// Sends requests using routes already resolved by the caller. + /// Cancellation stops future retries while allowing in-flight requests to complete. pub(crate) async fn execute( &self, writes: Vec<(PeerRef, RegionRequest)>, histogram: Option, + cancellation: Option<&CancellationToken>, ) -> Result { let histogram = histogram.as_ref(); let writes = if self.retry.is_some() { @@ -194,7 +197,7 @@ impl WriteDispatcher { .map(|(peer, request)| self.send(peer, request, histogram)), ) .await; - self.finish(attempts, histogram).await + self.finish(attempts, histogram, cancellation).await } /// Sends the first attempt immediately so bulk encoding can overlap other sends. @@ -231,10 +234,11 @@ impl WriteDispatcher { &self, attempts: Vec, histogram: Option<&Histogram>, + cancellation: Option<&CancellationToken>, ) -> Result { if let Some(retry) = &self.retry { if retry.should_retry(&attempts, 0) { - return retry.retry(attempts, histogram).await; + return retry.retry(attempts, histogram, cancellation).await; } let mut response = RegionResponse::new(0); let failures = merge_attempts(attempts, &mut response)?; @@ -271,7 +275,7 @@ impl WriteDispatcher { .await .map_err(BoxedError::new) .context(ExternalSnafu)?; - self.execute(writes, histogram).await + self.execute(writes, histogram, None).await } } @@ -327,6 +331,7 @@ impl WriteRetry { &self, mut attempts: Vec, histogram: Option<&Histogram>, + cancellation: Option<&CancellationToken>, ) -> Result { let mut response = RegionResponse::new(0); let mut backoff = ExponentialBuilder::default() @@ -341,9 +346,20 @@ impl WriteRetry { let Some(delay) = backoff.next() else { break; }; + if let Some(cancellation) = cancellation { + tokio::select! { + biased; + _ = cancellation.cancelled() => break, + _ = tokio::time::sleep(delay) => {} + } + } else { + tokio::time::sleep(delay).await; + } + if cancellation.is_some_and(CancellationToken::is_cancelled) { + break; + } let failures = merge_attempts(attempts, &mut response)?; - tokio::time::sleep(delay).await; - + let mut first_error = None; let mut refresh = Vec::new(); let mut pending = Vec::new(); for (write, error) in failures { @@ -352,10 +368,16 @@ impl WriteRetry { } else { pending.push(write); } + first_error.get_or_insert(error); } if !refresh.is_empty() { pending.extend(self.reroute_failed_writes(refresh).await?); } + if let Some(error) = + first_error.filter(|_| cancellation.is_some_and(CancellationToken::is_cancelled)) + { + return Err(error); + } attempts = future::join_all(pending.into_iter().map(|write| self.send(write, histogram))) .await; diff --git a/src/operator/src/write_dispatcher/tests.rs b/src/operator/src/write_dispatcher/tests.rs index 250fe87f353..4a5691f2eb3 100644 --- a/src/operator/src/write_dispatcher/tests.rs +++ b/src/operator/src/write_dispatcher/tests.rs @@ -56,7 +56,7 @@ use partition::expr::{Operand, PartitionExpr, RestrictedOp}; use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef}; use store_api::storage::RegionId; use tokio::net::TcpListener; -use tokio::sync::{Barrier, Notify}; +use tokio::sync::{Barrier, Notify, Semaphore}; use tokio::task::JoinHandle; use tokio_stream::wrappers::TcpListenerStream; use tonic::{Request, Response, Status}; @@ -69,6 +69,7 @@ use crate::test_util; struct RecordingRegionService { requests: Arc>>, rejections: Arc>>>, + release: Arc>>>, } #[async_trait::async_trait] @@ -80,6 +81,10 @@ impl Region for RecordingRegionService { let request = request.into_inner(); let region_id = write_region_id(&request).unwrap().as_u64(); self.requests.lock().unwrap().push(request); + let release = self.release.lock().unwrap().clone(); + if let Some(release) = release { + release.acquire().await.unwrap().forget(); + } if let Some(hint) = self .rejections .lock() @@ -402,6 +407,96 @@ async fn test_should_retry_requires_a_retained_request_and_budget() { } } +#[tokio::test] +async fn test_write_retry_cancellation_stops_backoff() { + for hint in [RetryHint::Retryable, RetryHint::RetryAfterRouteRefresh] { + let server = RegionTestServer::start().await; + server.fail_region(1, [hint]); + let mut dispatcher = server.create_dispatcher(2).await; + let options = &mut dispatcher.retry.as_mut().unwrap().options; + options.initial_backoff = Duration::from_secs(60); + options.max_backoff = Duration::from_secs(60); + let attempt = dispatcher + .send( + Arc::new(server.peer.clone()), + build_insert_request(&[1]), + None, + ) + .await; + let cancellation = CancellationToken::new(); + let guard = cancellation.clone().drop_guard(); + let finish = dispatcher.finish(vec![attempt], None, Some(&cancellation)); + tokio::pin!(finish); + futures::future::poll_fn(|cx| { + assert!(finish.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + drop(guard); + let error = tokio::time::timeout(Duration::from_secs(1), finish) + .await + .unwrap() + .unwrap_err(); + assert_eq!(error.retry_hint(), hint); + assert_eq!( + server.attempt_counts(), + HashMap::from([(RegionId::new(1, 1).as_u64(), 1)]) + ); + assert!(server.invalidator.calls.lock().unwrap().is_empty()); + } +} + +#[tokio::test] +async fn test_write_retry_cancellation_drains_inflight_requests() { + let server = RegionTestServer::start().await; + server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]); + let release = Arc::new(Semaphore::new(0)); + *server.service.release.lock().unwrap() = Some(release.clone()); + let dispatcher = server.create_dispatcher(2).await; + let writes = dispatcher + .prepare(vec![build_insert_request(&[1, 2])]) + .await + .unwrap(); + let cancellation = CancellationToken::new(); + let guard = cancellation.clone().drop_guard(); + let execute = dispatcher.execute(writes, None, Some(&cancellation)); + tokio::pin!(execute); + tokio::time::timeout( + Duration::from_secs(1), + futures::future::poll_fn(|cx| { + assert!(execute.as_mut().poll(cx).is_pending()); + if server.service.requests.lock().unwrap().len() == 2 { + Poll::Ready(()) + } else { + cx.waker().wake_by_ref(); + Poll::Pending + } + }), + ) + .await + .unwrap(); + drop(guard); + futures::future::poll_fn(|cx| { + assert!(execute.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + release.add_permits(2); + let error = tokio::time::timeout(Duration::from_secs(1), execute) + .await + .unwrap() + .unwrap_err(); + assert_eq!(error.retry_hint(), RetryHint::RetryAfterRouteRefresh); + assert_eq!( + server.attempt_counts(), + HashMap::from([ + (RegionId::new(1, 1).as_u64(), 1), + (RegionId::new(1, 2).as_u64(), 1), + ]) + ); + assert!(server.invalidator.calls.lock().unwrap().is_empty()); +} + #[tokio::test] async fn test_write_retry_refreshes_routes_once_per_failed_round() { let server = RegionTestServer::start().await; @@ -469,7 +564,7 @@ async fn test_finish_retries_initial_failures_in_one_round() { (RegionId::new(1, 3).as_u64(), 1), ]) ); - let response = dispatcher.finish(attempts, None).await.unwrap(); + let response = dispatcher.finish(attempts, None, None).await.unwrap(); assert_eq!(response.affected_rows, 3); assert_eq!( server.attempt_counts(), @@ -929,7 +1024,7 @@ async fn test_dispatch_preserves_write_payloads_and_envelopes() { }); let writes = dispatcher.prepare(requests.to_vec()).await.unwrap(); assert!(received.try_recv().is_err()); - let response = dispatcher.execute(writes, None).await.unwrap(); + let response = dispatcher.execute(writes, None, None).await.unwrap(); assert_eq!(response.affected_rows, 2); for request in requests { assert_eq!(received.try_recv().unwrap(), (Peer::new(1, ""), request));