From 9e9346d2ff2ba754fb4d6740ec1be22b7c2e1800 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 22 Oct 2026 03:36:02 +0000 Subject: [PATCH] fix: preserve datanode write batches when retrying (#9720) * fix: preserve datanode write batches when retrying Signed-off-by: WenyXu * fix: isolate invalid entries in metric write batches Signed-off-by: WenyXu * refactor: adapt batch writes to the unified metric API Signed-off-by: WenyXu * fix: explicitly accept partial results for batched retries Signed-off-by: WenyXu * test: assert terminal retry hint instead of masked error text Signed-off-by: WenyXu * test: preserve application status in region failure fixtures Signed-off-by: WenyXu * fix: refresh routes after region close Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/common/meta/src/error.rs | 5 + src/mito2/src/error.rs | 16 ++ src/mito2/src/flush.rs | 29 +- src/operator/src/write_dispatcher.rs | 109 ++++++-- src/operator/src/write_dispatcher/tests.rs | 302 ++++++++++++++------- 5 files changed, 329 insertions(+), 132 deletions(-) diff --git a/src/common/meta/src/error.rs b/src/common/meta/src/error.rs index c048f499dff..7dd7b9b0a62 100644 --- a/src/common/meta/src/error.rs +++ b/src/common/meta/src/error.rs @@ -41,6 +41,9 @@ pub use retry_hint::{retry_hint_from_postgres_error, retry_hint_from_postgres_po #[snafu(visibility(pub))] #[stack_trace_debug] pub enum Error { + #[snafu(display("Failed to execute batch request"))] + BatchRequest { source: Arc }, + #[snafu(display( "Region request has {} failed entries; {}", response.failures.len(), @@ -1216,6 +1219,7 @@ impl ErrorExt for Error { Error::PartialRegion { response } => response .first_failure() .map_or(StatusCode::Internal, |(_, error)| error.status_code()), + Error::BatchRequest { source } => source.status_code(), IllegalServerState { .. } | EtcdTxnOpResponse { .. } | EtcdFailed { .. } @@ -1386,6 +1390,7 @@ impl ErrorExt for Error { match self { Error::PartialRegion { .. } => RetryHint::Retryable, + Error::BatchRequest { source } => source.retry_hint(), RetryLater { source, .. } => match source.retry_hint() { RetryHint::RetryAfterRouteRefresh => RetryHint::RetryAfterRouteRefresh, RetryHint::Retryable | RetryHint::NonRetryable => RetryHint::Retryable, diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index d7ecdd88b1b..71253e78bea 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -1679,6 +1679,7 @@ impl ErrorExt for Error { match self { PartitionExprVersionMismatch { .. } | RegionNotFound { .. } + | RegionClosed { .. } | RejectWriteInStaging { .. } | RegionState { state: RegionRoleState::Follower, @@ -1819,6 +1820,21 @@ mod tests { StatusCode::RegionNotFound, RetryHint::RetryAfterRouteRefresh, ), + ( + RegionClosedSnafu { region_id }.build(), + StatusCode::Cancelled, + RetryHint::RetryAfterRouteRefresh, + ), + ( + RegionDroppedSnafu { region_id }.build(), + StatusCode::Cancelled, + RetryHint::NonRetryable, + ), + ( + RegionTruncatedSnafu { region_id }.build(), + StatusCode::Cancelled, + RetryHint::NonRetryable, + ), ] { assert_eq!(error.status_code(), status); assert_eq!(error.retry_hint(), hint); diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index 52eb311de73..fba0607235b 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -1679,7 +1679,7 @@ impl FlushStatus { #[cfg(test)] mod tests { use api::v1::{OpType, Rows}; - use common_error::ext::ErrorExt; + use common_error::ext::{ErrorExt, RetryHint}; use common_error::status_code::StatusCode; use datatypes::arrow::datatypes::Schema; use datatypes::arrow::record_batch::RecordBatch; @@ -2122,28 +2122,39 @@ mod tests { } #[tokio::test] - async fn test_region_closed_notifies_pending_bulk_writes() { + async fn test_region_closed_notifies_pending_requests() { let region_id = RegionId::new(1, 1); let version_control = Arc::new(VersionControlBuilder::new().build()); - let (bulk_req, output_rx) = new_test_bulk_request(region_id); + let (write_req, write_rx) = new_test_write_request(region_id); + let (bulk_req, bulk_rx) = new_test_bulk_request(region_id); + let (ddl_tx, ddl_rx) = oneshot::channel(); let status = FlushStatus { region_id, version_control, state: CancellableTaskState::new(), pending_task: None, closing: false, - pending_ddls: Vec::new(), - pending_writes: Vec::new(), + pending_ddls: vec![SenderDdlRequest { + region_id, + sender: OptionOutputTx::from(ddl_tx), + request: DdlRequest::Truncate( + store_api::region_request::RegionTruncateRequest::All, + ), + }], + pending_writes: vec![write_req], pending_bulk_writes: vec![bulk_req], }; status.on_region_closed(Arc::new(RegionClosedSnafu { region_id }.build())); - let err = output_rx - .await - .expect("pending bulk write must receive explicit error") - .unwrap_err(); + for rx in [write_rx, bulk_rx] { + let err = rx.await.unwrap().unwrap_err(); + assert_eq!(err.status_code(), StatusCode::Cancelled); + assert_eq!(err.retry_hint(), RetryHint::RetryAfterRouteRefresh); + } + let err = ddl_rx.await.unwrap().unwrap_err(); assert_eq!(err.status_code(), StatusCode::Cancelled); + assert_eq!(err.retry_hint(), RetryHint::RetryAfterRouteRefresh); } #[tokio::test] diff --git a/src/operator/src/write_dispatcher.rs b/src/operator/src/write_dispatcher.rs index d8a6ac9905e..2d12feb74c5 100644 --- a/src/operator/src/write_dispatcher.rs +++ b/src/operator/src/write_dispatcher.rs @@ -25,7 +25,9 @@ use backon::{BackoffBuilder, ExponentialBuilder}; use catalog::CatalogManagerRef; use common_error::ext::{BoxedError, ErrorExt, RetryHint}; use common_meta::cache_invalidator::{CacheInvalidatorRef, Context}; -use common_meta::error::{Error, ExternalSnafu, Result, TableInfoNotFoundSnafu, UnexpectedSnafu}; +use common_meta::error::{ + BatchRequestSnafu, Error, ExternalSnafu, Result, TableInfoNotFoundSnafu, UnexpectedSnafu, +}; use common_meta::instruction::CacheIdent; use common_meta::node_manager::NodeManagerRef; use common_meta::peer::{Peer, PeerRef}; @@ -33,7 +35,7 @@ use common_meta::rpc::router; use futures::future; use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef}; use prometheus::Histogram; -use snafu::{OptionExt, ResultExt, ensure}; +use snafu::{IntoError, OptionExt, ResultExt, ensure}; use store_api::storage::RegionId; use table::table_name::TableName; use tokio_util::sync::CancellationToken; @@ -178,24 +180,17 @@ impl WriteDispatcher { cancellation: Option<&CancellationToken>, ) -> Result { let histogram = histogram.as_ref(); - let writes = if self.retry.is_some() { - writes - .into_iter() - .flat_map(|(peer, request)| { - split_request(request) - .into_iter() - .map(move |request| (Arc::clone(&peer), request)) - }) - .collect() - } else { - writes - }; - let attempts = future::join_all( - writes - .into_iter() - .map(|(peer, request)| self.send(peer, request, histogram)), - ) - .await; + let attempts = future::join_all(writes.into_iter().map(|(peer, request)| async move { + if self.retry.is_some() { + self.send_batch(peer, request, histogram).await + } else { + vec![self.send(peer, request, histogram).await] + } + })) + .await + .into_iter() + .flatten() + .collect(); self.finish(attempts, histogram, cancellation).await } @@ -267,6 +262,80 @@ impl WriteDispatcher { resolve_peers(&self.partition_manager, regions).await } + async fn send_batch( + &self, + peer: PeerRef, + mut request: RegionRequest, + histogram: Option<&Histogram>, + ) -> Vec { + if !matches!(&request.body, Some(Body::Inserts(_) | Body::Deletes(_))) { + return vec![self.send(peer, request, histogram).await]; + } + // Upgrade datanodes first and roll back frontends first: old datanodes + // cannot report which entries succeeded, so retrying a batch may replay writes. + api::region::set_accept_partial_result(request.header.get_or_insert_default()); + let request = Arc::new(request); + let requester = self.node_manager.datanode(&peer).await; + let result = { + let _timer = histogram.map(Histogram::start_timer); + requester.handle(Arc::clone(&request)).await + }; + let error = match result { + Ok(response) => return vec![WriteAttempt::Success(response)], + Err(error) => error, + }; + let request = match Arc::try_unwrap(request) { + Ok(request) => request, + Err(_) => { + return vec![WriteAttempt::Failed { + error: UnexpectedSnafu { + err_msg: "Region request is still shared after its RPC completed", + } + .build(), + request: None, + }]; + } + }; + let requests = split_request(request); + match error { + Error::PartialRegion { response } => { + let mut failures: HashMap<_, _> = response.failures.into_iter().collect(); + let mut attempts = vec![WriteAttempt::Success(RegionResponse { + affected_rows: response.affected_rows, + extensions: response.extensions, + metadata: Vec::new(), + })]; + for (index, request) in requests.into_iter().enumerate() { + if let Some(error) = failures.remove(&index) { + let request = error.retry_hint().is_retryable().then(|| RoutedWrite { + peer: Arc::clone(&peer), + request: Arc::new(request), + }); + attempts.push(WriteAttempt::Failed { error, request }); + } + } + attempts + } + error => { + // A whole-RPC failure cannot confirm success for any entry. + let error = Arc::new(error); + requests + .into_iter() + .map(|request| { + let request = error.retry_hint().is_retryable().then(|| RoutedWrite { + peer: Arc::clone(&peer), + request: Arc::new(request), + }); + WriteAttempt::Failed { + error: BatchRequestSnafu.into_error(Arc::clone(&error)), + request, + } + }) + .collect() + } + } + } + async fn dispatch_inner( &self, requests: Vec, diff --git a/src/operator/src/write_dispatcher/tests.rs b/src/operator/src/write_dispatcher/tests.rs index a249a8dfe8e..2607a3f194d 100644 --- a/src/operator/src/write_dispatcher/tests.rs +++ b/src/operator/src/write_dispatcher/tests.rs @@ -21,8 +21,8 @@ use std::time::{Duration, Instant}; use api::v1::region::region_server::{Region, RegionServer}; use api::v1::region::{ - AlignedSchemaVersion, BulkInsertRequest, DeleteRequest, InsertRequest, RegionRequestHeader, - RegionResponse as PbRegionResponse, bulk_insert_request, + AlignedSchemaVersion, BulkInsertRequest, DeleteRequest, DeleteRequests, InsertRequest, + InsertRequests, RegionRequestHeader, RegionResponse as PbRegionResponse, bulk_insert_request, }; use api::v1::value::ValueData; use api::v1::{ @@ -78,6 +78,8 @@ use crate::test_util; #[derive(Clone, Default)] struct RecordingRegionService { + request_sizes: Arc>>, + rpc_rejections: Arc>>, requests: Arc>>, rejections: Arc>>>, release: Arc>>>, @@ -90,6 +92,70 @@ impl Region for RecordingRegionService { request: Request, ) -> std::result::Result, Status> { let request = request.into_inner(); + if !matches!(&request.body, Some(Body::Inserts(_) | Body::Deletes(_))) { + return self.handle_entry(request).await; + } + let accepts_partial = request + .header + .as_ref() + .and_then(|header| header.query_context.as_ref()) + .and_then(|context| context.extensions.get(api::region::ACCEPT_PARTIAL_RESULT)) + .is_some_and(|value| value == "true"); + let entries = split_request(request); + let return_tonic_error = !accepts_partial || entries.len() == 1; + self.request_sizes.lock().unwrap().push(entries.len()); + if let Some(status) = self.rpc_rejections.lock().unwrap().pop_front() { + self.requests.lock().unwrap().extend(entries); + return Err(status); + } + let mut response = PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus::default()), + }), + ..Default::default() + }; + let results = + future::join_all(entries.into_iter().map(|entry| self.handle_entry(entry))).await; + for (index, result) in results.into_iter().enumerate() { + match result { + Ok(result) => response.affected_rows += result.into_inner().affected_rows, + Err(status) => { + if return_tonic_error { + return Err(status); + } + let error = client::Error::from(status); + response + .failures + .push(api::v1::region::RegionRequestFailure { + request_index: index as u32, + status: Some(PbStatus { + status_code: error.status_code() as u32, + err_msg: error.output_msg(), + retry_hint: error.retry_hint().as_str().to_string(), + }), + }); + } + } + } + if !response.failures.is_empty() { + response + .header + .as_mut() + .unwrap() + .status + .as_mut() + .unwrap() + .status_code = StatusCode::PartialFailure as u32; + } + Ok(Response::new(response)) + } +} + +impl RecordingRegionService { + async fn handle_entry( + &self, + request: RegionRequest, + ) -> std::result::Result, Status> { let region_id = write_region_id(&request).unwrap().as_u64(); self.requests.lock().unwrap().push(request); let release = self.release.lock().unwrap().clone(); @@ -237,9 +303,19 @@ impl RegionTestServer { hints .into_iter() .map(|hint| { - let mut status = Status::failed_precondition(format!( - "region {region_id} rejected the write" - )); + let code = match hint { + RetryHint::Retryable => StatusCode::RegionBusy, + RetryHint::RetryAfterRouteRefresh => StatusCode::Cancelled, + RetryHint::NonRetryable => StatusCode::InvalidArguments, + }; + let mut status = Status::new( + common_error::status_code::status_to_tonic_code(code), + format!("region {region_id} rejected the write"), + ); + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_CODE, + (code as u32).to_string().parse().unwrap(), + ); status.metadata_mut().insert( common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT, hint.as_str().parse().unwrap(), @@ -343,18 +419,26 @@ fn build_insert_request(regions: &[u32]) -> RegionRequest { #[tokio::test] async fn test_write_retry_disabled_and_enabled_ablation() { common_telemetry::init_default_ut_logging(); - for run in 0..3 { + for regions in [vec![1], vec![1, 2]] { for max_retries in [0, 1] { let server = RegionTestServer::start().await; server.fail_region(1, [RetryHint::Retryable]); let dispatcher = server.create_dispatcher(max_retries).await; - let result = dispatcher.dispatch(vec![build_insert_request(&[1])]).await; + let result = dispatcher + .dispatch(vec![build_insert_request(®ions)]) + .await; let attempts = server.attempt_counts()[&RegionId::new(1, 1).as_u64()]; common_telemetry::info!( - "write_retry_ablation run={run} max_retries={max_retries} success={} attempts={attempts}", + "write_retry_ablation regions={regions:?} max_retries={max_retries} success={} attempts={attempts}", result.is_ok() ); - assert_eq!(result.is_ok(), max_retries > 0); + if max_retries == 0 { + let error = result.unwrap_err(); + assert_eq!(error.status_code(), StatusCode::RegionBusy); + assert_eq!(error.retry_hint(), RetryHint::Retryable); + } else { + assert_eq!(result.unwrap().affected_rows, regions.len()); + } assert_eq!(attempts, max_retries + 1); for request in server.service.requests.lock().unwrap().iter() { let accepts_partial = request @@ -372,22 +456,29 @@ async fn test_write_retry_disabled_and_enabled_ablation() { #[tokio::test] async fn test_write_retry_does_not_replay_successful_regions() { - let server = RegionTestServer::start().await; - server.fail_region(2, [RetryHint::Retryable]); - let dispatcher = server.create_dispatcher(2).await; - let response = dispatcher - .dispatch(vec![build_insert_request(&[1, 2])]) - .await - .unwrap(); - assert_eq!(response.affected_rows, 2); - assert_eq!( - server.attempt_counts(), - HashMap::from([ - (RegionId::new(1, 1).as_u64(), 1), - (RegionId::new(1, 2).as_u64(), 2), - ]) - ); - assert!(server.invalidator.calls.lock().unwrap().is_empty()); + for regions in [vec![1, 2], vec![1, 1, 1]] { + let server = RegionTestServer::start().await; + server.fail_region(1, [RetryHint::Retryable]); + let dispatcher = server.create_dispatcher(2).await; + let response = dispatcher + .dispatch(vec![build_insert_request(®ions)]) + .await + .unwrap(); + assert_eq!(response.affected_rows, regions.len()); + let mut expected = HashMap::new(); + for region in ®ions { + *expected + .entry(RegionId::new(1, *region).as_u64()) + .or_insert(0) += 1; + } + *expected.get_mut(&RegionId::new(1, 1).as_u64()).unwrap() += 1; + assert_eq!(server.attempt_counts(), expected); + assert_eq!( + *server.service.request_sizes.lock().unwrap(), + vec![regions.len(), 1] + ); + assert!(server.invalidator.calls.lock().unwrap().is_empty()); + } } #[tokio::test] @@ -429,10 +520,10 @@ async fn test_write_retry_transport_errors() { } server .service - .rejections + .rpc_rejections .lock() .unwrap() - .insert(region_id, VecDeque::from([status])); + .push_back(status); let dispatcher = server.create_dispatcher(2).await; let attempt = dispatcher .send( @@ -462,12 +553,12 @@ async fn test_write_retry_transport_attempt_budget() { for failures in [1, 3] { let server = RegionTestServer::start().await; let region_id = RegionId::new(1, 1).as_u64(); - server.service.rejections.lock().unwrap().insert( - region_id, - (0..failures) - .map(|_| Status::new(code, "write failed")) - .collect(), - ); + server + .service + .rpc_rejections + .lock() + .unwrap() + .extend((0..failures).map(|_| Status::new(code, "write failed"))); let dispatcher = server.create_dispatcher(2).await; let response = dispatcher.dispatch(vec![build_insert_request(&[1])]).await; let expected_attempts = if failures == 1 { @@ -500,8 +591,11 @@ async fn test_write_retry_stops_when_a_sibling_is_nonretryable() { .await .unwrap_err(); let terminal_region = RegionId::new(1, 2).as_u64(); + assert_eq!(error.retry_hint(), RetryHint::NonRetryable); + assert_eq!(error.status_code(), StatusCode::InvalidArguments); assert!( - format!("{error:?}") + error + .output_msg() .contains(&format!("region {terminal_region} rejected the write")), "{error:?}" ); @@ -736,81 +830,83 @@ async fn test_finish_retries_initial_failures_in_one_round() { #[tokio::test] async fn test_write_retry_refreshes_stale_route_and_sends_only_failed_region_to_new_leader() { - let old_server = RegionTestServer::start().await; - let new_server = RegionTestServer::start().await; - old_server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]); - let backend = test_util::prepare_mocked_backend().await; - // Reuse the operator fixture to populate table metadata and partition rules. - test_util::create_partition_rule_manager(backend.clone()).await; - update_table_leader(&backend, &old_server.peer).await; + for regions in [vec![2], vec![1, 2]] { + let old_server = RegionTestServer::start().await; + let new_server = RegionTestServer::start().await; + old_server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]); + let backend = test_util::prepare_mocked_backend().await; + // Reuse the operator fixture to populate table metadata and partition rules. + test_util::create_partition_rule_manager(backend.clone()).await; + update_table_leader(&backend, &old_server.peer).await; - let (partition_manager, table_info_cache, registry, catalog_manager) = - create_routing_caches(backend.clone()); - let failed_region = RegionId::new(1, 2); - assert_eq!( - partition_manager - .find_region_leader(failed_region) + let (partition_manager, table_info_cache, registry, catalog_manager) = + create_routing_caches(backend.clone()); + let failed_region = RegionId::new(1, 2); + assert_eq!( + partition_manager + .find_region_leader(failed_region) + .await + .unwrap(), + old_server.peer + ); + let table_info = table_info_cache.get(1).await.unwrap().unwrap(); + let (_, versions) = partition_manager + .find_table_partition_rule(&table_info) .await - .unwrap(), - old_server.peer - ); - let table_info = table_info_cache.get(1).await.unwrap().unwrap(); - let (_, versions) = partition_manager - .find_table_partition_rule(&table_info) - .await - .unwrap(); + .unwrap(); - // Metadata has migrated while the frontend still holds the old route. - update_table_leader(&backend, &new_server.peer).await; - assert_eq!( - partition_manager - .find_region_leader(failed_region) - .await - .unwrap(), - old_server.peer - ); + // Metadata has migrated while the frontend still holds the old route. + update_table_leader(&backend, &new_server.peer).await; + assert_eq!( + partition_manager + .find_region_leader(failed_region) + .await + .unwrap(), + old_server.peer + ); - let invalidator = Arc::new(RecordingCacheInvalidator { - calls: Mutex::default(), - inner: Some(registry), - ..Default::default() - }); - let clients = Arc::new(NodeClients::new(ChannelConfig::default())); - let dispatcher = WriteDispatcher::new(clients, partition_manager).with_retry( - build_retry_options(2), - catalog_manager, - invalidator.clone(), - ); - let mut request = build_insert_request(&[1, 2]); - if let Some(Body::Inserts(inserts)) = &mut request.body { - for insert in &mut inserts.requests { - insert.partition_expr_version = versions - [&RegionId::from_u64(insert.region_id).region_number()] - .map(|value| PartitionExprVersion { value }); + let invalidator = Arc::new(RecordingCacheInvalidator { + calls: Mutex::default(), + inner: Some(registry), + ..Default::default() + }); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + let dispatcher = WriteDispatcher::new(clients, partition_manager).with_retry( + build_retry_options(2), + catalog_manager, + invalidator.clone(), + ); + let mut request = build_insert_request(®ions); + if let Some(Body::Inserts(inserts)) = &mut request.body { + for insert in &mut inserts.requests { + insert.partition_expr_version = versions + [&RegionId::from_u64(insert.region_id).region_number()] + .map(|value| PartitionExprVersion { value }); + } } + let response = dispatcher.dispatch(vec![request]).await.unwrap(); + assert_eq!(response.affected_rows, regions.len()); + assert_eq!( + old_server.attempt_counts(), + regions + .iter() + .map(|region| (RegionId::new(1, *region).as_u64(), 1)) + .collect() + ); + assert_eq!( + new_server.attempt_counts(), + HashMap::from([(failed_region.as_u64(), 1)]) + ); + assert_eq!( + *invalidator.calls.lock().unwrap(), + vec![ + vec![CacheIdent::TableId(1)], + vec![CacheIdent::TableName(TableName::new( + "greptime", "public", "table_1" + ))], + ] + ); } - let response = dispatcher.dispatch(vec![request]).await.unwrap(); - assert_eq!(response.affected_rows, 2); - assert_eq!( - old_server.attempt_counts(), - HashMap::from([ - (RegionId::new(1, 1).as_u64(), 1), - (failed_region.as_u64(), 1), - ]) - ); - assert_eq!( - new_server.attempt_counts(), - HashMap::from([(failed_region.as_u64(), 1)]) - ); - assert_eq!( - *invalidator.calls.lock().unwrap(), - vec![ - vec![CacheIdent::TableId(1)], - vec![CacheIdent::TableName(TableName::new( - "greptime", "public", "table_1" - ))], - ] - ); } #[derive(Clone, Copy, Debug)]