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..b2fcb59e4ff 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -14,22 +14,20 @@ use std::collections::HashSet; -use ahash::{HashMap, HashMapExt}; 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}; @@ -40,7 +38,8 @@ 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 +53,26 @@ impl Inserter { batch: RecordBatch, ctx: QueryContextRef, ) -> Result { - let (rule, versions) = self - .partition_manager - .find_table_partition_rule(&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 { + let requests = TablePartitioner::try_new(&self.partition_manager, &table_info) + .await? + .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(); + Ok(self + .write_dispatcher + .dispatch(writes) + .await + .context(error::RequestInsertsSnafu)? + .affected_rows) } /// Handle bulk insert request. @@ -208,12 +161,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 +180,12 @@ 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 result = self + .write_dispatcher + .dispatch_with_metrics( + vec![request], + metrics::HANDLE_BULK_INSERT_ELAPSED.with_label_values(&["datanode_handle"]), + ) .await .context(error::RequestRegionSnafu) .map(|r| r.affected_rows); @@ -250,121 +197,110 @@ 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 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 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, + })), + })), + }; + + Ok(request) + }); + handles.push(handle); } - let region_responses = try_join_all(handles).await.context(error::JoinTaskSnafu)?; + let dispatcher = self.write_dispatcher.clone(); + let rows_inserted = common_runtime::spawn_global(async move { + let writes = future::try_join_all(handles) + .await + .context(error::JoinTaskSnafu)? + .into_iter() + .collect::>>()?; + dispatcher + .dispatch_with_metrics( + writes, + metrics::HANDLE_BULK_INSERT_ELAPSED.with_label_values(&["datanode_handle"]), + ) + .await + .context(error::RequestRegionSnafu) + }) + .await + .context(error::JoinTaskSnafu)?? + .affected_rows; 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 +329,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..f2647af195b 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, 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,20 @@ 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 affected_rows = + common_runtime::spawn_global(async move { dispatcher.dispatch(vec![request]).await }) + .await + .context(crate::error::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..29da4e6d4b8 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, + 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,14 @@ 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 affected_rows = + common_runtime::spawn_global(async move { dispatcher.dispatch(vec![request]).await }) + .await + .context(crate::error::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 +849,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 +1673,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..47eaebaf142 100644 --- a/src/operator/src/req_convert/common/partitioner.rs +++ b/src/operator/src/req_convert/common/partitioner.rs @@ -12,72 +12,171 @@ // 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) - .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) + TablePartitioner::try_new(self.partition_manager, table_info) + .await? + .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(); + TablePartitioner::try_new(self.partition_manager, table_info) + .await? + .partition_delete_requests(rows) + } +} - let requests = self - .partition_manager - .split_rows(table_info, 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, + ) -> Result { + let (rule, versions) = partition_manager + .find_table_partition_rule(table_info) .await + .context(error::InvalidPartitionSnafu)?; + Ok(Self { + table_id: table_info.table_id(), + rule, + versions, + }) + } + + /// 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, 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_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)| 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() + .map(|value| PartitionExprVersion { value }) } } diff --git a/src/operator/src/write_dispatcher.rs b/src/operator/src/write_dispatcher.rs new file mode 100644 index 00000000000..d334f96e9d6 --- /dev/null +++ b/src/operator/src/write_dispatcher.rs @@ -0,0 +1,236 @@ +// 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}; +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::{self, Result}; +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; + +/// 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 + } + + async fn dispatch_inner( + &self, + requests: Vec, + histogram: Option, + ) -> Result { + self.dispatch_once(requests, histogram.as_ref()).await + } + + async fn dispatch_once( + &self, + requests: Vec, + histogram: Option<&Histogram>, + ) -> Result { + let requests = requests.into_iter().map(split_request).collect::>(); + let regions = requests + .iter() + .flatten() + .map(write_region_id) + .collect::>>()?; + let peers = self.resolve_peers(®ions).await?; + 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)?)?; + match by_peer.entry(peer) { + std::collections::hash_map::Entry::Vacant(entry) => { + entry.insert(request); + } + std::collections::hash_map::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 error::UnexpectedSnafu { + err_msg: "Cannot combine different write kinds", + } + .fail(); + } + } + } + } + } + writes.extend(by_peer); + } + let results = future::join_all(writes.into_iter().map(|(peer, request)| async move { + let datanode = self.node_manager.datanode(&peer).await; + let _timer = histogram.map(Histogram::start_timer); + datanode.handle(request).await + })) + .await; + let mut response = RegionResponse::new(0); + for result in results { + merge_response(&mut response, result?); + } + Ok(response) + } + + async fn resolve_peers(&self, regions: &[RegionId]) -> Result> { + let table_ids = regions + .iter() + .map(|region| region.table_id()) + .collect::>(); + let mut physical_tables = HashMap::new(); + let mut leaders = HashMap::new(); + for table_id in table_ids { + let (physical_id, route) = self + .partition_manager + .find_physical_table_route_with_id(table_id) + .await + .map_err(BoxedError::new) + .context(error::ExternalSnafu)?; + physical_tables.insert(table_id, physical_id); + leaders.entry(physical_id).or_insert_with(|| { + router::convert_to_region_leader_map(&route.region_routes) + .into_iter() + .map(|(region, peer)| (region, peer.clone())) + .collect::>() + }); + } + 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| leaders.get(physical_id)) + .and_then(|leaders| leaders.get(®ion.region_number())) + .cloned() + .context(partition::error::FindLeaderSnafu { + region_id: *region, + table_id: region.table_id(), + }) + .map_err(BoxedError::new) + .context(error::ExternalSnafu)?; + peers.insert(*region, peer); + } + Ok(peers) + } +} + +fn get_peer(peers: &HashMap, region: RegionId) -> Result { + peers.get(®ion).cloned().context(error::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(|| { + error::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..e20b18d302c --- /dev/null +++ b/src/operator/src/write_dispatcher/tests.rs @@ -0,0 +1,255 @@ +// 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::InsertRequest; +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::key::table_route::TableRouteValue; +use common_meta::node_manager::{ + Datanode, DatanodeManager, DatanodeRef, FlownodeManager, FlownodeRef, +}; +use common_query::request::QueryRequest; +use common_recordbatch::SendableRecordBatchStream; +use partition::manager::PartitionRuleManager; +use tokio::sync::{Barrier, Notify}; + +use super::*; +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; + common_meta::key::TableMetadataManager::new(backend) + .create_table_metadata( + crate::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()); +} + +struct GatedNodeManager { + entered: Arc, + rejected: Arc, + completed: Arc, + release: Arc, +} + +struct GatedDatanode { + reject: bool, + entered: Arc, + rejected: Arc, + completed: Arc, + release: Arc, +} + +#[async_trait::async_trait] +impl DatanodeManager for GatedNodeManager { + async fn datanode(&self, peer: &Peer) -> DatanodeRef { + Arc::new(GatedDatanode { + reject: peer.id == 1, + entered: self.entered.clone(), + rejected: self.rejected.clone(), + completed: self.completed.clone(), + release: self.release.clone(), + }) + } +} + +#[async_trait::async_trait] +impl FlownodeManager for GatedNodeManager { + async fn flownode(&self, _peer: &Peer) -> FlownodeRef { + unimplemented!() + } +} + +#[async_trait::async_trait] +impl Datanode for GatedDatanode { + async fn handle(&self, _request: RegionRequest) -> Result { + self.entered.wait().await; + if self.reject { + self.rejected.store(true, Ordering::SeqCst); + return error::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, _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 nodes = Arc::new(GatedNodeManager { + entered: Arc::new(Barrier::new(2)), + rejected: Arc::default(), + completed: Arc::default(), + release: Arc::default(), + }); + let dispatcher = WriteDispatcher::new(nodes.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 nodes.rejected.load(Ordering::SeqCst) { + Poll::Ready(()) + } else { + Poll::Pending + } + }), + ) + .await + .expect("both region sends must reach the barrier concurrently"); + assert!(!nodes.completed.load(Ordering::SeqCst)); + nodes.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!(nodes.completed.load(Ordering::SeqCst)); +} + +#[tokio::test] +async fn test_resolve_peers_cache_lookup_ablation() { + let backend = test_util::prepare_mocked_backend().await; + test_util::create_partition_rule_manager(backend.clone()).await; + common_meta::key::TableMetadataManager::new(backend.clone()) + .create_table_metadata( + crate::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 { + [false, true] + } else { + [true, false] + }; + for batched in modes { + let before = counter.get(); + let started = Instant::now(); + let peers = if batched { + dispatcher.resolve_peers(®ions).await.unwrap() + } else { + let mut peers = HashMap::new(); + for region in ®ions { + 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); + assert_eq!(gets, if batched { 3 } else { 384 }); + eprintln!( + "route_resolution_ablation run={run} batched={batched} regions={} cache_gets={gets} elapsed_ns={}", + regions.len(), + elapsed.as_nanos() + ); + } + } +}