From 344db17dd499e2ad1539889c4fa6a343152df04b Mon Sep 17 00:00:00 2001 From: WenyXu Date: Mon, 12 Oct 2026 11:06:07 +0000 Subject: [PATCH] feat(operator): retry failed writes with refreshed partitions Signed-off-by: WenyXu --- Cargo.lock | 4 + src/operator/AGENTS.md | 4 +- src/operator/Cargo.toml | 4 + src/operator/src/bulk_insert.rs | 43 +- .../src/req_convert/common/partitioner.rs | 436 ++++++++++- src/operator/src/write_dispatcher.rs | 310 +++++++- src/operator/src/write_dispatcher/options.rs | 106 +++ src/operator/src/write_dispatcher/tests.rs | 707 +++++++++++++++++- 8 files changed, 1576 insertions(+), 38 deletions(-) create mode 100644 src/operator/src/write_dispatcher/options.rs diff --git a/Cargo.lock b/Cargo.lock index f75bde16328..8678104f5d4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10095,7 +10095,9 @@ dependencies = [ "async-stream", "async-trait", "axum 0.8.4", + "backon", "bytes", + "cache", "catalog", "chrono", "client", @@ -10126,6 +10128,7 @@ dependencies = [ "futures", "futures-util", "humantime", + "humantime-serde", "itertools 0.14.0", "jsonb", "lazy_static", @@ -10152,6 +10155,7 @@ dependencies = [ "substrait 1.4.0-alpha.0", "table", "tokio", + "tokio-stream", "tokio-util", "tonic 0.14.2", "tracing", diff --git a/src/operator/AGENTS.md b/src/operator/AGENTS.md index 46359eb5c50..bb0ebb2b2d4 100644 --- a/src/operator/AGENTS.md +++ b/src/operator/AGENTS.md @@ -14,7 +14,7 @@ insert/delete routing. | Statements | `src/operator/src/statement.rs`, `src/operator/src/statement/` | `StatementExecutor` and per-statement handlers | | Inserts | `src/operator/src/insert.rs`, `src/operator/src/bulk_insert.rs` | Insert, auto-create/alter, routing, dispatch | | Deletes | `src/operator/src/delete.rs` | Delete conversion and dispatch | -| Write dispatch | `src/operator/src/write_dispatcher.rs` | Shared insert/delete/bulk dispatch and batched peer resolution | +| Write dispatch | `src/operator/src/write_dispatcher.rs` | Shared insert/delete/bulk dispatch, bounded retries, route refresh and repartition | | Request conversion | `src/operator/src/req_convert.rs`, `src/operator/src/req_convert/` | Table/row/column requests to region requests | | Region requests | `src/operator/src/region_req_factory.rs` | Region-level DDL/DML request construction | | Procedures | `src/operator/src/procedure.rs` | Metasrv procedure administration | @@ -39,6 +39,8 @@ cargo sqlness bare -t ``` Check standalone and distributed routing when changing region dispatch. +Write retries operate on prepared region requests in `write_dispatcher`; keep +admission and Flow side effects outside retries, and preserve successful regions. Write admission uses fallible `write_meter!` calls in `Inserter::do_request` and `handle_bulk_insert`, before flow side effects and datanode dispatch. Bulk diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml index 72f953e0f89..7f4233e633e 100644 --- a/src/operator/Cargo.toml +++ b/src/operator/Cargo.toml @@ -24,6 +24,7 @@ api.workspace = true arrow.workspace = true async-stream.workspace = true async-trait.workspace = true +backon.workspace = true bytes.workspace = true catalog.workspace = true chrono.workspace = true @@ -54,6 +55,7 @@ file-engine.workspace = true futures.workspace = true futures-util.workspace = true humantime.workspace = true +humantime-serde.workspace = true itertools.workspace = true jsonb.workspace = true lazy_static.workspace = true @@ -86,8 +88,10 @@ url.workspace = true [dev-dependencies] axum.workspace = true +cache.workspace = true catalog = { workspace = true, features = ["testing"] } common-meta = { workspace = true, features = ["testing"] } common-test-util.workspace = true object-store = { workspace = true, features = ["testing"] } path-slash = "0.2" +tokio-stream = { workspace = true, features = ["net"] } diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index d5aa7168b44..73f3792e782 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -14,7 +14,6 @@ use std::collections::HashSet; -use api::region::RegionResponse; use api::v1::flow::{DirtyWindowRequest, DirtyWindowRequests}; use api::v1::region::{ BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request, @@ -41,6 +40,7 @@ use crate::error::Result; use crate::insert::Inserter; use crate::req_convert::TablePartitioner; use crate::req_convert::insert; +use crate::write_dispatcher::WriteOutcome; use crate::{error, metrics}; impl Inserter { @@ -246,7 +246,7 @@ impl Inserter { } else { None }; - let handle: common_runtime::JoinHandle> = + let handle: common_runtime::JoinHandle> = common_runtime::spawn_global(async move { let (header, payload) = if mask.select_all() { // SAFETY: raw data must be present, we can avoid re-encoding. @@ -301,27 +301,40 @@ impl Inserter { })), }; - dispatcher - .dispatch_prepared( - vec![(peer, request)], + Ok(dispatcher + .dispatch_initial( + peer, + request, Some( - metrics::HANDLE_BULK_INSERT_ELAPSED + &metrics::HANDLE_BULK_INSERT_ELAPSED .with_label_values(&["datanode_handle"]), ), ) - .await - .context(error::RequestRegionSnafu) + .await) }); handles.push(handle); } - let results = future::try_join_all(handles) - .await - .context(error::JoinTaskSnafu)?; - let rows_inserted = results - .into_iter() - .map(|result| result.map(|response| response.affected_rows)) - .sum::>()?; + let dispatcher = self.write_dispatcher.clone(); + let rows_inserted = common_runtime::spawn_global(async move { + let results = future::try_join_all(handles) + .await + .context(error::JoinTaskSnafu)?; + let outcomes = results.into_iter().collect::>>()?; + dispatcher + .finish_dispatch( + outcomes, + Some( + &metrics::HANDLE_BULK_INSERT_ELAPSED + .with_label_values(&["datanode_handle"]), + ), + ) + .await + .context(error::RequestRegionSnafu) + .map(|response| response.affected_rows) + }) + .await + .context(error::JoinTaskSnafu)??; wait_all_datanode_timer.observe_duration(); crate::metrics::DIST_INGEST_ROW_COUNT .with_label_values(&[db_name.as_str()]) diff --git a/src/operator/src/req_convert/common/partitioner.rs b/src/operator/src/req_convert/common/partitioner.rs index 044d3b97db5..ad806d04f9e 100644 --- a/src/operator/src/req_convert/common/partitioner.rs +++ b/src/operator/src/req_convert/common/partitioner.rs @@ -14,7 +14,10 @@ use std::collections::HashMap; -use api::v1::region::{BulkInsertRequest, DeleteRequest, InsertRequest, bulk_insert_request}; +use api::v1::region::{ + BulkInsertRequest, DeleteRequest, DeleteRequests, InsertRequest, InsertRequests, RegionRequest, + bulk_insert_request, region_request::Body, +}; use api::v1::{ArrowIpc, PartitionExprVersion, Rows}; use common_error::ext::BoxedError; use common_grpc::flight; @@ -23,7 +26,7 @@ use datatypes::arrow::record_batch::RecordBatch; use partition::PartitionRuleRef; use partition::manager::PartitionRuleManager; use partition::splitter::RowSplitter; -use snafu::ResultExt; +use snafu::{OptionExt, ResultExt}; use store_api::storage::{RegionId, RegionNumber, TableId}; use table::metadata::TableInfo; @@ -94,6 +97,94 @@ impl TablePartitioner { }) } + /// Repartitions a rejected write after the caller refreshes the routing caches. + /// + /// Takes ownership so row payloads can move into their new regions. Requests whose + /// region and partition version still match retain their original payloads. + pub fn repartition_request(&self, request: RegionRequest) -> Result> { + let RegionRequest { header, body } = request; + let bodies = match body.context(error::UnexpectedSnafu { + violated: "Cannot reroute a request without a body", + })? { + Body::Inserts(inserts) => { + let mut bodies = Vec::new(); + for insert in inserts.requests { + let region_id = RegionId::from_u64(insert.region_id); + if self.matches_region(region_id, insert.partition_expr_version) { + bodies.push(Body::Inserts(InsertRequests { + requests: vec![insert], + })); + continue; + } + let rows = insert.rows.context(error::UnexpectedSnafu { + violated: "Cannot reroute an insert without rows", + })?; + for request in self.partition_insert_requests(rows, insert.skip_wal)? { + bodies.push(Body::Inserts(InsertRequests { + requests: vec![request], + })); + } + } + bodies + } + Body::Deletes(deletes) => { + let mut bodies = Vec::new(); + for delete in deletes.requests { + let region_id = RegionId::from_u64(delete.region_id); + if self.matches_region(region_id, delete.partition_expr_version) { + bodies.push(Body::Deletes(DeleteRequests { + requests: vec![delete], + })); + continue; + } + let rows = delete.rows.context(error::UnexpectedSnafu { + violated: "Cannot reroute a delete without rows", + })?; + for request in self.partition_delete_requests(rows)? { + bodies.push(Body::Deletes(DeleteRequests { + requests: vec![request], + })); + } + } + bodies + } + Body::BulkInsert(bulk) => { + let region_id = RegionId::from_u64(bulk.region_id); + if self.matches_region(region_id, bulk.partition_expr_version) { + vec![Body::BulkInsert(bulk)] + } else { + let bulk_insert_request::Body::ArrowIpc(ipc) = + bulk.body.context(error::UnexpectedSnafu { + violated: "Cannot reroute a bulk insert without Arrow IPC data", + })?; + let batch = flight::FlightDecoder::try_from_schema_bytes(&ipc.schema) + .and_then(|mut decoder| { + decoder.try_decode_record_batch(&ipc.data_header, &ipc.payload) + }) + .map_err(BoxedError::new) + .context(error::ExternalSnafu)?; + self.partition_bulk_insert_requests(batch, Some(ipc), bulk.skip_wal)? + .into_iter() + .map(Body::BulkInsert) + .collect() + } + } + _ => { + return error::UnexpectedSnafu { + violated: "Only writes can be rerouted", + } + .fail(); + } + }; + Ok(bodies + .into_iter() + .map(|body| RegionRequest { + header: header.clone(), + body: Some(body), + }) + .collect()) + } + /// Moves insert rows into requests for this snapshot's destination regions. pub fn partition_insert_requests( &self, @@ -178,6 +269,12 @@ impl TablePartitioner { Ok(requests) } + /// Whether an earlier request still matches this partition snapshot. + fn matches_region(&self, region_id: RegionId, version: Option) -> bool { + region_id.table_id() == self.table_id + && self.versions.get(®ion_id.region_number()) == Some(&version) + } + fn version(&self, region: RegionNumber) -> Option { self.versions.get(®ion).copied().flatten() } @@ -187,12 +284,26 @@ impl TablePartitioner { mod tests { use std::sync::Arc; + use api::v1::region::{ + AlignedSchemaVersion, BulkInsertRequest, DeleteRequest, InsertRequest, RegionRequestHeader, + }; + use api::v1::value::ValueData; + use api::v1::{ArrowIpc, PartitionExprVersion}; + use api::v1::{ColumnDataType, ColumnSchema, Row, Rows, Value}; use common_error::ext::{ErrorExt, RetryHint}; use common_error::status_code::StatusCode; - use common_meta::cache; + use common_grpc::flight; + use common_meta::cache::{self, TableInfoCache}; + use common_meta::key::TableMetadataManager; + use common_meta::key::table_route::TableRouteValue; + use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::kv_backend::test_util::MockKvBackendBuilder; use common_meta::rpc::store::RangeResponse; + use datatypes::arrow::array::Int32Array; + use datatypes::arrow::datatypes::{DataType, Field, Schema}; + use datatypes::arrow::record_batch::RecordBatch; use moka::future::Cache; + use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef}; use super::*; use crate::error::Error; @@ -249,4 +360,323 @@ mod tests { } } } + + async fn create_fixture() -> (PartitionRuleManagerRef, TableInfoCache) { + let backend = Arc::new(MemoryKvBackend::new()); + let manager = test_util::create_partition_rule_manager(backend.clone()).await; + TableMetadataManager::new(backend.clone()) + .create_table_metadata( + test_util::new_test_table_info(42, "logical_reroute", [1, 2, 3].into_iter()), + TableRouteValue::logical(1), + Default::default(), + ) + .await + .unwrap(); + let cache = cache::new_table_info_cache( + "reroute_test_table_info".to_string(), + moka::future::Cache::new(16), + backend, + ); + (manager, cache) + } + + fn build_request(body: Body) -> RegionRequest { + RegionRequest { + header: Some(RegionRequestHeader { + dbname: "greptime.public".to_string(), + ..Default::default() + }), + body: Some(body), + } + } + + fn build_rows(values: &[i32]) -> Rows { + Rows { + schema: vec![ColumnSchema { + column_name: "a".to_string(), + datatype: ColumnDataType::Int32 as i32, + ..Default::default() + }], + rows: values + .iter() + .map(|value| Row { + values: vec![Value { + value_data: Some(ValueData::I32Value(*value)), + }], + }) + .collect(), + } + } + + fn build_row_request( + rows: Rows, + insert: bool, + region: u32, + version: Option, + ) -> RegionRequest { + let region_id = RegionId::new(1, region).as_u64(); + let partition_expr_version = version.map(|value| PartitionExprVersion { value }); + let body = if insert { + Body::Inserts(InsertRequests { + requests: vec![InsertRequest { + region_id, + partition_expr_version, + rows: Some(rows), + skip_wal: true, + }], + }) + } else { + Body::Deletes(DeleteRequests { + requests: vec![DeleteRequest { + region_id, + partition_expr_version, + rows: Some(rows), + }], + }) + }; + build_request(body) + } + + fn row_request_parts(request: &RegionRequest) -> (u64, Option, &Rows) { + match request.body.as_ref().unwrap() { + Body::Inserts(inserts) => { + assert_eq!(inserts.requests.len(), 1); + let insert = &inserts.requests[0]; + assert!(insert.skip_wal); + ( + insert.region_id, + insert.partition_expr_version.map(|v| v.value), + insert.rows.as_ref().unwrap(), + ) + } + Body::Deletes(deletes) => { + assert_eq!(deletes.requests.len(), 1); + let delete = &deletes.requests[0]; + ( + delete.region_id, + delete.partition_expr_version.map(|v| v.value), + delete.rows.as_ref().unwrap(), + ) + } + _ => panic!("expected row write"), + } + } + + #[tokio::test] + async fn test_repartition_rows_updates_partitions_without_losing_rows() { + let (manager, cache) = create_fixture().await; + let table_info = cache.get(1).await.unwrap().unwrap(); + let (_, versions) = manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + for insert in [true, false] { + let request = build_row_request(build_rows(&[1, 20, 70]), insert, 99, Some(0)); + let header = request.header.clone(); + let rerouted = repartition_with_cached_table(&manager, &cache, request) + .await + .unwrap(); + assert_eq!(rerouted.len(), 3); + let mut actual = Vec::new(); + for request in rerouted { + assert_eq!(request.header, header); + let (id, version, rows) = row_request_parts(&request); + let region_id = RegionId::from_u64(id); + assert_eq!(region_id.table_id(), 1); + assert_eq!(version, versions[®ion_id.region_number()]); + assert_eq!(rows.schema, build_rows(&[]).schema); + assert_eq!(rows.rows.len(), 1); + actual.push((region_id.region_number(), rows.rows[0].clone())); + } + actual.sort_by_key(|(region, _)| *region); + assert_eq!( + actual, + vec![ + (1, build_rows(&[70]).rows.remove(0)), + (2, build_rows(&[20]).rows.remove(0)), + (3, build_rows(&[1]).rows.remove(0)) + ] + ); + } + } + + #[tokio::test] + async fn test_unchanged_row_routes_preserve_allocations() { + let (manager, cache) = create_fixture().await; + let table_info = cache.get(1).await.unwrap().unwrap(); + let (_, versions) = manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + for insert in [true, false] { + let request = build_row_request(build_rows(&[1]), insert, 3, versions[&3]); + let original_rows = row_request_parts(&request).2; + let row_ptr = original_rows.rows.as_ptr(); + let value_ptr = original_rows.rows[0].values.as_ptr(); + let rerouted = repartition_with_cached_table(&manager, &cache, request) + .await + .unwrap(); + assert_eq!(rerouted.len(), 1); + let rows = row_request_parts(&rerouted[0]).2; + assert_eq!(rows.rows.as_ptr(), row_ptr); + assert_eq!(rows.rows[0].values.as_ptr(), value_ptr); + assert_eq!(rows, &build_rows(&[1])); + } + } + + fn build_bulk_request(region: u32, version: Option, values: Vec) -> RegionRequest { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(values))], + ) + .unwrap(); + let (schema, data_header, payload) = flight::record_batch_to_ipc(batch).unwrap(); + build_request(Body::BulkInsert(BulkInsertRequest { + region_id: RegionId::new(1, region).as_u64(), + partition_expr_version: version.map(|value| PartitionExprVersion { value }), + aligned_schema_version: Some(AlignedSchemaVersion { + schema_version: 123, + }), + skip_wal: true, + body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { + schema, + data_header, + payload, + })), + })) + } + + #[tokio::test] + async fn test_repartition_bulk_reencodes_only_changed_partitions() { + let (manager, cache) = create_fixture().await; + let table_info = cache.get(1).await.unwrap().unwrap(); + let (_, versions) = manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + let request = build_bulk_request(99, Some(0), vec![1, 20, 70]); + let header = request.header.clone(); + let rerouted = repartition_with_cached_table(&manager, &cache, request) + .await + .unwrap(); + assert_eq!(rerouted.len(), 3); + let mut actual = Vec::new(); + for request in rerouted { + assert_eq!(request.header, header); + let Some(Body::BulkInsert(bulk)) = request.body else { + panic!("expected bulk insert") + }; + let region = RegionId::from_u64(bulk.region_id); + assert_eq!(region.table_id(), 1); + assert!(bulk.skip_wal); + assert_eq!(bulk.aligned_schema_version, None); + assert_eq!( + bulk.partition_expr_version.map(|v| v.value), + versions[®ion.region_number()] + ); + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = bulk.body else { + panic!("expected Arrow IPC") + }; + let batch = flight::FlightDecoder::try_from_schema_bytes(&ipc.schema) + .unwrap() + .try_decode_record_batch(&ipc.data_header, &ipc.payload) + .unwrap(); + assert_eq!(batch.num_rows(), 1); + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + actual.push((region.region_number(), values.value(0))); + } + actual.sort_unstable(); + assert_eq!(actual, [(1, 70), (2, 20), (3, 1)]); + + // A leader-only change retains the original IPC buffers and schema hint. + let request = build_bulk_request(3, versions[&3], vec![1]); + let Some(Body::BulkInsert(bulk)) = request.body.as_ref() else { + unreachable!() + }; + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &bulk.body else { + unreachable!() + }; + let pointers = ( + ipc.schema.as_ptr(), + ipc.data_header.as_ptr(), + ipc.payload.as_ptr(), + ); + let rerouted = repartition_with_cached_table(&manager, &cache, request) + .await + .unwrap(); + assert_eq!(rerouted.len(), 1); + let Some(Body::BulkInsert(bulk)) = rerouted[0].body.as_ref() else { + unreachable!() + }; + assert_eq!( + bulk.aligned_schema_version, + Some(AlignedSchemaVersion { + schema_version: 123 + }) + ); + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &bulk.body else { + unreachable!() + }; + assert_eq!( + ( + ipc.schema.as_ptr(), + ipc.data_header.as_ptr(), + ipc.payload.as_ptr() + ), + pointers + ); + } + #[tokio::test] + async fn test_repartition_logical_table_preserves_logical_region_ids() { + let (manager, cache) = create_fixture().await; + for insert in [true, false] { + let mut request = build_row_request(build_rows(&[1, 20, 70]), insert, 99, Some(0)); + match request.body.as_mut().unwrap() { + Body::Inserts(inserts) => { + inserts.requests[0].region_id = RegionId::new(42, 99).as_u64() + } + Body::Deletes(deletes) => { + deletes.requests[0].region_id = RegionId::new(42, 99).as_u64() + } + _ => unreachable!(), + } + let rerouted = repartition_with_cached_table(&manager, &cache, request) + .await + .unwrap(); + assert_eq!(rerouted.len(), 3); + let mut regions = rerouted + .iter() + .map(|request| { + let region = RegionId::from_u64(row_request_parts(request).0); + assert_eq!(region.table_id(), 42); + region.region_number() + }) + .collect::>(); + regions.sort_unstable(); + assert_eq!(regions, [1, 2, 3]); + } + } + + async fn repartition_with_cached_table( + manager: &PartitionRuleManager, + cache: &TableInfoCache, + request: RegionRequest, + ) -> Result> { + let region_id = match request.body.as_ref().unwrap() { + Body::Inserts(inserts) => inserts.requests[0].region_id, + Body::Deletes(deletes) => deletes.requests[0].region_id, + Body::BulkInsert(bulk) => bulk.region_id, + _ => unreachable!(), + }; + let table_id = RegionId::from_u64(region_id).table_id(); + let table_info = cache.get(table_id).await.unwrap().unwrap(); + let partitioner = TablePartitioner::try_new(manager, &table_info) + .await + .context(error::InvalidPartitionSnafu)?; + partitioner.repartition_request(request) + } } diff --git a/src/operator/src/write_dispatcher.rs b/src/operator/src/write_dispatcher.rs index 4af4fbfafdf..4d4065bbcbe 100644 --- a/src/operator/src/write_dispatcher.rs +++ b/src/operator/src/write_dispatcher.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod options; #[cfg(test)] mod tests; @@ -20,8 +21,13 @@ use std::sync::Arc; use api::region::RegionResponse; use api::v1::region::{DeleteRequests, InsertRequests, RegionRequest, region_request::Body}; -use common_error::ext::BoxedError; -use common_meta::error::{ExternalSnafu, Result, UnexpectedSnafu}; +use backon::{BackoffBuilder, ExponentialBuilder}; +use client::client_manager::NodeClients; +use common_error::ext::{BoxedError, ErrorExt, RetryHint}; +use common_meta::cache::TableInfoCacheRef; +use common_meta::cache_invalidator::{CacheInvalidatorRef, Context}; +use common_meta::error::{ExternalSnafu, Result, TableInfoNotFoundSnafu, UnexpectedSnafu}; +use common_meta::instruction::CacheIdent; use common_meta::node_manager::NodeManagerRef; use common_meta::peer::Peer; use common_meta::rpc::router; @@ -31,32 +37,63 @@ use prometheus::Histogram; use snafu::{OptionExt, ResultExt}; use store_api::storage::RegionId; +pub use self::options::WriteRetryOptions; use crate::error::{self, FindRegionLeaderSnafu, RequestRegionSnafu}; +use crate::req_convert::TablePartitioner; -/// Dispatches already partitioned insert, delete, and bulk requests to their peers. +/// Dispatches already partitioned writes, retaining only failed payloads for retries. pub struct WriteDispatcher { node_manager: NodeManagerRef, partition_manager: PartitionRuleManagerRef, + remote: Option, } /// Shared reference to a [`WriteDispatcher`]. pub type WriteDispatcherRef = Arc; +/// A completed send, retaining its request only when a retry is possible. +pub(crate) struct WriteOutcome { + result: Result, + pending: Option, +} + impl WriteDispatcher { /// Creates an owned-request dispatcher without retries, including for standalone. pub fn new(node_manager: NodeManagerRef, partition_manager: PartitionRuleManagerRef) -> Self { Self { node_manager, partition_manager, + remote: None, } } + /// Enables shared remote sends and retries using frontend configuration and caches. + pub fn with_retry( + mut self, + clients: Arc, + options: WriteRetryOptions, + table_info_cache: TableInfoCacheRef, + cache_invalidator: CacheInvalidatorRef, + ) -> Self { + if options.max_retries > 0 { + self.remote = Some(RemoteRetry { + clients, + options, + table_info_cache, + cache_invalidator, + }); + } else { + self.remote = None; + } + self + } + /// Writes planned requests without repeating admission, default values, or Flow side effects. pub async fn dispatch(&self, requests: Vec) -> Result { self.dispatch_inner(requests, None).await } - /// Records each datanode send in the supplied histogram. + /// Records each datanode send attempt, including retries, in the supplied histogram. pub async fn dispatch_with_metrics( &self, requests: Vec, @@ -124,15 +161,65 @@ impl WriteDispatcher { histogram: Option, ) -> Result { let histogram = histogram.as_ref(); - let results = future::join_all(writes.into_iter().map(|(peer, request)| async move { - let _timer = histogram.map(Histogram::start_timer); - let datanode = self.node_manager.datanode(&peer).await; - datanode.handle(request).await - })) + let writes = if self.remote.is_some() { + writes + .into_iter() + .flat_map(|(peer, request)| { + split_request(request) + .into_iter() + .map(move |request| (peer.clone(), request)) + }) + .collect() + } else { + writes + }; + let outcomes = future::join_all( + writes + .into_iter() + .map(|(peer, request)| self.dispatch_initial(peer, request, histogram)), + ) .await; + self.finish_dispatch(outcomes, histogram).await + } + + /// Sends the first attempt immediately so bulk encoding can overlap other sends. + pub(crate) async fn dispatch_initial( + &self, + peer: Peer, + request: RegionRequest, + histogram: Option<&Histogram>, + ) -> WriteOutcome { + if let Some(remote) = &self.remote { + return Self::send_write( + remote, + PendingWrite { + peer, + request: Arc::new(request), + }, + histogram, + ) + .await; + } + let _timer = histogram.map(Histogram::start_timer); + let datanode = self.node_manager.datanode(&peer).await; + WriteOutcome { + result: datanode.handle(request).await, + pending: None, + } + } + + /// Collects first attempts and retries their failures together, refreshing routes once per round. + pub(crate) async fn finish_dispatch( + &self, + outcomes: Vec, + histogram: Option<&Histogram>, + ) -> Result { + if let Some(remote) = &self.remote { + return self.dispatch_with_retry(remote, outcomes, histogram).await; + } let mut response = RegionResponse::new(0); - for result in results { - merge_response(&mut response, result?); + for outcome in outcomes { + merge_response(&mut response, outcome.result?); } Ok(response) } @@ -176,6 +263,169 @@ impl WriteDispatcher { Ok(peers) } + async fn send_write( + remote: &RemoteRetry, + write: PendingWrite, + histogram: Option<&Histogram>, + ) -> WriteOutcome { + let _timer = histogram.map(Histogram::start_timer); + let requester = remote.clients.make_region_requester(&write.peer).await; + let result = requester.handle(Arc::clone(&write.request)).await; + let pending = if result + .as_ref() + .is_err_and(|error| error.retry_hint().is_retryable()) + { + Some(write) + } else { + None + }; + WriteOutcome { + result: result.map_err(Into::into), + pending, + } + } + + async fn dispatch_with_retry( + &self, + remote: &RemoteRetry, + mut outcomes: Vec, + histogram: Option<&Histogram>, + ) -> Result { + let mut response = RegionResponse::new(0); + let mut backoff = ExponentialBuilder::default() + .with_min_delay(remote.options.initial_backoff) + .with_max_delay(remote.options.max_backoff) + .with_factor(remote.options.backoff_multiplier) + .with_max_times(remote.options.max_retries) + .with_jitter() + .build(); + + loop { + let mut failures = Vec::new(); + let mut terminal_error = None; + for outcome in outcomes { + match (outcome.pending, outcome.result) { + (_, Ok(result)) => { + merge_response(&mut response, result); + } + (Some(write), Err(error)) => { + failures.push((write, error)); + } + (None, Err(error)) => { + terminal_error.get_or_insert(error); + } + } + } + // All in-flight requests finish before returning an error. Dropping a + // sibling future does not cancel a write already accepted by a datanode. + if let Some(error) = terminal_error { + return Err(error); + } + if failures.is_empty() { + return Ok(response); + } + let Some(delay) = backoff.next() else { + let (_, error) = failures.into_iter().next().context(UnexpectedSnafu { + err_msg: "Write retry budget exhausted without a failed request", + })?; + return Err(error); + }; + tokio::time::sleep(delay).await; + + let mut refresh = Vec::new(); + let mut pending = Vec::new(); + for (write, error) in failures { + if error.retry_hint() == RetryHint::RetryAfterRouteRefresh { + refresh.push(write); + } else { + pending.push(write); + } + } + if !refresh.is_empty() { + pending.extend(self.reroute_failed_writes(remote, refresh).await?); + } + outcomes = future::join_all( + pending + .into_iter() + .map(|write| Self::send_write(remote, write, histogram)), + ) + .await; + } + } + + async fn reroute_failed_writes( + &self, + remote: &RemoteRetry, + writes: Vec, + ) -> Result> { + let regions = writes + .iter() + .map(|write| write_region_id(&write.request)) + .collect::>>()?; + self.invalidate_routes(®ions, &remote.cache_invalidator) + .await?; + let mut partitioners = HashMap::new(); + for table_id in regions + .iter() + .map(|region| region.table_id()) + .collect::>() + { + let table_info = remote + .table_info_cache + .get(table_id) + .await? + .with_context(|| TableInfoNotFoundSnafu { + table: table_id.to_string(), + })?; + let partitioner = TablePartitioner::try_new(&self.partition_manager, &table_info) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + partitioners.insert(table_id, partitioner); + } + let mut repartitioned = Vec::new(); + for write in writes { + let request = Arc::try_unwrap(write.request).map_err(|_| { + UnexpectedSnafu { + err_msg: "Write request is still shared after its RPC completed", + } + .build() + })?; + let table_id = write_region_id(&request)?.table_id(); + let partitioner = partitioners.get(&table_id).context(UnexpectedSnafu { + err_msg: "Refreshed write has no partition snapshot", + })?; + let requests = partitioner + .repartition_request(request) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + repartitioned.extend(requests); + } + self.prepare_writes(repartitioned).await + } + + async fn prepare_writes(&self, requests: Vec) -> Result> { + let regions = requests + .iter() + .map(write_region_id) + .collect::>>()?; + let peers = self + .resolve_peers(®ions) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + requests + .into_iter() + .zip(regions) + .map(|(request, region)| { + Ok(PendingWrite { + peer: get_peer(&peers, region)?, + request: Arc::new(request), + }) + }) + .collect() + } + async fn dispatch_inner( &self, requests: Vec, @@ -188,6 +438,44 @@ impl WriteDispatcher { .context(ExternalSnafu)?; self.dispatch_prepared(writes, histogram).await } + async fn invalidate_routes( + &self, + regions: &[RegionId], + invalidator: &CacheInvalidatorRef, + ) -> Result<()> { + let mut table_ids = regions + .iter() + .map(|region| region.table_id()) + .collect::>(); + let mut physical_ids = HashSet::new(); + for table_id in &table_ids { + let (physical_id, _) = self + .partition_manager + .find_physical_table_route_with_id(*table_id) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + physical_ids.insert(physical_id); + } + table_ids.extend(physical_ids); + let caches = table_ids + .into_iter() + .map(CacheIdent::TableId) + .collect::>(); + invalidator.invalidate(&Context::default(), &caches).await + } +} + +struct RemoteRetry { + clients: Arc, + options: WriteRetryOptions, + table_info_cache: TableInfoCacheRef, + cache_invalidator: CacheInvalidatorRef, +} + +struct PendingWrite { + peer: Peer, + request: Arc, } fn get_peer(peers: &HashMap, region: RegionId) -> Result { diff --git a/src/operator/src/write_dispatcher/options.rs b/src/operator/src/write_dispatcher/options.rs new file mode 100644 index 00000000000..20c8ba535f3 --- /dev/null +++ b/src/operator/src/write_dispatcher/options.rs @@ -0,0 +1,106 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::time::Duration; + +use serde::{Deserialize, Serialize}; + +/// Exponential backoff for retryable distributed writes. +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] +#[serde(default)] +pub struct WriteRetryOptions { + /// Maximum additional attempts after the original write. Zero disables retries. + pub max_retries: usize, + /// Initial backoff before jitter is applied. + #[serde(with = "humantime_serde")] + pub initial_backoff: Duration, + /// Maximum backoff before jitter is applied. + #[serde(with = "humantime_serde")] + pub max_backoff: Duration, + /// Multiplier applied to the backoff after each retry. + pub backoff_multiplier: f32, +} + +impl Default for WriteRetryOptions { + fn default() -> Self { + Self { + max_retries: 3, + initial_backoff: Duration::from_millis(1500), + max_backoff: Duration::from_secs(10), + backoff_multiplier: 2.0, + } + } +} + +impl WriteRetryOptions { + /// Checks that the configured backoff is positive, bounded, and nondecreasing. + pub fn validate(&self) -> Result<(), &'static str> { + if self.initial_backoff.is_zero() { + return Err("write_retry.initial_backoff must be positive"); + } + if self.max_backoff < self.initial_backoff { + return Err("write_retry.max_backoff must be at least initial_backoff"); + } + if !self.backoff_multiplier.is_finite() || self.backoff_multiplier < 1.0 { + return Err("write_retry.backoff_multiplier must be finite and at least 1"); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_validate_write_retry_options() { + for options in [ + WriteRetryOptions::default(), + WriteRetryOptions { + max_retries: 0, + ..Default::default() + }, + WriteRetryOptions { + backoff_multiplier: 1.0, + ..Default::default() + }, + ] { + options.validate().unwrap(); + } + for options in [ + WriteRetryOptions { + initial_backoff: Duration::ZERO, + ..Default::default() + }, + WriteRetryOptions { + max_backoff: Duration::from_secs(1), + ..Default::default() + }, + WriteRetryOptions { + backoff_multiplier: 0.5, + ..Default::default() + }, + WriteRetryOptions { + backoff_multiplier: f32::NAN, + ..Default::default() + }, + WriteRetryOptions { + backoff_multiplier: f32::INFINITY, + ..Default::default() + }, + ] { + assert!(options.validate().is_err(), "{options:?}"); + } + } +} diff --git a/src/operator/src/write_dispatcher/tests.rs b/src/operator/src/write_dispatcher/tests.rs index 1191f73d55a..ac505e9f631 100644 --- a/src/operator/src/write_dispatcher/tests.rs +++ b/src/operator/src/write_dispatcher/tests.rs @@ -12,36 +12,249 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashMap; +use std::collections::{HashMap, VecDeque}; use std::future::Future; -use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; use std::task::Poll; use std::time::{Duration, Instant}; +use api::v1::region::region_server::{Region, RegionServer}; use api::v1::region::{ - BulkInsertRequest, DeleteRequest, InsertRequest, RegionRequestHeader, bulk_insert_request, + AlignedSchemaVersion, BulkInsertRequest, DeleteRequest, InsertRequest, RegionRequestHeader, + RegionResponse as PbRegionResponse, bulk_insert_request, }; +use api::v1::value::ValueData; use api::v1::{ - ArrowIpc, ColumnDataType, PartitionExprVersion, Row, Rows, helper, value::ValueData, + ArrowIpc, ColumnDataType, ColumnSchema, PartitionExprVersion, Row, Rows, Value, helper, }; +use api::v1::{ResponseHeader, Status as PbStatus}; use client::client_manager::NodeClients; -use common_error::ext::ErrorExt; +use common_error::ext::{ErrorExt, RetryHint}; +use common_error::status_code::StatusCode; use common_grpc::channel_manager::ChannelConfig; -use common_meta::cache as meta_cache; +use common_grpc::flight::{self, FlightDecoder}; +use common_meta::cache::{ + self as meta_cache, LayeredCacheRegistryBuilder, TableInfoCacheRef, TableRouteCacheRef, +}; +use common_meta::cache_invalidator::{CacheInvalidator, CacheInvalidatorRef, Context}; use common_meta::ddl::test_util::datanode_handler::DatanodeWatcher; +use common_meta::error::{Result, UnexpectedSnafu}; +use common_meta::instruction::CacheIdent; use common_meta::key::TableMetadataManager; -use common_meta::key::table_route::TableRouteValue; +use common_meta::key::table_route::{TableRouteManager, TableRouteValue}; +use common_meta::kv_backend::KvBackendRef; +use common_meta::peer::Peer; use common_meta::test_util::{MockDatanodeHandler, MockDatanodeManager}; use common_query::request::QueryRequest; use common_recordbatch::SendableRecordBatchStream; -use partition::manager::PartitionRuleManager; +use datatypes::arrow::array::Int32Array; +use datatypes::arrow::datatypes::{DataType, Field, Schema}; +use datatypes::arrow::record_batch::RecordBatch; +use partition::cache::PartitionInfoCacheRef; +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::task::JoinHandle; +use tokio_stream::wrappers::TcpListenerStream; +use tonic::{Request, Response, Status}; use super::*; use crate::error::Error; use crate::test_util; +#[derive(Clone, Default)] +struct RecordingRegionService { + requests: Arc>>, + failures: Arc>>>, +} + +#[async_trait::async_trait] +impl Region for RecordingRegionService { + async fn handle( + &self, + request: Request, + ) -> std::result::Result, Status> { + let request = request.into_inner(); + let region_id = write_region_id(&request).unwrap().as_u64(); + self.requests.lock().unwrap().push(request); + if let Some(hint) = self + .failures + .lock() + .unwrap() + .get_mut(®ion_id) + .and_then(VecDeque::pop_front) + { + let mut status = Status::unavailable("scripted write rejection"); + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT, + hint.as_str().parse().unwrap(), + ); + return Err(status); + } + Ok(Response::new(PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus { + status_code: StatusCode::Success as u32, + ..Default::default() + }), + }), + affected_rows: 1, + ..Default::default() + })) + } +} + +#[derive(Default)] +struct RecordingCacheInvalidator { + calls: Mutex>>, + inner: Option, + route_update: Mutex>, +} + +#[async_trait::async_trait] +impl CacheInvalidator for RecordingCacheInvalidator { + async fn invalidate(&self, ctx: &Context, caches: &[CacheIdent]) -> Result<()> { + self.calls.lock().unwrap().push(caches.to_vec()); + let route_update = self.route_update.lock().unwrap().take(); + if let Some((backend, updated)) = route_update { + let manager = TableRouteManager::new(backend.clone()); + let storage = manager.table_route_storage(); + let current = storage.get_with_raw_bytes(1).await?.unwrap(); + let (txn, _) = storage.build_update_txn(1, ¤t, &updated)?; + assert!(backend.txn(txn).await?.succeeded); + } + if let Some(inner) = &self.inner { + inner.invalidate(ctx, caches).await + } else { + UnexpectedSnafu { + err_msg: "stop after recording route refresh", + } + .fail() + } + } + + fn invalidate_all(&self) -> Result<()> { + // This recorder owns no cache entries. + Ok(()) + } +} + +struct RegionTestServer { + peer: Peer, + service: Arc, + invalidator: Arc, + task: JoinHandle<()>, +} + +impl RegionTestServer { + async fn start() -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let peer = Peer::new(1, listener.local_addr().unwrap().to_string()); + let service = Arc::new(RecordingRegionService::default()); + let server_service = service.clone(); + let task = tokio::spawn(async move { + tonic::transport::Server::builder() + .add_service(RegionServer::new(server_service.as_ref().clone())) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + .unwrap(); + }); + Self { + peer, + service, + invalidator: Arc::default(), + task, + } + } + + async fn create_dispatcher(&self, max_retries: usize) -> WriteDispatcher { + let backend = test_util::prepare_mocked_backend().await; + let partition_manager = test_util::create_partition_rule_manager(backend.clone()).await; + update_table_leader(&backend, &self.peer).await; + let registry = cache::build_fundamental_cache_registry(backend); + let table_info_cache = registry.get::().unwrap(); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + WriteDispatcher::new(clients.clone(), partition_manager).with_retry( + clients, + build_retry_options(max_retries), + table_info_cache, + self.invalidator.clone(), + ) + } + + fn fail_region(&self, region_number: u32, hints: impl IntoIterator) { + self.service.failures.lock().unwrap().insert( + RegionId::new(1, region_number).as_u64(), + hints.into_iter().collect(), + ); + } + + fn attempt_counts(&self) -> HashMap { + let mut counts = HashMap::new(); + for request in self.service.requests.lock().unwrap().iter() { + *counts + .entry(write_region_id(request).unwrap().as_u64()) + .or_default() += 1; + } + counts + } +} + +impl Drop for RegionTestServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +fn build_retry_options(max_retries: usize) -> WriteRetryOptions { + WriteRetryOptions { + max_retries, + initial_backoff: Duration::from_millis(1), + max_backoff: Duration::from_millis(2), + ..Default::default() + } +} + +fn create_routing_caches( + backend: KvBackendRef, +) -> ( + PartitionRuleManagerRef, + TableInfoCacheRef, + CacheInvalidatorRef, +) { + let registry = Arc::new( + cache::with_default_composite_cache_registry( + LayeredCacheRegistryBuilder::default() + .add_cache_registry(cache::build_fundamental_cache_registry(backend.clone())), + ) + .unwrap() + .build(), + ); + let partition_manager = Arc::new(PartitionRuleManager::new( + backend.clone(), + registry.get::().unwrap(), + registry.get::().unwrap(), + )); + let table_info_cache = registry.get::().unwrap(); + (partition_manager, table_info_cache, registry) +} + +async fn update_table_leader(backend: &KvBackendRef, peer: &Peer) { + let manager = TableRouteManager::new(backend.clone()); + let storage = manager.table_route_storage(); + let current = storage.get_with_raw_bytes(1).await.unwrap().unwrap(); + let mut routes = current.region_routes().unwrap().clone(); + for route in &mut routes { + route.leader_peer = Some(peer.clone()); + } + let updated = current.update(routes).unwrap(); + let (txn, _) = storage.build_update_txn(1, ¤t, &updated).unwrap(); + assert!(backend.txn(txn).await.unwrap().succeeded); +} + fn build_insert_request(regions: &[u32]) -> RegionRequest { RegionRequest { body: Some(Body::Inserts(InsertRequests { @@ -57,6 +270,459 @@ 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 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 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}", + result.is_ok() + ); + assert_eq!(result.is_ok(), max_retries > 0); + assert_eq!(attempts, max_retries + 1); + assert!(server.invalidator.calls.lock().unwrap().is_empty()); + } + } +} + +#[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()); +} + +#[tokio::test] +async fn test_write_retry_attempt_budget_and_nonretryable_errors() { + for (hint, expected_attempts) in [(RetryHint::Retryable, 3), (RetryHint::NonRetryable, 1)] { + let server = RegionTestServer::start().await; + server.fail_region(1, [hint; 4]); + let dispatcher = server.create_dispatcher(2).await; + let error = dispatcher + .dispatch(vec![build_insert_request(&[1])]) + .await + .unwrap_err(); + // The dispatcher uses the explicit client hint for its attempt budget. + // The existing common-meta conversion promotes Unavailable to RetryLater. + assert_eq!(error.retry_hint(), RetryHint::Retryable); + assert_eq!( + server.attempt_counts()[&RegionId::new(1, 1).as_u64()], + expected_attempts + ); + 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; + for region in [1, 2] { + server.fail_region(region, [RetryHint::RetryAfterRouteRefresh]); + } + let dispatcher = server.create_dispatcher(2).await; + let error = dispatcher + .dispatch(vec![build_insert_request(&[1, 2])]) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("stop after recording route refresh") + ); + let calls = server.invalidator.calls.lock().unwrap(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0], vec![CacheIdent::TableId(1)]); + assert_eq!(server.service.requests.lock().unwrap().len(), 2); +} + +#[tokio::test] +async fn test_finish_dispatch_retries_initial_failures_in_one_round() { + let server = RegionTestServer::start().await; + for region in [2, 3] { + server.fail_region(region, [RetryHint::RetryAfterRouteRefresh]); + } + let backend = test_util::prepare_mocked_backend().await; + test_util::create_partition_rule_manager(backend.clone()).await; + update_table_leader(&backend, &server.peer).await; + let (partition_manager, table_info_cache, registry) = create_routing_caches(backend); + let table_info = table_info_cache.get(1).await.unwrap().unwrap(); + let (_, versions) = partition_manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + let invalidator = Arc::new(RecordingCacheInvalidator { + inner: Some(registry), + ..Default::default() + }); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + let dispatcher = WriteDispatcher::new(clients.clone(), partition_manager).with_retry( + clients, + build_retry_options(2), + table_info_cache, + invalidator.clone(), + ); + let mut outcomes = Vec::new(); + for (region, value) in [(1, 101), (2, 11), (3, 1)] { + let request = + build_partitioned_request(WriteKind::Bulk, region, versions[®ion], &[value]); + outcomes.push( + dispatcher + .dispatch_initial(server.peer.clone(), request, None) + .await, + ); + } + assert!(invalidator.calls.lock().unwrap().is_empty()); + assert_eq!( + server.attempt_counts(), + HashMap::from([ + (RegionId::new(1, 1).as_u64(), 1), + (RegionId::new(1, 2).as_u64(), 1), + (RegionId::new(1, 3).as_u64(), 1), + ]) + ); + let response = dispatcher.finish_dispatch(outcomes, None).await.unwrap(); + assert_eq!(response.affected_rows, 3); + assert_eq!( + server.attempt_counts(), + HashMap::from([ + (RegionId::new(1, 1).as_u64(), 1), + (RegionId::new(1, 2).as_u64(), 2), + (RegionId::new(1, 3).as_u64(), 2), + ]) + ); + assert_eq!( + *invalidator.calls.lock().unwrap(), + vec![vec![CacheIdent::TableId(1)]] + ); +} + +#[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; + + let (partition_manager, table_info_cache, registry) = 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(); + + // 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.clone(), partition_manager).with_retry( + clients, + build_retry_options(2), + table_info_cache, + 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| api::v1::PartitionExprVersion { value }); + } + } + 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)]] + ); +} + +#[derive(Clone, Copy, Debug)] +enum WriteKind { + Insert, + Delete, + Bulk, +} + +fn build_partitioned_request( + kind: WriteKind, + region: u32, + version: Option, + values: &[i32], +) -> RegionRequest { + let region_id = RegionId::new(1, region).as_u64(); + let partition_expr_version = version.map(|value| PartitionExprVersion { value }); + let rows = Rows { + schema: vec![ColumnSchema { + column_name: "a".to_string(), + datatype: ColumnDataType::Int32 as i32, + ..Default::default() + }], + rows: values + .iter() + .map(|value| Row { + values: vec![Value { + value_data: Some(ValueData::I32Value(*value)), + }], + }) + .collect(), + }; + let body = match kind { + WriteKind::Insert => Body::Inserts(InsertRequests { + requests: vec![InsertRequest { + region_id, + partition_expr_version, + rows: Some(rows), + skip_wal: true, + }], + }), + WriteKind::Delete => Body::Deletes(DeleteRequests { + requests: vec![DeleteRequest { + region_id, + partition_expr_version, + rows: Some(rows), + }], + }), + WriteKind::Bulk => { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(values.to_vec()))], + ) + .unwrap(); + let (schema, data_header, payload) = flight::record_batch_to_ipc(batch).unwrap(); + Body::BulkInsert(BulkInsertRequest { + region_id, + partition_expr_version, + skip_wal: true, + aligned_schema_version: Some(AlignedSchemaVersion { + schema_version: 123, + }), + body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { + schema, + data_header, + payload, + })), + }) + } + }; + RegionRequest { + header: None, + body: Some(body), + } +} + +fn request_values(request: &RegionRequest) -> (u32, Option, Vec) { + let region = write_region_id(request).unwrap().region_number(); + let (version, rows) = match request.body.as_ref().unwrap() { + Body::Inserts(inserts) => { + assert!(inserts.requests[0].skip_wal); + ( + inserts.requests[0].partition_expr_version, + inserts.requests[0].rows.as_ref().unwrap(), + ) + } + Body::Deletes(deletes) => ( + deletes.requests[0].partition_expr_version, + deletes.requests[0].rows.as_ref().unwrap(), + ), + Body::BulkInsert(bulk) => { + assert!(bulk.skip_wal); + assert_eq!(bulk.aligned_schema_version, None); + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &bulk.body else { + panic!("expected IPC") + }; + let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema) + .unwrap() + .try_decode_record_batch(&ipc.data_header, &ipc.payload) + .unwrap(); + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + return ( + region, + bulk.partition_expr_version.map(|v| v.value), + values.values().to_vec(), + ); + } + _ => panic!("expected write request"), + }; + ( + region, + version.map(|v| v.value), + rows.rows + .iter() + .map(|row| { + let Some(ValueData::I32Value(value)) = row.values[0].value_data else { + panic!("expected int32") + }; + value + }) + .collect(), + ) +} + +#[tokio::test] +async fn test_write_retry_repartitions_only_failed_rows_after_route_refresh() { + for kind in [WriteKind::Insert, WriteKind::Delete, WriteKind::Bulk] { + 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; + test_util::create_partition_rule_manager(backend.clone()).await; + update_table_leader(&backend, &old_server.peer).await; + let (partition_manager, table_info_cache, registry) = + create_routing_caches(backend.clone()); + let table_info = table_info_cache.get(1).await.unwrap().unwrap(); + let (_, old_versions) = partition_manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + let manager = TableRouteManager::new(backend.clone()); + let current = manager + .table_route_storage() + .get_with_raw_bytes(1) + .await + .unwrap() + .unwrap(); + let old_routes = current.region_routes().unwrap(); + let replaced = old_routes + .iter() + .find(|route| route.region.id.region_number() == 2) + .unwrap(); + let mut routes = old_routes + .iter() + .filter(|route| route.region.id.region_number() != 2) + .cloned() + .collect::>(); + for (region, lower, upper) in [(4, 10, 30), (5, 30, 50)] { + let mut route = replaced.clone(); + route.region.id = RegionId::new(1, region); + route.leader_peer = Some(new_server.peer.clone()); + route.region.partition_expr = PartitionExpr::new( + Operand::Expr(PartitionExpr::new( + Operand::Column("a".to_string()), + RestrictedOp::GtEq, + Operand::Value(datatypes::value::Value::Int32(lower)), + )), + RestrictedOp::And, + Operand::Expr(PartitionExpr::new( + Operand::Column("a".to_string()), + RestrictedOp::Lt, + Operand::Value(datatypes::value::Value::Int32(upper)), + )), + ) + .as_json_str() + .unwrap(); + routes.push(route); + } + let invalidator = Arc::new(RecordingCacheInvalidator { + calls: Mutex::default(), + inner: Some(registry), + route_update: Mutex::new(Some((backend, current.update(routes).unwrap()))), + }); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + let dispatcher = WriteDispatcher::new(clients.clone(), partition_manager.clone()) + .with_retry( + clients, + build_retry_options(2), + table_info_cache, + invalidator.clone(), + ); + let response = dispatcher + .dispatch(vec![ + build_partitioned_request(kind, 1, old_versions[&1], &[70]), + build_partitioned_request(kind, 2, old_versions[&2], &[20, 40]), + ]) + .await + .unwrap(); + assert_eq!(response.affected_rows, 3, "{kind:?}"); + assert_eq!( + old_server.attempt_counts(), + HashMap::from([ + (RegionId::new(1, 1).as_u64(), 1), + (RegionId::new(1, 2).as_u64(), 1), + ]), + "{kind:?}" + ); + let (_, versions) = partition_manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + let mut writes = new_server + .service + .requests + .lock() + .unwrap() + .iter() + .map(request_values) + .collect::>(); + writes.sort_by_key(|write| write.0); + assert_eq!( + writes, + vec![(4, versions[&4], vec![20]), (5, versions[&5], vec![40])], + "{kind:?}" + ); + assert_eq!( + *invalidator.calls.lock().unwrap(), + vec![vec![CacheIdent::TableId(1)]] + ); + } +} + #[tokio::test] async fn test_resolve_peers_batches_logical_and_physical_regions() { let backend = test_util::prepare_mocked_backend().await; @@ -348,3 +1014,28 @@ async fn test_resolve_peers_cache_lookup_ablation() { } } } + +#[tokio::test] +async fn test_with_retry_can_disable_previously_enabled_retries() { + let server = RegionTestServer::start().await; + let dispatcher = server.create_dispatcher(1).await; + let remote = dispatcher.remote.as_ref().unwrap(); + let clients = remote.clients.clone(); + let table_info_cache = remote.table_info_cache.clone(); + let cache_invalidator = remote.cache_invalidator.clone(); + let dispatcher = dispatcher.with_retry( + clients, + build_retry_options(0), + table_info_cache, + cache_invalidator, + ); + assert!(dispatcher.remote.is_none()); + server.fail_region(1, [RetryHint::Retryable]); + assert!( + dispatcher + .dispatch(vec![build_insert_request(&[1])]) + .await + .is_err() + ); + assert_eq!(server.attempt_counts()[&RegionId::new(1, 1).as_u64()], 1); +}