fix(operator): stop detached write retries when callers cancel

Signed-off-by: WenyXu <wenymedia@gmail.com>

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)

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);

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);

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<Histogram>,
+        cancellation: Option<&CancellationToken>,
     ) -> Result<RegionResponse> {
         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<WriteAttempt>,
         histogram: Option<&Histogram>,
+        cancellation: Option<&CancellationToken>,
     ) -> Result<RegionResponse> {
         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<WriteAttempt>,
         histogram: Option<&Histogram>,
+        cancellation: Option<&CancellationToken>,
     ) -> Result<RegionResponse> {
         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;

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<Mutex<Vec<RegionRequest>>>,
     rejections: Arc<Mutex<HashMap<u64, VecDeque<RetryHint>>>>,
+    release: Arc<Mutex<Option<Arc<Semaphore>>>>,
 }
 
 #[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));