feat(operator): retry failed writes with refreshed partitions (#9542)

* feat(operator): retry failed writes with refreshed partitions

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

* style(operator): simplify retry test type references

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

* refactor(operator): clarify write request rebuilding names

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

* refactor(operator): encapsulate write retry attempts

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

* style(operator): group inserter configuration methods

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

* refactor(operator): share routed peers and separate retry decisions

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

* refactor(operator): avoid temporary write routing allocations

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

---------

Signed-off-by: WenyXu <wenymedia@gmail.com>
This commit is contained in:
Weny Xu
2026-10-13 06:02:43 +00:00
committed by GitHub
parent 18634e17bb
commit d86741d0ad
13 changed files with 1958 additions and 117 deletions
Generated
+4
View File
@@ -10095,7 +10095,9 @@ dependencies = [
"async-stream",
"async-trait",
"axum 0.8.4",
"backon",
"bytes",
"cache",
"catalog",
"chrono",
"client",
@@ -10126,6 +10128,7 @@ dependencies = [
"futures",
"futures-util",
"humantime",
"humantime-serde",
"itertools 0.14.0",
"jsonb",
"lazy_static",
@@ -10152,6 +10155,7 @@ dependencies = [
"substrait 1.4.0-alpha.0",
"table",
"tokio",
"tokio-stream",
"tokio-util",
"tonic 0.14.2",
"tracing",
+3
View File
@@ -22,6 +22,9 @@ use crate::cluster::NodeInfo;
use crate::error::Error;
use crate::{DatanodeId, FlownodeId};
/// Shared reference to a [`Peer`].
pub type PeerRef = Arc<Peer>;
/// [`PeerResolver`] provides methods to look up peer information by node ID.
///
/// This trait allows resolving both datanode and flownode peers, regardless of their current activity status.
+32 -1
View File
@@ -114,15 +114,25 @@ pub fn convert_to_region_leader_map(region_routes: &[RegionRoute]) -> HashMap<Re
.collect::<HashMap<_, _>>()
}
/// Returns the leader of the first route matching the region number.
pub fn find_region_leader(
region_routes: &[RegionRoute],
region_number: RegionNumber,
) -> Option<Peer> {
find_region_leader_ref(region_routes, region_number).cloned()
}
/// Borrows the leader of the first route matching the region number.
///
/// Returns `None` if no route matches or the first matching route has no leader.
pub fn find_region_leader_ref(
region_routes: &[RegionRoute],
region_number: RegionNumber,
) -> Option<&Peer> {
region_routes
.iter()
.find(|x| x.region.id.region_number() == region_number)
.and_then(|r| r.leader_peer.as_ref())
.cloned()
}
/// Returns the region numbers of the leader regions on the target datanode.
@@ -442,6 +452,27 @@ mod tests {
}
}
#[test]
fn test_find_region_leader() {
let mut routes = vec![
new_test_region_route(RegionId::new(1, 1)),
new_test_region_route(RegionId::new(1, 1)),
];
routes[1].leader_peer = Some(Peer::new(2, "a2"));
assert_eq!(
find_region_leader_ref(&routes, 1),
routes[0].leader_peer.as_ref()
);
assert_eq!(find_region_leader(&routes, 1), routes[0].leader_peer);
assert_eq!(find_region_leader_ref(&routes, 2), None);
assert_eq!(find_region_leader(&routes, 2), None);
routes[0].leader_peer = None;
assert_eq!(find_region_leader_ref(&routes, 1), None);
assert_eq!(find_region_leader(&routes, 1), None);
}
#[test]
fn test_leader_is_downgraded() {
let mut region_route = RegionRoute {
+3 -1
View File
@@ -14,7 +14,7 @@ insert/delete routing.
| Statements | `src/operator/src/statement.rs`, `src/operator/src/statement/` | `StatementExecutor` and per-statement handlers |
| Inserts | `src/operator/src/insert.rs`, `src/operator/src/bulk_insert.rs` | Insert, auto-create/alter, routing, dispatch |
| Deletes | `src/operator/src/delete.rs` | Delete conversion and dispatch |
| Write dispatch | `src/operator/src/write_dispatcher.rs` | Shared insert/delete/bulk dispatch and batched peer resolution |
| Write dispatch | `src/operator/src/write_dispatcher.rs` | Shared insert/delete/bulk dispatch, bounded retries, route refresh and write request rebuilding |
| Request conversion | `src/operator/src/req_convert.rs`, `src/operator/src/req_convert/` | Table/row/column requests to region requests |
| Region requests | `src/operator/src/region_req_factory.rs` | Region-level DDL/DML request construction |
| Procedures | `src/operator/src/procedure.rs` | Metasrv procedure administration |
@@ -39,6 +39,8 @@ cargo sqlness bare -t <case>
```
Check standalone and distributed routing when changing region dispatch.
Write retries operate on prepared region requests in `write_dispatcher`; keep
admission and Flow side effects outside retries, and preserve successful regions.
Write admission uses fallible `write_meter!` calls in `Inserter::do_request`
and `handle_bulk_insert`, before flow side effects and datanode dispatch. Bulk
+4
View File
@@ -24,6 +24,7 @@ api.workspace = true
arrow.workspace = true
async-stream.workspace = true
async-trait.workspace = true
backon.workspace = true
bytes.workspace = true
catalog.workspace = true
chrono.workspace = true
@@ -54,6 +55,7 @@ file-engine.workspace = true
futures.workspace = true
futures-util.workspace = true
humantime.workspace = true
humantime-serde.workspace = true
itertools.workspace = true
jsonb.workspace = true
lazy_static.workspace = true
@@ -86,8 +88,10 @@ url.workspace = true
[dev-dependencies]
axum.workspace = true
cache.workspace = true
catalog = { workspace = true, features = ["testing"] }
common-meta = { workspace = true, features = ["testing"] }
common-test-util.workspace = true
object-store = { workspace = true, features = ["testing"] }
path-slash = "0.2"
tokio-stream = { workspace = true, features = ["net"] }
+30 -17
View File
@@ -14,7 +14,6 @@
use std::collections::HashSet;
use api::region::RegionResponse;
use api::v1::flow::{DirtyWindowRequest, DirtyWindowRequests};
use api::v1::region::{
BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request,
@@ -41,6 +40,7 @@ use crate::error::Result;
use crate::insert::Inserter;
use crate::req_convert::TablePartitioner;
use crate::req_convert::insert;
use crate::write_dispatcher::WriteAttempt;
use crate::{error, metrics};
impl Inserter {
@@ -72,7 +72,7 @@ impl Inserter {
let writes = self.write_dispatcher.prepare(writes).await?;
Ok(self
.write_dispatcher
.dispatch_prepared(writes, None)
.execute(writes, None)
.await
.context(error::RequestInsertsSnafu)?
.affected_rows)
@@ -186,7 +186,7 @@ impl Inserter {
let writes = self.write_dispatcher.prepare(vec![request]).await?;
let result = self
.write_dispatcher
.dispatch_prepared(
.execute(
writes,
Some(
metrics::HANDLE_BULK_INSERT_ELAPSED.with_label_values(&["datanode_handle"]),
@@ -246,7 +246,7 @@ impl Inserter {
} else {
None
};
let handle: common_runtime::JoinHandle<Result<RegionResponse>> =
let handle: common_runtime::JoinHandle<Result<WriteAttempt>> =
common_runtime::spawn_global(async move {
let (header, payload) = if mask.select_all() {
// SAFETY: raw data must be present, we can avoid re-encoding.
@@ -301,27 +301,40 @@ impl Inserter {
})),
};
dispatcher
.dispatch_prepared(
vec![(peer, request)],
Ok(dispatcher
.send(
peer,
request,
Some(
metrics::HANDLE_BULK_INSERT_ELAPSED
&metrics::HANDLE_BULK_INSERT_ELAPSED
.with_label_values(&["datanode_handle"]),
),
)
.await
.context(error::RequestRegionSnafu)
.await)
});
handles.push(handle);
}
let results = future::try_join_all(handles)
.await
.context(error::JoinTaskSnafu)?;
let rows_inserted = results
.into_iter()
.map(|result| result.map(|response| response.affected_rows))
.sum::<Result<usize>>()?;
let dispatcher = self.write_dispatcher.clone();
let rows_inserted = common_runtime::spawn_global(async move {
let results = future::try_join_all(handles)
.await
.context(error::JoinTaskSnafu)?;
let attempts = results.into_iter().collect::<Result<Vec<_>>>()?;
dispatcher
.finish(
attempts,
Some(
&metrics::HANDLE_BULK_INSERT_ELAPSED
.with_label_values(&["datanode_handle"]),
),
)
.await
.context(error::RequestRegionSnafu)
.map(|response| response.affected_rows)
})
.await
.context(error::JoinTaskSnafu)??;
wait_all_datanode_timer.observe_duration();
crate::metrics::DIST_INGEST_ROW_COUNT
.with_label_values(&[db_name.as_str()])
+6 -7
View File
@@ -141,13 +141,12 @@ impl Deleter {
let dispatcher = self.write_dispatcher.clone();
let request = request_factory.build_delete(requests);
let writes = dispatcher.prepare(vec![request]).await?;
let affected_rows = common_runtime::spawn_global(async move {
dispatcher.dispatch_prepared(writes, None).await
})
.await
.context(JoinTaskSnafu)?
.context(RequestDeletesSnafu)?
.affected_rows;
let affected_rows =
common_runtime::spawn_global(async move { dispatcher.execute(writes, None).await })
.await
.context(JoinTaskSnafu)?
.context(RequestDeletesSnafu)?
.affected_rows;
crate::metrics::DIST_DELETE_ROW_COUNT
.with_label_values(&[ctx.get_db_string().as_str()])
.inc_by(affected_rows as u64);
+12 -13
View File
@@ -316,6 +316,12 @@ impl Inserter {
}
}
/// Installs the shared dispatcher used by direct writes and batch flushes.
pub fn with_write_dispatcher(mut self, dispatcher: WriteDispatcherRef) -> Self {
self.write_dispatcher = dispatcher;
self
}
/// Installs the shared batcher; callers explicitly select its ingestion entry point.
pub fn with_pending_rows_batcher(
mut self,
@@ -835,13 +841,12 @@ impl Inserter {
let dispatcher = self.write_dispatcher.clone();
let request = request_factory.build_insert(normal_requests);
let writes = dispatcher.prepare(vec![request]).await?;
let affected_rows = common_runtime::spawn_global(async move {
dispatcher.dispatch_prepared(writes, None).await
})
.await
.context(JoinTaskSnafu)?
.context(RequestInsertsSnafu)?
.affected_rows;
let affected_rows =
common_runtime::spawn_global(async move { dispatcher.execute(writes, None).await })
.await
.context(JoinTaskSnafu)?
.context(RequestInsertsSnafu)?
.affected_rows;
crate::metrics::DIST_INGEST_ROW_COUNT
.with_label_values(&[ctx.get_db_string().as_str()])
.inc_by(affected_rows as u64);
@@ -1675,12 +1680,6 @@ impl Inserter {
}
}
/// Installs the shared dispatcher used by direct writes and batch flushes.
pub fn with_write_dispatcher(mut self, dispatcher: WriteDispatcherRef) -> Self {
self.write_dispatcher = dispatcher;
self
}
/// Returns the dispatcher shared with batching execution paths.
pub fn write_dispatcher(&self) -> &WriteDispatcherRef {
&self.write_dispatcher
-1
View File
@@ -16,5 +16,4 @@ pub(crate) mod common;
pub mod delete;
pub mod insert;
pub use common::partitioner::Partitioner;
pub(crate) use common::partitioner::TablePartitioner;
@@ -14,16 +14,19 @@
use std::collections::HashMap;
use api::v1::region::{BulkInsertRequest, DeleteRequest, InsertRequest, bulk_insert_request};
use api::v1::region::{
BulkInsertRequest, DeleteRequest, DeleteRequests, InsertRequest, InsertRequests, RegionRequest,
bulk_insert_request, region_request::Body,
};
use api::v1::{ArrowIpc, PartitionExprVersion, Rows};
use common_error::ext::BoxedError;
use common_grpc::flight;
use common_grpc::flight::{self, FlightDecoder};
use datatypes::arrow::compute;
use datatypes::arrow::record_batch::RecordBatch;
use partition::PartitionRuleRef;
use partition::manager::PartitionRuleManager;
use partition::splitter::RowSplitter;
use snafu::ResultExt;
use snafu::{OptionExt, ResultExt};
use store_api::storage::{RegionId, RegionNumber, TableId};
use table::metadata::TableInfo;
@@ -94,6 +97,37 @@ impl TablePartitioner {
})
}
/// Rebuilds write requests using this partition snapshot after a route refresh.
///
/// Takes ownership so row payloads can move into their new regions. Requests whose
/// region and partition version still match retain their original payloads.
pub(crate) fn rebuild_write_requests(
&self,
request: RegionRequest,
) -> Result<Vec<RegionRequest>> {
let RegionRequest { header, body } = request;
let bodies = match body.context(error::UnexpectedSnafu {
violated: "Cannot rebuild a request without a body",
})? {
Body::Inserts(inserts) => self.rebuild_insert_requests(inserts)?,
Body::Deletes(deletes) => self.rebuild_delete_requests(deletes)?,
Body::BulkInsert(bulk) => self.rebuild_bulk_insert_request(bulk)?,
_ => {
return error::UnexpectedSnafu {
violated: "Only writes can be rebuilt",
}
.fail();
}
};
Ok(bodies
.into_iter()
.map(|body| RegionRequest {
header: header.clone(),
body: Some(body),
})
.collect())
}
/// Moves insert rows into requests for this snapshot's destination regions.
pub fn partition_insert_requests(
&self,
@@ -178,6 +212,84 @@ impl TablePartitioner {
Ok(requests)
}
fn rebuild_insert_requests(&self, inserts: InsertRequests) -> Result<Vec<Body>> {
let mut bodies = Vec::new();
for insert in inserts.requests {
let region_id = RegionId::from_u64(insert.region_id);
if self.is_partition_expr_unchanged(region_id, insert.partition_expr_version) {
bodies.push(Body::Inserts(InsertRequests {
requests: vec![insert],
}));
continue;
}
let rows = insert.rows.context(error::UnexpectedSnafu {
violated: "Cannot rebuild an insert without rows",
})?;
for request in self.partition_insert_requests(rows, insert.skip_wal)? {
bodies.push(Body::Inserts(InsertRequests {
requests: vec![request],
}));
}
}
Ok(bodies)
}
fn rebuild_delete_requests(&self, deletes: DeleteRequests) -> Result<Vec<Body>> {
let mut bodies = Vec::new();
for delete in deletes.requests {
let region_id = RegionId::from_u64(delete.region_id);
if self.is_partition_expr_unchanged(region_id, delete.partition_expr_version) {
bodies.push(Body::Deletes(DeleteRequests {
requests: vec![delete],
}));
continue;
}
let rows = delete.rows.context(error::UnexpectedSnafu {
violated: "Cannot rebuild a delete without rows",
})?;
for request in self.partition_delete_requests(rows)? {
bodies.push(Body::Deletes(DeleteRequests {
requests: vec![request],
}));
}
}
Ok(bodies)
}
fn rebuild_bulk_insert_request(&self, bulk: BulkInsertRequest) -> Result<Vec<Body>> {
let region_id = RegionId::from_u64(bulk.region_id);
if self.is_partition_expr_unchanged(region_id, bulk.partition_expr_version) {
return Ok(vec![Body::BulkInsert(bulk)]);
}
let bulk_insert_request::Body::ArrowIpc(ipc) =
bulk.body.context(error::UnexpectedSnafu {
violated: "Cannot rebuild a bulk insert without Arrow IPC data",
})?;
let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema)
.and_then(|mut decoder| decoder.try_decode_record_batch(&ipc.data_header, &ipc.payload))
.map_err(BoxedError::new)
.context(error::ExternalSnafu)?;
Ok(self
.partition_bulk_insert_requests(batch, Some(ipc), bulk.skip_wal)?
.into_iter()
.map(Body::BulkInsert)
.collect())
}
/// Whether the request's region still exists with the same partition expression version.
///
/// A leader change leaves the partition unchanged, so the original payload can be
/// sent to the new peer. A removed region or changed partition expression version
/// requires rebuilding the request using this snapshot's partition rule.
fn is_partition_expr_unchanged(
&self,
region_id: RegionId,
version: Option<PartitionExprVersion>,
) -> bool {
region_id.table_id() == self.table_id
&& self.versions.get(&region_id.region_number()) == Some(&version)
}
fn version(&self, region: RegionNumber) -> Option<PartitionExprVersion> {
self.versions.get(&region).copied().flatten()
}
@@ -187,12 +299,26 @@ impl TablePartitioner {
mod tests {
use std::sync::Arc;
use api::v1::region::{
AlignedSchemaVersion, BulkInsertRequest, DeleteRequest, InsertRequest, RegionRequestHeader,
};
use api::v1::value::ValueData;
use api::v1::{ArrowIpc, PartitionExprVersion};
use api::v1::{ColumnDataType, ColumnSchema, Row, Rows, Value};
use common_error::ext::{ErrorExt, RetryHint};
use common_error::status_code::StatusCode;
use common_meta::cache;
use common_grpc::flight;
use common_meta::cache::{self, TableInfoCache};
use common_meta::key::TableMetadataManager;
use common_meta::key::table_route::TableRouteValue;
use common_meta::kv_backend::memory::MemoryKvBackend;
use common_meta::kv_backend::test_util::MockKvBackendBuilder;
use common_meta::rpc::store::RangeResponse;
use datatypes::arrow::array::Int32Array;
use datatypes::arrow::datatypes::{DataType, Field, Schema};
use datatypes::arrow::record_batch::RecordBatch;
use moka::future::Cache;
use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef};
use super::*;
use crate::error::Error;
@@ -249,4 +375,324 @@ mod tests {
}
}
}
async fn create_fixture() -> (PartitionRuleManagerRef, TableInfoCache) {
let backend = Arc::new(MemoryKvBackend::new());
let manager = test_util::create_partition_rule_manager(backend.clone()).await;
TableMetadataManager::new(backend.clone())
.create_table_metadata(
test_util::new_test_table_info(42, "logical_reroute", [1, 2, 3].into_iter()),
TableRouteValue::logical(1),
Default::default(),
)
.await
.unwrap();
let cache = cache::new_table_info_cache(
"reroute_test_table_info".to_string(),
Cache::new(16),
backend,
);
(manager, cache)
}
fn build_request(body: Body) -> RegionRequest {
RegionRequest {
header: Some(RegionRequestHeader {
dbname: "greptime.public".to_string(),
..Default::default()
}),
body: Some(body),
}
}
fn build_rows(values: &[i32]) -> Rows {
Rows {
schema: vec![ColumnSchema {
column_name: "a".to_string(),
datatype: ColumnDataType::Int32 as i32,
..Default::default()
}],
rows: values
.iter()
.map(|value| Row {
values: vec![Value {
value_data: Some(ValueData::I32Value(*value)),
}],
})
.collect(),
}
}
fn build_row_request(
rows: Rows,
insert: bool,
region: u32,
version: Option<u64>,
) -> RegionRequest {
let region_id = RegionId::new(1, region).as_u64();
let partition_expr_version = version.map(|value| PartitionExprVersion { value });
let body = if insert {
Body::Inserts(InsertRequests {
requests: vec![InsertRequest {
region_id,
partition_expr_version,
rows: Some(rows),
skip_wal: true,
}],
})
} else {
Body::Deletes(DeleteRequests {
requests: vec![DeleteRequest {
region_id,
partition_expr_version,
rows: Some(rows),
}],
})
};
build_request(body)
}
fn row_request_parts(request: &RegionRequest) -> (u64, Option<u64>, &Rows) {
match request.body.as_ref().unwrap() {
Body::Inserts(inserts) => {
assert_eq!(inserts.requests.len(), 1);
let insert = &inserts.requests[0];
assert!(insert.skip_wal);
(
insert.region_id,
insert.partition_expr_version.map(|v| v.value),
insert.rows.as_ref().unwrap(),
)
}
Body::Deletes(deletes) => {
assert_eq!(deletes.requests.len(), 1);
let delete = &deletes.requests[0];
(
delete.region_id,
delete.partition_expr_version.map(|v| v.value),
delete.rows.as_ref().unwrap(),
)
}
_ => panic!("expected row write"),
}
}
#[tokio::test]
async fn test_rebuild_rows_updates_partitions_without_losing_rows() {
let (manager, cache) = create_fixture().await;
let table_info = cache.get(1).await.unwrap().unwrap();
let (_, versions) = manager
.find_table_partition_rule(&table_info)
.await
.unwrap();
for insert in [true, false] {
let request = build_row_request(build_rows(&[1, 20, 70]), insert, 99, Some(0));
let header = request.header.clone();
let rerouted = rebuild_with_cached_table(&manager, &cache, request)
.await
.unwrap();
assert_eq!(rerouted.len(), 3);
let mut actual = Vec::new();
for request in rerouted {
assert_eq!(request.header, header);
let (id, version, rows) = row_request_parts(&request);
let region_id = RegionId::from_u64(id);
assert_eq!(region_id.table_id(), 1);
assert_eq!(version, versions[&region_id.region_number()]);
assert_eq!(rows.schema, build_rows(&[]).schema);
assert_eq!(rows.rows.len(), 1);
actual.push((region_id.region_number(), rows.rows[0].clone()));
}
actual.sort_by_key(|(region, _)| *region);
assert_eq!(
actual,
vec![
(1, build_rows(&[70]).rows.remove(0)),
(2, build_rows(&[20]).rows.remove(0)),
(3, build_rows(&[1]).rows.remove(0))
]
);
}
}
#[tokio::test]
async fn test_unchanged_row_routes_preserve_allocations() {
let (manager, cache) = create_fixture().await;
let table_info = cache.get(1).await.unwrap().unwrap();
let (_, versions) = manager
.find_table_partition_rule(&table_info)
.await
.unwrap();
for insert in [true, false] {
let request = build_row_request(build_rows(&[1]), insert, 3, versions[&3]);
let original_rows = row_request_parts(&request).2;
let row_ptr = original_rows.rows.as_ptr();
let value_ptr = original_rows.rows[0].values.as_ptr();
let rerouted = rebuild_with_cached_table(&manager, &cache, request)
.await
.unwrap();
assert_eq!(rerouted.len(), 1);
let rows = row_request_parts(&rerouted[0]).2;
assert_eq!(rows.rows.as_ptr(), row_ptr);
assert_eq!(rows.rows[0].values.as_ptr(), value_ptr);
assert_eq!(rows, &build_rows(&[1]));
}
}
fn build_bulk_request(region: u32, version: Option<u64>, values: Vec<i32>) -> RegionRequest {
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(values))],
)
.unwrap();
let (schema, data_header, payload) = flight::record_batch_to_ipc(batch).unwrap();
build_request(Body::BulkInsert(BulkInsertRequest {
region_id: RegionId::new(1, region).as_u64(),
partition_expr_version: version.map(|value| PartitionExprVersion { value }),
aligned_schema_version: Some(AlignedSchemaVersion {
schema_version: 123,
}),
skip_wal: true,
body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc {
schema,
data_header,
payload,
})),
}))
}
#[tokio::test]
async fn test_rebuild_bulk_reencodes_only_changed_partitions() {
let (manager, cache) = create_fixture().await;
let table_info = cache.get(1).await.unwrap().unwrap();
let (_, versions) = manager
.find_table_partition_rule(&table_info)
.await
.unwrap();
let request = build_bulk_request(99, Some(0), vec![1, 20, 70]);
let header = request.header.clone();
let rerouted = rebuild_with_cached_table(&manager, &cache, request)
.await
.unwrap();
assert_eq!(rerouted.len(), 3);
let mut actual = Vec::new();
for request in rerouted {
assert_eq!(request.header, header);
let Some(Body::BulkInsert(bulk)) = request.body else {
panic!("expected bulk insert")
};
let region = RegionId::from_u64(bulk.region_id);
assert_eq!(region.table_id(), 1);
assert!(bulk.skip_wal);
assert_eq!(bulk.aligned_schema_version, None);
assert_eq!(
bulk.partition_expr_version.map(|v| v.value),
versions[&region.region_number()]
);
let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = bulk.body else {
panic!("expected Arrow IPC")
};
let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema)
.unwrap()
.try_decode_record_batch(&ipc.data_header, &ipc.payload)
.unwrap();
assert_eq!(batch.num_rows(), 1);
let values = batch
.column(0)
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
actual.push((region.region_number(), values.value(0)));
}
actual.sort_unstable();
assert_eq!(actual, [(1, 70), (2, 20), (3, 1)]);
// A leader-only change retains the original IPC buffers and schema hint.
let request = build_bulk_request(3, versions[&3], vec![1]);
let Some(Body::BulkInsert(bulk)) = request.body.as_ref() else {
unreachable!()
};
let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &bulk.body else {
unreachable!()
};
let pointers = (
ipc.schema.as_ptr(),
ipc.data_header.as_ptr(),
ipc.payload.as_ptr(),
);
let rerouted = rebuild_with_cached_table(&manager, &cache, request)
.await
.unwrap();
assert_eq!(rerouted.len(), 1);
let Some(Body::BulkInsert(bulk)) = rerouted[0].body.as_ref() else {
unreachable!()
};
assert_eq!(
bulk.aligned_schema_version,
Some(AlignedSchemaVersion {
schema_version: 123
})
);
let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &bulk.body else {
unreachable!()
};
assert_eq!(
(
ipc.schema.as_ptr(),
ipc.data_header.as_ptr(),
ipc.payload.as_ptr()
),
pointers
);
}
#[tokio::test]
async fn test_rebuild_logical_table_preserves_logical_region_ids() {
let (manager, cache) = create_fixture().await;
for insert in [true, false] {
let mut request = build_row_request(build_rows(&[1, 20, 70]), insert, 99, Some(0));
match request.body.as_mut().unwrap() {
Body::Inserts(inserts) => {
inserts.requests[0].region_id = RegionId::new(42, 99).as_u64()
}
Body::Deletes(deletes) => {
deletes.requests[0].region_id = RegionId::new(42, 99).as_u64()
}
_ => unreachable!(),
}
let rerouted = rebuild_with_cached_table(&manager, &cache, request)
.await
.unwrap();
assert_eq!(rerouted.len(), 3);
let mut regions = rerouted
.iter()
.map(|request| {
let region = RegionId::from_u64(row_request_parts(request).0);
assert_eq!(region.table_id(), 42);
region.region_number()
})
.collect::<Vec<_>>();
regions.sort_unstable();
assert_eq!(regions, [1, 2, 3]);
}
}
async fn rebuild_with_cached_table(
manager: &PartitionRuleManager,
cache: &TableInfoCache,
request: RegionRequest,
) -> Result<Vec<RegionRequest>> {
let region_id = match request.body.as_ref().unwrap() {
Body::Inserts(inserts) => inserts.requests[0].region_id,
Body::Deletes(deletes) => deletes.requests[0].region_id,
Body::BulkInsert(bulk) => bulk.region_id,
_ => unreachable!(),
};
let table_id = RegionId::from_u64(region_id).table_id();
let table_info = cache.get(table_id).await.unwrap().unwrap();
let partitioner = TablePartitioner::try_new(manager, &table_info)
.await
.context(error::InvalidPartitionSnafu)?;
partitioner.rebuild_write_requests(request)
}
}
+395 -56
View File
@@ -12,6 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
mod options;
#[cfg(test)]
mod tests;
@@ -20,43 +21,85 @@ use std::sync::Arc;
use api::region::RegionResponse;
use api::v1::region::{DeleteRequests, InsertRequests, RegionRequest, region_request::Body};
use common_error::ext::BoxedError;
use common_meta::error::{ExternalSnafu, Result, UnexpectedSnafu};
use backon::{BackoffBuilder, ExponentialBuilder};
use client::client_manager::NodeClients;
use common_error::ext::{BoxedError, ErrorExt, RetryHint};
use common_meta::cache::TableInfoCacheRef;
use common_meta::cache_invalidator::{CacheInvalidatorRef, Context};
use common_meta::error::{Error, ExternalSnafu, Result, TableInfoNotFoundSnafu, UnexpectedSnafu};
use common_meta::instruction::CacheIdent;
use common_meta::node_manager::NodeManagerRef;
use common_meta::peer::Peer;
use common_meta::peer::{Peer, PeerRef};
use common_meta::rpc::router;
use futures::future;
use partition::manager::PartitionRuleManagerRef;
use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef};
use prometheus::Histogram;
use snafu::{OptionExt, ResultExt};
use store_api::storage::RegionId;
pub use self::options::WriteRetryOptions;
use crate::error::{self, FindRegionLeaderSnafu, RequestRegionSnafu};
use crate::req_convert::TablePartitioner;
/// Dispatches already partitioned insert, delete, and bulk requests to their peers.
/// Dispatches already partitioned writes, retaining only failed payloads for retries.
pub struct WriteDispatcher {
node_manager: NodeManagerRef,
partition_manager: PartitionRuleManagerRef,
retry: Option<WriteRetry>,
}
/// Shared reference to a [`WriteDispatcher`].
pub type WriteDispatcherRef = Arc<WriteDispatcher>;
/// The result of one write RPC attempt.
pub(crate) enum WriteAttempt {
Success(RegionResponse),
Failed {
error: Error,
/// Retained when retries are enabled and the original client error allows retry.
/// The remaining retry budget is checked separately.
request: Option<RoutedWrite>,
},
}
impl WriteDispatcher {
/// Creates an owned-request dispatcher without retries, including for standalone.
pub fn new(node_manager: NodeManagerRef, partition_manager: PartitionRuleManagerRef) -> Self {
Self {
node_manager,
partition_manager,
retry: None,
}
}
/// Enables shared remote sends and retries using frontend configuration and caches.
pub fn with_retry(
mut self,
clients: Arc<NodeClients>,
options: WriteRetryOptions,
table_info_cache: TableInfoCacheRef,
cache_invalidator: CacheInvalidatorRef,
) -> Self {
if options.max_retries > 0 {
self.retry = Some(WriteRetry {
clients,
partition_manager: self.partition_manager.clone(),
options,
table_info_cache,
cache_invalidator,
});
} else {
self.retry = None;
}
self
}
/// Writes planned requests without repeating admission, default values, or Flow side effects.
pub async fn dispatch(&self, requests: Vec<RegionRequest>) -> Result<RegionResponse> {
self.dispatch_inner(requests, None).await
}
/// Records each datanode send in the supplied histogram.
/// Records each datanode send attempt, including retries, in the supplied histogram.
pub async fn dispatch_with_metrics(
&self,
requests: Vec<RegionRequest>,
@@ -69,11 +112,17 @@ impl WriteDispatcher {
pub(crate) async fn prepare(
&self,
requests: Vec<RegionRequest>,
) -> error::Result<Vec<(Peer, RegionRequest)>> {
let requests = requests.into_iter().map(split_request).collect::<Vec<_>>();
) -> error::Result<Vec<(PeerRef, RegionRequest)>> {
let requests = requests
.into_iter()
.map(|mut request| {
let header = request.header.take();
(header, split_request(request))
})
.collect::<Vec<_>>();
let regions = requests
.iter()
.flatten()
.flat_map(|(_, requests)| requests)
.map(write_region_id)
.collect::<Result<Vec<_>>>()
.context(RequestRegionSnafu)?;
@@ -83,8 +132,8 @@ impl WriteDispatcher {
.context(FindRegionLeaderSnafu)?;
let mut writes = Vec::new();
// Preserve each original envelope and peer batching when retries are disabled.
for requests in requests {
let mut by_peer: HashMap<Peer, RegionRequest> = HashMap::new();
for (header, requests) in requests {
let mut by_peer: HashMap<PeerRef, RegionRequest> = HashMap::new();
for request in requests {
let peer = get_peer(
&peers,
@@ -112,27 +161,94 @@ impl WriteDispatcher {
},
}
}
writes.extend(by_peer);
for (peer, mut request) in by_peer {
request.header = header.clone();
writes.push((peer, request));
}
}
Ok(writes)
}
/// Sends requests using routes already resolved by the caller.
pub(crate) async fn dispatch_prepared(
pub(crate) async fn execute(
&self,
writes: Vec<(Peer, RegionRequest)>,
writes: Vec<(PeerRef, RegionRequest)>,
histogram: Option<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
}))
let writes = if self.retry.is_some() {
writes
.into_iter()
.flat_map(|(peer, request)| {
split_request(request)
.into_iter()
.map(move |request| (Arc::clone(&peer), request))
})
.collect()
} else {
writes
};
let attempts = future::join_all(
writes
.into_iter()
.map(|(peer, request)| self.send(peer, request, histogram)),
)
.await;
self.finish(attempts, histogram).await
}
/// Sends the first attempt immediately so bulk encoding can overlap other sends.
pub(crate) async fn send(
&self,
peer: PeerRef,
request: RegionRequest,
histogram: Option<&Histogram>,
) -> WriteAttempt {
if let Some(retry) = &self.retry {
return retry
.send(
RoutedWrite {
peer,
request: Arc::new(request),
},
histogram,
)
.await;
}
let _timer = histogram.map(Histogram::start_timer);
let datanode = self.node_manager.datanode(&peer).await;
match datanode.handle(request).await {
Ok(response) => WriteAttempt::Success(response),
Err(error) => WriteAttempt::Failed {
error,
request: None,
},
}
}
/// Collects first attempts and retries their failures together, refreshing routes once per round.
pub(crate) async fn finish(
&self,
attempts: Vec<WriteAttempt>,
histogram: Option<&Histogram>,
) -> Result<RegionResponse> {
if let Some(retry) = &self.retry {
if retry.should_retry(&attempts, 0) {
return retry.retry(attempts, histogram).await;
}
let mut response = RegionResponse::new(0);
let failures = merge_attempts(attempts, &mut response)?;
return match failures.into_iter().next() {
Some((_, error)) => Err(error),
None => Ok(response),
};
}
let mut response = RegionResponse::new(0);
for result in results {
merge_response(&mut response, result?);
for attempt in attempts {
match attempt {
WriteAttempt::Success(result) => merge_response(&mut response, result),
WriteAttempt::Failed { error, .. } => return Err(error),
}
}
Ok(response)
}
@@ -141,39 +257,8 @@ impl WriteDispatcher {
pub(crate) async fn resolve_peers(
&self,
regions: &[RegionId],
) -> partition::error::Result<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)
) -> partition::error::Result<HashMap<RegionId, PeerRef>> {
resolve_peers(&self.partition_manager, regions).await
}
async fn dispatch_inner(
@@ -186,16 +271,270 @@ impl WriteDispatcher {
.await
.map_err(BoxedError::new)
.context(ExternalSnafu)?;
self.dispatch_prepared(writes, histogram).await
self.execute(writes, histogram).await
}
}
fn get_peer(peers: &HashMap<RegionId, Peer>, region: RegionId) -> Result<Peer> {
/// Sends shared requests and coordinates retries with refreshed routes.
struct WriteRetry {
clients: Arc<NodeClients>,
partition_manager: PartitionRuleManagerRef,
options: WriteRetryOptions,
table_info_cache: TableInfoCacheRef,
cache_invalidator: CacheInvalidatorRef,
}
/// A single-region write request paired with its destination peer.
pub(crate) struct RoutedWrite {
peer: PeerRef,
request: Arc<RegionRequest>,
}
impl WriteRetry {
async fn send(&self, write: RoutedWrite, histogram: Option<&Histogram>) -> WriteAttempt {
let _timer = histogram.map(Histogram::start_timer);
let requester = self.clients.make_region_requester(&write.peer).await;
match requester.handle(Arc::clone(&write.request)).await {
Ok(response) => WriteAttempt::Success(response),
Err(error) => {
let request = error.retry_hint().is_retryable().then_some(write);
WriteAttempt::Failed {
error: error.into(),
request,
}
}
}
}
fn should_retry(&self, attempts: &[WriteAttempt], retries: usize) -> bool {
if retries >= self.options.max_retries {
return false;
}
let mut retryable = false;
for attempt in attempts {
match attempt {
WriteAttempt::Success(_) => {}
WriteAttempt::Failed {
request: Some(_), ..
} => retryable = true,
WriteAttempt::Failed { request: None, .. } => return false,
}
}
retryable
}
async fn retry(
&self,
mut attempts: Vec<WriteAttempt>,
histogram: Option<&Histogram>,
) -> Result<RegionResponse> {
let mut response = RegionResponse::new(0);
let mut backoff = ExponentialBuilder::default()
.with_min_delay(self.options.initial_backoff)
.with_max_delay(self.options.max_backoff)
.with_factor(self.options.backoff_multiplier)
.with_max_times(self.options.max_retries)
.with_jitter()
.build();
let mut retries = 0;
while self.should_retry(&attempts, retries) {
let Some(delay) = backoff.next() else {
break;
};
let failures = merge_attempts(attempts, &mut response)?;
tokio::time::sleep(delay).await;
let mut refresh = Vec::new();
let mut pending = Vec::new();
for (write, error) in failures {
if error.retry_hint() == RetryHint::RetryAfterRouteRefresh {
refresh.push(write);
} else {
pending.push(write);
}
}
if !refresh.is_empty() {
pending.extend(self.reroute_failed_writes(refresh).await?);
}
attempts =
future::join_all(pending.into_iter().map(|write| self.send(write, histogram)))
.await;
retries += 1;
}
match merge_attempts(attempts, &mut response)?.into_iter().next() {
Some((_, error)) => Err(error),
None => Ok(response),
}
}
async fn reroute_failed_writes(&self, writes: Vec<RoutedWrite>) -> Result<Vec<RoutedWrite>> {
let regions = writes
.iter()
.map(|write| write_region_id(&write.request))
.collect::<Result<Vec<_>>>()?;
self.invalidate_routes(&regions).await?;
let mut partitioners = HashMap::new();
for table_id in regions
.iter()
.map(|region| region.table_id())
.collect::<HashSet<_>>()
{
let table_info = self
.table_info_cache
.get(table_id)
.await?
.with_context(|| TableInfoNotFoundSnafu {
table: table_id.to_string(),
})?;
let partitioner = TablePartitioner::try_new(&self.partition_manager, &table_info)
.await
.map_err(BoxedError::new)
.context(ExternalSnafu)?;
partitioners.insert(table_id, partitioner);
}
let mut rebuilt_requests = Vec::new();
for write in writes {
let request = Arc::try_unwrap(write.request).map_err(|_| {
UnexpectedSnafu {
err_msg: "Write request is still shared after its RPC completed",
}
.build()
})?;
let table_id = write_region_id(&request)?.table_id();
let partitioner = partitioners.get(&table_id).context(UnexpectedSnafu {
err_msg: "Refreshed write has no partition snapshot",
})?;
let requests = partitioner
.rebuild_write_requests(request)
.map_err(BoxedError::new)
.context(ExternalSnafu)?;
rebuilt_requests.extend(requests);
}
self.prepare(rebuilt_requests).await
}
async fn prepare(&self, requests: Vec<RegionRequest>) -> Result<Vec<RoutedWrite>> {
let regions = requests
.iter()
.map(write_region_id)
.collect::<Result<Vec<_>>>()?;
let peers = resolve_peers(&self.partition_manager, &regions)
.await
.map_err(BoxedError::new)
.context(ExternalSnafu)?;
requests
.into_iter()
.zip(regions)
.map(|(request, region)| {
Ok(RoutedWrite {
peer: get_peer(&peers, region)?,
request: Arc::new(request),
})
})
.collect()
}
async fn invalidate_routes(&self, regions: &[RegionId]) -> Result<()> {
let mut table_ids = regions
.iter()
.map(|region| region.table_id())
.collect::<HashSet<_>>();
let mut physical_ids = HashSet::new();
for table_id in &table_ids {
let (physical_id, _) = self
.partition_manager
.find_physical_table_route_with_id(*table_id)
.await
.map_err(BoxedError::new)
.context(ExternalSnafu)?;
physical_ids.insert(physical_id);
}
table_ids.extend(physical_ids);
let caches = table_ids
.into_iter()
.map(CacheIdent::TableId)
.collect::<Vec<_>>();
self.cache_invalidator
.invalidate(&Context::default(), &caches)
.await
}
}
async fn resolve_peers(
partition_manager: &PartitionRuleManager,
regions: &[RegionId],
) -> partition::error::Result<HashMap<RegionId, PeerRef>> {
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) = partition_manager
.find_physical_table_route_with_id(table_id)
.await?;
physical_tables.insert(table_id, physical_id);
routes.entry(physical_id).or_insert(route);
}
let mut shared_peers: HashMap<&Peer, PeerRef> = HashMap::new();
let mut peers = HashMap::with_capacity(regions.len());
for region in regions {
if peers.contains_key(region) {
continue;
}
let peer = physical_tables
.get(&region.table_id())
.and_then(|physical_id| routes.get(physical_id))
.and_then(|route| {
router::find_region_leader_ref(&route.region_routes, region.region_number())
})
.context(partition::error::FindLeaderSnafu {
region_id: *region,
table_id: region.table_id(),
})?;
let peer = shared_peers
.entry(peer)
.or_insert_with(|| Arc::new(peer.clone()));
peers.insert(*region, Arc::clone(peer));
}
Ok(peers)
}
fn get_peer(peers: &HashMap<RegionId, PeerRef>, region: RegionId) -> Result<PeerRef> {
peers.get(&region).cloned().context(UnexpectedSnafu {
err_msg: "Planned write has no resolved peer",
})
}
fn merge_attempts(
attempts: Vec<WriteAttempt>,
response: &mut RegionResponse,
) -> Result<Vec<(RoutedWrite, Error)>> {
let mut failures = Vec::new();
let mut terminal_error = None;
for attempt in attempts {
match attempt {
WriteAttempt::Success(result) => merge_response(response, result),
WriteAttempt::Failed {
error,
request: Some(write),
} => failures.push((write, error)),
WriteAttempt::Failed {
error,
request: None,
} => {
terminal_error.get_or_insert(error);
}
}
}
// A terminal error takes precedence over retained retryable failures.
match terminal_error {
Some(error) => Err(error),
None => Ok(failures),
}
}
fn merge_response(response: &mut RegionResponse, result: RegionResponse) {
response.affected_rows += result.affected_rows;
response.extensions.extend(result.extensions);
@@ -0,0 +1,106 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::time::Duration;
use serde::{Deserialize, Serialize};
/// Exponential backoff for retryable distributed writes.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
#[serde(default)]
pub struct WriteRetryOptions {
/// Maximum additional attempts after the original write. Zero disables retries.
pub max_retries: usize,
/// Initial backoff before jitter is applied.
#[serde(with = "humantime_serde")]
pub initial_backoff: Duration,
/// Maximum backoff before jitter is applied.
#[serde(with = "humantime_serde")]
pub max_backoff: Duration,
/// Multiplier applied to the backoff after each retry.
pub backoff_multiplier: f32,
}
impl Default for WriteRetryOptions {
fn default() -> Self {
Self {
max_retries: 3,
initial_backoff: Duration::from_millis(1500),
max_backoff: Duration::from_secs(10),
backoff_multiplier: 2.0,
}
}
}
impl WriteRetryOptions {
/// Checks that the configured backoff is positive, bounded, and nondecreasing.
pub fn validate(&self) -> Result<(), &'static str> {
if self.initial_backoff.is_zero() {
return Err("write_retry.initial_backoff must be positive");
}
if self.max_backoff < self.initial_backoff {
return Err("write_retry.max_backoff must be at least initial_backoff");
}
if !self.backoff_multiplier.is_finite() || self.backoff_multiplier < 1.0 {
return Err("write_retry.backoff_multiplier must be finite and at least 1");
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_validate_write_retry_options() {
for options in [
WriteRetryOptions::default(),
WriteRetryOptions {
max_retries: 0,
..Default::default()
},
WriteRetryOptions {
backoff_multiplier: 1.0,
..Default::default()
},
] {
options.validate().unwrap();
}
for options in [
WriteRetryOptions {
initial_backoff: Duration::ZERO,
..Default::default()
},
WriteRetryOptions {
max_backoff: Duration::from_secs(1),
..Default::default()
},
WriteRetryOptions {
backoff_multiplier: 0.5,
..Default::default()
},
WriteRetryOptions {
backoff_multiplier: f32::NAN,
..Default::default()
},
WriteRetryOptions {
backoff_multiplier: f32::INFINITY,
..Default::default()
},
] {
assert!(options.validate().is_err(), "{options:?}");
}
}
}
File diff suppressed because it is too large Load Diff