rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-15 16:32:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(operator): centralize write partitioning and dispatch

Browse Source 
Signed-off-by: WenyXu <wenymedia@gmail.com>
This commit is contained in:
WenyXu
2026-10-12 13:11:24 +00:00
parent 1627e6c06b
commit d3d395f6f0
9 changed files with 805 additions and 330 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/operator/AGENTS.md
+1
View File
@@ -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 |
src/operator/src/bulk_insert.rs
+122 -186
View File
@@ -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<AffectedRows> {
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(&region_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::<Result<usize>>()?;
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(&region_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(&region_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<Result<api::region::RegionResponse>> =
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(&region_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<Result<RegionRequest>> =
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::<Result<Vec<_>>>()?;
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
src/operator/src/delete.rs
+24 -47
View File
@@ -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<Deleter>;
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::<Result<AffectedRows>>()?;
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<HashMap<Peer, RegionDeleteRequests>> {
let mut deletes: HashMap<Peer, RegionDeleteRequests> = 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,
src/operator/src/insert.rs
+29 -62
View File
@@ -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::<Result<AffectedRows>>()?;
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<HashMap<Peer, RegionInsertRequests>> {
// 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<RegionId, RegionInsertRequests> = 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<Peer, RegionInsertRequests> = 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
}
src/operator/src/lib.rs
+1
View File
@@ -29,3 +29,4 @@ pub mod table;
#[cfg(any(test, feature = "testing"))]
pub mod test_util;
pub mod utils;
pub mod write_dispatcher;
src/operator/src/req_convert.rs
+3
View File
@@ -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;
src/operator/src/req_convert/common/partitioner.rs
+134 -35
View File
@@ -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<RegionNumber, Option<u64>>,
}
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<Vec<InsertRequest>> {
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<Vec<DeleteRequest>> {
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<Self> {
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<Vec<InsertRequest>> {
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<Vec<DeleteRequest>> {
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<ArrowIpc>,
skip_wal: bool,
) -> Result<Vec<BulkInsertRequest>> {
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<PartitionExprVersion> {
self.versions
.get(&region)
.copied()
.flatten()
.map(|value| PartitionExprVersion { value })
}
}
src/operator/src/write_dispatcher.rs
+236
View File
@@ -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<WriteDispatcher>;
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<RegionRequest>) -> Result<RegionResponse> {
self.dispatch_inner(requests, None).await
}
/// Records each datanode send in the supplied histogram.
pub async fn dispatch_with_metrics(
&self,
requests: Vec<RegionRequest>,
histogram: Histogram,
) -> Result<RegionResponse> {
self.dispatch_inner(requests, Some(histogram)).await
}
async fn dispatch_inner(
&self,
requests: Vec<RegionRequest>,
histogram: Option<Histogram>,
) -> Result<RegionResponse> {
self.dispatch_once(requests, histogram.as_ref()).await
}
async fn dispatch_once(
&self,
requests: Vec<RegionRequest>,
histogram: Option<&Histogram>,
) -> Result<RegionResponse> {
let requests = requests.into_iter().map(split_request).collect::<Vec<_>>();
let regions = requests
.iter()
.flatten()
.map(write_region_id)
.collect::<Result<Vec<_>>>()?;
let peers = self.resolve_peers(&regions).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<Peer, RegionRequest> = 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<HashMap<RegionId, Peer>> {
let table_ids = regions
.iter()
.map(|region| region.table_id())
.collect::<HashSet<_>>();
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::<HashMap<_, _>>()
});
}
let mut peers = HashMap::with_capacity(regions.len());
for region in regions {
if peers.contains_key(region) {
continue;
}
let peer = physical_tables
.get(&region.table_id())
.and_then(|physical_id| leaders.get(physical_id))
.and_then(|leaders| leaders.get(&region.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<RegionId, Peer>, region: RegionId) -> Result<Peer> {
peers.get(&region).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<RegionRequest> {
let RegionRequest { header, body } = request;
let bodies: Vec<Body> = 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<RegionId> {
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()
})
}
src/operator/src/write_dispatcher/tests.rs
+255
View File
@@ -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(&regions).await.unwrap();
assert_eq!(peers.len(), 4);
for region in regions {
assert_eq!(
peers[&region],
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<Barrier>,
rejected: Arc<AtomicBool>,
completed: Arc<AtomicBool>,
release: Arc<Notify>,
}
struct GatedDatanode {
reject: bool,
entered: Arc<Barrier>,
rejected: Arc<AtomicBool>,
completed: Arc<AtomicBool>,
release: Arc<Notify>,
}
#[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<RegionResponse> {
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<SendableRecordBatchStream> {
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::<Vec<_>>();
let expected = dispatcher.resolve_peers(&regions).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(&regions).await.unwrap()
} else {
let mut peers = HashMap::new();
for region in &regions {
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()
);
}
}
}