diff --git a/src/operator/AGENTS.md b/src/operator/AGENTS.md index 49cbd617567..46359eb5c50 100644 --- a/src/operator/AGENTS.md +++ b/src/operator/AGENTS.md @@ -14,6 +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 | | 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 | diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index f8d58f5d173..d5aa7168b44 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -14,33 +14,33 @@ use std::collections::HashSet; -use ahash::{HashMap, HashMapExt}; +use api::region::RegionResponse; use api::v1::flow::{DirtyWindowRequest, DirtyWindowRequests}; use api::v1::region::{ BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request, }; use api::v1::{ArrowIpc, PartitionExprVersion}; -use arrow::compute::filter_record_batch; +use arrow::compute; use arrow::record_batch::RecordBatch; use bytes::Bytes; use common_base::AffectedRows; -use common_error::ext::BoxedError; use common_grpc::FlightData; -use common_grpc::flight::{FlightEncoder, FlightMessage, record_batch_to_ipc}; +use common_grpc::flight::{FlightEncoder, FlightMessage}; use common_telemetry::error; use common_telemetry::tracing_context::TracingContext; -use futures::future::{join_all, try_join_all}; +use futures::future; use meter_core::data::MeterRecord; use meter_macros::write_meter; use session::context::{Channel, QueryContextRef}; -use snafu::{ResultExt, ensure}; +use snafu::{OptionExt, ResultExt, ensure}; use store_api::storage::RegionId; use table::TableRef; use table::metadata::TableInfoRef; use crate::error::Result; use crate::insert::Inserter; -use crate::req_convert::insert::extract_timestamps; +use crate::req_convert::TablePartitioner; +use crate::req_convert::insert; use crate::{error, metrics}; impl Inserter { @@ -54,72 +54,28 @@ impl Inserter { batch: RecordBatch, ctx: QueryContextRef, ) -> Result { - let (rule, versions) = self - .partition_manager - .find_table_partition_rule(&table_info) + let requests = TablePartitioner::try_new(&self.partition_manager, &table_info) .await - .context(error::InvalidPartitionSnafu)?; - let masks = rule - .split_record_batch(&batch) - .context(error::SplitInsertSnafu)?; - let mut writes = Vec::with_capacity(masks.len()); - for (region_number, mask) in masks { - if mask.select_none() { - continue; - } - let region_id = RegionId::new(table_info.table_id(), region_number); - let selected = if mask.select_all() { - batch.clone() - } else { - filter_record_batch(&batch, mask.array()).context(error::ComputeArrowSnafu)? - }; - let (schema, data_header, payload) = record_batch_to_ipc(selected) - .map_err(BoxedError::new) - .context(error::ExternalSnafu)?; - let peer = self - .partition_manager - .find_region_leader(region_id) - .await - .context(error::FindRegionLeaderSnafu)?; - let request = RegionRequest { + .context(error::InvalidPartitionSnafu)? + .partition_bulk_insert_requests(batch, None, ctx.skip_wal())?; + let writes = requests + .into_iter() + .map(|request| RegionRequest { header: Some(RegionRequestHeader { dbname: ctx.get_db_string(), tracing_context: TracingContext::from_current_span().to_w3c(), ..Default::default() }), - body: Some(region_request::Body::BulkInsert(BulkInsertRequest { - skip_wal: ctx.skip_wal(), - region_id: region_id.as_u64(), - partition_expr_version: versions - .get(®ion_number) - .copied() - .flatten() - .map(|value| PartitionExprVersion { value }), - // Let the datanode revalidate against its current schema. - aligned_schema_version: None, - body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { - schema, - data_header, - payload, - })), - })), - }; - writes.push((peer, request)); - } - let results = join_all(writes.into_iter().map(|(peer, request)| async move { - self.node_manager - .datanode(&peer) - .await - .handle(request) - .await - .context(error::RequestInsertsSnafu) - })) - .await; - let affected_rows = results - .into_iter() - .map(|result| result.map(|response| response.affected_rows)) - .sum::>()?; - Ok(affected_rows) + body: Some(region_request::Body::BulkInsert(request)), + }) + .collect(); + let writes = self.write_dispatcher.prepare(writes).await?; + Ok(self + .write_dispatcher + .dispatch_prepared(writes, None) + .await + .context(error::RequestInsertsSnafu)? + .affected_rows) } /// Handle bulk insert request. @@ -208,12 +164,6 @@ impl Inserter { .get(®ion_number) .copied() .unwrap_or_default(); - let datanode = self - .partition_manager - .find_region_leader(region_id) - .await - .context(error::FindRegionLeaderSnafu)?; - let request = RegionRequest { header: Some(RegionRequestHeader { tracing_context: TracingContext::from_current_span().to_w3c(), @@ -233,12 +183,15 @@ impl Inserter { })), }; - let _datanode_handle_timer = metrics::HANDLE_BULK_INSERT_ELAPSED - .with_label_values(&["datanode_handle"]) - .start_timer(); - let datanode = self.node_manager.datanode(&datanode).await; - let result = datanode - .handle(request) + let writes = self.write_dispatcher.prepare(vec![request]).await?; + let result = self + .write_dispatcher + .dispatch_prepared( + writes, + Some( + metrics::HANDLE_BULK_INSERT_ELAPSED.with_label_values(&["datanode_handle"]), + ), + ) .await .context(error::RequestRegionSnafu) .map(|r| r.affected_rows); @@ -250,121 +203,126 @@ impl Inserter { return result; } - let mut mask_per_datanode = HashMap::with_capacity(region_masks.len()); - for (region_number, mask) in region_masks { - let region_id = RegionId::new(table_id, region_number); - let datanode = self - .partition_manager - .find_region_leader(region_id) - .await - .context(error::FindRegionLeaderSnafu)?; - mask_per_datanode - .entry(datanode) - .or_insert_with(Vec::new) - .push((region_id, mask)); - } + let regions = region_masks + .keys() + .map(|region| RegionId::new(table_id, *region)) + .collect::>(); + let mut peers = self + .write_dispatcher + .resolve_peers(®ions) + .await + .context(error::FindRegionLeaderSnafu)?; let wait_all_datanode_timer = metrics::HANDLE_BULK_INSERT_ELAPSED .with_label_values(&["wait_all_datanode"]) .start_timer(); - let mut handles = Vec::with_capacity(mask_per_datanode.len()); + let mut handles = Vec::with_capacity(region_masks.len()); - for (peer, masks) in mask_per_datanode { - for (region_id, mask) in masks { - if mask.select_none() { - continue; - } - let partition_expr_version = partition_versions - .get(®ion_id.region_number()) - .copied() - .unwrap_or_default(); - let rb = record_batch.clone(); - let schema_bytes = schema_bytes.clone(); - let node_manager = self.node_manager.clone(); - let peer = peer.clone(); - let raw_header_and_data = if mask.select_all() { - Some(( - raw_flight_data.data_header.clone(), - raw_flight_data.data_body.clone(), - )) - } else { - None - }; - 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. - raw_header_and_data.unwrap() - } else { - let filter_timer = metrics::HANDLE_BULK_INSERT_ELAPSED - .with_label_values(&["filter"]) - .start_timer(); - let batch = filter_record_batch(&rb, mask.array()) - .context(error::ComputeArrowSnafu)?; - filter_timer.observe_duration(); - metrics::BULK_REQUEST_ROWS - .with_label_values(&["rows_per_region"]) - .observe(batch.num_rows() as f64); - - let encode_timer = metrics::HANDLE_BULK_INSERT_ELAPSED - .with_label_values(&["encode"]) - .start_timer(); - let mut iter = FlightEncoder::default() - .encode(FlightMessage::RecordBatch(batch)) - .into_iter(); - let Some(flight_data) = iter.next() else { - // Safety: `iter` on a type of `Vec1`, which is guaranteed to have - // at least one element. - unreachable!() - }; - ensure!( - iter.next().is_none(), - error::NotSupportedSnafu { - feat: "bulk insert RecordBatch with dictionary arrays", - } - ); - encode_timer.observe_duration(); - (flight_data.data_header, flight_data.data_body) - }; - let _datanode_handle_timer = metrics::HANDLE_BULK_INSERT_ELAPSED - .with_label_values(&["datanode_handle"]) - .start_timer(); - let request = RegionRequest { - header: Some(RegionRequestHeader { - tracing_context: TracingContext::from_current_span().to_w3c(), - ..Default::default() - }), - body: Some(region_request::Body::BulkInsert(BulkInsertRequest { - skip_wal, - region_id: region_id.as_u64(), - partition_expr_version: partition_expr_version - .map(|value| PartitionExprVersion { value }), - aligned_schema_version: None, - body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { - schema: schema_bytes, - data_header: header, - payload, - })), - })), - }; - - let datanode = node_manager.datanode(&peer).await; - datanode - .handle(request) - .await - .context(error::RequestRegionSnafu) - }); - handles.push(handle); + for (region_number, mask) in region_masks { + let region_id = RegionId::new(table_id, region_number); + if mask.select_none() { + continue; } + let partition_expr_version = partition_versions + .get(®ion_id.region_number()) + .copied() + .unwrap_or_default(); + let peer = peers + .remove(®ion_id) + .context(partition::error::FindLeaderSnafu { + region_id, + table_id, + }) + .context(error::FindRegionLeaderSnafu)?; + let dispatcher = self.write_dispatcher.clone(); + let rb = record_batch.clone(); + let schema_bytes = schema_bytes.clone(); + let raw_header_and_data = if mask.select_all() { + Some(( + raw_flight_data.data_header.clone(), + raw_flight_data.data_body.clone(), + )) + } else { + None + }; + 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. + raw_header_and_data.unwrap() + } else { + let filter_timer = metrics::HANDLE_BULK_INSERT_ELAPSED + .with_label_values(&["filter"]) + .start_timer(); + let batch = compute::filter_record_batch(&rb, mask.array()) + .context(error::ComputeArrowSnafu)?; + filter_timer.observe_duration(); + metrics::BULK_REQUEST_ROWS + .with_label_values(&["rows_per_region"]) + .observe(batch.num_rows() as f64); + + let encode_timer = metrics::HANDLE_BULK_INSERT_ELAPSED + .with_label_values(&["encode"]) + .start_timer(); + let mut iter = FlightEncoder::default() + .encode(FlightMessage::RecordBatch(batch)) + .into_iter(); + let Some(flight_data) = iter.next() else { + // Safety: `iter` on a type of `Vec1`, which is guaranteed to have + // at least one element. + unreachable!() + }; + ensure!( + iter.next().is_none(), + error::NotSupportedSnafu { + feat: "bulk insert RecordBatch with dictionary arrays", + } + ); + encode_timer.observe_duration(); + (flight_data.data_header, flight_data.data_body) + }; + let request = RegionRequest { + header: Some(RegionRequestHeader { + tracing_context: TracingContext::from_current_span().to_w3c(), + ..Default::default() + }), + body: Some(region_request::Body::BulkInsert(BulkInsertRequest { + skip_wal, + region_id: region_id.as_u64(), + partition_expr_version: partition_expr_version + .map(|value| PartitionExprVersion { value }), + aligned_schema_version: None, + body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { + schema: schema_bytes, + data_header: header, + payload, + })), + })), + }; + + dispatcher + .dispatch_prepared( + vec![(peer, request)], + Some( + metrics::HANDLE_BULK_INSERT_ELAPSED + .with_label_values(&["datanode_handle"]), + ), + ) + .await + .context(error::RequestRegionSnafu) + }); + handles.push(handle); } - let region_responses = try_join_all(handles).await.context(error::JoinTaskSnafu)?; + 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::>()?; wait_all_datanode_timer.observe_duration(); - let mut rows_inserted: usize = 0; - for res in region_responses { - rows_inserted += res?.affected_rows; - } crate::metrics::DIST_INGEST_ROW_COUNT .with_label_values(&[db_name.as_str()]) .inc_by(rows_inserted as u64); @@ -393,7 +351,7 @@ impl Inserter { return; } - let Ok(timestamps) = extract_timestamps( + let Ok(timestamps) = insert::extract_timestamps( &record_batch, &table_info .meta diff --git a/src/operator/src/delete.rs b/src/operator/src/delete.rs index 9641e0e27d4..abd8302c965 100644 --- a/src/operator/src/delete.rs +++ b/src/operator/src/delete.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use std::sync::Arc; use std::{iter, mem}; @@ -20,10 +20,8 @@ use api::v1::region::{DeleteRequests as RegionDeleteRequests, RegionRequestHeade use api::v1::{DeleteRequests, RowDeleteRequests}; use catalog::CatalogManagerRef; use common_meta::node_manager::{AffectedRows, NodeManagerRef}; -use common_meta::peer::Peer; use common_query::Output; use common_telemetry::tracing_context::TracingContext; -use futures_util::future; use partition::manager::PartitionRuleManagerRef; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt, ensure}; @@ -31,34 +29,45 @@ use table::TableRef; use table::requests::DeleteRequest as TableDeleteRequest; use crate::error::{ - CatalogSnafu, FindRegionLeaderSnafu, InvalidDeleteRequestSnafu, JoinTaskSnafu, - MissingTimeIndexColumnSnafu, RequestDeletesSnafu, Result, TableNotFoundSnafu, + CatalogSnafu, InvalidDeleteRequestSnafu, JoinTaskSnafu, MissingTimeIndexColumnSnafu, + RequestDeletesSnafu, Result, TableNotFoundSnafu, }; use crate::region_req_factory::RegionRequestFactory; use crate::req_convert::common::preprocess_row_delete_requests; use crate::req_convert::delete::{ColumnToRow, RowToRegion, TableToRegion}; +use crate::write_dispatcher::{WriteDispatcher, WriteDispatcherRef}; pub struct Deleter { catalog_manager: CatalogManagerRef, partition_manager: PartitionRuleManagerRef, - node_manager: NodeManagerRef, + write_dispatcher: WriteDispatcherRef, } pub type DeleterRef = Arc; impl Deleter { + /// Creates a row deleter with the default write dispatcher. pub fn new( catalog_manager: CatalogManagerRef, partition_manager: PartitionRuleManagerRef, node_manager: NodeManagerRef, ) -> Self { Self { + write_dispatcher: Arc::new(WriteDispatcher::new( + node_manager, + partition_manager.clone(), + )), catalog_manager, partition_manager, - node_manager, } } + /// Installs the dispatcher shared with the other write paths. + pub fn with_write_dispatcher(mut self, dispatcher: WriteDispatcherRef) -> Self { + self.write_dispatcher = dispatcher; + self + } + pub async fn handle_column_deletes( &self, requests: DeleteRequests, @@ -129,52 +138,22 @@ impl Deleter { ..Default::default() }); - let tasks = self - .group_requests_by_peer(requests) - .await? - .into_iter() - .map(|(peer, deletes)| { - let request = request_factory.build_delete(deletes); - let node_manager = self.node_manager.clone(); - common_runtime::spawn_global(async move { - node_manager - .datanode(&peer) - .await - .handle(request) - .await - .context(RequestDeletesSnafu) - }) - }); - let results = future::try_join_all(tasks).await.context(JoinTaskSnafu)?; - - let affected_rows = results - .into_iter() - .map(|resp| resp.map(|r| r.affected_rows)) - .sum::>()?; + 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; crate::metrics::DIST_DELETE_ROW_COUNT .with_label_values(&[ctx.get_db_string().as_str()]) .inc_by(affected_rows as u64); Ok(affected_rows) } - async fn group_requests_by_peer( - &self, - requests: RegionDeleteRequests, - ) -> Result> { - let mut deletes: HashMap = HashMap::new(); - - for req in requests.requests { - let peer = self - .partition_manager - .find_region_leader(req.region_id.into()) - .await - .context(FindRegionLeaderSnafu)?; - deletes.entry(peer).or_default().requests.push(req); - } - - Ok(deletes) - } - async fn trim_columns( &self, mut requests: RowDeleteRequests, diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index c7f60c0883c..89c1195cda4 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -42,7 +42,7 @@ use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME; use common_grpc_expr::util::ColumnExpr; use common_meta::cache::TableFlownodeSetCacheRef; use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME; -use common_meta::node_manager::{AffectedRows, NodeManagerRef}; +use common_meta::node_manager::NodeManagerRef; use common_meta::peer::Peer; use common_meta::rpc::ddl::TriggerReason; use common_query::Output; @@ -83,9 +83,8 @@ use table::table_reference::TableReference; use crate::batcher::PendingRowsBatcher; use crate::error::{ - CatalogSnafu, ColumnOptionsSnafu, CreatePartitionRulesSnafu, FindRegionLeaderSnafu, - InvalidInsertRequestSnafu, JoinTaskSnafu, RequestInsertsSnafu, Result, TableNotFoundSnafu, - WriteRejectedSnafu, + CatalogSnafu, ColumnOptionsSnafu, CreatePartitionRulesSnafu, InvalidInsertRequestSnafu, + JoinTaskSnafu, RequestInsertsSnafu, Result, TableNotFoundSnafu, WriteRejectedSnafu, }; use crate::expr_helper; use crate::region_req_factory::RegionRequestFactory; @@ -95,11 +94,13 @@ use crate::req_convert::insert::{ fill_reqs_with_impure_default, rows_to_record_batch, }; use crate::statement::StatementExecutor; +use crate::write_dispatcher::{WriteDispatcher, WriteDispatcherRef}; pub struct Inserter { catalog_manager: CatalogManagerRef, pub(crate) partition_manager: PartitionRuleManagerRef, pub(crate) node_manager: NodeManagerRef, + pub(crate) write_dispatcher: WriteDispatcherRef, pub(crate) table_flownode_set_cache: TableFlownodeSetCacheRef, /// Server-side upper bound for auto table creation on write. /// When `false`, missing tables are never auto-created regardless of the @@ -292,6 +293,7 @@ impl Inserter { cost } + /// Creates a row inserter with the default write dispatcher. pub fn new( catalog_manager: CatalogManagerRef, partition_manager: PartitionRuleManagerRef, @@ -300,6 +302,10 @@ impl Inserter { auto_create_table: bool, ) -> Self { Self { + write_dispatcher: Arc::new(WriteDispatcher::new( + node_manager.clone(), + partition_manager.clone(), + )), catalog_manager, partition_manager, node_manager, @@ -826,30 +832,16 @@ impl Inserter { has_instant_rows, )?; - // Write requests to datanode and wait for response - let write_tasks = self - .group_requests_by_peer(normal_requests) - .await? - .into_iter() - .map(|(peer, inserts)| { - let node_manager = self.node_manager.clone(); - let request = request_factory.build_insert(inserts); - common_runtime::spawn_global(async move { - node_manager - .datanode(&peer) - .await - .handle(request) - .await - .context(RequestInsertsSnafu) - }) - }); - let results = future::try_join_all(write_tasks) - .await - .context(JoinTaskSnafu)?; - let affected_rows = results - .into_iter() - .map(|resp| resp.map(|r| r.affected_rows)) - .sum::>()?; + 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; crate::metrics::DIST_INGEST_ROW_COUNT .with_label_values(&[ctx.get_db_string().as_str()]) .inc_by(affected_rows as u64); @@ -859,40 +851,6 @@ impl Inserter { )) } - async fn group_requests_by_peer( - &self, - requests: RegionInsertRequests, - ) -> Result> { - // group by region ids first to reduce repeatedly call `find_region_leader` - // TODO(discord9): determine if a addition clone is worth it - let mut requests_per_region: HashMap = HashMap::new(); - for req in requests.requests { - let region_id = RegionId::from_u64(req.region_id); - requests_per_region - .entry(region_id) - .or_default() - .requests - .push(req); - } - - let mut inserts: HashMap = HashMap::new(); - - for (region_id, reqs) in requests_per_region { - let peer = self - .partition_manager - .find_region_leader(region_id) - .await - .context(FindRegionLeaderSnafu)?; - inserts - .entry(peer) - .or_default() - .requests - .extend(reqs.requests); - } - - Ok(inserts) - } - /// Returns `Some(reason)` if the config or request hint disables automatic /// table creation. Exempt private system tables are handled by /// [`Self::is_auto_create_exempt_private_table`]. @@ -1717,6 +1675,17 @@ 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 + } + pub fn node_manager(&self) -> &NodeManagerRef { &self.node_manager } diff --git a/src/operator/src/lib.rs b/src/operator/src/lib.rs index a971ca41f70..901f351e22c 100644 --- a/src/operator/src/lib.rs +++ b/src/operator/src/lib.rs @@ -29,3 +29,4 @@ pub mod table; #[cfg(any(test, feature = "testing"))] pub mod test_util; pub mod utils; +pub mod write_dispatcher; diff --git a/src/operator/src/req_convert.rs b/src/operator/src/req_convert.rs index 456ff0b328d..1161f62795d 100644 --- a/src/operator/src/req_convert.rs +++ b/src/operator/src/req_convert.rs @@ -15,3 +15,6 @@ 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 3353e2fcda7..044d3b97db5 100644 --- a/src/operator/src/req_convert/common/partitioner.rs +++ b/src/operator/src/req_convert/common/partitioner.rs @@ -12,72 +12,241 @@ // See the License for the specific language governing permissions and // limitations under the License. -use api::v1::region::{DeleteRequest, InsertRequest}; -use api::v1::{PartitionExprVersion, Rows}; +use std::collections::HashMap; + +use api::v1::region::{BulkInsertRequest, DeleteRequest, InsertRequest, bulk_insert_request}; +use api::v1::{ArrowIpc, PartitionExprVersion, Rows}; +use common_error::ext::BoxedError; +use common_grpc::flight; +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 store_api::storage::RegionId; +use store_api::storage::{RegionId, RegionNumber, TableId}; use table::metadata::TableInfo; -use crate::error::{Result, SplitDeleteSnafu, SplitInsertSnafu}; +use crate::error::{self, Result, SplitDeleteSnafu, SplitInsertSnafu}; +/// Loads table routing rules and converts rows to region write requests. pub struct Partitioner<'a> { partition_manager: &'a PartitionRuleManager, } +/// A single snapshot of a table's partition rule and expression versions. +pub(crate) struct TablePartitioner { + table_id: TableId, + rule: PartitionRuleRef, + versions: HashMap>, +} + impl<'a> Partitioner<'a> { + /// Creates a partitioner backed by the shared routing caches. pub fn new(partition_manager: &'a PartitionRuleManager) -> Self { Self { partition_manager } } + /// Routes insert rows using the current table partitions. pub async fn partition_insert_requests( &self, table_info: &TableInfo, rows: Rows, skip_wal: bool, ) -> Result> { - let table_id = table_info.table_id(); - let requests = self - .partition_manager - .split_rows(table_info, rows) + TablePartitioner::try_new(self.partition_manager, table_info) .await .context(SplitInsertSnafu)? - .into_iter() - .map( - |(region_number, (rows, partition_expr_version))| InsertRequest { - skip_wal, - region_id: RegionId::new(table_id, region_number).into(), - rows: Some(rows), - partition_expr_version: partition_expr_version - .map(|value| PartitionExprVersion { value }), - }, - ) - .collect(); - Ok(requests) + .partition_insert_requests(rows, skip_wal) } + /// Routes delete rows using the current table partitions. pub async fn partition_delete_requests( &self, table_info: &TableInfo, rows: Rows, ) -> Result> { - let table_id = table_info.table_id(); - - let requests = self - .partition_manager - .split_rows(table_info, rows) + TablePartitioner::try_new(self.partition_manager, table_info) .await .context(SplitDeleteSnafu)? + .partition_delete_requests(rows) + } +} + +impl TablePartitioner { + /// Reads a table's partition rule and expression versions into one routing snapshot. + pub async fn try_new( + partition_manager: &PartitionRuleManager, + table_info: &TableInfo, + ) -> partition::error::Result { + let (rule, versions) = partition_manager + .find_table_partition_rule(table_info) + .await?; + Ok(Self { + table_id: table_info.table_id(), + rule, + versions: versions + .into_iter() + .map(|(region, version)| { + (region, version.map(|value| PartitionExprVersion { value })) + }) + .collect(), + }) + } + + /// Moves insert rows into requests for this snapshot's destination regions. + pub fn partition_insert_requests( + &self, + rows: Rows, + skip_wal: bool, + ) -> Result> { + Ok(RowSplitter::new(self.rule.clone()) + .split(rows) + .context(SplitInsertSnafu)? .into_iter() - .map( - |(region_number, (rows, partition_expr_version))| DeleteRequest { - region_id: RegionId::new(table_id, region_number).into(), - rows: Some(rows), - partition_expr_version: partition_expr_version - .map(|value| PartitionExprVersion { value }), - }, - ) - .collect(); + .map(|(region, rows)| InsertRequest { + region_id: RegionId::new(self.table_id, region).as_u64(), + rows: Some(rows), + partition_expr_version: self.version(region), + skip_wal, + }) + .collect()) + } + + /// Moves delete rows into requests for this snapshot's destination regions. + pub fn partition_delete_requests(&self, rows: Rows) -> Result> { + Ok(RowSplitter::new(self.rule.clone()) + .split(rows) + .context(SplitDeleteSnafu)? + .into_iter() + .map(|(region, rows)| DeleteRequest { + region_id: RegionId::new(self.table_id, region).as_u64(), + rows: Some(rows), + partition_expr_version: self.version(region), + }) + .collect()) + } + + /// Splits an Arrow batch, reusing the original IPC when one region selects all rows. + pub fn partition_bulk_insert_requests( + &self, + batch: RecordBatch, + mut original_ipc: Option, + skip_wal: bool, + ) -> Result> { + let masks = self + .rule + .split_record_batch(&batch) + .context(SplitInsertSnafu)?; + let mut requests = Vec::with_capacity(masks.len()); + for (region, mask) in masks { + if mask.select_none() { + continue; + } + let original = if mask.select_all() { + original_ipc.take() + } else { + None + }; + let ipc = if let Some(ipc) = original { + ipc + } else { + let selected = if mask.select_all() { + batch.clone() + } else { + compute::filter_record_batch(&batch, mask.array()) + .context(error::ComputeArrowSnafu)? + }; + let (schema, data_header, payload) = flight::record_batch_to_ipc(selected) + .map_err(BoxedError::new) + .context(error::ExternalSnafu)?; + ArrowIpc { + schema, + data_header, + payload, + } + }; + requests.push(BulkInsertRequest { + region_id: RegionId::new(self.table_id, region).as_u64(), + partition_expr_version: self.version(region), + // The destination datanode validates against its current schema. + aligned_schema_version: None, + skip_wal, + body: Some(bulk_insert_request::Body::ArrowIpc(ipc)), + }); + } Ok(requests) + } + + fn version(&self, region: RegionNumber) -> Option { + self.versions.get(®ion).copied().flatten() + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_error::ext::{ErrorExt, RetryHint}; + use common_error::status_code::StatusCode; + use common_meta::cache; + use common_meta::kv_backend::test_util::MockKvBackendBuilder; + use common_meta::rpc::store::RangeResponse; + use moka::future::Cache; + + use super::*; + use crate::error::Error; + use crate::test_util; + + #[tokio::test] + async fn test_partition_requests_preserve_routing_errors() { + let table_info = test_util::new_test_table_info(1, "table_1", [1].into_iter()); + for (backend_unavailable, expected_status, expected_hint) in [ + (false, StatusCode::TableNotFound, RetryHint::NonRetryable), + (true, StatusCode::Internal, RetryHint::Retryable), + ] { + let backend = Arc::new( + MockKvBackendBuilder::default() + .range_fn(Arc::new(move |_| { + if backend_unavailable { + common_meta::error::NoLeaderSnafu { table_id: 1u32 }.fail() + } else { + Ok(RangeResponse { + kvs: Vec::new(), + more: false, + }) + } + })) + .build() + .unwrap(), + ); + let routes = Arc::new(cache::new_table_route_cache( + "partitioner_test_routes".to_string(), + Cache::new(16), + backend.clone(), + )); + let partitions = Arc::new(partition::cache::new_partition_info_cache( + "partitioner_test_partitions".to_string(), + Cache::new(16), + routes.clone(), + )); + let manager = PartitionRuleManager::new(backend, routes, partitions); + let partitioner = Partitioner::new(&manager); + let insert_error = partitioner + .partition_insert_requests(&table_info, Rows::default(), false) + .await + .unwrap_err(); + let delete_error = partitioner + .partition_delete_requests(&table_info, Rows::default()) + .await + .unwrap_err(); + + assert!(matches!(insert_error, Error::SplitInsert { .. })); + assert!(matches!(delete_error, Error::SplitDelete { .. })); + for error in [insert_error, delete_error] { + assert_eq!(error.status_code(), expected_status); + assert_eq!(error.retry_hint(), expected_hint); + } + } } } diff --git a/src/operator/src/write_dispatcher.rs b/src/operator/src/write_dispatcher.rs new file mode 100644 index 00000000000..4af4fbfafdf --- /dev/null +++ b/src/operator/src/write_dispatcher.rs @@ -0,0 +1,256 @@ +// 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. + +#[cfg(test)] +mod tests; + +use std::collections::{HashMap, HashSet, hash_map::Entry}; +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 common_meta::node_manager::NodeManagerRef; +use common_meta::peer::Peer; +use common_meta::rpc::router; +use futures::future; +use partition::manager::PartitionRuleManagerRef; +use prometheus::Histogram; +use snafu::{OptionExt, ResultExt}; +use store_api::storage::RegionId; + +use crate::error::{self, FindRegionLeaderSnafu, RequestRegionSnafu}; + +/// Dispatches already partitioned insert, delete, and bulk requests to their peers. +pub struct WriteDispatcher { + node_manager: NodeManagerRef, + partition_manager: PartitionRuleManagerRef, +} + +/// Shared reference to a [`WriteDispatcher`]. +pub type WriteDispatcherRef = Arc; + +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, + } + } + + /// 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. + pub async fn dispatch_with_metrics( + &self, + requests: Vec, + histogram: Histogram, + ) -> Result { + self.dispatch_inner(requests, Some(histogram)).await + } + + /// Resolves routes before writes are detached from the caller. + pub(crate) async fn prepare( + &self, + requests: Vec, + ) -> error::Result> { + let requests = requests.into_iter().map(split_request).collect::>(); + let regions = requests + .iter() + .flatten() + .map(write_region_id) + .collect::>>() + .context(RequestRegionSnafu)?; + let peers = self + .resolve_peers(®ions) + .await + .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 request in requests { + let peer = get_peer( + &peers, + write_region_id(&request).context(RequestRegionSnafu)?, + ) + .context(RequestRegionSnafu)?; + match by_peer.entry(peer) { + Entry::Vacant(entry) => { + entry.insert(request); + } + Entry::Occupied(mut entry) => match (&mut entry.get_mut().body, request.body) { + (Some(Body::Inserts(existing)), Some(Body::Inserts(next))) => { + existing.requests.extend(next.requests) + } + (Some(Body::Deletes(existing)), Some(Body::Deletes(next))) => { + existing.requests.extend(next.requests) + } + _ => { + return UnexpectedSnafu { + err_msg: "Cannot combine different write kinds", + } + .fail() + .context(RequestRegionSnafu); + } + }, + } + } + writes.extend(by_peer); + } + Ok(writes) + } + + /// Sends requests using routes already resolved by the caller. + pub(crate) async fn dispatch_prepared( + &self, + writes: Vec<(Peer, 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 + })) + .await; + let mut response = RegionResponse::new(0); + for result in results { + merge_response(&mut response, result?); + } + Ok(response) + } + + /// Resolves only the requested peers, loading each table route once. + 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) + } + + async fn dispatch_inner( + &self, + requests: Vec, + histogram: Option, + ) -> Result { + let writes = self + .prepare(requests) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + self.dispatch_prepared(writes, histogram).await + } +} + +fn get_peer(peers: &HashMap, region: RegionId) -> Result { + peers.get(®ion).cloned().context(UnexpectedSnafu { + err_msg: "Planned write has no resolved peer", + }) +} + +fn merge_response(response: &mut RegionResponse, result: RegionResponse) { + response.affected_rows += result.affected_rows; + response.extensions.extend(result.extensions); + response.metadata = result.metadata; +} + +fn split_request(request: RegionRequest) -> Vec { + let RegionRequest { header, body } = request; + let bodies: Vec = match body { + Some(Body::Inserts(inserts)) => inserts + .requests + .into_iter() + .map(|request| { + Body::Inserts(InsertRequests { + requests: vec![request], + }) + }) + .collect(), + Some(Body::Deletes(deletes)) => deletes + .requests + .into_iter() + .map(|request| { + Body::Deletes(DeleteRequests { + requests: vec![request], + }) + }) + .collect(), + body => return vec![RegionRequest { header, body }], + }; + bodies + .into_iter() + .map(|body| RegionRequest { + header: header.clone(), + body: Some(body), + }) + .collect() +} + +fn write_region_id(request: &RegionRequest) -> Result { + let region_id = match &request.body { + Some(Body::Inserts(inserts)) => match inserts.requests.as_slice() { + [request] => Some(request.region_id), + _ => None, + }, + Some(Body::Deletes(deletes)) => match deletes.requests.as_slice() { + [request] => Some(request.region_id), + _ => None, + }, + Some(Body::BulkInsert(bulk)) => Some(bulk.region_id), + _ => None, + }; + region_id.map(RegionId::from_u64).ok_or_else(|| { + UnexpectedSnafu { + err_msg: "Expected a single-region write request", + } + .build() + }) +} diff --git a/src/operator/src/write_dispatcher/tests.rs b/src/operator/src/write_dispatcher/tests.rs new file mode 100644 index 00000000000..1191f73d55a --- /dev/null +++ b/src/operator/src/write_dispatcher/tests.rs @@ -0,0 +1,350 @@ +// 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::collections::HashMap; +use std::future::Future; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::task::Poll; +use std::time::{Duration, Instant}; + +use api::v1::region::{ + BulkInsertRequest, DeleteRequest, InsertRequest, RegionRequestHeader, bulk_insert_request, +}; +use api::v1::{ + ArrowIpc, ColumnDataType, PartitionExprVersion, Row, Rows, helper, value::ValueData, +}; +use client::client_manager::NodeClients; +use common_error::ext::ErrorExt; +use common_grpc::channel_manager::ChannelConfig; +use common_meta::cache as meta_cache; +use common_meta::ddl::test_util::datanode_handler::DatanodeWatcher; +use common_meta::key::TableMetadataManager; +use common_meta::key::table_route::TableRouteValue; +use common_meta::test_util::{MockDatanodeHandler, MockDatanodeManager}; +use common_query::request::QueryRequest; +use common_recordbatch::SendableRecordBatchStream; +use partition::manager::PartitionRuleManager; +use tokio::sync::{Barrier, Notify}; + +use super::*; +use crate::error::Error; +use crate::test_util; + +fn build_insert_request(regions: &[u32]) -> RegionRequest { + RegionRequest { + body: Some(Body::Inserts(InsertRequests { + requests: regions + .iter() + .map(|region| InsertRequest { + region_id: RegionId::new(1, *region).as_u64(), + ..Default::default() + }) + .collect(), + })), + ..Default::default() + } +} + +#[tokio::test] +async fn test_resolve_peers_batches_logical_and_physical_regions() { + let backend = test_util::prepare_mocked_backend().await; + let partition_manager = test_util::create_partition_rule_manager(backend.clone()).await; + TableMetadataManager::new(backend) + .create_table_metadata( + test_util::new_test_table_info(42, "dispatcher_logical", [1, 2, 3].into_iter()), + TableRouteValue::logical(1), + Default::default(), + ) + .await + .unwrap(); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + let dispatcher = WriteDispatcher::new(clients, partition_manager.clone()); + let regions = [ + RegionId::new(1, 1), + RegionId::new(1, 2), + RegionId::new(42, 1), + RegionId::new(42, 2), + RegionId::new(42, 1), + ]; + let peers = dispatcher.resolve_peers(®ions).await.unwrap(); + assert_eq!(peers.len(), 4); + for region in regions { + assert_eq!( + peers[®ion], + partition_manager.find_region_leader(region).await.unwrap() + ); + } + assert_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) + .await + .unwrap_err(); + let actual = dispatcher.resolve_peers(&[missing]).await.unwrap_err(); + assert_eq!(actual.status_code(), expected.status_code()); + assert_eq!(actual.retry_hint(), expected.retry_hint()); + assert!(dispatcher.resolve_peers(&[]).await.unwrap().is_empty()); +} + +#[tokio::test] +async fn test_dispatch_preserves_write_payloads_and_envelopes() { + let backend = test_util::prepare_mocked_backend().await; + let partition_manager = test_util::create_partition_rule_manager(backend.clone()).await; + TableMetadataManager::new(backend) + .create_table_metadata( + test_util::new_test_table_info(42, "dispatcher_logical", [1, 2, 3].into_iter()), + TableRouteValue::logical(1), + Default::default(), + ) + .await + .unwrap(); + let (sender, mut received) = tokio::sync::mpsc::channel(8); + let watcher = DatanodeWatcher::new(sender).with_handler(|_, _| Ok(RegionResponse::new(1))); + let dispatcher = WriteDispatcher::new( + Arc::new(MockDatanodeManager::new(watcher)), + partition_manager, + ); + // Physical and logical regions share a peer but retain distinct region IDs. + let regions = [RegionId::new(1, 1), RegionId::new(42, 1)]; + let rows = Rows { + schema: vec![helper::tag_column_schema("a", ColumnDataType::Int32)], + rows: vec![Row { + values: vec![ValueData::I32Value(7).into()], + }], + }; + let version = Some(PartitionExprVersion { value: 17 }); + let bodies = [ + Body::Inserts(InsertRequests { + requests: regions + .iter() + .map(|region| InsertRequest { + region_id: region.as_u64(), + rows: Some(rows.clone()), + partition_expr_version: version, + skip_wal: true, + }) + .collect(), + }), + Body::Deletes(DeleteRequests { + requests: regions + .iter() + .map(|region| DeleteRequest { + region_id: region.as_u64(), + rows: Some(rows.clone()), + partition_expr_version: version, + }) + .collect(), + }), + ]; + for body in bodies { + let requests = ["first", "second"].map(|dbname| RegionRequest { + header: Some(RegionRequestHeader { + dbname: dbname.to_string(), + ..Default::default() + }), + body: Some(body.clone()), + }); + let writes = dispatcher.prepare(requests.to_vec()).await.unwrap(); + assert!(received.try_recv().is_err()); + let response = dispatcher.dispatch_prepared(writes, None).await.unwrap(); + assert_eq!(response.affected_rows, 2); + for request in requests { + assert_eq!(received.try_recv().unwrap(), (Peer::new(1, ""), request)); + } + assert!(received.try_recv().is_err()); + } + let requests = regions.map(|region| RegionRequest { + header: Some(RegionRequestHeader { + dbname: "bulk".to_string(), + ..Default::default() + }), + body: Some(Body::BulkInsert(BulkInsertRequest { + region_id: region.as_u64(), + partition_expr_version: version, + aligned_schema_version: None, + skip_wal: true, + body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { + schema: "schema".into(), + data_header: "header".into(), + payload: "payload".into(), + })), + })), + }); + let response = dispatcher.dispatch(requests.to_vec()).await.unwrap(); + assert_eq!(response.affected_rows, 2); + for request in requests { + assert_eq!(received.try_recv().unwrap(), (Peer::new(1, ""), request)); + } + assert!(received.try_recv().is_err()); + + let error = dispatcher + .prepare(vec![build_insert_request(&[1, 999])]) + .await + .unwrap_err(); + assert!(matches!(error, Error::FindRegionLeader { .. })); + assert!(received.try_recv().is_err()); +} + +#[derive(Clone)] +struct GatedDatanodeHandler { + entered: Arc, + rejected: Arc, + completed: Arc, + release: Arc, +} + +#[async_trait::async_trait] +impl MockDatanodeHandler for GatedDatanodeHandler { + async fn handle(&self, peer: &Peer, _request: RegionRequest) -> Result { + self.entered.wait().await; + if peer.id == 1 { + self.rejected.store(true, Ordering::SeqCst); + return UnexpectedSnafu { + err_msg: "first region rejected", + } + .fail(); + } + self.release.notified().await; + self.completed.store(true, Ordering::SeqCst); + Ok(RegionResponse::new(1)) + } + + async fn handle_query( + &self, + _peer: &Peer, + _request: QueryRequest, + ) -> Result { + unimplemented!() + } +} + +#[tokio::test] +async fn test_dispatch_sends_concurrently_and_waits_for_siblings_after_error() { + let backend = test_util::prepare_mocked_backend().await; + let partition_manager = test_util::create_partition_rule_manager(backend).await; + let handler = GatedDatanodeHandler { + entered: Arc::new(Barrier::new(2)), + rejected: Arc::default(), + completed: Arc::default(), + release: Arc::default(), + }; + let dispatcher = WriteDispatcher::new( + Arc::new(MockDatanodeManager::new(handler.clone())), + partition_manager, + ); + let dispatch = dispatcher.dispatch(vec![build_insert_request(&[1, 2])]); + tokio::pin!(dispatch); + tokio::time::timeout( + Duration::from_secs(1), + futures::future::poll_fn(|cx| { + assert!( + dispatch.as_mut().poll(cx).is_pending(), + "dispatch returned while a sibling remained blocked" + ); + if handler.rejected.load(Ordering::SeqCst) { + Poll::Ready(()) + } else { + Poll::Pending + } + }), + ) + .await + .expect("both region sends must reach the barrier concurrently"); + assert!(!handler.completed.load(Ordering::SeqCst)); + handler.release.notify_one(); + let error = tokio::time::timeout(Duration::from_secs(1), dispatch) + .await + .unwrap() + .unwrap_err(); + assert!(error.to_string().contains("first region rejected")); + assert!(handler.completed.load(Ordering::SeqCst)); +} + +#[tokio::test] +async fn test_resolve_peers_cache_lookup_ablation() { + common_telemetry::init_default_ut_logging(); + 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, "resolution_ablation", [1, 2, 3].into_iter()), + TableRouteValue::logical(1), + Default::default(), + ) + .await + .unwrap(); + // A dedicated cache label isolates counters from other concurrently running tests. + let cache_name = "write_dispatcher_resolution_ablation_table_route"; + let routes = Arc::new(meta_cache::new_table_route_cache( + cache_name.to_string(), + moka::future::Cache::new(16), + backend.clone(), + )); + let partitions = Arc::new(partition::cache::new_partition_info_cache( + "write_dispatcher_resolution_ablation_partitions".to_string(), + moka::future::Cache::new(16), + routes.clone(), + )); + let partition_manager = Arc::new(PartitionRuleManager::new(backend, routes, partitions)); + let dispatcher = + WriteDispatcher::new(Arc::new(NodeClients::default()), partition_manager.clone()); + let regions = (0..256) + .map(|index| RegionId::new(if index % 2 == 0 { 1 } else { 42 }, index % 3 + 1)) + .collect::>(); + let expected = dispatcher.resolve_peers(®ions).await.unwrap(); + let counter = common_meta::metrics::CACHE_CONTAINER_CACHE_GET.with_label_values(&[cache_name]); + for run in 0..6 { + // Alternate order; compare the same warm cache, routes and deterministic workload. + let modes = if run % 2 == 0 { + ["delete", "insert", "batched"] + } else { + ["batched", "insert", "delete"] + }; + for mode in modes { + let before = counter.get(); + let started = Instant::now(); + let peers = if mode == "batched" { + dispatcher.resolve_peers(®ions).await.unwrap() + } else { + let mut peers = HashMap::new(); + for region in ®ions { + if mode == "insert" && peers.contains_key(region) { + continue; + } + peers.insert( + *region, + partition_manager.find_region_leader(*region).await.unwrap(), + ); + } + peers + }; + let elapsed = started.elapsed(); + let gets = counter.get() - before; + assert_eq!(peers, expected); + let expected_gets = match mode { + "delete" => 384, + "insert" => 9, + "batched" => 3, + _ => unreachable!(), + }; + assert_eq!(gets, expected_gets); + common_telemetry::info!( + "route_resolution_ablation run={run} mode={mode} regions={} cache_gets={gets} elapsed_ns={}", + regions.len(), + elapsed.as_nanos() + ); + } + } +}