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

refactor(operator): centralize write partitioning and dispatch (#9541)

Browse Source 
* refactor(operator): centralize write partitioning and dispatch

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix(operator): preserve write dispatch behavior during refactor

Signed-off-by: WenyXu <wenymedia@gmail.com>

* style(operator): import write dispatch error paths

Signed-off-by: WenyXu <wenymedia@gmail.com>

* style(operator): align dispatch imports and test utilities

Signed-off-by: WenyXu <wenymedia@gmail.com>

* style(operator): use imported types in partition tests

Signed-off-by: WenyXu <wenymedia@gmail.com>

---------

Signed-off-by: WenyXu <wenymedia@gmail.com>
This commit is contained in:
Weny Xu
2026-10-13 02:30:14 +00:00
committed by GitHub
parent 684c8ff291
commit a2d28e6060
9 changed files with 1014 additions and 328 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
+143 -185
View File
@@ -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<AffectedRows> {
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(&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();
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(&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 +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::<Vec<_>>();
let mut peers = self
.write_dispatcher
.resolve_peers(&regions)
.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(&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 peer = peers
.remove(&region_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<Result<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 = 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::<Result<usize>>()?;
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
src/operator/src/delete.rs
+26 -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, 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<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,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::<Result<AffectedRows>>()?;
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<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
+31 -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,
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::<Result<AffectedRows>>()?;
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<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 +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
}
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
+203 -34
View File
@@ -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<RegionNumber, Option<PartitionExprVersion>>,
}
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)
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<Vec<DeleteRequest>> {
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<Self> {
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<Vec<InsertRequest>> {
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<Vec<DeleteRequest>> {
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<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()
}
}
#[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);
}
}
}
}
src/operator/src/write_dispatcher.rs
+256
View File
@@ -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<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
}
/// Resolves routes before writes are detached from the caller.
pub(crate) async fn prepare(
&self,
requests: Vec<RegionRequest>,
) -> error::Result<Vec<(Peer, RegionRequest)>> {
let requests = requests.into_iter().map(split_request).collect::<Vec<_>>();
let regions = requests
.iter()
.flatten()
.map(write_region_id)
.collect::<Result<Vec<_>>>()
.context(RequestRegionSnafu)?;
let peers = self
.resolve_peers(&regions)
.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<Peer, RegionRequest> = 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<Histogram>,
) -> Result<RegionResponse> {
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<HashMap<RegionId, Peer>> {
let table_ids = regions
.iter()
.map(|region| region.table_id())
.collect::<HashSet<_>>();
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(&region.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<RegionRequest>,
histogram: Option<Histogram>,
) -> Result<RegionResponse> {
let writes = self
.prepare(requests)
.await
.map_err(BoxedError::new)
.context(ExternalSnafu)?;
self.dispatch_prepared(writes, histogram).await
}
}
fn get_peer(peers: &HashMap<RegionId, Peer>, region: RegionId) -> Result<Peer> {
peers.get(&region).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<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(|| {
UnexpectedSnafu {
err_msg: "Expected a single-region write request",
}
.build()
})
}
src/operator/src/write_dispatcher/tests.rs
+350
View File
@@ -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(&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());
}
#[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<Barrier>,
rejected: Arc<AtomicBool>,
completed: Arc<AtomicBool>,
release: Arc<Notify>,
}
#[async_trait::async_trait]
impl MockDatanodeHandler for GatedDatanodeHandler {
async fn handle(&self, peer: &Peer, _request: RegionRequest) -> Result<RegionResponse> {
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<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 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::<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 {
["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(&regions).await.unwrap()
} else {
let mut peers = HashMap::new();
for region in &regions {
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()
);
}
}
}