From d86741d0ad92b98cff2cc0cff101734fd02e270a Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 13 Oct 2026 06:02:43 +0000 Subject: [PATCH] feat(operator): retry failed writes with refreshed partitions (#9542) * feat(operator): retry failed writes with refreshed partitions Signed-off-by: WenyXu * style(operator): simplify retry test type references Signed-off-by: WenyXu * refactor(operator): clarify write request rebuilding names Signed-off-by: WenyXu * refactor(operator): encapsulate write retry attempts Signed-off-by: WenyXu * style(operator): group inserter configuration methods Signed-off-by: WenyXu * refactor(operator): share routed peers and separate retry decisions Signed-off-by: WenyXu * refactor(operator): avoid temporary write routing allocations Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- Cargo.lock | 4 + src/common/meta/src/peer.rs | 3 + src/common/meta/src/rpc/router.rs | 33 +- src/operator/AGENTS.md | 4 +- src/operator/Cargo.toml | 4 + src/operator/src/bulk_insert.rs | 47 +- src/operator/src/delete.rs | 13 +- src/operator/src/insert.rs | 25 +- src/operator/src/req_convert.rs | 1 - .../src/req_convert/common/partitioner.rs | 454 ++++++++- src/operator/src/write_dispatcher.rs | 451 +++++++-- src/operator/src/write_dispatcher/options.rs | 106 ++ src/operator/src/write_dispatcher/tests.rs | 930 +++++++++++++++++- 13 files changed, 1958 insertions(+), 117 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/common/meta/src/peer.rs b/src/common/meta/src/peer.rs index 924c293bbe0..18ffd92f8bb 100644 --- a/src/common/meta/src/peer.rs +++ b/src/common/meta/src/peer.rs @@ -22,6 +22,9 @@ use crate::cluster::NodeInfo; use crate::error::Error; use crate::{DatanodeId, FlownodeId}; +/// Shared reference to a [`Peer`]. +pub type PeerRef = Arc; + /// [`PeerResolver`] provides methods to look up peer information by node ID. /// /// This trait allows resolving both datanode and flownode peers, regardless of their current activity status. diff --git a/src/common/meta/src/rpc/router.rs b/src/common/meta/src/rpc/router.rs index 20fcd955abb..e089318b272 100644 --- a/src/common/meta/src/rpc/router.rs +++ b/src/common/meta/src/rpc/router.rs @@ -114,15 +114,25 @@ pub fn convert_to_region_leader_map(region_routes: &[RegionRoute]) -> HashMap>() } +/// Returns the leader of the first route matching the region number. pub fn find_region_leader( region_routes: &[RegionRoute], region_number: RegionNumber, ) -> Option { + find_region_leader_ref(region_routes, region_number).cloned() +} + +/// Borrows the leader of the first route matching the region number. +/// +/// Returns `None` if no route matches or the first matching route has no leader. +pub fn find_region_leader_ref( + region_routes: &[RegionRoute], + region_number: RegionNumber, +) -> Option<&Peer> { region_routes .iter() .find(|x| x.region.id.region_number() == region_number) .and_then(|r| r.leader_peer.as_ref()) - .cloned() } /// Returns the region numbers of the leader regions on the target datanode. @@ -442,6 +452,27 @@ mod tests { } } + #[test] + fn test_find_region_leader() { + let mut routes = vec![ + new_test_region_route(RegionId::new(1, 1)), + new_test_region_route(RegionId::new(1, 1)), + ]; + routes[1].leader_peer = Some(Peer::new(2, "a2")); + + assert_eq!( + find_region_leader_ref(&routes, 1), + routes[0].leader_peer.as_ref() + ); + assert_eq!(find_region_leader(&routes, 1), routes[0].leader_peer); + assert_eq!(find_region_leader_ref(&routes, 2), None); + assert_eq!(find_region_leader(&routes, 2), None); + + routes[0].leader_peer = None; + assert_eq!(find_region_leader_ref(&routes, 1), None); + assert_eq!(find_region_leader(&routes, 1), None); + } + #[test] fn test_leader_is_downgraded() { let mut region_route = RegionRoute { diff --git a/src/operator/AGENTS.md b/src/operator/AGENTS.md index 46359eb5c50..4d31f44de96 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 write request rebuilding | | 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..29537959df5 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::WriteAttempt; use crate::{error, metrics}; impl Inserter { @@ -72,7 +72,7 @@ impl Inserter { let writes = self.write_dispatcher.prepare(writes).await?; Ok(self .write_dispatcher - .dispatch_prepared(writes, None) + .execute(writes, None) .await .context(error::RequestInsertsSnafu)? .affected_rows) @@ -186,7 +186,7 @@ impl Inserter { let writes = self.write_dispatcher.prepare(vec![request]).await?; let result = self .write_dispatcher - .dispatch_prepared( + .execute( writes, Some( metrics::HANDLE_BULK_INSERT_ELAPSED.with_label_values(&["datanode_handle"]), @@ -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 + .send( + 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 attempts = results.into_iter().collect::>>()?; + dispatcher + .finish( + attempts, + 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/delete.rs b/src/operator/src/delete.rs index abd8302c965..c389374b7ac 100644 --- a/src/operator/src/delete.rs +++ b/src/operator/src/delete.rs @@ -141,13 +141,12 @@ 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.dispatch_prepared(writes, None).await - }) - .await - .context(JoinTaskSnafu)? - .context(RequestDeletesSnafu)? - .affected_rows; + let affected_rows = + common_runtime::spawn_global(async move { dispatcher.execute(writes, None).await }) + .await + .context(JoinTaskSnafu)? + .context(RequestDeletesSnafu)? + .affected_rows; crate::metrics::DIST_DELETE_ROW_COUNT .with_label_values(&[ctx.get_db_string().as_str()]) .inc_by(affected_rows as u64); diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 89c1195cda4..2c01a35b8d9 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -316,6 +316,12 @@ impl Inserter { } } + /// Installs the shared dispatcher used by direct writes and batch flushes. + pub fn with_write_dispatcher(mut self, dispatcher: WriteDispatcherRef) -> Self { + self.write_dispatcher = dispatcher; + self + } + /// Installs the shared batcher; callers explicitly select its ingestion entry point. pub fn with_pending_rows_batcher( mut self, @@ -835,13 +841,12 @@ 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.dispatch_prepared(writes, None).await - }) - .await - .context(JoinTaskSnafu)? - .context(RequestInsertsSnafu)? - .affected_rows; + let affected_rows = + common_runtime::spawn_global(async move { dispatcher.execute(writes, None).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); @@ -1675,12 +1680,6 @@ impl Inserter { } } - /// Installs the shared dispatcher used by direct writes and batch flushes. - pub fn with_write_dispatcher(mut self, dispatcher: WriteDispatcherRef) -> Self { - self.write_dispatcher = dispatcher; - self - } - /// Returns the dispatcher shared with batching execution paths. pub fn write_dispatcher(&self) -> &WriteDispatcherRef { &self.write_dispatcher diff --git a/src/operator/src/req_convert.rs b/src/operator/src/req_convert.rs index 1161f62795d..82681e09968 100644 --- a/src/operator/src/req_convert.rs +++ b/src/operator/src/req_convert.rs @@ -16,5 +16,4 @@ pub(crate) mod common; pub mod delete; pub mod insert; -pub use common::partitioner::Partitioner; pub(crate) use common::partitioner::TablePartitioner; diff --git a/src/operator/src/req_convert/common/partitioner.rs b/src/operator/src/req_convert/common/partitioner.rs index 044d3b97db5..f1c3eff4fe3 100644 --- a/src/operator/src/req_convert/common/partitioner.rs +++ b/src/operator/src/req_convert/common/partitioner.rs @@ -14,16 +14,19 @@ 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; +use common_grpc::flight::{self, FlightDecoder}; use datatypes::arrow::compute; 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,37 @@ impl TablePartitioner { }) } + /// Rebuilds write requests using this partition snapshot after a route refresh. + /// + /// Takes ownership so row payloads can move into their new regions. Requests whose + /// region and partition version still match retain their original payloads. + pub(crate) fn rebuild_write_requests( + &self, + request: RegionRequest, + ) -> Result> { + let RegionRequest { header, body } = request; + let bodies = match body.context(error::UnexpectedSnafu { + violated: "Cannot rebuild a request without a body", + })? { + Body::Inserts(inserts) => self.rebuild_insert_requests(inserts)?, + Body::Deletes(deletes) => self.rebuild_delete_requests(deletes)?, + Body::BulkInsert(bulk) => self.rebuild_bulk_insert_request(bulk)?, + _ => { + return error::UnexpectedSnafu { + violated: "Only writes can be rebuilt", + } + .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 +212,84 @@ impl TablePartitioner { Ok(requests) } + fn rebuild_insert_requests(&self, inserts: InsertRequests) -> Result> { + let mut bodies = Vec::new(); + for insert in inserts.requests { + let region_id = RegionId::from_u64(insert.region_id); + if self.is_partition_expr_unchanged(region_id, insert.partition_expr_version) { + bodies.push(Body::Inserts(InsertRequests { + requests: vec![insert], + })); + continue; + } + let rows = insert.rows.context(error::UnexpectedSnafu { + violated: "Cannot rebuild an insert without rows", + })?; + for request in self.partition_insert_requests(rows, insert.skip_wal)? { + bodies.push(Body::Inserts(InsertRequests { + requests: vec![request], + })); + } + } + Ok(bodies) + } + + fn rebuild_delete_requests(&self, deletes: DeleteRequests) -> Result> { + let mut bodies = Vec::new(); + for delete in deletes.requests { + let region_id = RegionId::from_u64(delete.region_id); + if self.is_partition_expr_unchanged(region_id, delete.partition_expr_version) { + bodies.push(Body::Deletes(DeleteRequests { + requests: vec![delete], + })); + continue; + } + let rows = delete.rows.context(error::UnexpectedSnafu { + violated: "Cannot rebuild a delete without rows", + })?; + for request in self.partition_delete_requests(rows)? { + bodies.push(Body::Deletes(DeleteRequests { + requests: vec![request], + })); + } + } + Ok(bodies) + } + + fn rebuild_bulk_insert_request(&self, bulk: BulkInsertRequest) -> Result> { + let region_id = RegionId::from_u64(bulk.region_id); + if self.is_partition_expr_unchanged(region_id, bulk.partition_expr_version) { + return Ok(vec![Body::BulkInsert(bulk)]); + } + let bulk_insert_request::Body::ArrowIpc(ipc) = + bulk.body.context(error::UnexpectedSnafu { + violated: "Cannot rebuild a bulk insert without Arrow IPC data", + })?; + let batch = 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)?; + Ok(self + .partition_bulk_insert_requests(batch, Some(ipc), bulk.skip_wal)? + .into_iter() + .map(Body::BulkInsert) + .collect()) + } + + /// Whether the request's region still exists with the same partition expression version. + /// + /// A leader change leaves the partition unchanged, so the original payload can be + /// sent to the new peer. A removed region or changed partition expression version + /// requires rebuilding the request using this snapshot's partition rule. + fn is_partition_expr_unchanged( + &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 +299,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 +375,324 @@ 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(), + 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_rebuild_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 = rebuild_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 = rebuild_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_rebuild_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 = rebuild_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 = 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 = rebuild_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_rebuild_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 = rebuild_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 rebuild_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.rebuild_write_requests(request) + } } diff --git a/src/operator/src/write_dispatcher.rs b/src/operator/src/write_dispatcher.rs index 4af4fbfafdf..c2ac6543411 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,43 +21,85 @@ 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::{Error, ExternalSnafu, Result, TableInfoNotFoundSnafu, UnexpectedSnafu}; +use common_meta::instruction::CacheIdent; use common_meta::node_manager::NodeManagerRef; -use common_meta::peer::Peer; +use common_meta::peer::{Peer, PeerRef}; use common_meta::rpc::router; use futures::future; -use partition::manager::PartitionRuleManagerRef; +use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef}; 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, + retry: Option, } /// Shared reference to a [`WriteDispatcher`]. pub type WriteDispatcherRef = Arc; +/// The result of one write RPC attempt. +pub(crate) enum WriteAttempt { + Success(RegionResponse), + Failed { + error: Error, + /// Retained when retries are enabled and the original client error allows retry. + /// The remaining retry budget is checked separately. + request: 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, + retry: 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.retry = Some(WriteRetry { + clients, + partition_manager: self.partition_manager.clone(), + options, + table_info_cache, + cache_invalidator, + }); + } else { + self.retry = 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, @@ -69,11 +112,17 @@ impl WriteDispatcher { pub(crate) async fn prepare( &self, requests: Vec, - ) -> error::Result> { - let requests = requests.into_iter().map(split_request).collect::>(); + ) -> error::Result> { + let requests = requests + .into_iter() + .map(|mut request| { + let header = request.header.take(); + (header, split_request(request)) + }) + .collect::>(); let regions = requests .iter() - .flatten() + .flat_map(|(_, requests)| requests) .map(write_region_id) .collect::>>() .context(RequestRegionSnafu)?; @@ -83,8 +132,8 @@ impl WriteDispatcher { .context(FindRegionLeaderSnafu)?; let mut writes = Vec::new(); // Preserve each original envelope and peer batching when retries are disabled. - for requests in requests { - let mut by_peer: HashMap = HashMap::new(); + for (header, requests) in requests { + let mut by_peer: HashMap = HashMap::new(); for request in requests { let peer = get_peer( &peers, @@ -112,27 +161,94 @@ impl WriteDispatcher { }, } } - writes.extend(by_peer); + for (peer, mut request) in by_peer { + request.header = header.clone(); + writes.push((peer, request)); + } } Ok(writes) } /// Sends requests using routes already resolved by the caller. - pub(crate) async fn dispatch_prepared( + pub(crate) async fn execute( &self, - writes: Vec<(Peer, RegionRequest)>, + writes: Vec<(PeerRef, RegionRequest)>, 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.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; + self.finish(attempts, histogram).await + } + + /// Sends the first attempt immediately so bulk encoding can overlap other sends. + pub(crate) async fn send( + &self, + peer: PeerRef, + request: RegionRequest, + histogram: Option<&Histogram>, + ) -> WriteAttempt { + if let Some(retry) = &self.retry { + return retry + .send( + RoutedWrite { + peer, + request: Arc::new(request), + }, + histogram, + ) + .await; + } + let _timer = histogram.map(Histogram::start_timer); + let datanode = self.node_manager.datanode(&peer).await; + match datanode.handle(request).await { + Ok(response) => WriteAttempt::Success(response), + Err(error) => WriteAttempt::Failed { + error, + request: None, + }, + } + } + + /// Collects first attempts and retries their failures together, refreshing routes once per round. + pub(crate) async fn finish( + &self, + attempts: Vec, + histogram: Option<&Histogram>, + ) -> Result { + if let Some(retry) = &self.retry { + if retry.should_retry(&attempts, 0) { + return retry.retry(attempts, histogram).await; + } + let mut response = RegionResponse::new(0); + let failures = merge_attempts(attempts, &mut response)?; + return match failures.into_iter().next() { + Some((_, error)) => Err(error), + None => Ok(response), + }; + } let mut response = RegionResponse::new(0); - for result in results { - merge_response(&mut response, result?); + for attempt in attempts { + match attempt { + WriteAttempt::Success(result) => merge_response(&mut response, result), + WriteAttempt::Failed { error, .. } => return Err(error), + } } Ok(response) } @@ -141,39 +257,8 @@ impl WriteDispatcher { pub(crate) async fn resolve_peers( &self, regions: &[RegionId], - ) -> partition::error::Result> { - let table_ids = regions - .iter() - .map(|region| region.table_id()) - .collect::>(); - let mut physical_tables = HashMap::new(); - let mut routes = HashMap::new(); - for table_id in table_ids { - let (physical_id, route) = self - .partition_manager - .find_physical_table_route_with_id(table_id) - .await?; - physical_tables.insert(table_id, physical_id); - routes.entry(physical_id).or_insert(route); - } - let mut peers = HashMap::with_capacity(regions.len()); - for region in regions { - if peers.contains_key(region) { - continue; - } - let peer = physical_tables - .get(®ion.table_id()) - .and_then(|physical_id| routes.get(physical_id)) - .and_then(|route| { - router::find_region_leader(&route.region_routes, region.region_number()) - }) - .context(partition::error::FindLeaderSnafu { - region_id: *region, - table_id: region.table_id(), - })?; - peers.insert(*region, peer); - } - Ok(peers) + ) -> partition::error::Result> { + resolve_peers(&self.partition_manager, regions).await } async fn dispatch_inner( @@ -186,16 +271,270 @@ impl WriteDispatcher { .await .map_err(BoxedError::new) .context(ExternalSnafu)?; - self.dispatch_prepared(writes, histogram).await + self.execute(writes, histogram).await } } -fn get_peer(peers: &HashMap, region: RegionId) -> Result { +/// Sends shared requests and coordinates retries with refreshed routes. +struct WriteRetry { + clients: Arc, + partition_manager: PartitionRuleManagerRef, + options: WriteRetryOptions, + table_info_cache: TableInfoCacheRef, + cache_invalidator: CacheInvalidatorRef, +} + +/// A single-region write request paired with its destination peer. +pub(crate) struct RoutedWrite { + peer: PeerRef, + request: Arc, +} + +impl WriteRetry { + async fn send(&self, write: RoutedWrite, histogram: Option<&Histogram>) -> WriteAttempt { + let _timer = histogram.map(Histogram::start_timer); + let requester = self.clients.make_region_requester(&write.peer).await; + match requester.handle(Arc::clone(&write.request)).await { + Ok(response) => WriteAttempt::Success(response), + Err(error) => { + let request = error.retry_hint().is_retryable().then_some(write); + WriteAttempt::Failed { + error: error.into(), + request, + } + } + } + } + + fn should_retry(&self, attempts: &[WriteAttempt], retries: usize) -> bool { + if retries >= self.options.max_retries { + return false; + } + let mut retryable = false; + for attempt in attempts { + match attempt { + WriteAttempt::Success(_) => {} + WriteAttempt::Failed { + request: Some(_), .. + } => retryable = true, + WriteAttempt::Failed { request: None, .. } => return false, + } + } + retryable + } + + async fn retry( + &self, + mut attempts: Vec, + histogram: Option<&Histogram>, + ) -> Result { + let mut response = RegionResponse::new(0); + let mut backoff = ExponentialBuilder::default() + .with_min_delay(self.options.initial_backoff) + .with_max_delay(self.options.max_backoff) + .with_factor(self.options.backoff_multiplier) + .with_max_times(self.options.max_retries) + .with_jitter() + .build(); + let mut retries = 0; + while self.should_retry(&attempts, retries) { + let Some(delay) = backoff.next() else { + break; + }; + let failures = merge_attempts(attempts, &mut response)?; + 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(refresh).await?); + } + attempts = + future::join_all(pending.into_iter().map(|write| self.send(write, histogram))) + .await; + retries += 1; + } + match merge_attempts(attempts, &mut response)?.into_iter().next() { + Some((_, error)) => Err(error), + None => Ok(response), + } + } + + async fn reroute_failed_writes(&self, writes: Vec) -> Result> { + let regions = writes + .iter() + .map(|write| write_region_id(&write.request)) + .collect::>>()?; + self.invalidate_routes(®ions).await?; + let mut partitioners = HashMap::new(); + for table_id in regions + .iter() + .map(|region| region.table_id()) + .collect::>() + { + let table_info = self + .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 rebuilt_requests = 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 + .rebuild_write_requests(request) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + rebuilt_requests.extend(requests); + } + self.prepare(rebuilt_requests).await + } + + async fn prepare(&self, requests: Vec) -> Result> { + let regions = requests + .iter() + .map(write_region_id) + .collect::>>()?; + let peers = resolve_peers(&self.partition_manager, ®ions) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + requests + .into_iter() + .zip(regions) + .map(|(request, region)| { + Ok(RoutedWrite { + peer: get_peer(&peers, region)?, + request: Arc::new(request), + }) + }) + .collect() + } + + async fn invalidate_routes(&self, regions: &[RegionId]) -> 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::>(); + self.cache_invalidator + .invalidate(&Context::default(), &caches) + .await + } +} + +async fn resolve_peers( + partition_manager: &PartitionRuleManager, + regions: &[RegionId], +) -> partition::error::Result> { + let table_ids = regions + .iter() + .map(|region| region.table_id()) + .collect::>(); + let mut physical_tables = HashMap::new(); + let mut routes = HashMap::new(); + for table_id in table_ids { + let (physical_id, route) = partition_manager + .find_physical_table_route_with_id(table_id) + .await?; + physical_tables.insert(table_id, physical_id); + routes.entry(physical_id).or_insert(route); + } + let mut shared_peers: HashMap<&Peer, PeerRef> = HashMap::new(); + let mut peers = HashMap::with_capacity(regions.len()); + for region in regions { + if peers.contains_key(region) { + continue; + } + let peer = physical_tables + .get(®ion.table_id()) + .and_then(|physical_id| routes.get(physical_id)) + .and_then(|route| { + router::find_region_leader_ref(&route.region_routes, region.region_number()) + }) + .context(partition::error::FindLeaderSnafu { + region_id: *region, + table_id: region.table_id(), + })?; + let peer = shared_peers + .entry(peer) + .or_insert_with(|| Arc::new(peer.clone())); + peers.insert(*region, Arc::clone(peer)); + } + Ok(peers) +} + +fn get_peer(peers: &HashMap, region: RegionId) -> Result { peers.get(®ion).cloned().context(UnexpectedSnafu { err_msg: "Planned write has no resolved peer", }) } +fn merge_attempts( + attempts: Vec, + response: &mut RegionResponse, +) -> Result> { + let mut failures = Vec::new(); + let mut terminal_error = None; + for attempt in attempts { + match attempt { + WriteAttempt::Success(result) => merge_response(response, result), + WriteAttempt::Failed { + error, + request: Some(write), + } => failures.push((write, error)), + WriteAttempt::Failed { + error, + request: None, + } => { + terminal_error.get_or_insert(error); + } + } + } + // A terminal error takes precedence over retained retryable failures. + match terminal_error { + Some(error) => Err(error), + None => Ok(failures), + } +} + fn merge_response(response: &mut RegionResponse, result: RegionResponse) { response.affected_rows += result.affected_rows; response.extensions.extend(result.extensions); 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..250fe87f353 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>>, + rejections: 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 + .rejections + .lock() + .unwrap() + .get_mut(®ion_id) + .and_then(VecDeque::pop_front) + { + let mut status = Status::unavailable(format!("region {region_id} rejected the write")); + 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.rejections.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,529 @@ 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_stops_when_a_sibling_is_nonretryable() { + for hint in [RetryHint::Retryable, RetryHint::RetryAfterRouteRefresh] { + for regions in [[1, 2], [2, 1]] { + let server = RegionTestServer::start().await; + server.fail_region(1, [hint]); + server.fail_region(2, [RetryHint::NonRetryable]); + let dispatcher = server.create_dispatcher(2).await; + let error = dispatcher + .dispatch(vec![build_insert_request(®ions)]) + .await + .unwrap_err(); + let terminal_region = RegionId::new(1, 2).as_u64(); + assert!( + format!("{error:?}") + .contains(&format!("region {terminal_region} rejected the write")), + "{error:?}" + ); + assert_eq!( + server.attempt_counts(), + HashMap::from([(RegionId::new(1, 1).as_u64(), 1), (terminal_region, 1),]) + ); + assert!(server.invalidator.calls.lock().unwrap().is_empty()); + } + } +} + +#[tokio::test] +async fn test_should_retry_requires_a_retained_request_and_budget() { + let backend = test_util::prepare_mocked_backend().await; + let (partition_manager, table_info_cache, cache_invalidator) = create_routing_caches(backend); + let retry = WriteRetry { + clients: Arc::new(NodeClients::default()), + partition_manager, + options: build_retry_options(2), + table_info_cache, + cache_invalidator, + }; + let failure = |request| { + // The converted error alone cannot distinguish an explicitly nonretryable attempt. + let error: client::Error = Status::unavailable("write rejected").into(); + WriteAttempt::Failed { + error: error.into(), + request, + } + }; + let retryable = || { + failure(Some(RoutedWrite { + peer: Arc::new(Peer::new(1, "peer")), + request: Arc::new(build_insert_request(&[1])), + })) + }; + for (name, attempts, retries, expected) in [ + ("empty", vec![], 0, false), + ( + "all succeeded", + vec![WriteAttempt::Success(RegionResponse::new(1))], + 0, + false, + ), + ("retryable", vec![retryable()], 0, true), + ("last retry", vec![retryable()], 1, true), + ("budget exhausted", vec![retryable()], 2, false), + ("terminal first", vec![failure(None), retryable()], 0, false), + ("terminal last", vec![retryable(), failure(None)], 0, false), + ] { + assert_eq!(retry.should_retry(&attempts, retries), expected, "{name}"); + } +} + +#[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_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 attempts = Vec::new(); + for (region, value) in [(1, 101), (2, 11), (3, 1)] { + let request = + build_partitioned_request(WriteKind::BulkInsert, region, versions[®ion], &[value]); + attempts.push( + dispatcher + .send(Arc::new(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(attempts, 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| 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, + BulkInsert, +} + +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::BulkInsert => { + 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_parts(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_rebuilds_only_failed_rows_after_route_refresh() { + for kind in [WriteKind::Insert, WriteKind::Delete, WriteKind::BulkInsert] { + 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_parts) + .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; @@ -82,11 +818,14 @@ async fn test_resolve_peers_batches_logical_and_physical_regions() { assert_eq!(peers.len(), 4); for region in regions { assert_eq!( - peers[®ion], - partition_manager.find_region_leader(region).await.unwrap() + peers[®ion].as_ref(), + &partition_manager.find_region_leader(region).await.unwrap() ); } - assert_eq!(peers[&RegionId::new(1, 1)], peers[&RegionId::new(42, 1)]); + assert!(Arc::ptr_eq( + &peers[&RegionId::new(1, 1)], + &peers[&RegionId::new(42, 1)], + )); let missing = RegionId::new(42, 999); let expected = partition_manager .find_region_leader(missing) @@ -98,6 +837,38 @@ async fn test_resolve_peers_batches_logical_and_physical_regions() { assert!(dispatcher.resolve_peers(&[]).await.unwrap().is_empty()); } +#[tokio::test] +async fn test_resolve_peers_shares_only_identical_peers() { + let backend = test_util::prepare_mocked_backend().await; + test_util::create_partition_rule_manager(backend.clone()).await; + 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(); + let peer = Peer::new(7, "same-address"); + let other = Peer::new(7, "other-address"); + for route in &mut routes { + route.leader_peer = Some(if route.region.id.region_number() == 3 { + other.clone() + } else { + 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); + let (partition_manager, _, _) = create_routing_caches(backend); + let dispatcher = WriteDispatcher::new(Arc::new(NodeClients::default()), partition_manager); + let regions = [1, 2, 3].map(|region| RegionId::new(1, region)); + let peers = dispatcher.resolve_peers(®ions).await.unwrap(); + assert_eq!(peers.len(), 3); + assert_eq!(peers[®ions[0]].as_ref(), &peer); + assert_eq!(peers[®ions[1]].as_ref(), &peer); + assert_eq!(peers[®ions[2]].as_ref(), &other); + assert!(Arc::ptr_eq(&peers[®ions[0]], &peers[®ions[1]])); + assert!(!Arc::ptr_eq(&peers[®ions[0]], &peers[®ions[2]])); +} + #[tokio::test] async fn test_dispatch_preserves_write_payloads_and_envelopes() { let backend = test_util::prepare_mocked_backend().await; @@ -158,7 +929,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.dispatch_prepared(writes, None).await.unwrap(); + let response = dispatcher.execute(writes, None).await.unwrap(); assert_eq!(response.affected_rows, 2); for request in requests { assert_eq!(received.try_recv().unwrap(), (Peer::new(1, ""), request)); @@ -315,8 +1086,9 @@ async fn test_resolve_peers_cache_lookup_ablation() { for mode in modes { let before = counter.get(); let started = Instant::now(); - let peers = if mode == "batched" { - dispatcher.resolve_peers(®ions).await.unwrap() + let (elapsed, peers_match) = if mode == "batched" { + let peers = dispatcher.resolve_peers(®ions).await.unwrap(); + (started.elapsed(), peers == expected) } else { let mut peers = HashMap::new(); for region in ®ions { @@ -328,11 +1100,17 @@ async fn test_resolve_peers_cache_lookup_ablation() { partition_manager.find_region_leader(*region).await.unwrap(), ); } - peers + let elapsed = started.elapsed(); + let peers_match = peers.len() == expected.len() + && peers.iter().all(|(region, peer)| { + expected + .get(region) + .is_some_and(|expected| expected.as_ref() == peer) + }); + (elapsed, peers_match) }; - let elapsed = started.elapsed(); let gets = counter.get() - before; - assert_eq!(peers, expected); + assert!(peers_match, "{mode}"); let expected_gets = match mode { "delete" => 384, "insert" => 9, @@ -348,3 +1126,121 @@ async fn test_resolve_peers_cache_lookup_ablation() { } } } + +#[tokio::test] +async fn test_write_retry_rebuilds_logical_rows_in_mixed_table_envelope() { + for kind in [WriteKind::Insert, WriteKind::Delete] { + let server = RegionTestServer::start().await; + let backend = test_util::prepare_mocked_backend().await; + test_util::create_partition_rule_manager(backend.clone()).await; + TableMetadataManager::new(backend.clone()) + .create_table_metadata( + test_util::new_test_table_info(42, "retry_logical", [1, 2, 3].into_iter()), + TableRouteValue::logical(1), + Default::default(), + ) + .await + .unwrap(); + 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 physical_id = RegionId::new(1, 1).as_u64(); + let stale_logical_id = RegionId::new(42, 2).as_u64(); + let rebuilt_logical_id = RegionId::new(42, 3).as_u64(); + server.service.rejections.lock().unwrap().insert( + stale_logical_id, + VecDeque::from([RetryHint::RetryAfterRouteRefresh]), + ); + + let header = Some(RegionRequestHeader { + dbname: "greptime.public".to_string(), + ..Default::default() + }); + let mut physical = build_partitioned_request(kind, 1, versions[&1], &[70]); + physical.header = header.clone(); + // The stale logical request must move to region 3 while keeping table 42. + let mut logical = build_partitioned_request(kind, 2, Some(0), &[1]); + let mut rebuilt = build_partitioned_request(kind, 3, versions[&3], &[1]); + for request in [&mut logical, &mut rebuilt] { + request.header = header.clone(); + let region_id = match request.body.as_mut().unwrap() { + Body::Inserts(inserts) => &mut inserts.requests[0].region_id, + Body::Deletes(deletes) => &mut deletes.requests[0].region_id, + _ => unreachable!(), + }; + *region_id = RegionId::new(42, RegionId::from_u64(*region_id).region_number()).as_u64(); + } + let mut envelope = physical.clone(); + match ( + envelope.body.as_mut().unwrap(), + logical.body.clone().unwrap(), + ) { + (Body::Inserts(physical), Body::Inserts(logical)) => { + physical.requests.extend(logical.requests); + } + (Body::Deletes(physical), Body::Deletes(logical)) => { + physical.requests.extend(logical.requests); + } + _ => unreachable!(), + } + let response = dispatcher.dispatch(vec![envelope]).await.unwrap(); + assert_eq!(response.affected_rows, 2, "{kind:?}"); + assert_eq!( + server.attempt_counts(), + HashMap::from([ + (physical_id, 1), + (stale_logical_id, 1), + (rebuilt_logical_id, 1) + ]), + "{kind:?}" + ); + let mut requests = server.service.requests.lock().unwrap().clone(); + requests.sort_by_key(|request| write_region_id(request).unwrap().as_u64()); + assert_eq!(requests, vec![physical, logical, rebuilt], "{kind:?}"); + let calls = invalidator.calls.lock().unwrap(); + assert_eq!(calls.len(), 1, "{kind:?}"); + assert_eq!(calls[0].len(), 2, "{kind:?}"); + assert!(calls[0].contains(&CacheIdent::TableId(1))); + assert!(calls[0].contains(&CacheIdent::TableId(42))); + } +} + +#[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 retry = dispatcher.retry.as_ref().unwrap(); + let clients = retry.clients.clone(); + let table_info_cache = retry.table_info_cache.clone(); + let cache_invalidator = retry.cache_invalidator.clone(); + let dispatcher = dispatcher.with_retry( + clients, + build_retry_options(0), + table_info_cache, + cache_invalidator, + ); + assert!(dispatcher.retry.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); +}