diff --git a/Cargo.lock b/Cargo.lock
index d84446150d6..e97dd203ee3 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -10136,6 +10136,7 @@ dependencies = [
"meta-client",
"meter-core",
"meter-macros",
+ "mito-codec",
"moka",
"object-store",
"partition",
diff --git a/config/config.md b/config/config.md
index 5d76c8f7445..381abde24c4 100644
--- a/config/config.md
+++ b/config/config.md
@@ -391,6 +391,11 @@
| `query.experimental_spill_path` | String | Unset | Spill directory path. Ignored unless mode is "custom". |
| `query.experimental_spill_max_temp_directory_size` | String | `1GiB` | Maximum total size of spill directory (default: "1GiB").
Ignored unless mode is "custom". |
| `query.experimental_spill_compression` | String | `uncompressed` | Compression for spilled data files: "uncompressed" (default), "lz4_frame", "zstd".
Ignored unless mode is "custom". |
+| `write_retry` | -- | -- | Retry retryable distributed write failures.
Backoff grows exponentially, with jitter to spread concurrent retries. |
+| `write_retry.max_retries` | Integer | `3` | Maximum retries after the initial write attempt. Set to 0 to disable retries. |
+| `write_retry.initial_backoff` | String | `1500ms` | Initial backoff before jitter is applied. Must be positive. |
+| `write_retry.max_backoff` | String | `10s` | Maximum backoff before jitter is applied. Must be at least initial_backoff. |
+| `write_retry.backoff_multiplier` | Float | `2.0` | Backoff multiplier. Must be finite and at least 1. |
| `datanode` | -- | -- | Datanode options. |
| `datanode.client` | -- | -- | Datanode client options. |
| `datanode.client.connect_timeout` | String | 1s | Connect server timeout. Defaults to `1s`; this example sets it to `10s`. |
diff --git a/config/frontend.example.toml b/config/frontend.example.toml
index 9a1c3db63fa..c42d254671c 100644
--- a/config/frontend.example.toml
+++ b/config/frontend.example.toml
@@ -398,6 +398,18 @@ memory_pool_size = "50%"
#+ experimental_spill_compression = "uncompressed"
# --- End spill-to-disk ---
+## Retry retryable distributed write failures.
+## Backoff grows exponentially, with jitter to spread concurrent retries.
+[write_retry]
+## Maximum retries after the initial write attempt. Set to 0 to disable retries.
+max_retries = 3
+## Initial backoff before jitter is applied. Must be positive.
+initial_backoff = "1500ms"
+## Maximum backoff before jitter is applied. Must be at least initial_backoff.
+max_backoff = "10s"
+## Backoff multiplier. Must be finite and at least 1.
+backoff_multiplier = 2.0
+
## Datanode options.
[datanode]
## Datanode client options.
diff --git a/src/client/src/client_manager.rs b/src/client/src/client_manager.rs
index ef26c107cd5..4d723361e80 100644
--- a/src/client/src/client_manager.rs
+++ b/src/client/src/client_manager.rs
@@ -52,7 +52,18 @@ impl Debug for NodeClients {
#[async_trait::async_trait]
impl DatanodeManager for NodeClients {
async fn datanode(&self, datanode: &Peer) -> DatanodeRef {
- Arc::new(self.make_region_requester(datanode).await)
+ let client = self.get_client(datanode).await;
+
+ let ChannelConfig {
+ send_compression,
+ accept_compression,
+ ..
+ } = self.control_channel_manager.config();
+ Arc::new(RegionRequester::new(
+ client,
+ *send_compression,
+ *accept_compression,
+ ))
}
}
@@ -78,18 +89,6 @@ impl NodeClients {
}
}
- /// Creates a [`RegionRequester`] with the cached client and control channel compression settings.
- pub async fn make_region_requester(&self, datanode: &Peer) -> RegionRequester {
- let client = self.get_client(datanode).await;
-
- let ChannelConfig {
- send_compression,
- accept_compression,
- ..
- } = self.control_channel_manager.config();
- RegionRequester::new(client, *send_compression, *accept_compression)
- }
-
pub async fn get_client(&self, datanode: &Peer) -> Client {
self.clients
.get_with_by_ref(datanode, async move {
diff --git a/src/client/src/error.rs b/src/client/src/error.rs
index 5c70e5fbe91..fdf0df9582f 100644
--- a/src/client/src/error.rs
+++ b/src/client/src/error.rs
@@ -191,9 +191,22 @@ impl ErrorExt for Error {
Error::Tonic { retry_hint, .. } => *retry_hint,
Error::FlightGet { source, .. }
| Error::FlightStream { source, .. }
- | Error::RegionServer { source, .. }
| Error::FlowServer { source, .. }
| Error::External { source, .. } => source.retry_hint(),
+ Error::RegionServer { code, source, .. } => {
+ let hint = source.retry_hint();
+ // Preserve the transport retry policy previously applied by should_retry().
+ if !hint.is_retryable()
+ && matches!(
+ code,
+ Code::Unavailable | Code::Cancelled | Code::DeadlineExceeded
+ )
+ {
+ RetryHint::Retryable
+ } else {
+ hint
+ }
+ }
Error::ConvertFlightData { source, .. }
| Error::CreateChannel { source, .. }
| Error::CreateTlsChannel { source, .. } => source.retry_hint(),
@@ -207,13 +220,9 @@ define_from_tonic_status!(Error, Tonic);
impl From for common_meta::error::Error {
fn from(error: Error) -> Self {
- if error.should_retry() {
- Self::retry_later(error)
- } else {
- Self::External {
- source: BoxedError::new(error),
- location: location!(),
- }
+ Self::External {
+ source: BoxedError::new(error),
+ location: location!(),
}
}
}
diff --git a/src/client/src/region.rs b/src/client/src/region.rs
index 47c477b06b0..f9ebbb0d0bf 100644
--- a/src/client/src/region.rs
+++ b/src/client/src/region.rs
@@ -67,7 +67,7 @@ pub struct RegionRequester {
#[async_trait]
impl Datanode for RegionRequester {
- async fn handle(&self, request: RegionRequest) -> MetaResult {
+ async fn handle_request(&self, request: Arc) -> MetaResult {
self.handle_inner(request).await.map_err(Into::into)
}
@@ -220,16 +220,7 @@ impl RegionRequester {
.await
}
.await
- .map_err(|e| {
- let code = e.code();
- // Uses `Error::RegionServer` instead of `Error::Server`
- error::Error::RegionServer {
- addr,
- code,
- source: BoxedError::new(error::Error::from(e)),
- location: location!(),
- }
- })?
+ .map_err(|status| map_region_error(status, addr))?
.into_inner();
check_response_header(&response.header)?;
@@ -238,6 +229,15 @@ impl RegionRequester {
}
}
+fn map_region_error(status: tonic::Status, addr: String) -> error::Error {
+ error::Error::RegionServer {
+ addr,
+ code: status.code(),
+ source: BoxedError::new(error::Error::from(status)),
+ location: location!(),
+ }
+}
+
async fn recordbatches_from_flight_message_stream(
addr: String,
flight_message_stream: S,
@@ -761,21 +761,104 @@ mod tests {
assert_eq!(error.retry_hint(), baseline.retry_hint());
assert_eq!(error.tonic_code(), baseline.tonic_code());
}
- let error = Datanode::handle(&requester, build_insert_request())
- .await
- .unwrap_err();
+ let datanode: &dyn Datanode = &requester;
+ let error = datanode.handle(build_insert_request()).await.unwrap_err();
assert_eq!(error.status_code(), StatusCode::RegionNotFound);
assert_eq!(error.retry_hint(), hint);
- match hint {
- RetryHint::Retryable | RetryHint::RetryAfterRouteRefresh => assert!(matches!(
- error,
- meta_error::Error::RetryLater {
- clean_poisons: false,
+ assert!(matches!(error, meta_error::Error::External { .. }));
+ }
+ }
+
+ #[tokio::test]
+ async fn test_datanode_transport_retry_hints() {
+ let server = RegionTestServer::start().await;
+ let requester = server.create_requester(ChannelConfig::default(), false, false);
+ let datanode: &dyn Datanode = &requester;
+ let statuses = [
+ tonic::Code::Unavailable,
+ tonic::Code::Cancelled,
+ tonic::Code::DeadlineExceeded,
+ tonic::Code::FailedPrecondition,
+ ]
+ .into_iter()
+ .flat_map(|code| {
+ [
+ None,
+ Some(RetryHint::NonRetryable),
+ Some(RetryHint::Retryable),
+ Some(RetryHint::RetryAfterRouteRefresh),
+ ]
+ .into_iter()
+ .map(move |hint| (code, hint))
+ });
+ for (code, hint) in statuses {
+ for body in [
+ region_request::Body::Inserts(Default::default()),
+ region_request::Body::Deletes(Default::default()),
+ region_request::Body::BulkInsert(Default::default()),
+ region_request::Body::Create(Default::default()),
+ region_request::Body::RemoteDynFilter(Default::default()),
+ ] {
+ let request_type = body.as_ref().to_string();
+ let mut status = Status::new(code, "request failed");
+ // Metadata status is deliberately different from the transport status.
+ status.metadata_mut().insert(
+ common_error::GREPTIME_DB_HEADER_ERROR_CODE,
+ (StatusCode::RegionNotReady as u32)
+ .to_string()
+ .parse()
+ .unwrap(),
+ );
+ if let Some(hint) = hint {
+ status.metadata_mut().insert(
+ common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT,
+ hint.as_str().parse().unwrap(),
+ );
+ }
+ let original = Error::from(status.clone());
+ let expected = if original.should_retry() && !original.retry_hint().is_retryable() {
+ RetryHint::Retryable
+ } else {
+ original.retry_hint()
+ };
+ server
+ .service
+ .responses
+ .lock()
+ .unwrap()
+ .extend([Err(status.clone()), Err(status)]);
+ let request = RegionRequest {
+ header: None,
+ body: Some(body),
+ };
+ let shared = Arc::new(request);
+ let shared_error = datanode.handle(Arc::clone(&shared)).await.unwrap_err();
+ assert_eq!(Arc::strong_count(&shared), 1);
+ let request = Arc::try_unwrap(shared).unwrap();
+ let owned_error = datanode.handle(request).await.unwrap_err();
+ for error in [owned_error, shared_error] {
+ assert_eq!(
+ error.retry_hint(),
+ expected,
+ "{code:?}, {hint:?}, {request_type}"
+ );
+ assert_eq!(error.retry_hint().is_retryable(), original.should_retry());
+ assert_eq!(error.status_code(), StatusCode::RegionNotReady);
+ let meta_error::Error::External { source, .. } = error else {
+ panic!("expected an external error");
+ };
+ assert_eq!(source.retry_hint(), expected);
+ let Error::RegionServer {
+ source,
+ code: actual_code,
..
- }
- )),
- RetryHint::NonRetryable => {
- assert!(matches!(error, meta_error::Error::External { .. }))
+ } = source.as_any().downcast_ref::().unwrap()
+ else {
+ panic!("expected a region server error");
+ };
+ assert_eq!(*actual_code, code);
+ assert_eq!(source.retry_hint(), original.retry_hint());
+ assert_eq!(source.to_string(), original.to_string());
}
}
}
@@ -846,8 +929,11 @@ mod tests {
baseline.status_code(),
"shared: {error:?}; generated: {baseline:?}"
);
+ let Error::RegionServer { source, .. } = &error else {
+ panic!("expected a region server error: {error:?}");
+ };
assert_eq!(
- error.retry_hint(),
+ source.retry_hint(),
baseline.retry_hint(),
"shared: {error:?}; generated: {baseline:?}"
);
diff --git a/src/common/meta/src/ddl/drop_table.rs b/src/common/meta/src/ddl/drop_table.rs
index b73bd95e1bc..622a27d804b 100644
--- a/src/common/meta/src/ddl/drop_table.rs
+++ b/src/common/meta/src/ddl/drop_table.rs
@@ -19,6 +19,8 @@ use std::collections::HashMap;
use async_trait::async_trait;
use common_error::ext::BoxedError;
+#[cfg(feature = "enterprise")]
+use common_error::ext::ErrorExt;
use common_event_recorder::Event;
use common_procedure::error::{ExternalSnafu, FromJsonSnafu, ToJsonSnafu};
use common_procedure::{
@@ -53,7 +55,7 @@ use crate::rpc::router::{RegionRoute, operating_leader_region_roles};
#[cfg(feature = "enterprise")]
fn ensure_retry_later(err: error::Error) -> error::Error {
- if err.is_retry_later() {
+ if err.retry_hint().is_retryable() {
err
} else {
error::Error::retry_later(err)
diff --git a/src/common/meta/src/ddl/undrop_table.rs b/src/common/meta/src/ddl/undrop_table.rs
index 73a361ffd0e..efb1d4824d7 100644
--- a/src/common/meta/src/ddl/undrop_table.rs
+++ b/src/common/meta/src/ddl/undrop_table.rs
@@ -17,8 +17,8 @@ use api::v1::region::{
use async_trait::async_trait;
use common_procedure::error::{FromJsonSnafu, ToJsonSnafu};
use common_procedure::{
- Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure,
- Result as ProcedureResult, Status,
+ Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey,
+ Procedure, Result as ProcedureResult, Status,
};
use common_telemetry::tracing_context::TracingContext;
use common_telemetry::warn;
@@ -156,6 +156,7 @@ impl UndropTableProcedure {
}
.await;
if let Err(err) = restore_result {
+ // Keep this classification aligned with map_restore_error().
let should_cleanup_opened_regions = !err.is_retry_later();
let err = self.map_restore_metadata_error(err);
if should_cleanup_opened_regions
@@ -274,6 +275,19 @@ impl UndropTableProcedure {
Ok(())
}
+ fn map_restore_error(err: error::Error) -> ProcedureError {
+ // Restore retries keep opened regions, while other failures close them. Preserve
+ // this policy even when the underlying metadata error carries a retry hint:
+ // retrying after cleanup would restore metadata without reopening the regions.
+ if err.is_retry_later() {
+ map_to_procedure_error(err)
+ } else if err.need_clean_poisons() {
+ ProcedureError::external_and_clean_poisons(err)
+ } else {
+ ProcedureError::external(err)
+ }
+ }
+
fn map_restore_metadata_error(&self, err: error::Error) -> error::Error {
match err {
error::Error::TombstoneTargetAlreadyExists { .. } => {
@@ -413,7 +427,12 @@ impl Procedure for UndropTableProcedure {
async fn execute(&mut self, _: &ProcedureContext) -> ProcedureResult {
match self.data.state {
UndropTableState::Prepare => self.on_prepare().await,
- UndropTableState::RestoreMetadata => self.on_restore_metadata().await,
+ UndropTableState::RestoreMetadata => {
+ return self
+ .on_restore_metadata()
+ .await
+ .map_err(Self::map_restore_error);
+ }
UndropTableState::OpenRegions => self.on_open_regions().await,
UndropTableState::InvalidateTableCache => self.on_broadcast().await,
}
@@ -606,7 +625,8 @@ mod tests {
use api::v1::region::region_request;
use async_trait::async_trait;
use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
- use common_error::ext::ErrorExt;
+ use common_error::ext::{BoxedError, ErrorExt};
+ use common_error::mock::MockError;
use common_error::status_code::StatusCode;
use store_api::storage::RegionId;
use table::table_name::TableName;
@@ -635,6 +655,38 @@ mod tests {
}
}
+ #[test]
+ fn test_restore_metadata_retry_policy() {
+ let source = error::Error::retry_later(MockError::new(StatusCode::Internal));
+ let err = Err::<(), _>(BoxedError::new(source))
+ .context(error::ExternalSnafu)
+ .unwrap_err();
+ assert!(err.retry_hint().is_retryable());
+ assert!(!err.is_retry_later());
+ let err = UndropTableProcedure::map_restore_error(err);
+ assert!(!err.is_retry_later());
+ assert!(!err.need_clean_poisons());
+
+ for retryable in [false, true] {
+ for clean_poisons in [false, true] {
+ let source = BoxedError::new(MockError::new(StatusCode::Internal));
+ let err = if retryable {
+ error::Error::RetryLater {
+ source,
+ clean_poisons,
+ }
+ } else {
+ Err::<(), _>(source)
+ .context(error::AbortProcedureSnafu { clean_poisons })
+ .unwrap_err()
+ };
+ let err = UndropTableProcedure::map_restore_error(err);
+ assert_eq!(err.is_retry_later(), retryable);
+ assert_eq!(err.need_clean_poisons(), clean_poisons);
+ }
+ }
+ }
+
#[test]
fn test_map_restore_metadata_error_without_table_name() {
let context = new_ddl_context(Arc::new(MockDatanodeManager::new(())));
diff --git a/src/common/meta/src/ddl/utils.rs b/src/common/meta/src/ddl/utils.rs
index ebfbd6fd522..a8058d89b94 100644
--- a/src/common/meta/src/ddl/utils.rs
+++ b/src/common/meta/src/ddl/utils.rs
@@ -28,7 +28,7 @@ use api::v1::region::{
region_request,
};
use common_catalog::consts::{METRIC_ENGINE, MITO_ENGINE};
-use common_error::ext::BoxedError;
+use common_error::ext::{BoxedError, ErrorExt};
use common_procedure::error::Error as ProcedureError;
use common_telemetry::tracing_context::TracingContext;
use common_telemetry::{error, info, warn};
@@ -61,7 +61,7 @@ use crate::wal_provider::RegionWalOptions;
pub fn add_peer_context_if_needed(datanode: Peer) -> impl FnOnce(Error) -> Error {
move |err| {
error!(err; "Failed to operate datanode, peer: {}", datanode);
- if !err.is_retry_later() {
+ if !err.retry_hint().is_retryable() {
return Err::<(), BoxedError>(BoxedError::new(err))
.context(OperateDatanodeSnafu { peer: datanode })
.unwrap_err();
@@ -84,7 +84,7 @@ pub(crate) fn is_metric_engine_logical_table(
/// This function determines whether the error should be retried and if poison cleanup is needed,
/// then maps it to the appropriate procedure error variant.
pub fn map_to_procedure_error(e: Error) -> ProcedureError {
- match (e.is_retry_later(), e.need_clean_poisons()) {
+ match (e.retry_hint().is_retryable(), e.need_clean_poisons()) {
(true, true) => ProcedureError::retry_later_and_clean_poisons(e),
(true, false) => ProcedureError::retry_later(e),
(false, true) => ProcedureError::external_and_clean_poisons(e),
@@ -256,7 +256,7 @@ pub fn handle_multiple_results(results: Vec>) -> MultipleRes
match result {
Ok(value) => ok_results.push(value),
Err(err) => {
- if err.is_retry_later() {
+ if err.retry_hint().is_retryable() {
retryable_results.push(err);
} else {
non_retryable_results.push(err);
@@ -475,8 +475,74 @@ pub fn extract_column_metadatas(
#[cfg(test)]
mod tests {
+ use common_error::mock::MockError;
+ use common_error::status_code::StatusCode;
+
use super::*;
+ fn external_error(retryable: bool) -> Error {
+ let source = MockError::new(StatusCode::Internal);
+ let source = if retryable {
+ BoxedError::new(Error::retry_later(source))
+ } else {
+ BoxedError::new(source)
+ };
+ Err::<(), _>(source)
+ .context(error::ExternalSnafu)
+ .unwrap_err()
+ }
+
+ #[test]
+ fn test_map_to_procedure_error() {
+ for retryable in [false, true] {
+ for clean_poisons in [false, true] {
+ let err = if retryable {
+ Error::RetryLater {
+ source: BoxedError::new(external_error(false)),
+ clean_poisons,
+ }
+ } else {
+ Err::<(), _>(BoxedError::new(external_error(true)))
+ .context(error::AbortProcedureSnafu { clean_poisons })
+ .unwrap_err()
+ };
+ let err = map_to_procedure_error(err);
+ assert_eq!(err.is_retry_later(), retryable);
+ assert_eq!(err.need_clean_poisons(), clean_poisons);
+ }
+ let err = map_to_procedure_error(external_error(retryable));
+ assert_eq!(err.is_retry_later(), retryable);
+ assert!(!err.need_clean_poisons());
+ }
+ }
+
+ #[test]
+ fn test_handle_multiple_results_retry_hint() {
+ assert!(matches!(
+ handle_multiple_results(vec![Ok(()), Err(external_error(true))]),
+ MultipleResults::PartialRetryable(_)
+ ));
+ assert!(matches!(
+ handle_multiple_results::<()>(vec![
+ Err(external_error(true)),
+ Err(external_error(true))
+ ]),
+ MultipleResults::AllRetryable(_)
+ ));
+ assert!(matches!(
+ handle_multiple_results(vec![
+ Ok(()),
+ Err(external_error(true)),
+ Err(external_error(false))
+ ]),
+ MultipleResults::PartialNonRetryable(_)
+ ));
+ assert!(matches!(
+ handle_multiple_results::<()>(vec![Err(external_error(false))]),
+ MultipleResults::AllNonRetryable(_)
+ ));
+ }
+
#[test]
fn test_get_catalog_and_schema() {
let test_catalog = "my_catalog";
diff --git a/src/common/meta/src/error.rs b/src/common/meta/src/error.rs
index 297af792799..adfa158adf0 100644
--- a/src/common/meta/src/error.rs
+++ b/src/common/meta/src/error.rs
@@ -1373,15 +1373,15 @@ impl ErrorExt for Error {
| StartProcedureManager { source, .. }
| StopProcedureManager { source, .. }
| RegisterProcedureLoader { source, .. }
- | PutPoison { source, .. }
- | ProcedureStateReceiver { source, .. } => source.retry_hint(),
+ | PutPoison { source, .. } => source.retry_hint(),
External { source, .. }
| ResponseExceededSizeLimit { source, .. }
| OperateDatanode { source, .. }
- | AbortProcedure { source, .. }
| RegisterRepartitionProcedureLoader { source, .. }
| CreateRepartitionProcedure { source, .. }
| PersistRepartitionGcRequirement { source, .. } => source.retry_hint(),
+ // These errors end an operation; their sources cannot make it retryable.
+ AbortProcedure { .. } | ProcedureStateReceiver { .. } => RetryHint::NonRetryable,
Table { source, .. } => source.retry_hint(),
ConvertAlterTableRequest { source, .. } => source.retry_hint(),
ConvertColumnDef { source, .. } => source.retry_hint(),
@@ -1466,6 +1466,7 @@ mod retry_hint_tests {
use std::sync::Arc;
use common_error::mock::MockError;
+ use snafu::ResultExt;
use super::*;
@@ -1483,6 +1484,42 @@ mod retry_hint_tests {
assert_eq!(err.retry_hint(), RetryHint::Retryable);
}
+ #[test]
+ fn test_leader_retry_hints() {
+ for (error, expected) in [
+ (
+ NoLeaderSnafu { table_id: 1_u32 }.build(),
+ RetryHint::Retryable,
+ ),
+ (ElectionNoLeaderSnafu.build(), RetryHint::Retryable),
+ (
+ ElectionLeaderLeaseExpiredSnafu.build(),
+ RetryHint::Retryable,
+ ),
+ (
+ ElectionLeaderLeaseChangedSnafu.build(),
+ RetryHint::Retryable,
+ ),
+ ] {
+ assert_eq!(error.retry_hint(), expected, "{error:?}");
+ }
+ }
+
+ #[test]
+ fn test_procedure_error_retry_hints() {
+ let source = common_procedure::Error::retry_later(MockError::new(StatusCode::Internal));
+ let err = Err::<(), _>(source).context(PutPoisonSnafu).unwrap_err();
+ assert_eq!(err.retry_hint(), RetryHint::Retryable);
+
+ let source = common_procedure::Error::retry_later(MockError::new(StatusCode::Internal));
+ let err = Err::<(), _>(source)
+ .context(ProcedureStateReceiverSnafu {
+ procedure_id: ProcedureId::random(),
+ })
+ .unwrap_err();
+ assert_eq!(err.retry_hint(), RetryHint::NonRetryable);
+ }
+
#[test]
fn test_get_cache_forwards_retry_hint() {
let source = Arc::new(Error::retry_later(MockError::new(StatusCode::Internal)));
diff --git a/src/common/meta/src/node_manager.rs b/src/common/meta/src/node_manager.rs
index 6cc6c3e1539..d18e2cc8471 100644
--- a/src/common/meta/src/node_manager.rs
+++ b/src/common/meta/src/node_manager.rs
@@ -28,12 +28,22 @@ use crate::peer::Peer;
#[async_trait::async_trait]
pub trait Datanode: Send + Sync {
/// Handles DML, and DDL requests.
- async fn handle(&self, request: RegionRequest) -> Result;
+ async fn handle_request(&self, request: Arc) -> Result;
/// Handles query requests
async fn handle_query(&self, request: QueryRequest) -> Result;
}
+impl dyn Datanode + '_ {
+ /// Handles an owned or shared request without cloning its payload.
+ pub async fn handle(&self, request: R) -> Result
+ where
+ R: Into> + Send,
+ {
+ self.handle_request(request.into()).await
+ }
+}
+
pub type DatanodeRef = Arc;
/// The trait for handling requests to flownode
diff --git a/src/common/meta/src/reconciliation/reconcile_catalog.rs b/src/common/meta/src/reconciliation/reconcile_catalog.rs
index e28e81d323d..f015a483b6e 100644
--- a/src/common/meta/src/reconciliation/reconcile_catalog.rs
+++ b/src/common/meta/src/reconciliation/reconcile_catalog.rs
@@ -16,6 +16,7 @@ use std::any::Any;
use std::fmt::Debug;
use std::time::Instant;
+use common_error::ext::ErrorExt;
use common_procedure::error::FromJsonSnafu;
use common_procedure::{
Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey,
@@ -191,7 +192,7 @@ impl Procedure for ReconcileCatalogProcedure {
Ok(status)
}
Err(e) => {
- if e.is_retry_later() {
+ if e.retry_hint().is_retryable() {
metrics::METRIC_META_RECONCILIATION_PROCEDURE_ERROR
.with_label_values(&[procedure_name, step, metrics::ERROR_TYPE_RETRYABLE])
.inc();
diff --git a/src/common/meta/src/reconciliation/reconcile_database.rs b/src/common/meta/src/reconciliation/reconcile_database.rs
index f5e2fe9f9a0..49443aca9db 100644
--- a/src/common/meta/src/reconciliation/reconcile_database.rs
+++ b/src/common/meta/src/reconciliation/reconcile_database.rs
@@ -23,6 +23,7 @@ use std::fmt::Debug;
use std::time::Instant;
use async_trait::async_trait;
+use common_error::ext::ErrorExt;
use common_procedure::error::{FromJsonSnafu, ToJsonSnafu};
use common_procedure::{
Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey,
@@ -231,7 +232,7 @@ impl Procedure for ReconcileDatabaseProcedure {
Ok(status)
}
Err(e) => {
- if e.is_retry_later() {
+ if e.retry_hint().is_retryable() {
metrics::METRIC_META_RECONCILIATION_PROCEDURE_ERROR
.with_label_values(&[procedure_name, step, metrics::ERROR_TYPE_RETRYABLE])
.inc();
diff --git a/src/common/meta/src/reconciliation/reconcile_logical_tables.rs b/src/common/meta/src/reconciliation/reconcile_logical_tables.rs
index 8c16e5a450f..5c5c654f6df 100644
--- a/src/common/meta/src/reconciliation/reconcile_logical_tables.rs
+++ b/src/common/meta/src/reconciliation/reconcile_logical_tables.rs
@@ -22,6 +22,7 @@ use std::any::Any;
use std::fmt::Debug;
use async_trait::async_trait;
+use common_error::ext::ErrorExt;
use common_procedure::error::{FromJsonSnafu, ToJsonSnafu};
use common_procedure::{
Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey,
@@ -207,7 +208,7 @@ impl Procedure for ReconcileLogicalTablesProcedure {
Ok(status)
}
Err(e) => {
- if e.is_retry_later() {
+ if e.retry_hint().is_retryable() {
metrics::METRIC_META_RECONCILIATION_PROCEDURE_ERROR
.with_label_values(&[procedure_name, step, metrics::ERROR_TYPE_RETRYABLE])
.inc();
diff --git a/src/common/meta/src/reconciliation/reconcile_table.rs b/src/common/meta/src/reconciliation/reconcile_table.rs
index d48999b895f..a14fe13a3db 100644
--- a/src/common/meta/src/reconciliation/reconcile_table.rs
+++ b/src/common/meta/src/reconciliation/reconcile_table.rs
@@ -20,6 +20,7 @@ pub(crate) mod update_table_info;
use std::fmt::Debug;
+use common_error::ext::ErrorExt;
use common_procedure::error::{FromJsonSnafu, ToJsonSnafu};
use common_procedure::{
Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey,
@@ -318,7 +319,7 @@ impl Procedure for ReconcileTableProcedure {
Ok(status)
}
Err(e) => {
- if e.is_retry_later() {
+ if e.retry_hint().is_retryable() {
metrics::METRIC_META_RECONCILIATION_PROCEDURE_ERROR
.with_label_values(&[procedure_name, step, metrics::ERROR_TYPE_RETRYABLE])
.inc();
diff --git a/src/common/meta/src/test_util.rs b/src/common/meta/src/test_util.rs
index d78dc815d09..618603c13e5 100644
--- a/src/common/meta/src/test_util.rs
+++ b/src/common/meta/src/test_util.rs
@@ -111,7 +111,8 @@ struct MockNode {
#[async_trait::async_trait]
impl Datanode for MockNode {
- async fn handle(&self, request: RegionRequest) -> Result {
+ async fn handle_request(&self, request: Arc) -> Result {
+ let request = Arc::try_unwrap(request).expect("mock request must be exclusively owned");
self.handler.handle(&self.peer, request).await
}
diff --git a/src/common/procedure/src/error.rs b/src/common/procedure/src/error.rs
index 67fab54b156..b3f4a83d5f4 100644
--- a/src/common/procedure/src/error.rs
+++ b/src/common/procedure/src/error.rs
@@ -306,7 +306,8 @@ impl ErrorExt for Error {
| Error::DeletePoison { source, .. }
| Error::GetPoison { source, .. }
| Error::CheckStatus { source, .. } => source.retry_hint(),
- Error::ProcedureExec { source, .. } => source.retry_hint(),
+ // The watcher reports this only after the child has failed or been poisoned.
+ Error::ProcedureExec { .. } => RetryHint::NonRetryable,
Error::StartRemoveOutdatedMetaTask { source, .. }
| Error::StopRemoveOutdatedMetaTask { source, .. } => source.retry_hint(),
Error::DeleteState { error, .. } => retry_hint_from_opendal_error(error),
@@ -379,8 +380,11 @@ mod tests {
use std::sync::Arc;
use common_error::mock::MockError;
+ use tokio::sync::watch;
use super::*;
+ use crate::procedure::ProcedureState;
+ use crate::watcher;
#[test]
fn test_retry_later_hint_is_retryable() {
@@ -397,6 +401,21 @@ mod tests {
assert_eq!(err.retry_hint(), RetryHint::Retryable);
}
+ #[tokio::test]
+ async fn test_finished_procedure_hint_is_non_retryable() {
+ let source = Arc::new(Error::retry_later(MockError::new(StatusCode::Internal)));
+ for state in [
+ ProcedureState::failed(source.clone()),
+ ProcedureState::poisoned(Default::default(), source),
+ ] {
+ let (sender, mut receiver) = watch::channel(ProcedureState::Running);
+ sender.send(state).unwrap();
+ let err = watcher::wait(&mut receiver).await.unwrap_err();
+
+ assert_eq!(err.retry_hint(), RetryHint::NonRetryable);
+ }
+ }
+
#[test]
fn test_retry_exceeded_hint_is_non_retryable() {
let source = Arc::new(Error::retry_later(MockError::new(StatusCode::Internal)));
diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md
index fdaa5dbcfe4..8d858242e8c 100644
--- a/src/frontend/AGENTS.md
+++ b/src/frontend/AGENTS.md
@@ -48,7 +48,9 @@ including caches owned by pipeline and protocol handlers.
- **Insert** (`instance/grpc.rs`): `handle_inserts` / `handle_row_inserts` →
`check_permission` → `operator`'s `Inserter` (schema validation, optional
auto-create, partition routing, meter admission) → local `RegionServer`
- (standalone) or RPC to datanodes (distributed). Arrow bulk inserts pass the
+ (standalone) or RPC to datanodes (distributed). `FrontendBuilder` configures
+ the operator write dispatcher with frontend retry options for distributed
+ inserts, deletes, and bulk writes. Arrow bulk inserts pass the
request channel to `Inserter` and check meter admission for each nonempty batch.
- Finite ingestion requests split internally admit their total rows per database
before dispatch (`operator::insert::admit_write` / `admit_row_insert_batches`).
@@ -113,7 +115,8 @@ cargo nextest run -p frontend
`servers`; permissions, planning, and routing live here.
- Standalone vs distributed diverge in datanode access (the `standalone` crate's
local `RegionServer` adapter vs `NodeClients` RPC), MetaClient usage, and
- whether heartbeat matters. In standalone, the cache invalidator is a no-op.
+ whether heartbeat matters. The builder defaults local invalidation to its
+ layered cache registry; callers may supply a custom invalidator.
## Maintenance contract
diff --git a/src/frontend/src/error.rs b/src/frontend/src/error.rs
index 9fd2312e736..451040c884a 100644
--- a/src/frontend/src/error.rs
+++ b/src/frontend/src/error.rs
@@ -323,6 +323,13 @@ pub enum Error {
name: String,
},
+ #[snafu(display("Invalid write retry config: {}", reason))]
+ InvalidWriteRetryConfig {
+ reason: &'static str,
+ #[snafu(implicit)]
+ location: Location,
+ },
+
#[snafu(display("Invalid tls config"))]
InvalidTlsConfig {
#[snafu(source)]
@@ -391,7 +398,8 @@ impl ErrorExt for Error {
| Error::IllegalAuthConfig { .. }
| Error::ColumnNoneDefaultValue { .. }
| Error::IncompleteGrpcRequest { .. }
- | Error::InvalidTlsConfig { .. } => StatusCode::InvalidArguments,
+ | Error::InvalidTlsConfig { .. }
+ | Error::InvalidWriteRetryConfig { .. } => StatusCode::InvalidArguments,
Error::NotSupported { .. } => StatusCode::Unsupported,
diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs
index e594ba9f7b1..2ccbc3915e5 100644
--- a/src/frontend/src/frontend.rs
+++ b/src/frontend/src/frontend.rs
@@ -36,9 +36,10 @@ use crate::heartbeat::HeartbeatTask;
use crate::instance::Instance;
use crate::service_config::{
InfluxdbOptions, JaegerOptions, MysqlOptions, OpentsdbOptions, OtlpOptions,
- PendingRowsBatcherOptions, PostgresOptions, PromStoreOptions,
+ PendingRowsBatcherOptions, PostgresOptions, PromStoreOptions, WriteRetryOptions,
};
+/// Configures frontend protocols, query execution, and distributed write retries.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
#[serde(default)]
pub struct FrontendOptions {
@@ -74,6 +75,8 @@ pub struct FrontendOptions {
pub meta_client: Option,
pub logging: LoggingOptions,
pub datanode: DatanodeClientOptions,
+ /// Retry policy for retryable distributed write failures.
+ pub write_retry: WriteRetryOptions,
pub user_provider: Option,
pub tracing: TracingOptions,
pub query: QueryOptions,
@@ -110,6 +113,7 @@ impl Default for FrontendOptions {
meta_client: None,
logging: LoggingOptions::default(),
datanode: DatanodeClientOptions::default(),
+ write_retry: WriteRetryOptions::default(),
user_provider: None,
tracing: TracingOptions::default(),
query: QueryOptions::default(),
@@ -225,6 +229,72 @@ mod tests {
type GrpcStream =
Pin> + Send + Sync + 'static>>;
+ #[test]
+ fn test_write_retry_config() {
+ let defaults: FrontendOptions = toml::from_str("").unwrap();
+ assert_eq!(defaults.write_retry, WriteRetryOptions::default());
+ let options: FrontendOptions = toml::from_str(
+ r#"
+[write_retry]
+max_retries = 5
+initial_backoff = "250ms"
+max_backoff = "5s"
+backoff_multiplier = 3.0
+"#,
+ )
+ .unwrap();
+ assert_eq!(
+ options.write_retry,
+ WriteRetryOptions {
+ max_retries: 5,
+ initial_backoff: Duration::from_millis(250),
+ max_backoff: Duration::from_secs(5),
+ backoff_multiplier: 3.0,
+ }
+ );
+ options.write_retry.validate().unwrap();
+ let serialized = toml::to_string(&options).unwrap();
+ let parsed: FrontendOptions = toml::from_str(&serialized).unwrap();
+ assert_eq!(options.write_retry, parsed.write_retry);
+ let disabled: FrontendOptions = toml::from_str("[write_retry]\nmax_retries = 0").unwrap();
+ assert_eq!(disabled.write_retry.max_retries, 0);
+ assert_eq!(
+ disabled.write_retry.initial_backoff,
+ defaults.write_retry.initial_backoff
+ );
+ }
+
+ #[test]
+ fn test_write_retry_config_from_env() {
+ temp_env::with_vars(
+ [
+ ("FRONTEND_RETRY_TEST__WRITE_RETRY__MAX_RETRIES", Some("4")),
+ (
+ "FRONTEND_RETRY_TEST__WRITE_RETRY__INITIAL_BACKOFF",
+ Some("500ms"),
+ ),
+ ("FRONTEND_RETRY_TEST__WRITE_RETRY__MAX_BACKOFF", Some("8s")),
+ (
+ "FRONTEND_RETRY_TEST__WRITE_RETRY__BACKOFF_MULTIPLIER",
+ Some("1.5"),
+ ),
+ ],
+ || {
+ let options =
+ FrontendOptions::load_layered_options(None, "FRONTEND_RETRY_TEST").unwrap();
+ assert_eq!(
+ options.write_retry,
+ WriteRetryOptions {
+ max_retries: 4,
+ initial_backoff: Duration::from_millis(500),
+ max_backoff: Duration::from_secs(8),
+ backoff_multiplier: 1.5,
+ }
+ );
+ },
+ );
+ }
+
#[test]
fn test_logical_batcher_protocols_from_env() {
temp_env::with_vars(
diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs
index 9db7b70c221..f93113d3426 100644
--- a/src/frontend/src/instance/builder.rs
+++ b/src/frontend/src/instance/builder.rs
@@ -25,8 +25,10 @@ use common_base::Plugins;
use common_batcher::request_limiter::RequestLimiter;
use common_datasource::object_store::LocalFileAccess;
use common_event_recorder::{EventRecorderImpl, EventRecorderRef};
-use common_meta::cache::{LayeredCacheRegistryRef, TableFlownodeSetCacheRef, TableRouteCacheRef};
-use common_meta::cache_invalidator::{CacheInvalidatorRef, DummyCacheInvalidator};
+use common_meta::cache::{
+ LayeredCacheRegistryRef, TableFlownodeSetCacheRef, TableInfoCacheRef, TableRouteCacheRef,
+};
+use common_meta::cache_invalidator::CacheInvalidatorRef;
use common_meta::key::TableMetadataManager;
use common_meta::key::flow::FlowMetadataManager;
use common_meta::kv_backend::KvBackendRef;
@@ -46,6 +48,7 @@ use operator::statement::{
StatementExecutor, StatementExecutorConfiguratorRef, StatementExecutorRef,
};
use operator::table::TableMutationOperator;
+use operator::write_dispatcher::WriteDispatcher;
use partition::cache::PartitionInfoCacheRef;
use partition::manager::PartitionRuleManager;
use pipeline::pipeline_operator::PipelineOperator;
@@ -70,7 +73,7 @@ pub struct FrontendBuilder {
options: FrontendOptions,
kv_backend: KvBackendRef,
layered_cache_registry: LayeredCacheRegistryRef,
- local_cache_invalidator: Option,
+ local_cache_invalidator: CacheInvalidatorRef,
catalog_manager: CatalogManagerRef,
node_manager: NodeManagerRef,
plugins: Option,
@@ -80,6 +83,9 @@ pub struct FrontendBuilder {
}
impl FrontendBuilder {
+ /// Creates a frontend builder from its configuration and shared services.
+ ///
+ /// Uses the layered cache registry for local invalidation unless overridden.
#[allow(clippy::too_many_arguments)]
pub fn new(
options: FrontendOptions,
@@ -93,8 +99,8 @@ impl FrontendBuilder {
Self {
options,
kv_backend,
+ local_cache_invalidator: layered_cache_registry.clone(),
layered_cache_registry,
- local_cache_invalidator: None,
catalog_manager,
node_manager,
plugins: None,
@@ -140,9 +146,10 @@ impl FrontendBuilder {
)
}
+ /// Overrides the default local cache invalidator.
pub fn with_local_cache_invalidator(self, cache_invalidator: CacheInvalidatorRef) -> Self {
Self {
- local_cache_invalidator: Some(cache_invalidator),
+ local_cache_invalidator: cache_invalidator,
..self
}
}
@@ -187,7 +194,11 @@ impl FrontendBuilder {
}
}
+ /// Validates configuration and assembles the frontend [`Instance`].
pub async fn try_build(self) -> Result {
+ if let Err(reason) = self.options.write_retry.validate() {
+ return error::InvalidWriteRetryConfigSnafu { reason }.fail();
+ }
let kv_backend = self.kv_backend;
let node_manager = self.node_manager;
let plugins = self.plugins.unwrap_or_default();
@@ -210,9 +221,24 @@ impl FrontendBuilder {
partition_info_cache.clone(),
));
- let local_cache_invalidator = self
- .local_cache_invalidator
- .unwrap_or_else(|| Arc::new(DummyCacheInvalidator));
+ let local_cache_invalidator = self.local_cache_invalidator;
+
+ let mut write_dispatcher =
+ WriteDispatcher::new(node_manager.clone(), partition_manager.clone());
+ if self.options.write_retry.max_retries > 0 {
+ let table_info_cache: TableInfoCacheRef =
+ self.layered_cache_registry
+ .get()
+ .context(error::CacheRequiredSnafu {
+ name: cache::TABLE_INFO_CACHE_NAME,
+ })?;
+ write_dispatcher = write_dispatcher.with_retry(
+ self.options.write_retry.clone(),
+ table_info_cache,
+ local_cache_invalidator.clone(),
+ );
+ }
+ let write_dispatcher = Arc::new(write_dispatcher);
let region_query_handler =
if let Some(factory) = plugins.get::() {
@@ -236,6 +262,7 @@ impl FrontendBuilder {
table_flownode_cache.clone(),
self.options.auto_create_table,
)
+ .with_write_dispatcher(write_dispatcher.clone())
};
// The execution-only inserter owns no batchers, avoiding an Arc cycle.
let bulk_inserter = Arc::new(create_inserter());
@@ -264,11 +291,14 @@ impl FrontendBuilder {
create_inserter()
.with_pending_rows_batcher(build_batcher(self.options.table_batcher_options())),
);
- let deleter = Arc::new(Deleter::new(
- self.catalog_manager.clone(),
- partition_manager.clone(),
- node_manager.clone(),
- ));
+ let deleter = Arc::new(
+ Deleter::new(
+ self.catalog_manager.clone(),
+ partition_manager.clone(),
+ node_manager.clone(),
+ )
+ .with_write_dispatcher(write_dispatcher),
+ );
let requester = Arc::new(Requester::new(
self.catalog_manager.clone(),
partition_manager.clone(),
@@ -421,3 +451,76 @@ impl FrontendBuilder {
})
}
}
+
+#[cfg(test)]
+mod tests {
+ use api::v1::meta::Role;
+ use common_meta::cache_invalidator::{Context, DummyCacheInvalidator};
+ use common_meta::instruction::CacheIdent;
+ use common_meta::key::table_route::{TableRouteKey, TableRouteValue};
+ use common_meta::key::{MetadataKey, MetadataValue};
+ use common_meta::rpc::router::{Region, RegionRoute};
+ use common_meta::rpc::store::PutRequest;
+ use meta_client::client::MetaClientBuilder;
+ use store_api::storage::RegionId;
+
+ use super::*;
+
+ #[tokio::test]
+ async fn test_local_cache_invalidator() {
+ for custom_invalidator in [false, true] {
+ let mut builder = FrontendBuilder::new_test(
+ &FrontendOptions::default(),
+ Arc::new(MetaClientBuilder::new(0, Role::Frontend).build()),
+ );
+ if custom_invalidator {
+ builder = builder.with_local_cache_invalidator(Arc::new(DummyCacheInvalidator));
+ }
+ let backend = builder.kv_backend.clone();
+ let routes: TableRouteCacheRef = builder.layered_cache_registry.get().unwrap();
+ let partitions: PartitionInfoCacheRef = builder.layered_cache_registry.get().unwrap();
+ let instance = builder.try_build().await.unwrap();
+ let table_id = 1024;
+ for region_number in [1, 2] {
+ let route = TableRouteValue::physical(vec![RegionRoute {
+ region: Region {
+ id: RegionId::new(table_id, region_number),
+ ..Default::default()
+ },
+ ..Default::default()
+ }]);
+ backend
+ .put(
+ PutRequest::new()
+ .with_key(TableRouteKey::new(table_id).to_bytes())
+ .with_value(route.try_as_raw_value().unwrap()),
+ )
+ .await
+ .unwrap();
+ if region_number == 1 {
+ routes.get(table_id).await.unwrap().unwrap();
+ partitions.get(table_id).await.unwrap().unwrap();
+ }
+ }
+ assert!(routes.contains_key(&table_id));
+ assert!(partitions.contains_key(&table_id));
+
+ instance
+ .cache_invalidator()
+ .invalidate(&Context::default(), &[CacheIdent::TableId(table_id)])
+ .await
+ .unwrap();
+
+ assert_eq!(custom_invalidator, routes.contains_key(&table_id));
+ assert_eq!(custom_invalidator, partitions.contains_key(&table_id));
+ let route = routes.get(table_id).await.unwrap().unwrap();
+ let region = route.as_physical_table_route_ref().unwrap().region_routes[0]
+ .region
+ .id;
+ assert_eq!(
+ if custom_invalidator { 1 } else { 2 },
+ region.region_number()
+ );
+ }
+ }
+}
diff --git a/src/frontend/src/instance/logical_batcher.rs b/src/frontend/src/instance/logical_batcher.rs
index 255766ac14d..423f2996cae 100644
--- a/src/frontend/src/instance/logical_batcher.rs
+++ b/src/frontend/src/instance/logical_batcher.rs
@@ -48,6 +48,7 @@ impl Instance {
LogicalTablePendingRowsBatcher::try_new(
self.partition_manager().clone(),
self.node_manager().clone(),
+ self.inserter.write_dispatcher().clone(),
self.catalog_manager().clone(),
self.table_flownode_set_cache().clone(),
true,
diff --git a/src/frontend/src/service_config.rs b/src/frontend/src/service_config.rs
index 2c761c4242a..535ae9b655b 100644
--- a/src/frontend/src/service_config.rs
+++ b/src/frontend/src/service_config.rs
@@ -25,6 +25,7 @@ pub use influxdb::{InfluxdbMergeMode, InfluxdbOptions};
pub use jaeger::JaegerOptions;
pub use mysql::MysqlOptions;
pub use opentsdb::OpentsdbOptions;
+pub use operator::write_dispatcher::WriteRetryOptions;
pub use otlp::OtlpOptions;
pub use pending_rows_batcher::{BatcherOptions, PendingRowsBatcherOptions};
pub use postgres::PostgresOptions;
diff --git a/src/metric-engine/src/engine/bulk_insert.rs b/src/metric-engine/src/engine/bulk_insert.rs
index dd20f34adf8..b7ad7646e3c 100644
--- a/src/metric-engine/src/engine/bulk_insert.rs
+++ b/src/metric-engine/src/engine/bulk_insert.rs
@@ -354,17 +354,32 @@ mod tests {
#[tokio::test]
async fn test_bulk_insert_physical_region_passthrough() {
- check_bulk_insert_physical_region_passthrough(false, false).await;
- check_bulk_insert_physical_region_passthrough(true, false).await;
+ check_bulk_insert_physical_region_passthrough(false, false, false).await;
+ check_bulk_insert_physical_region_passthrough(true, false, false).await;
}
#[tokio::test]
async fn test_bulk_insert_physical_region_close_recovery() {
- check_bulk_insert_physical_region_passthrough(false, true).await;
- check_bulk_insert_physical_region_passthrough(true, true).await;
+ check_bulk_insert_physical_region_passthrough(false, true, false).await;
+ check_bulk_insert_physical_region_passthrough(true, true, false).await;
}
- async fn check_bulk_insert_physical_region_passthrough(skip_wal: bool, close: bool) {
+ #[tokio::test]
+ async fn test_physical_bulk_partition_columns_preserve_stored_rows() {
+ for skip_wal in [false, true] {
+ let stripped =
+ check_bulk_insert_physical_region_passthrough(skip_wal, true, false).await;
+ let retained =
+ check_bulk_insert_physical_region_passthrough(skip_wal, true, true).await;
+ assert_eq!(retained, stripped, "skip_wal={skip_wal}");
+ }
+ }
+
+ async fn check_bulk_insert_physical_region_passthrough(
+ skip_wal: bool,
+ close: bool,
+ retain_partition_columns: bool,
+ ) -> String {
// Use flat format so that BulkMemtable is used (supports write_bulk).
let mito_config = MitoConfig {
default_flat_format: true,
@@ -394,7 +409,11 @@ mod tests {
index: 2,
column_id: 2, // column_id for "job" in the physical table
}];
- let non_tag_indices = vec![0, 1]; // timestamp, value
+ let non_tag_indices = if retain_partition_columns {
+ vec![0, 1, 2] // timestamp, value, and the job partition column
+ } else {
+ vec![0, 1] // timestamp, value
+ };
let second_batch = build_logical_batch(3, 3);
let physical_batch = modify_batch_sparse(
second_batch,
@@ -403,6 +422,10 @@ mod tests {
&non_tag_indices,
)
.unwrap();
+ assert_eq!(
+ physical_batch.column_by_name("job").is_some(),
+ retain_partition_columns
+ );
let request = build_bulk_request(physical_region_id, physical_batch, skip_wal);
let response = env
.metric()
@@ -469,6 +492,7 @@ mod tests {
batches.iter().map(|b| b.num_rows()).sum::(),
if skip_wal && !close { 0 } else { 6 },
);
+ batches.pretty_print().unwrap()
}
#[tokio::test]
diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml
index 7f4233e633e..b89bf4b13d3 100644
--- a/src/operator/Cargo.toml
+++ b/src/operator/Cargo.toml
@@ -62,6 +62,7 @@ lazy_static.workspace = true
meta-client.workspace = true
meter-core.workspace = true
meter-macros.workspace = true
+mito-codec.workspace = true
moka = { workspace = true, features = ["future"] }
object-store.workspace = true
partition.workspace = true
diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs
index 29537959df5..d4a1ed0497c 100644
--- a/src/operator/src/bulk_insert.rs
+++ b/src/operator/src/bulk_insert.rs
@@ -35,6 +35,7 @@ use snafu::{OptionExt, ResultExt, ensure};
use store_api::storage::RegionId;
use table::TableRef;
use table::metadata::TableInfoRef;
+use tokio_util::sync::CancellationToken;
use crate::error::Result;
use crate::insert::Inserter;
@@ -72,7 +73,7 @@ impl Inserter {
let writes = self.write_dispatcher.prepare(writes).await?;
Ok(self
.write_dispatcher
- .execute(writes, None)
+ .execute(writes, None, None)
.await
.context(error::RequestInsertsSnafu)?
.affected_rows)
@@ -191,6 +192,7 @@ impl Inserter {
Some(
metrics::HANDLE_BULK_INSERT_ELAPSED.with_label_values(&["datanode_handle"]),
),
+ None,
)
.await
.context(error::RequestRegionSnafu)
@@ -217,6 +219,9 @@ impl Inserter {
.with_label_values(&["wait_all_datanode"])
.start_timer();
+ // Dropping the caller stops retries without abandoning in-flight writes.
+ let cancellation = CancellationToken::new();
+ let _guard = cancellation.clone().drop_guard();
let mut handles = Vec::with_capacity(region_masks.len());
for (region_number, mask) in region_masks {
@@ -328,6 +333,7 @@ impl Inserter {
&metrics::HANDLE_BULK_INSERT_ELAPSED
.with_label_values(&["datanode_handle"]),
),
+ Some(&cancellation),
)
.await
.context(error::RequestRegionSnafu)
diff --git a/src/operator/src/delete.rs b/src/operator/src/delete.rs
index c389374b7ac..9bf14471696 100644
--- a/src/operator/src/delete.rs
+++ b/src/operator/src/delete.rs
@@ -27,6 +27,7 @@ use session::context::QueryContextRef;
use snafu::{OptionExt, ResultExt, ensure};
use table::TableRef;
use table::requests::DeleteRequest as TableDeleteRequest;
+use tokio_util::sync::CancellationToken;
use crate::error::{
CatalogSnafu, InvalidDeleteRequestSnafu, JoinTaskSnafu, MissingTimeIndexColumnSnafu,
@@ -141,12 +142,16 @@ 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.execute(writes, None).await })
- .await
- .context(JoinTaskSnafu)?
- .context(RequestDeletesSnafu)?
- .affected_rows;
+ // Dropping the caller stops retries without abandoning in-flight writes.
+ let cancellation = CancellationToken::new();
+ let _guard = cancellation.clone().drop_guard();
+ let affected_rows = common_runtime::spawn_global(async move {
+ dispatcher.execute(writes, None, Some(&cancellation)).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);
diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs
index 2b382dbdb5f..3f67d4fc576 100644
--- a/src/operator/src/insert.rs
+++ b/src/operator/src/insert.rs
@@ -80,6 +80,7 @@ use table::requests::{
validate_semantic_option,
};
use table::table_reference::TableReference;
+use tokio_util::sync::CancellationToken;
use crate::batcher::PendingRowsBatcher;
use crate::error::{
@@ -841,12 +842,16 @@ 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.execute(writes, None).await })
- .await
- .context(JoinTaskSnafu)?
- .context(RequestInsertsSnafu)?
- .affected_rows;
+ // Dropping the caller stops retries without abandoning in-flight writes.
+ let cancellation = CancellationToken::new();
+ let _guard = cancellation.clone().drop_guard();
+ let affected_rows = common_runtime::spawn_global(async move {
+ dispatcher.execute(writes, None, Some(&cancellation)).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);
@@ -3599,9 +3604,9 @@ mod tests {
#[async_trait::async_trait]
impl common_meta::node_manager::Datanode for RecordingNode {
- async fn handle(
+ async fn handle_request(
&self,
- _request: api::v1::region::RegionRequest,
+ _request: Arc,
) -> common_meta::error::Result {
let rows = match _request.body.as_ref() {
Some(api::v1::region::region_request::Body::Inserts(requests)) => requests
diff --git a/src/operator/src/req_convert/common/partitioner.rs b/src/operator/src/req_convert/common/partitioner.rs
index f1c3eff4fe3..6df73d7d1b9 100644
--- a/src/operator/src/req_convert/common/partitioner.rs
+++ b/src/operator/src/req_convert/common/partitioner.rs
@@ -13,6 +13,7 @@
// limitations under the License.
use std::collections::HashMap;
+use std::sync::Arc;
use api::v1::region::{
BulkInsertRequest, DeleteRequest, DeleteRequests, InsertRequest, InsertRequests, RegionRequest,
@@ -21,12 +22,17 @@ use api::v1::region::{
use api::v1::{ArrowIpc, PartitionExprVersion, Rows};
use common_error::ext::BoxedError;
use common_grpc::flight::{self, FlightDecoder};
+use datatypes::arrow::array::{Array, BinaryArray, StringBuilder};
use datatypes::arrow::compute;
+use datatypes::arrow::datatypes::{DataType, Field, Schema};
use datatypes::arrow::record_batch::RecordBatch;
+use mito_codec::row_converter::sparse::{SparseOffsetsCache, SparsePrimaryKeyView};
use partition::PartitionRuleRef;
use partition::manager::PartitionRuleManager;
+use partition::partition::RegionMask;
use partition::splitter::RowSplitter;
-use snafu::{OptionExt, ResultExt};
+use snafu::{OptionExt, ResultExt, ensure};
+use store_api::storage::consts;
use store_api::storage::{RegionId, RegionNumber, TableId};
use table::metadata::TableInfo;
@@ -101,9 +107,11 @@ impl TablePartitioner {
///
/// Takes ownership so row payloads can move into their new regions. Requests whose
/// region and partition version still match retain their original payloads.
+ /// The table metadata comes from the same refresh that created this partitioner.
pub(crate) fn rebuild_write_requests(
&self,
request: RegionRequest,
+ table_info: &TableInfo,
) -> Result> {
let RegionRequest { header, body } = request;
let bodies = match body.context(error::UnexpectedSnafu {
@@ -111,7 +119,7 @@ impl TablePartitioner {
})? {
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)?,
+ Body::BulkInsert(bulk) => self.rebuild_bulk_insert_request(bulk, table_info)?,
_ => {
return error::UnexpectedSnafu {
violated: "Only writes can be rebuilt",
@@ -165,13 +173,23 @@ impl TablePartitioner {
pub fn partition_bulk_insert_requests(
&self,
batch: RecordBatch,
- mut original_ipc: Option,
+ original_ipc: Option,
skip_wal: bool,
) -> Result> {
let masks = self
.rule
.split_record_batch(&batch)
.context(SplitInsertSnafu)?;
+ self.build_bulk_insert_requests(batch, masks, original_ipc, skip_wal)
+ }
+
+ fn build_bulk_insert_requests(
+ &self,
+ batch: RecordBatch,
+ masks: HashMap,
+ mut original_ipc: Option,
+ skip_wal: bool,
+ ) -> Result> {
let mut requests = Vec::with_capacity(masks.len());
for (region, mask) in masks {
if mask.select_none() {
@@ -256,7 +274,11 @@ impl TablePartitioner {
Ok(bodies)
}
- fn rebuild_bulk_insert_request(&self, bulk: BulkInsertRequest) -> Result> {
+ fn rebuild_bulk_insert_request(
+ &self,
+ bulk: BulkInsertRequest,
+ table_info: &TableInfo,
+ ) -> Result> {
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)]);
@@ -269,13 +291,92 @@ impl TablePartitioner {
.and_then(|mut decoder| decoder.try_decode_record_batch(&ipc.data_header, &ipc.payload))
.map_err(BoxedError::new)
.context(error::ExternalSnafu)?;
+ let routing_batch = self.restore_partition_tags(&batch, table_info)?;
+ let masks = self
+ .rule
+ .split_record_batch(routing_batch.as_ref().unwrap_or(&batch))
+ .context(SplitInsertSnafu)?;
Ok(self
- .partition_bulk_insert_requests(batch, Some(ipc), bulk.skip_wal)?
+ .build_bulk_insert_requests(batch, masks, Some(ipc), bulk.skip_wal)?
.into_iter()
.map(Body::BulkInsert)
.collect())
}
+ /// Restores omitted physical-table tags for routing without changing the write payload.
+ fn restore_partition_tags(
+ &self,
+ batch: &RecordBatch,
+ table_info: &TableInfo,
+ ) -> Result> {
+ let meta = &table_info.meta;
+ if !table_info.is_physical_table() || meta.column_ids.len() != meta.schema.num_columns() {
+ return Ok(None);
+ }
+ let missing = self
+ .rule
+ .partition_columns()
+ .iter()
+ .filter(|name| batch.column_by_name(name).is_none())
+ .filter_map(|name| {
+ let index = meta.schema.column_index_by_name(name)?;
+ (meta.primary_key_indices.contains(&index)
+ && meta.schema.column_schemas()[index].data_type.is_string())
+ .then_some((name, meta.column_ids[index]))
+ })
+ .collect::>();
+ if missing.is_empty() {
+ return Ok(None);
+ }
+ let Some(primary_key) = batch.column_by_name(consts::PRIMARY_KEY_COLUMN_NAME) else {
+ return Ok(None);
+ };
+ let primary_key =
+ primary_key
+ .as_any()
+ .downcast_ref::()
+ .context(error::UnexpectedSnafu {
+ violated: "Sparse primary key column must be Binary",
+ })?;
+ let mut builders = missing
+ .iter()
+ .map(|_| StringBuilder::new())
+ .collect::>();
+ let mut cache = SparseOffsetsCache::new();
+ let mut buffer = Vec::new();
+ for row in 0..batch.num_rows() {
+ ensure!(
+ !primary_key.is_null(row),
+ error::UnexpectedSnafu {
+ violated: "Sparse primary key must not be null",
+ }
+ );
+ let mut key = SparsePrimaryKeyView::new(primary_key.value(row), &mut cache)
+ .map_err(BoxedError::new)
+ .context(error::ExternalSnafu)?;
+ for ((_, column_id), builder) in missing.iter().zip(&mut builders) {
+ let label = key
+ .label(*column_id, &mut buffer)
+ .map_err(BoxedError::new)
+ .context(error::ExternalSnafu)?;
+ builder.append_option(label);
+ }
+ }
+ let schema = batch.schema();
+ let mut fields = schema.fields().to_vec();
+ let mut columns = batch.columns().to_vec();
+ for ((name, _), mut builder) in missing.into_iter().zip(builders) {
+ fields.push(Arc::new(Field::new(name, DataType::Utf8, true)));
+ columns.push(Arc::new(builder.finish()));
+ }
+ RecordBatch::try_new(
+ Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())),
+ columns,
+ )
+ .map(Some)
+ .context(error::ComputeArrowSnafu)
+ }
+
/// 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
@@ -313,16 +414,305 @@ mod tests {
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::router::{Region, RegionRoute};
use common_meta::rpc::store::RangeResponse;
- use datatypes::arrow::array::Int32Array;
+ use datatypes::arrow::array::{Int32Array, StringArray};
use datatypes::arrow::datatypes::{DataType, Field, Schema};
use datatypes::arrow::record_batch::RecordBatch;
+ use datatypes::data_type::ConcreteDataType;
+ use datatypes::schema;
+ use datatypes::value::{Value as DataValue, ValueRef};
+ use mito_codec::row_converter::sparse;
+ use mito_codec::row_converter::{PrimaryKeyCodec, SparsePrimaryKeyCodec};
use moka::future::Cache;
+ use partition::expr;
use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef};
+ use partition::multi_dim::MultiDimPartitionRule;
+ use store_api::metric_engine_consts;
use super::*;
use crate::error::Error;
use crate::test_util;
+ use crate::test_util::{test_new_partition_info_cache, test_new_table_route_cache};
+
+ fn build_sparse_partitioner() -> TablePartitioner {
+ TablePartitioner {
+ table_id: 1,
+ rule: Arc::new(
+ MultiDimPartitionRule::try_new(
+ vec!["zone".to_string()],
+ vec![1, 2],
+ vec![
+ expr::col("zone").lt(DataValue::String("m".into())),
+ expr::col("zone").gt_eq(DataValue::String("m".into())),
+ ],
+ true,
+ )
+ .unwrap(),
+ ),
+ versions: [
+ (1, Some(PartitionExprVersion { value: 2 })),
+ (2, Some(PartitionExprVersion { value: 2 })),
+ ]
+ .into_iter()
+ .collect(),
+ }
+ }
+
+ fn build_sparse_table_info() -> TableInfo {
+ let mut table_info = test_util::new_test_table_info(1, "physical_tags", [1, 2].into_iter());
+ table_info.meta.schema = Arc::new(schema::Schema::new(vec![
+ schema::ColumnSchema::new("host", ConcreteDataType::string_datatype(), true),
+ schema::ColumnSchema::new("zone", ConcreteDataType::string_datatype(), true),
+ ]));
+ table_info.meta.primary_key_indices = vec![0, 1];
+ table_info.meta.partition_key_indices = vec![1];
+ table_info.meta.column_ids = vec![10, 20];
+ table_info.meta.options.extra_options.insert(
+ metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY.to_string(),
+ String::new(),
+ );
+ table_info
+ }
+
+ fn build_sparse_batch() -> RecordBatch {
+ let codec = SparsePrimaryKeyCodec::schemaless();
+ let keys = [Some("east"), Some("west"), Some(""), None, None]
+ .into_iter()
+ .enumerate()
+ .map(|(index, zone)| {
+ let mut key = Vec::new();
+ let values = [
+ (
+ sparse::RESERVED_COLUMN_ID_TABLE_ID,
+ ValueRef::UInt32(42 + index as u32),
+ ),
+ (
+ sparse::RESERVED_COLUMN_ID_TSID,
+ ValueRef::UInt64(index as u64),
+ ),
+ (10, ValueRef::String("node")),
+ ];
+ codec.encode_value_refs(&values, &mut key).unwrap();
+ if let Some(zone) = zone {
+ codec
+ .encode_raw_tag_value([(20, zone.as_bytes())].into_iter(), &mut key)
+ .unwrap();
+ }
+ if index == 4 {
+ key.extend_from_slice(&20_u32.to_be_bytes());
+ key.push(0);
+ }
+ key
+ })
+ .collect::>();
+ RecordBatch::try_new(
+ Arc::new(Schema::new(vec![
+ Field::new("host", DataType::Utf8, false),
+ Field::new(consts::PRIMARY_KEY_COLUMN_NAME, DataType::Binary, false),
+ ])),
+ vec![
+ Arc::new(StringArray::from(vec!["node"; keys.len()])),
+ Arc::new(BinaryArray::from_iter_values(keys.iter())),
+ ],
+ )
+ .unwrap()
+ }
+
+ #[tokio::test]
+ async fn test_partition_tags_use_physical_table_metadata() {
+ let backend = Arc::new(MemoryKvBackend::new());
+ let table_route_cache = test_new_table_route_cache(backend.clone());
+ let partition_info_cache = test_new_partition_info_cache(table_route_cache.clone());
+ let manager =
+ PartitionRuleManager::new(backend.clone(), table_route_cache, partition_info_cache);
+ let table_info = build_sparse_table_info();
+ let routes = [
+ expr::col("zone").lt(DataValue::String("m".into())),
+ expr::col("zone").gt_eq(DataValue::String("m".into())),
+ ]
+ .into_iter()
+ .enumerate()
+ .map(|(index, expr)| RegionRoute {
+ region: Region {
+ id: RegionId::new(1, index as u32 + 1),
+ partition_expr: expr.as_json_str().unwrap(),
+ ..Default::default()
+ },
+ ..Default::default()
+ })
+ .collect();
+ TableMetadataManager::new(backend)
+ .create_table_metadata(
+ table_info.clone(),
+ TableRouteValue::physical(routes),
+ Default::default(),
+ )
+ .await
+ .unwrap();
+ let batch = build_sparse_batch();
+ let partitioner = TablePartitioner::try_new(&manager, &table_info)
+ .await
+ .unwrap();
+ let restored = partitioner
+ .restore_partition_tags(&batch, &table_info)
+ .unwrap()
+ .unwrap();
+ let zones = restored
+ .column_by_name("zone")
+ .unwrap()
+ .as_any()
+ .downcast_ref::()
+ .unwrap();
+ assert_eq!(
+ zones.iter().collect::>(),
+ vec![Some("east"), Some("west"), Some(""), None, None]
+ );
+
+ for (physical, tag, valid_ids) in [
+ (false, true, true),
+ (true, false, true),
+ (true, true, false),
+ ] {
+ let mut info = table_info.clone();
+ if !physical {
+ info.meta
+ .options
+ .extra_options
+ .remove(metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY);
+ }
+ if !tag {
+ info.meta.primary_key_indices = vec![0];
+ }
+ if !valid_ids {
+ info.meta.column_ids.clear();
+ }
+ let partitioner = TablePartitioner::try_new(&manager, &info).await.unwrap();
+ assert!(
+ partitioner
+ .restore_partition_tags(&batch, &info)
+ .unwrap()
+ .is_none()
+ );
+ assert!(
+ partitioner
+ .partition_bulk_insert_requests(batch.clone(), None, false)
+ .is_err()
+ );
+ }
+ }
+
+ #[test]
+ fn test_rebuild_sparse_bulk_with_changed_partition_column() {
+ let partitioner = build_sparse_partitioner();
+ let table_info = build_sparse_table_info();
+ let batch = build_sparse_batch();
+ let restored = partitioner
+ .restore_partition_tags(&batch, &table_info)
+ .unwrap()
+ .unwrap();
+ let zones = restored
+ .column_by_name("zone")
+ .unwrap()
+ .as_any()
+ .downcast_ref::()
+ .unwrap();
+ assert_eq!(
+ zones.iter().collect::>(),
+ vec![Some("east"), Some("west"), Some(""), None, None]
+ );
+ let (schema, data_header, payload) = flight::record_batch_to_ipc(batch.clone()).unwrap();
+ let rebuilt = partitioner
+ .rebuild_write_requests(
+ build_request(Body::BulkInsert(BulkInsertRequest {
+ region_id: RegionId::new(1, 1).as_u64(),
+ partition_expr_version: Some(PartitionExprVersion { value: 1 }),
+ skip_wal: true,
+ body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc {
+ schema,
+ data_header,
+ payload,
+ })),
+ ..Default::default()
+ })),
+ &table_info,
+ )
+ .unwrap();
+ assert_eq!(rebuilt.len(), 2);
+ for request in rebuilt {
+ let Some(Body::BulkInsert(bulk)) = request.body else {
+ unreachable!()
+ };
+ assert!(bulk.skip_wal);
+ assert_eq!(
+ bulk.partition_expr_version,
+ Some(PartitionExprVersion { value: 2 })
+ );
+ let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = bulk.body else {
+ unreachable!()
+ };
+ let actual = FlightDecoder::try_from_schema_bytes(&ipc.schema)
+ .unwrap()
+ .try_decode_record_batch(&ipc.data_header, &ipc.payload)
+ .unwrap();
+ let mask = if RegionId::from_u64(bulk.region_id).region_number() == 1 {
+ vec![true, false, true, true, true]
+ } else {
+ vec![false, true, false, false, false]
+ };
+ let expected = compute::filter_record_batch(&batch, &mask.into()).unwrap();
+ assert_eq!(actual, expected);
+ }
+ }
+
+ #[test]
+ fn test_restore_partition_tags_checks_only_missing_tags() {
+ let partitioner = build_sparse_partitioner();
+ let table_info = build_sparse_table_info();
+ let batch = build_sparse_batch();
+ let invalid = RecordBatch::try_new(
+ batch.schema(),
+ vec![
+ batch.column(0).clone(),
+ Arc::new(BinaryArray::from_iter_values([b"invalid"; 5])),
+ ],
+ )
+ .unwrap();
+ assert!(
+ partitioner
+ .restore_partition_tags(&invalid, &table_info)
+ .is_err()
+ );
+ let mut fields = invalid.schema().fields().to_vec();
+ fields.push(Arc::new(Field::new("zone", DataType::Utf8, true)));
+ let mut columns = invalid.columns().to_vec();
+ columns.push(Arc::new(StringArray::from(vec![Some("east"); 5])));
+ let complete = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap();
+ assert!(
+ partitioner
+ .restore_partition_tags(&complete, &table_info)
+ .unwrap()
+ .is_none()
+ );
+ let (schema, data_header, payload) = flight::record_batch_to_ipc(invalid).unwrap();
+ let request = build_request(Body::BulkInsert(BulkInsertRequest {
+ region_id: RegionId::new(1, 1).as_u64(),
+ partition_expr_version: Some(PartitionExprVersion { value: 2 }),
+ body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc {
+ schema,
+ data_header,
+ payload,
+ })),
+ ..Default::default()
+ }));
+ let expected = request.clone();
+ assert_eq!(
+ partitioner
+ .rebuild_write_requests(request, &table_info)
+ .unwrap(),
+ vec![expected]
+ );
+ }
#[tokio::test]
async fn test_partition_requests_preserve_routing_errors() {
@@ -693,6 +1083,6 @@ mod tests {
let partitioner = TablePartitioner::try_new(manager, &table_info)
.await
.context(error::InvalidPartitionSnafu)?;
- partitioner.rebuild_write_requests(request)
+ partitioner.rebuild_write_requests(request, &table_info)
}
}
diff --git a/src/operator/src/test_util.rs b/src/operator/src/test_util.rs
index aa4a2f30ab8..436c1bbd7b5 100644
--- a/src/operator/src/test_util.rs
+++ b/src/operator/src/test_util.rs
@@ -17,3 +17,8 @@ mod partition_manager;
pub use crate::test_util::kv_backend::prepare_mocked_backend;
pub use crate::test_util::partition_manager::{create_partition_rule_manager, new_test_table_info};
+
+#[cfg(test)]
+pub(crate) use self::partition_manager::{
+ test_new_partition_info_cache, test_new_table_route_cache,
+};
diff --git a/src/operator/src/test_util/partition_manager.rs b/src/operator/src/test_util/partition_manager.rs
index f80e2cc4fe2..260252d2da5 100644
--- a/src/operator/src/test_util/partition_manager.rs
+++ b/src/operator/src/test_util/partition_manager.rs
@@ -113,7 +113,7 @@ fn new_test_region_wal_options(regions: Vec) -> RegionWalOptions {
HashMap::default()
}
-fn test_new_table_route_cache(kv_backend: KvBackendRef) -> TableRouteCacheRef {
+pub(crate) fn test_new_table_route_cache(kv_backend: KvBackendRef) -> TableRouteCacheRef {
let cache = CacheBuilder::new(128).build();
Arc::new(new_table_route_cache(
"table_route_cache".to_string(),
@@ -122,7 +122,9 @@ fn test_new_table_route_cache(kv_backend: KvBackendRef) -> TableRouteCacheRef {
))
}
-fn test_new_partition_info_cache(table_route_cache: TableRouteCacheRef) -> PartitionInfoCacheRef {
+pub(crate) fn test_new_partition_info_cache(
+ table_route_cache: TableRouteCacheRef,
+) -> PartitionInfoCacheRef {
let cache = CacheBuilder::new(128).build();
Arc::new(new_partition_info_cache(
"partition_info_cache".to_string(),
diff --git a/src/operator/src/write_dispatcher.rs b/src/operator/src/write_dispatcher.rs
index c2ac6543411..da8a02eddcc 100644
--- a/src/operator/src/write_dispatcher.rs
+++ b/src/operator/src/write_dispatcher.rs
@@ -22,7 +22,6 @@ use std::sync::Arc;
use api::region::RegionResponse;
use api::v1::region::{DeleteRequests, InsertRequests, RegionRequest, region_request::Body};
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};
@@ -36,6 +35,7 @@ use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef};
use prometheus::Histogram;
use snafu::{OptionExt, ResultExt};
use store_api::storage::RegionId;
+use tokio_util::sync::CancellationToken;
pub use self::options::WriteRetryOptions;
use crate::error::{self, FindRegionLeaderSnafu, RequestRegionSnafu};
@@ -75,14 +75,13 @@ impl WriteDispatcher {
/// Enables shared remote sends and retries using frontend configuration and caches.
pub fn with_retry(
mut self,
- clients: Arc,
options: WriteRetryOptions,
table_info_cache: TableInfoCacheRef,
cache_invalidator: CacheInvalidatorRef,
) -> Self {
if options.max_retries > 0 {
self.retry = Some(WriteRetry {
- clients,
+ node_manager: self.node_manager.clone(),
partition_manager: self.partition_manager.clone(),
options,
table_info_cache,
@@ -170,10 +169,12 @@ impl WriteDispatcher {
}
/// Sends requests using routes already resolved by the caller.
+ /// Cancellation stops future retries while allowing in-flight requests to complete.
pub(crate) async fn execute(
&self,
writes: Vec<(PeerRef, RegionRequest)>,
histogram: Option,
+ cancellation: Option<&CancellationToken>,
) -> Result {
let histogram = histogram.as_ref();
let writes = if self.retry.is_some() {
@@ -194,7 +195,7 @@ impl WriteDispatcher {
.map(|(peer, request)| self.send(peer, request, histogram)),
)
.await;
- self.finish(attempts, histogram).await
+ self.finish(attempts, histogram, cancellation).await
}
/// Sends the first attempt immediately so bulk encoding can overlap other sends.
@@ -231,10 +232,11 @@ impl WriteDispatcher {
&self,
attempts: Vec,
histogram: Option<&Histogram>,
+ cancellation: Option<&CancellationToken>,
) -> Result {
if let Some(retry) = &self.retry {
if retry.should_retry(&attempts, 0) {
- return retry.retry(attempts, histogram).await;
+ return retry.retry(attempts, histogram, cancellation).await;
}
let mut response = RegionResponse::new(0);
let failures = merge_attempts(attempts, &mut response)?;
@@ -271,13 +273,13 @@ impl WriteDispatcher {
.await
.map_err(BoxedError::new)
.context(ExternalSnafu)?;
- self.execute(writes, histogram).await
+ self.execute(writes, histogram, None).await
}
}
/// Sends shared requests and coordinates retries with refreshed routes.
struct WriteRetry {
- clients: Arc,
+ node_manager: NodeManagerRef,
partition_manager: PartitionRuleManagerRef,
options: WriteRetryOptions,
table_info_cache: TableInfoCacheRef,
@@ -293,15 +295,12 @@ pub(crate) struct RoutedWrite {
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;
+ let requester = self.node_manager.datanode(&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,
- }
+ WriteAttempt::Failed { error, request }
}
}
}
@@ -327,6 +326,7 @@ impl WriteRetry {
&self,
mut attempts: Vec,
histogram: Option<&Histogram>,
+ cancellation: Option<&CancellationToken>,
) -> Result {
let mut response = RegionResponse::new(0);
let mut backoff = ExponentialBuilder::default()
@@ -341,9 +341,20 @@ impl WriteRetry {
let Some(delay) = backoff.next() else {
break;
};
+ if let Some(cancellation) = cancellation {
+ tokio::select! {
+ biased;
+ _ = cancellation.cancelled() => break,
+ _ = tokio::time::sleep(delay) => {}
+ }
+ } else {
+ tokio::time::sleep(delay).await;
+ }
+ if cancellation.is_some_and(CancellationToken::is_cancelled) {
+ break;
+ }
let failures = merge_attempts(attempts, &mut response)?;
- tokio::time::sleep(delay).await;
-
+ let mut first_error = None;
let mut refresh = Vec::new();
let mut pending = Vec::new();
for (write, error) in failures {
@@ -352,10 +363,16 @@ impl WriteRetry {
} else {
pending.push(write);
}
+ first_error.get_or_insert(error);
}
if !refresh.is_empty() {
pending.extend(self.reroute_failed_writes(refresh).await?);
}
+ if let Some(error) =
+ first_error.filter(|_| cancellation.is_some_and(CancellationToken::is_cancelled))
+ {
+ return Err(error);
+ }
attempts =
future::join_all(pending.into_iter().map(|write| self.send(write, histogram)))
.await;
@@ -390,7 +407,7 @@ impl WriteRetry {
.await
.map_err(BoxedError::new)
.context(ExternalSnafu)?;
- partitioners.insert(table_id, partitioner);
+ partitioners.insert(table_id, (table_info, partitioner));
}
let mut rebuilt_requests = Vec::new();
for write in writes {
@@ -401,11 +418,12 @@ impl WriteRetry {
.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 (table_info, partitioner) =
+ partitioners.get(&table_id).context(UnexpectedSnafu {
+ err_msg: "Refreshed write has no partition snapshot",
+ })?;
let requests = partitioner
- .rebuild_write_requests(request)
+ .rebuild_write_requests(request, table_info)
.map_err(BoxedError::new)
.context(ExternalSnafu)?;
rebuilt_requests.extend(requests);
diff --git a/src/operator/src/write_dispatcher/tests.rs b/src/operator/src/write_dispatcher/tests.rs
index 250fe87f353..c8daf73f089 100644
--- a/src/operator/src/write_dispatcher/tests.rs
+++ b/src/operator/src/write_dispatcher/tests.rs
@@ -44,22 +44,30 @@ use common_meta::instruction::CacheIdent;
use common_meta::key::TableMetadataManager;
use common_meta::key::table_route::{TableRouteManager, TableRouteValue};
use common_meta::kv_backend::KvBackendRef;
+use common_meta::kv_backend::memory::MemoryKvBackend;
use common_meta::peer::Peer;
+use common_meta::rpc::router::{self, RegionRoute};
use common_meta::test_util::{MockDatanodeHandler, MockDatanodeManager};
use common_query::request::QueryRequest;
use common_recordbatch::SendableRecordBatchStream;
-use datatypes::arrow::array::Int32Array;
+use datatypes::arrow::array::{BinaryArray, Int32Array, StringArray};
use datatypes::arrow::datatypes::{DataType, Field, Schema};
use datatypes::arrow::record_batch::RecordBatch;
+use datatypes::data_type::ConcreteDataType;
+use datatypes::schema;
+use datatypes::value::{self, ValueRef};
+use mito_codec::row_converter::sparse;
+use mito_codec::row_converter::{PrimaryKeyCodec, SparsePrimaryKeyCodec};
use partition::cache::PartitionInfoCacheRef;
-use partition::expr::{Operand, PartitionExpr, RestrictedOp};
+use partition::expr::{self, Operand, PartitionExpr, RestrictedOp};
use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef};
-use store_api::storage::RegionId;
+use store_api::metric_engine_consts;
+use store_api::storage::{RegionId, consts};
use tokio::net::TcpListener;
-use tokio::sync::{Barrier, Notify};
+use tokio::sync::{Barrier, Notify, Semaphore};
use tokio::task::JoinHandle;
use tokio_stream::wrappers::TcpListenerStream;
-use tonic::{Request, Response, Status};
+use tonic::{Code, Request, Response, Status};
use super::*;
use crate::error::Error;
@@ -68,7 +76,8 @@ use crate::test_util;
#[derive(Clone, Default)]
struct RecordingRegionService {
requests: Arc>>,
- rejections: Arc>>>,
+ rejections: Arc>>>,
+ release: Arc>>>,
}
#[async_trait::async_trait]
@@ -80,18 +89,17 @@ impl Region for RecordingRegionService {
let request = request.into_inner();
let region_id = write_region_id(&request).unwrap().as_u64();
self.requests.lock().unwrap().push(request);
- if let Some(hint) = self
+ let release = self.release.lock().unwrap().clone();
+ if let Some(release) = release {
+ release.acquire().await.unwrap().forget();
+ }
+ if let Some(status) = self
.rejections
.lock()
.unwrap()
.get_mut(®ion_id)
.and_then(VecDeque::pop_front)
{
- let mut status = Status::unavailable(format!("region {region_id} rejected the write"));
- status.metadata_mut().insert(
- common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT,
- hint.as_str().parse().unwrap(),
- );
return Err(status);
}
Ok(Response::new(PbRegionResponse {
@@ -177,8 +185,7 @@ impl RegionTestServer {
let registry = cache::build_fundamental_cache_registry(backend);
let table_info_cache = registry.get::().unwrap();
let clients = Arc::new(NodeClients::new(ChannelConfig::default()));
- WriteDispatcher::new(clients.clone(), partition_manager).with_retry(
- clients,
+ WriteDispatcher::new(clients, partition_manager).with_retry(
build_retry_options(max_retries),
table_info_cache,
self.invalidator.clone(),
@@ -186,9 +193,22 @@ impl RegionTestServer {
}
fn fail_region(&self, region_number: u32, hints: impl IntoIterator- ) {
+ let region_id = RegionId::new(1, region_number).as_u64();
self.service.rejections.lock().unwrap().insert(
- RegionId::new(1, region_number).as_u64(),
- hints.into_iter().collect(),
+ region_id,
+ hints
+ .into_iter()
+ .map(|hint| {
+ let mut status = Status::failed_precondition(format!(
+ "region {region_id} rejected the write"
+ ));
+ status.metadata_mut().insert(
+ common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT,
+ hint.as_str().parse().unwrap(),
+ );
+ status
+ })
+ .collect(),
);
}
@@ -321,9 +341,7 @@ async fn test_write_retry_attempt_budget_and_nonretryable_errors() {
.dispatch(vec![build_insert_request(&[1])])
.await
.unwrap_err();
- // The dispatcher uses the explicit client hint for its attempt budget.
- // The existing common-meta conversion promotes Unavailable to RetryLater.
- assert_eq!(error.retry_hint(), RetryHint::Retryable);
+ assert_eq!(error.retry_hint(), hint);
assert_eq!(
server.attempt_counts()[&RegionId::new(1, 1).as_u64()],
expected_attempts
@@ -332,6 +350,84 @@ async fn test_write_retry_attempt_budget_and_nonretryable_errors() {
}
}
+#[tokio::test]
+async fn test_write_retry_transport_errors() {
+ for code in [Code::Unavailable, Code::Cancelled, Code::DeadlineExceeded] {
+ for hint in [
+ None,
+ Some(RetryHint::NonRetryable),
+ Some(RetryHint::Retryable),
+ Some(RetryHint::RetryAfterRouteRefresh),
+ ] {
+ let server = RegionTestServer::start().await;
+ let region_id = RegionId::new(1, 1).as_u64();
+ let mut status = Status::new(code, "write failed");
+ if let Some(hint) = hint {
+ status.metadata_mut().insert(
+ common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT,
+ hint.as_str().parse().unwrap(),
+ );
+ }
+ server
+ .service
+ .rejections
+ .lock()
+ .unwrap()
+ .insert(region_id, VecDeque::from([status]));
+ let dispatcher = server.create_dispatcher(2).await;
+ let attempt = dispatcher
+ .send(
+ Arc::new(server.peer.clone()),
+ build_insert_request(&[1]),
+ None,
+ )
+ .await;
+ let WriteAttempt::Failed { error, request } = attempt else {
+ panic!("expected a failed write");
+ };
+ assert!(request.is_some(), "{code:?}, {hint:?}");
+ assert_eq!(
+ error.retry_hint(),
+ hint.filter(|hint| hint.is_retryable())
+ .unwrap_or(RetryHint::Retryable)
+ );
+ assert_eq!(server.attempt_counts(), HashMap::from([(region_id, 1)]));
+ assert!(server.invalidator.calls.lock().unwrap().is_empty());
+ }
+ }
+}
+
+#[tokio::test]
+async fn test_write_retry_transport_attempt_budget() {
+ for code in [Code::Unavailable, Code::Cancelled, Code::DeadlineExceeded] {
+ for failures in [1, 3] {
+ let server = RegionTestServer::start().await;
+ let region_id = RegionId::new(1, 1).as_u64();
+ server.service.rejections.lock().unwrap().insert(
+ region_id,
+ (0..failures)
+ .map(|_| Status::new(code, "write failed"))
+ .collect(),
+ );
+ let dispatcher = server.create_dispatcher(2).await;
+ let response = dispatcher.dispatch(vec![build_insert_request(&[1])]).await;
+ let expected_attempts = if failures == 1 {
+ assert_eq!(response.unwrap().affected_rows, 1);
+ 2
+ } else {
+ assert_eq!(response.unwrap_err().retry_hint(), RetryHint::Retryable);
+ 3
+ };
+ assert_eq!(
+ server.attempt_counts(),
+ HashMap::from([(region_id, expected_attempts)]),
+ "{code:?}, failures: {failures}"
+ );
+ assert!(server.invalidator.calls.lock().unwrap().is_empty());
+ }
+ }
+}
+
#[tokio::test]
async fn test_write_retry_stops_when_a_sibling_is_nonretryable() {
for hint in [RetryHint::Retryable, RetryHint::RetryAfterRouteRefresh] {
@@ -364,14 +460,13 @@ async fn test_should_retry_requires_a_retained_request_and_budget() {
let backend = test_util::prepare_mocked_backend().await;
let (partition_manager, table_info_cache, cache_invalidator) = create_routing_caches(backend);
let retry = WriteRetry {
- clients: Arc::new(NodeClients::default()),
+ node_manager: Arc::new(NodeClients::default()),
partition_manager,
options: build_retry_options(2),
table_info_cache,
cache_invalidator,
};
let failure = |request| {
- // The converted error alone cannot distinguish an explicitly nonretryable attempt.
let error: client::Error = Status::unavailable("write rejected").into();
WriteAttempt::Failed {
error: error.into(),
@@ -402,6 +497,96 @@ async fn test_should_retry_requires_a_retained_request_and_budget() {
}
}
+#[tokio::test]
+async fn test_write_retry_cancellation_stops_backoff() {
+ for hint in [RetryHint::Retryable, RetryHint::RetryAfterRouteRefresh] {
+ let server = RegionTestServer::start().await;
+ server.fail_region(1, [hint]);
+ let mut dispatcher = server.create_dispatcher(2).await;
+ let options = &mut dispatcher.retry.as_mut().unwrap().options;
+ options.initial_backoff = Duration::from_secs(60);
+ options.max_backoff = Duration::from_secs(60);
+ let attempt = dispatcher
+ .send(
+ Arc::new(server.peer.clone()),
+ build_insert_request(&[1]),
+ None,
+ )
+ .await;
+ let cancellation = CancellationToken::new();
+ let guard = cancellation.clone().drop_guard();
+ let finish = dispatcher.finish(vec![attempt], None, Some(&cancellation));
+ tokio::pin!(finish);
+ futures::future::poll_fn(|cx| {
+ assert!(finish.as_mut().poll(cx).is_pending());
+ Poll::Ready(())
+ })
+ .await;
+ drop(guard);
+ let error = tokio::time::timeout(Duration::from_secs(1), finish)
+ .await
+ .unwrap()
+ .unwrap_err();
+ assert_eq!(error.retry_hint(), hint);
+ assert_eq!(
+ server.attempt_counts(),
+ HashMap::from([(RegionId::new(1, 1).as_u64(), 1)])
+ );
+ assert!(server.invalidator.calls.lock().unwrap().is_empty());
+ }
+}
+
+#[tokio::test]
+async fn test_write_retry_cancellation_drains_inflight_requests() {
+ let server = RegionTestServer::start().await;
+ server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]);
+ let release = Arc::new(Semaphore::new(0));
+ *server.service.release.lock().unwrap() = Some(release.clone());
+ let dispatcher = server.create_dispatcher(2).await;
+ let writes = dispatcher
+ .prepare(vec![build_insert_request(&[1, 2])])
+ .await
+ .unwrap();
+ let cancellation = CancellationToken::new();
+ let guard = cancellation.clone().drop_guard();
+ let execute = dispatcher.execute(writes, None, Some(&cancellation));
+ tokio::pin!(execute);
+ tokio::time::timeout(
+ Duration::from_secs(1),
+ futures::future::poll_fn(|cx| {
+ assert!(execute.as_mut().poll(cx).is_pending());
+ if server.service.requests.lock().unwrap().len() == 2 {
+ Poll::Ready(())
+ } else {
+ cx.waker().wake_by_ref();
+ Poll::Pending
+ }
+ }),
+ )
+ .await
+ .unwrap();
+ drop(guard);
+ futures::future::poll_fn(|cx| {
+ assert!(execute.as_mut().poll(cx).is_pending());
+ Poll::Ready(())
+ })
+ .await;
+ release.add_permits(2);
+ let error = tokio::time::timeout(Duration::from_secs(1), execute)
+ .await
+ .unwrap()
+ .unwrap_err();
+ assert_eq!(error.retry_hint(), RetryHint::RetryAfterRouteRefresh);
+ assert_eq!(
+ server.attempt_counts(),
+ HashMap::from([
+ (RegionId::new(1, 1).as_u64(), 1),
+ (RegionId::new(1, 2).as_u64(), 1),
+ ])
+ );
+ assert!(server.invalidator.calls.lock().unwrap().is_empty());
+}
+
#[tokio::test]
async fn test_write_retry_refreshes_routes_once_per_failed_round() {
let server = RegionTestServer::start().await;
@@ -444,8 +629,7 @@ async fn test_finish_retries_initial_failures_in_one_round() {
..Default::default()
});
let clients = Arc::new(NodeClients::new(ChannelConfig::default()));
- let dispatcher = WriteDispatcher::new(clients.clone(), partition_manager).with_retry(
- clients,
+ let dispatcher = WriteDispatcher::new(clients, partition_manager).with_retry(
build_retry_options(2),
table_info_cache,
invalidator.clone(),
@@ -469,7 +653,7 @@ async fn test_finish_retries_initial_failures_in_one_round() {
(RegionId::new(1, 3).as_u64(), 1),
])
);
- let response = dispatcher.finish(attempts, None).await.unwrap();
+ let response = dispatcher.finish(attempts, None, None).await.unwrap();
assert_eq!(response.affected_rows, 3);
assert_eq!(
server.attempt_counts(),
@@ -526,8 +710,7 @@ async fn test_write_retry_refreshes_stale_route_and_sends_only_failed_region_to_
..Default::default()
});
let clients = Arc::new(NodeClients::new(ChannelConfig::default()));
- let dispatcher = WriteDispatcher::new(clients.clone(), partition_manager).with_retry(
- clients,
+ let dispatcher = WriteDispatcher::new(clients, partition_manager).with_retry(
build_retry_options(2),
table_info_cache,
invalidator.clone(),
@@ -745,13 +928,11 @@ async fn test_write_retry_rebuilds_only_failed_rows_after_route_refresh() {
route_update: Mutex::new(Some((backend, current.update(routes).unwrap()))),
});
let clients = Arc::new(NodeClients::new(ChannelConfig::default()));
- let dispatcher = WriteDispatcher::new(clients.clone(), partition_manager.clone())
- .with_retry(
- clients,
- build_retry_options(2),
- table_info_cache,
- invalidator.clone(),
- );
+ let dispatcher = WriteDispatcher::new(clients, partition_manager.clone()).with_retry(
+ build_retry_options(2),
+ table_info_cache,
+ invalidator.clone(),
+ );
let response = dispatcher
.dispatch(vec![
build_partitioned_request(kind, 1, old_versions[&1], &[70]),
@@ -793,6 +974,188 @@ async fn test_write_retry_rebuilds_only_failed_rows_after_route_refresh() {
}
}
+fn build_sparse_bulk_request(
+ region: u32,
+ version: Option,
+ tags: &[(&str, &str)],
+) -> (RegionRequest, RecordBatch) {
+ let codec = SparsePrimaryKeyCodec::schemaless();
+ let keys = tags
+ .iter()
+ .enumerate()
+ .map(|(index, (host, zone))| {
+ let mut key = Vec::new();
+ codec
+ .encode_value_refs(
+ &[
+ (sparse::RESERVED_COLUMN_ID_TABLE_ID, ValueRef::UInt32(42)),
+ (
+ sparse::RESERVED_COLUMN_ID_TSID,
+ ValueRef::UInt64(index as u64),
+ ),
+ (10, ValueRef::String(host)),
+ (20, ValueRef::String(zone)),
+ ],
+ &mut key,
+ )
+ .unwrap();
+ key
+ })
+ .collect::>();
+ let batch = RecordBatch::try_new(
+ Arc::new(Schema::new(vec![
+ Field::new(consts::PRIMARY_KEY_COLUMN_NAME, DataType::Binary, false),
+ Field::new("host", DataType::Utf8, false),
+ ])),
+ vec![
+ Arc::new(BinaryArray::from_iter_values(keys)),
+ Arc::new(StringArray::from_iter_values(
+ tags.iter().map(|(host, _)| *host),
+ )),
+ ],
+ )
+ .unwrap();
+ let (schema, data_header, payload) = flight::record_batch_to_ipc(batch.clone()).unwrap();
+ let request = RegionRequest {
+ body: Some(Body::BulkInsert(BulkInsertRequest {
+ region_id: RegionId::new(1, region).as_u64(),
+ partition_expr_version: version.map(|value| PartitionExprVersion { value }),
+ body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc {
+ schema,
+ data_header,
+ payload,
+ })),
+ ..Default::default()
+ })),
+ ..Default::default()
+ };
+ (request, batch)
+}
+
+#[tokio::test]
+async fn test_write_retry_restores_changed_partition_tags_for_failed_regions() {
+ let old_server = RegionTestServer::start().await;
+ let new_server = RegionTestServer::start().await;
+ old_server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]);
+ let backend: KvBackendRef = Arc::new(MemoryKvBackend::new());
+ let metadata = TableMetadataManager::new(backend.clone());
+ let mut table_info = test_util::new_test_table_info(1, "physical", [1, 2].into_iter());
+ table_info.meta.schema = Arc::new(schema::Schema::new(vec![
+ schema::ColumnSchema::new("host", ConcreteDataType::string_datatype(), true),
+ schema::ColumnSchema::new("zone", ConcreteDataType::string_datatype(), true),
+ ]));
+ table_info.meta.primary_key_indices = vec![0, 1];
+ table_info.meta.partition_key_indices = vec![0];
+ table_info.meta.column_ids = vec![10, 20];
+ table_info.meta.options.extra_options.insert(
+ metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY.to_string(),
+ String::new(),
+ );
+ let routes = [
+ expr::col("host").lt(value::Value::String("m".into())),
+ expr::col("host").gt_eq(value::Value::String("m".into())),
+ ]
+ .into_iter()
+ .enumerate()
+ .map(|(index, expr)| RegionRoute {
+ region: router::Region {
+ id: RegionId::new(1, index as u32 + 1),
+ partition_expr: expr.as_json_str().unwrap(),
+ ..Default::default()
+ },
+ leader_peer: Some(old_server.peer.clone()),
+ ..Default::default()
+ })
+ .collect();
+ metadata
+ .create_table_metadata(
+ table_info,
+ TableRouteValue::physical(routes),
+ Default::default(),
+ )
+ .await
+ .unwrap();
+ let (partition_manager, table_info_cache, registry) = create_routing_caches(backend.clone());
+ let cached_info = table_info_cache.get(1).await.unwrap().unwrap();
+ let (_, old_versions) = partition_manager
+ .find_table_partition_rule(&cached_info)
+ .await
+ .unwrap();
+ let current_info = metadata.table_info_manager().get(1).await.unwrap().unwrap();
+ let mut updated_info = current_info.table_info.clone();
+ updated_info.meta.partition_key_indices = vec![1];
+ metadata
+ .update_table_info(¤t_info, None, updated_info)
+ .await
+ .unwrap();
+ let current_route = metadata
+ .table_route_manager()
+ .table_route_storage()
+ .get_with_raw_bytes(1)
+ .await
+ .unwrap()
+ .unwrap();
+ let mut routes = current_route.region_routes().unwrap().clone();
+ for (route, expr) in routes.iter_mut().zip([
+ expr::col("zone").lt(value::Value::String("m".into())),
+ expr::col("zone").gt_eq(value::Value::String("m".into())),
+ ]) {
+ route.region.partition_expr = expr.as_json_str().unwrap();
+ route.leader_peer = Some(new_server.peer.clone());
+ }
+ let invalidator = Arc::new(RecordingCacheInvalidator {
+ calls: Mutex::default(),
+ inner: Some(registry),
+ route_update: Mutex::new(Some((backend, current_route.update(routes).unwrap()))),
+ });
+ let clients = Arc::new(NodeClients::new(ChannelConfig::default()));
+ let dispatcher = WriteDispatcher::new(clients, partition_manager.clone()).with_retry(
+ build_retry_options(1),
+ table_info_cache.clone(),
+ invalidator.clone(),
+ );
+ // The successful row would move under the new rule, but must not be replayed.
+ let (successful, _) = build_sparse_bulk_request(1, old_versions[&1], &[("a", "west")]);
+ let (failed, batch) =
+ build_sparse_bulk_request(2, old_versions[&2], &[("z", "east"), ("z", "west")]);
+ dispatcher.dispatch(vec![successful, failed]).await.unwrap();
+ let counts = HashMap::from([
+ (RegionId::new(1, 1).as_u64(), 1),
+ (RegionId::new(1, 2).as_u64(), 1),
+ ]);
+ assert_eq!(old_server.attempt_counts(), counts);
+ assert_eq!(new_server.attempt_counts(), counts);
+ assert_eq!(
+ *invalidator.calls.lock().unwrap(),
+ vec![vec![CacheIdent::TableId(1)]]
+ );
+ let refreshed_info = table_info_cache.get(1).await.unwrap().unwrap();
+ assert_eq!(refreshed_info.meta.partition_key_indices, vec![1]);
+ let (_, versions) = partition_manager
+ .find_table_partition_rule(&refreshed_info)
+ .await
+ .unwrap();
+ for request in new_server.service.requests.lock().unwrap().iter() {
+ let Some(Body::BulkInsert(bulk)) = &request.body else {
+ unreachable!()
+ };
+ let region = RegionId::from_u64(bulk.region_id).region_number();
+ assert_eq!(
+ bulk.partition_expr_version.map(|version| version.value),
+ versions[®ion]
+ );
+ let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &bulk.body else {
+ unreachable!()
+ };
+ let actual = FlightDecoder::try_from_schema_bytes(&ipc.schema)
+ .unwrap()
+ .try_decode_record_batch(&ipc.data_header, &ipc.payload)
+ .unwrap();
+ assert_eq!(actual, batch.slice(region as usize - 1, 1));
+ assert!(actual.column_by_name("zone").is_none());
+ }
+}
+
#[tokio::test]
async fn test_resolve_peers_batches_logical_and_physical_regions() {
let backend = test_util::prepare_mocked_backend().await;
@@ -929,7 +1292,7 @@ async fn test_dispatch_preserves_write_payloads_and_envelopes() {
});
let writes = dispatcher.prepare(requests.to_vec()).await.unwrap();
assert!(received.try_recv().is_err());
- let response = dispatcher.execute(writes, None).await.unwrap();
+ let response = dispatcher.execute(writes, None, None).await.unwrap();
assert_eq!(response.affected_rows, 2);
for request in requests {
assert_eq!(received.try_recv().unwrap(), (Peer::new(1, ""), request));
@@ -1153,8 +1516,7 @@ async fn test_write_retry_rebuilds_logical_rows_in_mixed_table_envelope() {
..Default::default()
});
let clients = Arc::new(NodeClients::new(ChannelConfig::default()));
- let dispatcher = WriteDispatcher::new(clients.clone(), partition_manager).with_retry(
- clients,
+ let dispatcher = WriteDispatcher::new(clients, partition_manager).with_retry(
build_retry_options(2),
table_info_cache,
invalidator.clone(),
@@ -1162,10 +1524,17 @@ async fn test_write_retry_rebuilds_logical_rows_in_mixed_table_envelope() {
let physical_id = RegionId::new(1, 1).as_u64();
let stale_logical_id = RegionId::new(42, 2).as_u64();
let rebuilt_logical_id = RegionId::new(42, 3).as_u64();
- server.service.rejections.lock().unwrap().insert(
- stale_logical_id,
- VecDeque::from([RetryHint::RetryAfterRouteRefresh]),
+ let mut status = Status::unavailable("partition changed");
+ status.metadata_mut().insert(
+ common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT,
+ RetryHint::RetryAfterRouteRefresh.as_str().parse().unwrap(),
);
+ server
+ .service
+ .rejections
+ .lock()
+ .unwrap()
+ .insert(stale_logical_id, VecDeque::from([status]));
let header = Some(RegionRequestHeader {
dbname: "greptime.public".to_string(),
@@ -1225,15 +1594,10 @@ async fn test_with_retry_can_disable_previously_enabled_retries() {
let server = RegionTestServer::start().await;
let dispatcher = server.create_dispatcher(1).await;
let retry = dispatcher.retry.as_ref().unwrap();
- let clients = retry.clients.clone();
let table_info_cache = retry.table_info_cache.clone();
let cache_invalidator = retry.cache_invalidator.clone();
- let dispatcher = dispatcher.with_retry(
- clients,
- build_retry_options(0),
- table_info_cache,
- cache_invalidator,
- );
+ let dispatcher =
+ dispatcher.with_retry(build_retry_options(0), table_info_cache, cache_invalidator);
assert!(dispatcher.retry.is_none());
server.fail_region(1, [RetryHint::Retryable]);
assert!(
diff --git a/src/servers/AGENTS.md b/src/servers/AGENTS.md
index c3b712d3924..6ee7675313c 100644
--- a/src/servers/AGENTS.md
+++ b/src/servers/AGENTS.md
@@ -43,6 +43,9 @@ cargo nextest run -p servers
Keep protocol translation here and permissions/database behavior in frontend.
Async table batching returns storage-completion futures with its early
acknowledgements. Keep completion reporting independent of the caller's lifetime.
+Logical-table physical flushes retain partition columns in Arrow RPC payloads and
+pass all region writes to the operator write dispatcher together; the dispatcher
+owns peer resolution and retries. See `batcher/logical_table/region_write.rs`.
`prom_remote_write/decode.rs` and `benches/prom_decode.rs` resolve
`benches/write_request.pb.data` using `CARGO_MANIFEST_DIR` at **run time**
diff --git a/src/servers/benches/flush_batch_physical.rs b/src/servers/benches/flush_batch_physical.rs
index be992692d13..d879486c2d3 100644
--- a/src/servers/benches/flush_batch_physical.rs
+++ b/src/servers/benches/flush_batch_physical.rs
@@ -15,8 +15,6 @@
use std::collections::HashMap;
use std::sync::Arc;
-use api::region::RegionResponse;
-use api::v1::meta::Peer;
use api::v1::region::RegionRequest;
use arrow::array::{Float64Array, StringArray, TimestampMillisecondArray};
use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema, TimeUnit};
@@ -29,11 +27,10 @@ use datatypes::schema::{ColumnSchema as DtColumnSchema, Schema as DtSchema};
use partition::error::Result as PartitionResult;
use partition::partition::{PartitionRule, PartitionRuleRef, RegionMask};
use servers::batcher::logical_table::{
- PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider,
+ PhysicalFlushCatalogProvider, PhysicalFlushPartitionProvider, PhysicalFlushWriteDispatcher,
PhysicalTableMetadata, RecordBatchWithTsIdx, TableBatch, flush_batch_physical,
};
-use servers::error::{self, Result};
-use store_api::storage::RegionId;
+use servers::error;
use table::test_util::table_info::test_table_info;
use tokio::runtime::Runtime;
@@ -95,24 +92,17 @@ impl PhysicalFlushPartitionProvider for BenchPartitionProvider {
async fn find_table_partition_rule(
&self,
_table_info: &table::metadata::TableInfo,
- ) -> PartitionResult {
- Ok(Arc::new(SingleRegionPartitionRule))
- }
-
- async fn find_region_leader(&self, _region_id: RegionId) -> Result {
- Ok(Peer {
- id: 1,
- addr: "bench-node".to_string(),
- })
+ ) -> PartitionResult<(PartitionRuleRef, HashMap>)> {
+ Ok((Arc::new(SingleRegionPartitionRule), HashMap::new()))
}
}
-struct BenchNodeRequester;
+struct BenchWriteDispatcher;
#[async_trait]
-impl PhysicalFlushNodeRequester for BenchNodeRequester {
- async fn handle(&self, _peer: &Peer, _request: RegionRequest) -> error::Result {
- Ok(RegionResponse::new(0))
+impl PhysicalFlushWriteDispatcher for BenchWriteDispatcher {
+ async fn dispatch(&self, _requests: Vec) -> error::Result {
+ Ok(0)
}
}
@@ -224,7 +214,7 @@ fn bench_flush_batch_physical(c: &mut Criterion) {
table: make_physical_table_metadata(num_tags),
};
let partition_provider = BenchPartitionProvider;
- let node_requester = BenchNodeRequester;
+ let write_dispatcher = BenchWriteDispatcher;
let mut group = c.benchmark_group("flush_batch_physical");
@@ -244,7 +234,7 @@ fn bench_flush_batch_physical(c: &mut Criterion) {
"phy",
&ctx,
&partition_provider,
- &node_requester,
+ &write_dispatcher,
&catalog_provider,
)
.await
@@ -271,7 +261,7 @@ fn bench_flush_batch_physical(c: &mut Criterion) {
"phy",
&ctx,
&partition_provider,
- &node_requester,
+ &write_dispatcher,
&catalog_provider,
)
.await
diff --git a/src/servers/src/batcher/logical_table.rs b/src/servers/src/batcher/logical_table.rs
index 7e0cea44b6a..f34b0ae406e 100644
--- a/src/servers/src/batcher/logical_table.rs
+++ b/src/servers/src/batcher/logical_table.rs
@@ -39,6 +39,7 @@ use common_query::prelude::GREPTIME_PHYSICAL_TABLE;
use common_time::timestamp::TimeUnit;
use meter_core::data::MeterRecord;
use meter_macros::write_meter;
+use operator::write_dispatcher::WriteDispatcherRef;
use partition::manager::PartitionRuleManagerRef;
use session::context::QueryContextRef;
use snafu::ResultExt;
@@ -52,7 +53,7 @@ use crate::batcher::logical_table::pending_worker::{
PendingWorker, WorkerCommand, remove_worker_if_same_channel, start_worker,
};
pub use crate::batcher::logical_table::region_write::{
- PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider,
+ PhysicalFlushCatalogProvider, PhysicalFlushPartitionProvider, PhysicalFlushWriteDispatcher,
PhysicalTableMetadata,
};
pub use crate::batcher::logical_table::tables::{
@@ -97,7 +98,7 @@ pub struct LogicalTablePendingRowsBatcher {
flush_interval: Duration,
flush_policy: TimingFlushPolicy,
partition_manager: PartitionRuleManagerRef,
- node_manager: NodeManagerRef,
+ write_dispatcher: WriteDispatcherRef,
catalog_manager: CatalogManagerRef,
flow_notification_tx: FlowNotifier,
flush_limiter: FlushLimiter,
@@ -110,10 +111,12 @@ pub struct LogicalTablePendingRowsBatcher {
}
impl LogicalTablePendingRowsBatcher {
+ /// Creates the batcher when metric-engine batching is enabled.
#[allow(clippy::too_many_arguments)]
pub fn try_new(
partition_manager: PartitionRuleManagerRef,
node_manager: NodeManagerRef,
+ write_dispatcher: WriteDispatcherRef,
catalog_manager: CatalogManagerRef,
table_flownode_set_cache: TableFlownodeSetCacheRef,
prom_store_with_metric_engine: bool,
@@ -153,7 +156,7 @@ impl LogicalTablePendingRowsBatcher {
flush_interval,
flush_policy,
partition_manager,
- node_manager,
+ write_dispatcher,
catalog_manager,
flow_notification_tx,
prom_store_with_metric_engine,
@@ -404,7 +407,7 @@ impl LogicalTablePendingRowsBatcher {
rx,
self.shutdown.clone(),
self.partition_manager.clone(),
- self.node_manager.clone(),
+ self.write_dispatcher.clone(),
self.catalog_manager.clone(),
self.flow_notification_tx.clone(),
worker_idle_timeout,
diff --git a/src/servers/src/batcher/logical_table/batch.rs b/src/servers/src/batcher/logical_table/batch.rs
index ac2db137de1..1ce90f39d01 100644
--- a/src/servers/src/batcher/logical_table/batch.rs
+++ b/src/servers/src/batcher/logical_table/batch.rs
@@ -18,45 +18,41 @@ use std::time::Instant;
use catalog::CatalogManagerRef;
use common_batcher::flush_limiter::FlushLimiter;
-use common_meta::node_manager::NodeManagerRef;
-use common_query::prelude::GREPTIME_PHYSICAL_TABLE;
+use common_query::prelude;
use common_telemetry::{debug, warn};
+use operator::write_dispatcher::WriteDispatcherRef;
use partition::manager::PartitionRuleManagerRef;
use session::context::QueryContextRef;
use snafu::OptionExt;
use crate::batcher::flow_notifier::FlowNotifier;
-use crate::batcher::logical_table::PHYSICAL_TABLE_KEY;
-use crate::batcher::logical_table::batch_convert::{
- TableBatch, concat_modified_batches, transform_logical_batches_to_physical,
-};
-use crate::batcher::logical_table::flow_notifier::enqueue_flow_notifications;
+use crate::batcher::logical_table;
+use crate::batcher::logical_table::batch_convert::{self, TableBatch};
+use crate::batcher::logical_table::flow_notifier;
use crate::batcher::logical_table::pending_worker::FlushWaiter;
use crate::batcher::logical_table::region_write::{
- CatalogManagerPhysicalFlushAdapter, NodeManagerPhysicalFlushAdapter,
- PartitionManagerPhysicalFlushAdapter, PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester,
- PhysicalFlushPartitionProvider, encode_region_write_requests, flush_region_writes_concurrently,
- plan_region_batches, resolve_region_targets,
+ self, CatalogManagerPhysicalFlushAdapter, PartitionManagerPhysicalFlushAdapter,
+ PhysicalFlushCatalogProvider, PhysicalFlushPartitionProvider, PhysicalFlushWriteDispatcher,
+ WriteDispatcherPhysicalFlushAdapter,
};
use crate::error;
use crate::error::Result;
-use crate::metrics::{
- FLUSH_DROPPED_ROWS, FLUSH_ELAPSED, FLUSH_FAILURES, FLUSH_ROWS, FLUSH_TOTAL,
- PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED,
-};
+use crate::metrics;
-pub(in crate::batcher::logical_table) struct Batch {
- pub(in crate::batcher::logical_table) table_batches: Vec,
- pub(in crate::batcher::logical_table) total_row_count: usize,
- pub(in crate::batcher::logical_table) db_string: String,
- pub(in crate::batcher::logical_table) ctx: QueryContextRef,
- pub(in crate::batcher::logical_table) waiters: Vec,
+/// Rows and completion waiters owned by one physical flush.
+pub(crate) struct Batch {
+ pub(crate) table_batches: Vec,
+ pub(crate) total_row_count: usize,
+ pub(crate) db_string: String,
+ pub(crate) ctx: QueryContextRef,
+ pub(crate) waiters: Vec,
}
-pub(in crate::batcher::logical_table) async fn spawn_flush(
+/// Spawns a physical flush with the configured concurrency limit.
+pub(crate) async fn spawn_flush(
flush: Batch,
partition_manager: PartitionRuleManagerRef,
- node_manager: NodeManagerRef,
+ write_dispatcher: WriteDispatcherRef,
catalog_manager: CatalogManagerRef,
flow_notification_tx: FlowNotifier,
flush_limiter: FlushLimiter,
@@ -68,7 +64,7 @@ pub(in crate::batcher::logical_table) async fn spawn_flush(
flush_batch_with_managers(
flush,
partition_manager,
- node_manager,
+ write_dispatcher,
catalog_manager,
flow_notification_tx,
)
@@ -80,7 +76,7 @@ pub(in crate::batcher::logical_table) async fn spawn_flush(
flush_batch_with_managers(
flush,
partition_manager,
- node_manager,
+ write_dispatcher,
catalog_manager,
flow_notification_tx,
)
@@ -89,32 +85,32 @@ pub(in crate::batcher::logical_table) async fn spawn_flush(
}
}
-pub(in crate::batcher::logical_table) async fn flush_batch_with_managers(
+/// Connects shared routing and dispatch dependencies to a physical flush.
+pub(crate) async fn flush_batch_with_managers(
flush: Batch,
partition_manager: PartitionRuleManagerRef,
- node_manager: NodeManagerRef,
+ write_dispatcher: WriteDispatcherRef,
catalog_manager: CatalogManagerRef,
flow_notification_tx: FlowNotifier,
) {
let partition_provider = PartitionManagerPhysicalFlushAdapter { partition_manager };
- let node_requester = NodeManagerPhysicalFlushAdapter {
- node_manager: node_manager.clone(),
- };
+ let write_dispatcher = WriteDispatcherPhysicalFlushAdapter { write_dispatcher };
let catalog_provider = CatalogManagerPhysicalFlushAdapter { catalog_manager };
flush_batch(
flush,
&partition_provider,
- &node_requester,
+ &write_dispatcher,
&catalog_provider,
flow_notification_tx,
)
.await;
}
-pub(in crate::batcher::logical_table) async fn flush_batch(
+/// Flushes a batch and reports completion to its waiters.
+pub(crate) async fn flush_batch(
flush: Batch,
partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized),
- node_manager: &(impl PhysicalFlushNodeRequester + ?Sized),
+ write_dispatcher: &(impl PhysicalFlushWriteDispatcher + ?Sized),
catalog_manager: &(impl PhysicalFlushCatalogProvider + ?Sized),
flow_notification_tx: FlowNotifier,
) {
@@ -130,21 +126,21 @@ pub(in crate::batcher::logical_table) async fn flush_batch(
// Physical-table-level flush: transform all logical table batches
// into physical format and write them together.
let physical_table_name = ctx
- .extension(PHYSICAL_TABLE_KEY)
- .unwrap_or(GREPTIME_PHYSICAL_TABLE)
+ .extension(logical_table::PHYSICAL_TABLE_KEY)
+ .unwrap_or(prelude::GREPTIME_PHYSICAL_TABLE)
.to_string();
let result = flush_batch_physical(
&table_batches,
&physical_table_name,
&ctx,
partition_manager,
- node_manager,
+ write_dispatcher,
catalog_manager,
)
.await;
let elapsed = start.elapsed().as_secs_f64();
- FLUSH_ELAPSED.observe(elapsed);
+ metrics::FLUSH_ELAPSED.observe(elapsed);
debug!(
"Pending rows batch flushed, total rows: {}, elapsed time: {}s",
@@ -153,18 +149,18 @@ pub(in crate::batcher::logical_table) async fn flush_batch(
match result {
Ok(affected_rows) => {
- FLUSH_TOTAL.inc();
- FLUSH_ROWS.observe(total_row_count as f64);
+ metrics::FLUSH_TOTAL.inc();
+ metrics::FLUSH_ROWS.observe(total_row_count as f64);
operator::metrics::DIST_INGEST_ROW_COUNT
.with_label_values(&[db_string.as_str()])
.inc_by(affected_rows as u64);
notify_waiters(waiters, Ok(()));
- enqueue_flow_notifications(table_batches, &flow_notification_tx);
+ flow_notifier::enqueue_flow_notifications(table_batches, &flow_notification_tx);
}
Err(err) => {
- FLUSH_FAILURES.inc();
- FLUSH_DROPPED_ROWS.inc_by(total_row_count as u64);
+ metrics::FLUSH_FAILURES.inc();
+ metrics::FLUSH_DROPPED_ROWS.inc_by(total_row_count as u64);
notify_waiters(waiters, Err(err));
}
}
@@ -185,12 +181,12 @@ pub async fn flush_batch_physical(
physical_table_name: &str,
ctx: &QueryContextRef,
partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized),
- node_manager: &(impl PhysicalFlushNodeRequester + ?Sized),
+ write_dispatcher: &(impl PhysicalFlushWriteDispatcher + ?Sized),
catalog_manager: &(impl PhysicalFlushCatalogProvider + ?Sized),
) -> Result {
// 1. Resolve the physical table and get column ID mapping
let physical_table = {
- let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
+ let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
.with_label_values(&["flush_physical_resolve_table"])
.start_timer();
catalog_manager
@@ -220,8 +216,8 @@ pub async fn flush_batch_physical(
})?;
// 2. Get the physical table's partition rule (one lookup instead of N)
- let partition_rule = {
- let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
+ let (partition_rule, partition_versions) = {
+ let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
.with_label_values(&["flush_physical_fetch_partition_rule"])
.start_timer();
partition_manager
@@ -233,24 +229,29 @@ pub async fn flush_batch_physical(
partition_columns.iter().map(String::as_str).collect();
// 3. Transform each logical table batch into physical format
- let modified_batches =
- transform_logical_batches_to_physical(table_batches, &name_to_ids, &partition_columns_set)?;
+ let modified_batches = batch_convert::transform_logical_batches_to_physical(
+ table_batches,
+ &name_to_ids,
+ &partition_columns_set,
+ )?;
// 4. Concatenate all modified batches (all share the same physical schema)
- let combined_batch = concat_modified_batches(&modified_batches)?;
+ let combined_batch = batch_convert::concat_modified_batches(&modified_batches)?;
// 5. Split by physical partition rule and send to regions
let physical_table_id = physical_table_info.table_id();
- let planned_batches = plan_region_batches(
+ let planned_batches = region_write::plan_region_batches(
combined_batch,
physical_table_id,
partition_rule.as_ref(),
- partition_columns,
)?;
- let resolved_batches = resolve_region_targets(planned_batches, partition_manager).await?;
- let region_writes = encode_region_write_requests(resolved_batches, ctx.skip_wal())?;
- flush_region_writes_concurrently(node_manager, region_writes).await
+ let region_writes = region_write::encode_region_write_requests(
+ planned_batches,
+ ctx.skip_wal(),
+ &partition_versions,
+ )?;
+ write_dispatcher.dispatch(region_writes).await
}
pub(in crate::batcher::logical_table) fn notify_waiters(
@@ -276,7 +277,6 @@ mod tests {
use std::task::Poll;
use std::time::Duration;
- use api::region::RegionResponse;
use api::v1::meta::Peer;
use api::v1::region::RegionRequest;
use arrow::record_batch::RecordBatch;
@@ -288,7 +288,6 @@ mod tests {
use datatypes::schema::{ColumnSchema as DtColumnSchema, Schema as DtSchema};
use partition::error::Result as PartitionResult;
use partition::partition::{PartitionRule, PartitionRuleRef, RegionMask};
- use store_api::storage::RegionId;
use table::metadata::TableId;
use table::test_util::table_info::test_table_info;
use tokio::sync::{mpsc, oneshot};
@@ -300,7 +299,7 @@ mod tests {
use crate::batcher::logical_table::batch_convert::TableBatch;
use crate::batcher::logical_table::pending_worker::FlushWaiter;
use crate::batcher::logical_table::region_write::{
- PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider,
+ PhysicalFlushCatalogProvider, PhysicalFlushPartitionProvider, PhysicalFlushWriteDispatcher,
PhysicalTableMetadata,
};
use crate::batcher::logical_table::test_util::{
@@ -344,9 +343,8 @@ mod tests {
},
&MockFlushPartitionProvider {
partition_rule_calls: Arc::new(AtomicUsize::new(0)),
- region_leader_calls: Arc::new(AtomicUsize::new(0)),
},
- &MockFlushNodeRequester {
+ &MockFlushWriteDispatcher {
writes: writes.clone(),
fail: false,
},
@@ -405,9 +403,8 @@ mod tests {
},
&MockFlushPartitionProvider {
partition_rule_calls: Arc::new(AtomicUsize::new(0)),
- region_leader_calls: Arc::new(AtomicUsize::new(0)),
},
- &MockFlushNodeRequester {
+ &MockFlushWriteDispatcher {
writes: writes.clone(),
fail: true,
},
@@ -451,8 +448,7 @@ mod tests {
row_count: 1,
}];
let partition_calls = Arc::new(AtomicUsize::new(0));
- let leader_calls = Arc::new(AtomicUsize::new(0));
- let node = MockFlushNodeRequester::default();
+ let dispatcher = MockFlushWriteDispatcher::default();
let ctx = session::context::QueryContext::arc();
flush_batch_physical(
@@ -461,9 +457,8 @@ mod tests {
&ctx,
&MockFlushPartitionProvider {
partition_rule_calls: partition_calls.clone(),
- region_leader_calls: leader_calls.clone(),
},
- &node,
+ &dispatcher,
&MockFlushCatalogProvider {
table: Some(mock_physical_table_metadata(1024)),
},
@@ -472,8 +467,7 @@ mod tests {
.unwrap();
assert_eq!(1, partition_calls.load(Ordering::SeqCst));
- assert_eq!(1, leader_calls.load(Ordering::SeqCst));
- assert_eq!(1, node.writes.load(Ordering::SeqCst));
+ assert_eq!(1, dispatcher.writes.load(Ordering::SeqCst));
}
#[tokio::test]
@@ -492,9 +486,8 @@ mod tests {
&ctx,
&MockFlushPartitionProvider {
partition_rule_calls: Arc::new(AtomicUsize::new(0)),
- region_leader_calls: Arc::new(AtomicUsize::new(0)),
},
- &AffectedRowsFlushNodeRequester { affected_rows: 7 },
+ &AffectedRowsFlushWriteDispatcher { affected_rows: 7 },
&MockFlushCatalogProvider {
table: Some(mock_physical_table_metadata(1024)),
},
@@ -506,7 +499,7 @@ mod tests {
}
#[tokio::test]
- async fn test_flush_batch_physical_stops_before_partition_and_node_when_table_missing() {
+ async fn test_flush_batch_physical_stops_before_partition_and_dispatch_when_table_missing() {
let table_batches = vec![TableBatch {
table_name: "t1".to_string(),
table_id: 11,
@@ -514,8 +507,7 @@ mod tests {
row_count: 1,
}];
let partition_calls = Arc::new(AtomicUsize::new(0));
- let leader_calls = Arc::new(AtomicUsize::new(0));
- let node = MockFlushNodeRequester::default();
+ let dispatcher = MockFlushWriteDispatcher::default();
let ctx = session::context::QueryContext::arc();
let err = flush_batch_physical(
@@ -524,9 +516,8 @@ mod tests {
&ctx,
&MockFlushPartitionProvider {
partition_rule_calls: partition_calls.clone(),
- region_leader_calls: leader_calls.clone(),
},
- &node,
+ &dispatcher,
&MockFlushCatalogProvider { table: None },
)
.await
@@ -537,8 +528,7 @@ mod tests {
.contains("Physical table 'missing_phy' not found")
);
assert_eq!(0, partition_calls.load(Ordering::SeqCst));
- assert_eq!(0, leader_calls.load(Ordering::SeqCst));
- assert_eq!(0, node.writes.load(Ordering::SeqCst));
+ assert_eq!(0, dispatcher.writes.load(Ordering::SeqCst));
}
#[tokio::test]
@@ -558,8 +548,7 @@ mod tests {
},
];
let partition_calls = Arc::new(AtomicUsize::new(0));
- let leader_calls = Arc::new(AtomicUsize::new(0));
- let node = MockFlushNodeRequester::default();
+ let dispatcher = MockFlushWriteDispatcher::default();
let ctx = session::context::QueryContext::arc();
let err = flush_batch_physical(
@@ -568,9 +557,8 @@ mod tests {
&ctx,
&MockFlushPartitionProvider {
partition_rule_calls: partition_calls.clone(),
- region_leader_calls: leader_calls.clone(),
},
- &node,
+ &dispatcher,
&MockFlushCatalogProvider {
table: Some(mock_physical_table_metadata(1024)),
},
@@ -580,8 +568,7 @@ mod tests {
assert!(err.to_string().contains("unknown_tag"));
assert_eq!(1, partition_calls.load(Ordering::SeqCst));
- assert_eq!(0, leader_calls.load(Ordering::SeqCst));
- assert_eq!(0, node.writes.load(Ordering::SeqCst));
+ assert_eq!(0, dispatcher.writes.load(Ordering::SeqCst));
}
fn mock_physical_table_metadata(table_id: TableId) -> PhysicalTableMetadata {
@@ -671,7 +658,6 @@ mod tests {
struct MockFlushPartitionProvider {
partition_rule_calls: Arc,
- region_leader_calls: Arc,
}
#[async_trait]
@@ -679,40 +665,28 @@ mod tests {
async fn find_table_partition_rule(
&self,
_table_info: &table::metadata::TableInfo,
- ) -> PartitionResult {
+ ) -> PartitionResult<(PartitionRuleRef, HashMap>)> {
self.partition_rule_calls.fetch_add(1, Ordering::SeqCst);
- Ok(Arc::new(SingleRegionPartitionRule))
- }
-
- async fn find_region_leader(&self, _region_id: RegionId) -> error::Result {
- self.region_leader_calls.fetch_add(1, Ordering::SeqCst);
- Ok(Peer {
- id: 1,
- addr: "node-1".to_string(),
- })
+ Ok((Arc::new(SingleRegionPartitionRule), HashMap::new()))
}
}
#[derive(Default)]
- struct MockFlushNodeRequester {
+ struct MockFlushWriteDispatcher {
writes: Arc,
fail: bool,
}
#[async_trait]
- impl PhysicalFlushNodeRequester for MockFlushNodeRequester {
- async fn handle(
- &self,
- _peer: &Peer,
- _request: RegionRequest,
- ) -> error::Result {
+ impl PhysicalFlushWriteDispatcher for MockFlushWriteDispatcher {
+ async fn dispatch(&self, _requests: Vec) -> error::Result {
self.writes.fetch_add(1, Ordering::SeqCst);
if self.fail {
return Err(Error::Internal {
err_msg: "physical write failed".to_string(),
});
}
- Ok(RegionResponse::new(0))
+ Ok(0)
}
}
@@ -726,18 +700,14 @@ mod tests {
}
#[derive(Default)]
- struct AffectedRowsFlushNodeRequester {
+ struct AffectedRowsFlushWriteDispatcher {
affected_rows: usize,
}
#[async_trait]
- impl PhysicalFlushNodeRequester for AffectedRowsFlushNodeRequester {
- async fn handle(
- &self,
- _peer: &Peer,
- _request: RegionRequest,
- ) -> error::Result {
- Ok(RegionResponse::new(self.affected_rows))
+ impl PhysicalFlushWriteDispatcher for AffectedRowsFlushWriteDispatcher {
+ async fn dispatch(&self, _requests: Vec) -> error::Result {
+ Ok(self.affected_rows)
}
}
}
diff --git a/src/servers/src/batcher/logical_table/batch_convert.rs b/src/servers/src/batcher/logical_table/batch_convert.rs
index 37080dbe258..65204ec038c 100644
--- a/src/servers/src/batcher/logical_table/batch_convert.rs
+++ b/src/servers/src/batcher/logical_table/batch_convert.rs
@@ -29,8 +29,6 @@ use crate::error;
use crate::error::Result;
use crate::metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED;
-pub(in crate::batcher::logical_table) const PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT: usize = 3;
-
/// An aligned logical record batch and its timestamp column index.
#[derive(Debug, Clone)]
pub struct RecordBatchWithTsIdx {
@@ -188,23 +186,6 @@ pub(in crate::batcher::logical_table) fn columns_taxonomy(
Ok((tag_columns, essential_column_indices))
}
-pub(in crate::batcher::logical_table) fn strip_partition_columns_from_batch(
- batch: RecordBatch,
-) -> Result {
- ensure!(
- batch.num_columns() >= PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT,
- error::InternalSnafu {
- err_msg: format!(
- "Expected at least {} columns in physical batch, got {}",
- PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT,
- batch.num_columns()
- ),
- }
- );
- let essential_indices: Vec = (0..PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT).collect();
- batch.project(&essential_indices).context(error::ArrowSnafu)
-}
-
/// Transforms logical table batches into physical format (sparse primary key encoding).
///
/// It identifies tag columns and essential columns (timestamp, value) for each logical batch
@@ -290,79 +271,18 @@ mod tests {
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
- use arrow::array::{Array, BinaryArray, StringArray, TimestampMillisecondArray};
+ use arrow::array::{Array, StringArray, TimestampMillisecondArray};
use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema};
use arrow::record_batch::RecordBatch;
use metric_engine::batch_modifier::modify_batch_sparse;
use smallvec::SmallVec;
use crate::batcher::logical_table::batch_convert::{
- TableBatch, columns_taxonomy, strip_partition_columns_from_batch,
- transform_logical_batches_to_physical,
+ TableBatch, columns_taxonomy, transform_logical_batches_to_physical,
};
use crate::batcher::logical_table::test_util::mock_aligned_tag_batch;
use crate::error::Error;
- #[test]
- fn test_strip_partition_columns_from_batch_removes_partition_tags() {
- let batch = RecordBatch::try_new(
- Arc::new(ArrowSchema::new(vec![
- Field::new("__primary_key", ArrowDataType::Binary, false),
- Field::new(
- "greptime_timestamp",
- ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
- false,
- ),
- Field::new("greptime_value", ArrowDataType::Float64, true),
- Field::new("host", ArrowDataType::Utf8, true),
- ])),
- vec![
- Arc::new(BinaryArray::from(vec![b"k1".as_slice()])),
- Arc::new(TimestampMillisecondArray::from(vec![1000_i64])),
- Arc::new(arrow::array::Float64Array::from(vec![42.0_f64])),
- Arc::new(StringArray::from(vec!["node-1"])),
- ],
- )
- .unwrap();
-
- let stripped = strip_partition_columns_from_batch(batch).unwrap();
-
- assert_eq!(3, stripped.num_columns());
- assert_eq!("__primary_key", stripped.schema().field(0).name());
- assert_eq!("greptime_timestamp", stripped.schema().field(1).name());
- assert_eq!("greptime_value", stripped.schema().field(2).name());
- }
-
- #[test]
- fn test_strip_partition_columns_from_batch_projects_essential_columns_without_lookup() {
- let batch = RecordBatch::try_new(
- Arc::new(ArrowSchema::new(vec![
- Field::new("__primary_key", ArrowDataType::Binary, false),
- Field::new(
- "greptime_timestamp",
- ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
- false,
- ),
- Field::new("greptime_value", ArrowDataType::Float64, true),
- Field::new("host", ArrowDataType::Utf8, true),
- ])),
- vec![
- Arc::new(BinaryArray::from(vec![b"k1".as_slice()])),
- Arc::new(TimestampMillisecondArray::from(vec![1000_i64])),
- Arc::new(arrow::array::Float64Array::from(vec![42.0_f64])),
- Arc::new(StringArray::from(vec!["node-1"])),
- ],
- )
- .unwrap();
-
- let stripped = strip_partition_columns_from_batch(batch).unwrap();
-
- assert_eq!(3, stripped.num_columns());
- assert_eq!("__primary_key", stripped.schema().field(0).name());
- assert_eq!("greptime_timestamp", stripped.schema().field(1).name());
- assert_eq!("greptime_value", stripped.schema().field(2).name());
- }
-
#[test]
fn test_collect_tag_columns_and_non_tag_indices_keeps_partition_tag_column() {
let schema = Arc::new(ArrowSchema::new(vec![
diff --git a/src/servers/src/batcher/logical_table/pending_worker.rs b/src/servers/src/batcher/logical_table/pending_worker.rs
index 4a83d8a87f4..ff363c7637c 100644
--- a/src/servers/src/batcher/logical_table/pending_worker.rs
+++ b/src/servers/src/batcher/logical_table/pending_worker.rs
@@ -22,8 +22,8 @@ use common_batcher::flush_policy::FlushTrigger;
use common_batcher::flush_policy::timing::TimingFlushPolicy;
use common_batcher::pending_worker::PendingWorker as PendingCore;
use common_batcher::worker_registry::WorkerRegistry;
-use common_meta::node_manager::NodeManagerRef;
use common_telemetry::debug;
+use operator::write_dispatcher::WriteDispatcherRef;
use partition::manager::PartitionRuleManagerRef;
use session::context::QueryContextRef;
use table::metadata::TableId;
@@ -31,7 +31,7 @@ use tokio::sync::{OwnedSemaphorePermit, broadcast, mpsc, oneshot};
use crate::batcher::flow_notifier::FlowNotifier;
use crate::batcher::logical_table::BatchKey;
-use crate::batcher::logical_table::batch::{Batch, flush_batch_with_managers, spawn_flush};
+use crate::batcher::logical_table::batch::{self, Batch};
use crate::batcher::logical_table::batch_convert::{RecordBatchWithTsIdx, TableBatch};
use crate::error::Error;
use crate::metrics::{PENDING_BATCHES, PENDING_ROWS, PENDING_WORKERS};
@@ -53,7 +53,8 @@ pub(in crate::batcher::logical_table) struct PendingWorker {
pub(in crate::batcher::logical_table) tx: mpsc::Sender,
}
-pub(in crate::batcher::logical_table) enum WorkerCommand {
+/// Commands accepted by a logical-table batching worker.
+pub(crate) enum WorkerCommand {
Submit {
table_batches: Vec<(String, u32, RecordBatchWithTsIdx)>,
total_rows: usize,
@@ -92,15 +93,16 @@ impl PendingBatch {
}
}
+/// Starts a worker that accumulates logical rows and dispatches physical flushes.
#[allow(clippy::too_many_arguments)]
-pub(in crate::batcher::logical_table) fn start_worker(
+pub(crate) fn start_worker(
key: BatchKey,
worker_tx: mpsc::Sender,
workers: Arc>,
mut rx: mpsc::Receiver,
shutdown: broadcast::Sender<()>,
partition_manager: PartitionRuleManagerRef,
- node_manager: NodeManagerRef,
+ write_dispatcher: WriteDispatcherRef,
catalog_manager: CatalogManagerRef,
flow_notification_tx: FlowNotifier,
worker_idle_timeout: Duration,
@@ -140,10 +142,10 @@ pub(in crate::batcher::logical_table) fn start_worker(
PENDING_ROWS.add(total_rows as i64);
if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, Some(FlushTrigger::Submission)) {
- spawn_flush(
+ batch::spawn_flush(
flush,
partition_manager.clone(),
- node_manager.clone(),
+ write_dispatcher.clone(),
catalog_manager.clone(),
flow_notification_tx.clone(),
flush_limiter.clone(),
@@ -152,10 +154,10 @@ pub(in crate::batcher::logical_table) fn start_worker(
}
None => {
if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, None) {
- flush_batch_with_managers(
+ batch::flush_batch_with_managers(
flush,
partition_manager.clone(),
- node_manager.clone(),
+ write_dispatcher.clone(),
catalog_manager.clone(),
flow_notification_tx.clone(),
).await;
@@ -189,10 +191,10 @@ pub(in crate::batcher::logical_table) fn start_worker(
}
_ = pending_flush.wait_flush() => {
if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, Some(FlushTrigger::Deadline)) {
- spawn_flush(
+ batch::spawn_flush(
flush,
partition_manager.clone(),
- node_manager.clone(),
+ write_dispatcher.clone(),
catalog_manager.clone(),
flow_notification_tx.clone(),
flush_limiter.clone(),
@@ -201,10 +203,10 @@ pub(in crate::batcher::logical_table) fn start_worker(
}
_ = shutdown_rx.recv() => {
if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, None) {
- flush_batch_with_managers(
+ batch::flush_batch_with_managers(
flush,
partition_manager.clone(),
- node_manager.clone(),
+ write_dispatcher.clone(),
catalog_manager.clone(),
flow_notification_tx.clone(),
).await;
@@ -286,6 +288,7 @@ mod tests {
use common_meta::kv_backend::memory::MemoryKvBackend;
use common_meta::node_manager::NodeManagerRef;
use moka::future::CacheBuilder;
+ use operator::write_dispatcher::WriteDispatcher;
use partition::cache::new_partition_info_cache;
use partition::manager::PartitionRuleManager;
use tokio::sync::{Semaphore, broadcast, mpsc, oneshot};
@@ -502,6 +505,10 @@ mod tests {
let node_manager: NodeManagerRef = Arc::new(ConcurrentMockNodeManager {
datanodes: Arc::new(HashMap::new()),
});
+ let write_dispatcher = Arc::new(WriteDispatcher::new(
+ node_manager,
+ partition_manager.clone(),
+ ));
let catalog_manager = MemoryCatalogManager::with_default_setup();
let (flow_notification_tx, _flow_notification_rx) =
FlowNotifier::try_new(1, FLOW_NOTIFICATION_DROPPED.clone()).unwrap();
@@ -515,7 +522,7 @@ mod tests {
worker_rx,
shutdown.clone(),
partition_manager,
- node_manager,
+ write_dispatcher,
catalog_manager,
flow_notification_tx,
worker_idle_timeout,
diff --git a/src/servers/src/batcher/logical_table/region_write.rs b/src/servers/src/batcher/logical_table/region_write.rs
index 71b558d1428..9511d21a998 100644
--- a/src/servers/src/batcher/logical_table/region_write.rs
+++ b/src/servers/src/batcher/logical_table/region_write.rs
@@ -14,29 +14,27 @@
use std::collections::HashMap;
-use api::v1::ArrowIpc;
-use api::v1::meta::Peer;
use api::v1::region::{
BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request,
};
-use arrow::compute::filter_record_batch;
+use api::v1::{ArrowIpc, PartitionExprVersion};
+use arrow::compute;
use arrow::record_batch::RecordBatch;
use async_trait::async_trait;
use catalog::CatalogManagerRef;
use common_grpc::error::Error as GrpcError;
-use common_grpc::flight::record_batch_to_ipc;
-use common_meta::node_manager::NodeManagerRef;
+use common_grpc::flight;
use common_telemetry::tracing_context::TracingContext;
+use operator::write_dispatcher::WriteDispatcherRef;
use partition::manager::PartitionRuleManagerRef;
use partition::partition::PartitionRuleRef;
use snafu::ResultExt;
use store_api::storage::RegionId;
use table::metadata::{TableId, TableInfo, TableInfoRef};
-use crate::batcher::logical_table::batch_convert::strip_partition_columns_from_batch;
use crate::error;
use crate::error::{Error, Result};
-use crate::metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED;
+use crate::metrics;
#[derive(Clone)]
pub struct PhysicalTableMetadata {
@@ -56,23 +54,21 @@ pub trait PhysicalFlushCatalogProvider: Send + Sync {
) -> catalog::error::Result>;
}
+/// Loads the physical partition rule and its expression versions.
#[async_trait]
pub trait PhysicalFlushPartitionProvider: Send + Sync {
+ /// Loads a routing snapshot for the physical table.
async fn find_table_partition_rule(
&self,
table_info: &TableInfo,
- ) -> partition::error::Result;
-
- async fn find_region_leader(&self, region_id: RegionId) -> Result;
+ ) -> partition::error::Result<(PartitionRuleRef, HashMap>)>;
}
+/// Dispatches the region requests produced by a physical flush.
#[async_trait]
-pub trait PhysicalFlushNodeRequester: Send + Sync {
- async fn handle(
- &self,
- peer: &Peer,
- request: RegionRequest,
- ) -> Result;
+pub trait PhysicalFlushWriteDispatcher: Send + Sync {
+ /// Dispatches one physical flush together, sharing routing refreshes across its writes.
+ async fn dispatch(&self, requests: Vec) -> Result;
}
#[derive(Clone)]
@@ -115,125 +111,61 @@ impl PhysicalFlushPartitionProvider for PartitionManagerPhysicalFlushAdapter {
async fn find_table_partition_rule(
&self,
table_info: &TableInfo,
- ) -> partition::error::Result {
+ ) -> partition::error::Result<(PartitionRuleRef, HashMap>)> {
self.partition_manager
.find_table_partition_rule(table_info)
.await
- .map(|(rule, _)| rule)
- }
-
- async fn find_region_leader(&self, region_id: RegionId) -> Result {
- let peer = self.partition_manager.find_region_leader(region_id).await?;
- Ok(peer)
}
}
+/// Sends a physical flush through the shared write dispatcher.
#[derive(Clone)]
-pub(in crate::batcher::logical_table) struct NodeManagerPhysicalFlushAdapter {
- pub(in crate::batcher::logical_table) node_manager: NodeManagerRef,
+pub(crate) struct WriteDispatcherPhysicalFlushAdapter {
+ pub(crate) write_dispatcher: WriteDispatcherRef,
}
#[async_trait]
-impl PhysicalFlushNodeRequester for NodeManagerPhysicalFlushAdapter {
- async fn handle(
- &self,
- peer: &Peer,
- request: RegionRequest,
- ) -> error::Result {
- let datanode = self.node_manager.datanode(peer).await;
- datanode
- .handle(request)
+impl PhysicalFlushWriteDispatcher for WriteDispatcherPhysicalFlushAdapter {
+ async fn dispatch(&self, requests: Vec) -> Result {
+ self.write_dispatcher
+ .dispatch_with_metrics(
+ requests,
+ metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
+ .with_label_values(&["flush_write_region"]),
+ )
.await
.context(error::CommonMetaSnafu)
+ .map(|response| response.affected_rows)
}
}
-pub(in crate::batcher::logical_table) struct FlushRegionWrite {
- pub(in crate::batcher::logical_table) datanode: Peer,
- pub(in crate::batcher::logical_table) request: RegionRequest,
-}
-
-pub(in crate::batcher::logical_table) struct PlannedRegionBatch {
- pub(in crate::batcher::logical_table) region_id: RegionId,
- pub(in crate::batcher::logical_table) batch: RecordBatch,
+/// Arrow rows selected for one physical region.
+pub(crate) struct PlannedRegionBatch {
+ pub(crate) region_id: RegionId,
+ pub(crate) batch: RecordBatch,
}
#[cfg(test)]
impl PlannedRegionBatch {
- pub(in crate::batcher::logical_table) fn num_rows(&self) -> usize {
+ pub(crate) fn num_rows(&self) -> usize {
self.batch.num_rows()
}
}
-pub(in crate::batcher::logical_table) struct ResolvedRegionBatch {
- pub(in crate::batcher::logical_table) planned: PlannedRegionBatch,
- pub(in crate::batcher::logical_table) datanode: Peer,
-}
-
-pub(in crate::batcher::logical_table) fn should_dispatch_concurrently(
- region_write_count: usize,
-) -> bool {
- region_write_count > 1
-}
-
-pub(in crate::batcher::logical_table) async fn flush_region_writes_concurrently(
- node_manager: &(impl PhysicalFlushNodeRequester + ?Sized),
- writes: Vec,
-) -> Result {
- let mut affected_rows = 0;
- if !should_dispatch_concurrently(writes.len()) {
- for write in writes {
- let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
- .with_label_values(&["flush_write_region"])
- .start_timer();
- affected_rows += node_manager
- .handle(&write.datanode, write.request)
- .await?
- .affected_rows;
- }
- return Ok(affected_rows);
- }
-
- let write_futures = writes.into_iter().map(|write| async move {
- let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
- .with_label_values(&["flush_write_region"])
- .start_timer();
-
- let response = node_manager.handle(&write.datanode, write.request).await?;
- Ok::<_, Error>(response.affected_rows)
- });
-
- // todo(hl): should be bounded.
- let affected_rows = futures::future::try_join_all(write_futures)
- .await?
- .into_iter()
- .sum();
- Ok(affected_rows)
-}
-
pub(in crate::batcher::logical_table) fn split_combined_batch_by_region(
combined_batch: &RecordBatch,
partition_rule: &dyn partition::partition::PartitionRule,
) -> Result> {
- let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
+ let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
.with_label_values(&["flush_physical_split_record_batch"])
.start_timer();
let map = partition_rule.split_record_batch(combined_batch)?;
Ok(map)
}
-pub(in crate::batcher::logical_table) fn prepare_physical_region_routing_batch(
- combined_batch: RecordBatch,
- partition_columns: &[String],
-) -> Result {
- if partition_columns.is_empty() {
- return Ok(combined_batch);
- }
- strip_partition_columns_from_batch(combined_batch)
-}
-
-pub(in crate::batcher::logical_table) fn plan_region_batch(
- stripped_batch: &RecordBatch,
+/// Selects the rows assigned to one physical region.
+pub(crate) fn plan_region_batch(
+ combined_batch: &RecordBatch,
physical_table_id: TableId,
region_number: u32,
mask: &partition::partition::RegionMask,
@@ -243,12 +175,12 @@ pub(in crate::batcher::logical_table) fn plan_region_batch(
}
let region_batch = if mask.select_all() {
- stripped_batch.clone()
+ combined_batch.clone()
} else {
- let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
+ let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
.with_label_values(&["flush_physical_filter_record_batch"])
.start_timer();
- filter_record_batch(stripped_batch, mask.array()).context(error::ArrowSnafu)?
+ compute::filter_record_batch(combined_batch, mask.array()).context(error::ArrowSnafu)?
};
let row_count = region_batch.num_rows();
@@ -262,19 +194,20 @@ pub(in crate::batcher::logical_table) fn plan_region_batch(
}))
}
-pub(in crate::batcher::logical_table) fn plan_region_batches(
+/// Partitions a physical flush while retaining columns needed for rerouting.
+pub(crate) fn plan_region_batches(
combined_batch: RecordBatch,
physical_table_id: TableId,
partition_rule: &dyn partition::partition::PartitionRule,
- partition_columns: &[String],
) -> Result> {
let region_masks = split_combined_batch_by_region(&combined_batch, partition_rule)?;
- let stripped_batch = prepare_physical_region_routing_batch(combined_batch, partition_columns)?;
+ // Retain partition columns in the RPC payload so routing can be rebuilt after repartition.
+ // Mito's sparse bulk conversion projects the stored fields and encoded primary key.
let mut planned_batches = Vec::new();
for (region_number, mask) in region_masks {
if let Some(planned_batch) =
- plan_region_batch(&stripped_batch, physical_table_id, region_number, &mask)?
+ plan_region_batch(&combined_batch, physical_table_id, region_number, &mask)?
{
planned_batches.push(planned_batch);
}
@@ -283,39 +216,20 @@ pub(in crate::batcher::logical_table) fn plan_region_batches(
Ok(planned_batches)
}
-pub(in crate::batcher::logical_table) async fn resolve_region_targets(
+/// Encodes each region batch with the partition version used to select its rows.
+pub(crate) fn encode_region_write_requests(
planned_batches: Vec,
- partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized),
-) -> Result> {
- let mut resolved_batches = Vec::with_capacity(planned_batches.len());
- for planned in planned_batches {
- let datanode = {
- let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
- .with_label_values(&["flush_physical_resolve_region_leader"])
- .start_timer();
- partition_manager
- .find_region_leader(planned.region_id)
- .await?
- };
-
- resolved_batches.push(ResolvedRegionBatch { planned, datanode });
- }
-
- Ok(resolved_batches)
-}
-
-pub(in crate::batcher::logical_table) fn encode_region_write_requests(
- resolved_batches: Vec,
skip_wal: bool,
-) -> Result> {
- let mut region_writes = Vec::with_capacity(resolved_batches.len());
- for resolved in resolved_batches {
- let region_id = resolved.planned.region_id;
+ partition_versions: &HashMap>,
+) -> Result> {
+ let mut region_writes = Vec::with_capacity(planned_batches.len());
+ for planned in planned_batches {
+ let region_id = planned.region_id;
let (schema_bytes, data_header, payload) = {
- let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
+ let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
.with_label_values(&["flush_physical_encode_ipc"])
.start_timer();
- record_batch_to_ipc(resolved.planned.batch).map_err(map_ipc_error)?
+ flight::record_batch_to_ipc(planned.batch).map_err(map_ipc_error)?
};
let request = RegionRequest {
@@ -326,7 +240,11 @@ pub(in crate::batcher::logical_table) fn encode_region_write_requests(
body: Some(region_request::Body::BulkInsert(BulkInsertRequest {
skip_wal,
region_id: region_id.as_u64(),
- partition_expr_version: None,
+ partition_expr_version: partition_versions
+ .get(®ion_id.region_number())
+ .copied()
+ .flatten()
+ .map(|value| PartitionExprVersion { value }),
// Set aligned_schema_version to None so that datanode will check the batch schema again to see if any
// column is missing.
aligned_schema_version: None,
@@ -338,10 +256,7 @@ pub(in crate::batcher::logical_table) fn encode_region_write_requests(
})),
};
- region_writes.push(FlushRegionWrite {
- datanode: resolved.datanode,
- request,
- });
+ region_writes.push(request);
}
Ok(region_writes)
@@ -359,89 +274,27 @@ pub(in crate::batcher::logical_table) fn map_ipc_error(error: GrpcError) -> Erro
#[cfg(test)]
mod tests {
- use std::collections::HashMap;
+ use std::collections::{HashMap, HashSet};
use std::sync::Arc;
- use std::sync::atomic::{AtomicUsize, Ordering};
- use std::time::Duration;
- use api::region::RegionResponse;
- use api::v1::meta::Peer;
- use api::v1::region::{RegionRequest, region_request};
- use arrow::array::{BinaryArray, BooleanArray, StringArray, TimestampMillisecondArray};
+ use api::v1::PartitionExprVersion;
+ use api::v1::region::{bulk_insert_request, region_request};
+ use arrow::array::{Array, BinaryArray, BooleanArray, StringArray, TimestampMillisecondArray};
use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema};
use arrow::record_batch::RecordBatch;
- use async_trait::async_trait;
- use common_meta::error::Result as MetaResult;
- use common_meta::node_manager::{Datanode, DatanodeRef};
- use common_query::request::QueryRequest;
- use common_recordbatch::SendableRecordBatchStream;
+ use common_grpc::flight::FlightDecoder;
use partition::partition::{PartitionRule, RegionMask};
use store_api::storage::RegionId;
- use tokio::time::sleep;
+ use crate::batcher::logical_table::batch_convert::{self, RecordBatchWithTsIdx, TableBatch};
use crate::batcher::logical_table::region_write::{
- FlushRegionWrite, PlannedRegionBatch, ResolvedRegionBatch, encode_region_write_requests,
- flush_region_writes_concurrently, plan_region_batches, should_dispatch_concurrently,
+ PlannedRegionBatch, encode_region_write_requests, plan_region_batches,
};
- use crate::batcher::logical_table::test_util::ConcurrentMockNodeManager;
-
- #[tokio::test]
- async fn test_flush_region_writes_concurrently_dispatches_multiple_datanodes() {
- let inflight = Arc::new(AtomicUsize::new(0));
- let max_inflight = Arc::new(AtomicUsize::new(0));
- let datanode1: DatanodeRef = Arc::new(ConcurrentMockDatanode {
- delay: Duration::from_millis(100),
- inflight: inflight.clone(),
- max_inflight: max_inflight.clone(),
- });
- let datanode2: DatanodeRef = Arc::new(ConcurrentMockDatanode {
- delay: Duration::from_millis(100),
- inflight,
- max_inflight: max_inflight.clone(),
- });
-
- let mut datanodes = HashMap::new();
- datanodes.insert(1, datanode1);
- datanodes.insert(2, datanode2);
- let node_manager = Arc::new(ConcurrentMockNodeManager {
- datanodes: Arc::new(datanodes),
- });
-
- let writes = vec![
- FlushRegionWrite {
- datanode: Peer {
- id: 1,
- addr: "node1".to_string(),
- },
- request: RegionRequest::default(),
- },
- FlushRegionWrite {
- datanode: Peer {
- id: 2,
- addr: "node2".to_string(),
- },
- request: RegionRequest::default(),
- },
- ];
-
- flush_region_writes_concurrently(node_manager.as_ref(), writes)
- .await
- .unwrap();
- assert!(max_inflight.load(Ordering::SeqCst) >= 2);
- }
#[test]
- fn test_should_dispatch_concurrently_by_region_count() {
- assert!(!should_dispatch_concurrently(0));
- assert!(!should_dispatch_concurrently(1));
- assert!(should_dispatch_concurrently(2));
- }
-
- #[test]
- fn test_plan_region_batches_splits_and_strips_partition_columns() {
- let combined_batch = RecordBatch::try_new(
+ fn test_plan_region_batches_preserves_partition_columns_through_encoding() {
+ let logical_batch = RecordBatch::try_new(
Arc::new(ArrowSchema::new(vec![
- Field::new("__primary_key", ArrowDataType::Binary, false),
Field::new(
"greptime_timestamp",
ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
@@ -449,22 +302,38 @@ mod tests {
),
Field::new("greptime_value", ArrowDataType::Float64, true),
Field::new("host", ArrowDataType::Utf8, true),
+ Field::new("zone", ArrowDataType::Utf8, true),
])),
vec![
- Arc::new(BinaryArray::from(vec![b"k1".as_slice(), b"k2".as_slice()])),
Arc::new(TimestampMillisecondArray::from(vec![1000_i64, 2000_i64])),
Arc::new(arrow::array::Float64Array::from(vec![1.0_f64, 2.0_f64])),
Arc::new(StringArray::from(vec!["node-1", "node-2"])),
+ Arc::new(StringArray::from(vec!["east", "west"])),
],
)
.unwrap();
+ let table_batches = [TableBatch {
+ table_name: "metric".to_string(),
+ table_id: 2048,
+ batches: vec![RecordBatchWithTsIdx::try_new(logical_batch, 0).unwrap()],
+ row_count: 2,
+ }];
+ let modified = batch_convert::transform_logical_batches_to_physical(
+ &table_batches,
+ &HashMap::from([("host".to_string(), 1), ("zone".to_string(), 2)]),
+ &HashSet::from(["host"]),
+ )
+ .unwrap();
+ let combined_batch = batch_convert::concat_modified_batches(&modified).unwrap();
+ assert!(combined_batch.column_by_name("host").is_some());
+ assert!(combined_batch.column_by_name("zone").is_none());
+ let expected = combined_batch.clone();
let mut planned_batches = plan_region_batches(
combined_batch,
1024,
&TwoRegionPartitionRule {
partition_columns: vec!["host".to_string()],
},
- &["host".to_string()],
)
.unwrap();
planned_batches.sort_by_key(|planned| planned.region_id.region_number());
@@ -472,10 +341,45 @@ mod tests {
assert_eq!(2, planned_batches.len());
assert_eq!(RegionId::new(1024, 1), planned_batches[0].region_id);
assert_eq!(1, planned_batches[0].num_rows());
- assert_eq!(3, planned_batches[0].batch.num_columns());
+ assert_eq!(4, planned_batches[0].batch.num_columns());
assert_eq!(RegionId::new(1024, 2), planned_batches[1].region_id);
assert_eq!(1, planned_batches[1].num_rows());
- assert_eq!(3, planned_batches[1].batch.num_columns());
+ assert_eq!(4, planned_batches[1].batch.num_columns());
+
+ let writes = encode_region_write_requests(
+ planned_batches,
+ false,
+ &HashMap::from([(1, Some(7)), (2, Some(8))]),
+ )
+ .unwrap();
+ for (index, write) in writes.into_iter().enumerate() {
+ let Some(region_request::Body::BulkInsert(request)) = write.body else {
+ panic!("expected bulk")
+ };
+ assert_eq!(
+ request.partition_expr_version.unwrap().value,
+ 7 + index as u64
+ );
+ let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = request.body else {
+ panic!("expected IPC")
+ };
+ let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema)
+ .unwrap()
+ .try_decode_record_batch(&ipc.data_header, &ipc.payload)
+ .unwrap();
+ // The partition-only Arrow projection must retain the complete sparse key,
+ // including the zone tag that a later partition rule can require.
+ assert!(batch.column_by_name("zone").is_none());
+ assert_eq!(batch, expected.slice(index, 1));
+ let tags = batch
+ .column_by_name("host")
+ .unwrap()
+ .as_any()
+ .downcast_ref::()
+ .unwrap();
+ assert_eq!(tags.len(), 1);
+ assert_eq!(tags.value(0), if index == 0 { "node-1" } else { "node-2" });
+ }
}
#[test]
@@ -520,44 +424,6 @@ mod tests {
}
}
- #[derive(Clone)]
- struct ConcurrentMockDatanode {
- delay: Duration,
- inflight: Arc,
- max_inflight: Arc,
- }
-
- #[async_trait]
- impl Datanode for ConcurrentMockDatanode {
- async fn handle(&self, _request: RegionRequest) -> MetaResult {
- let now = self.inflight.fetch_add(1, Ordering::SeqCst) + 1;
- loop {
- let max = self.max_inflight.load(Ordering::SeqCst);
- if now <= max {
- break;
- }
- if self
- .max_inflight
- .compare_exchange(max, now, Ordering::SeqCst, Ordering::SeqCst)
- .is_ok()
- {
- break;
- }
- }
-
- sleep(self.delay).await;
- self.inflight.fetch_sub(1, Ordering::SeqCst);
- Ok(RegionResponse::new(0))
- }
-
- async fn handle_query(
- &self,
- _request: QueryRequest,
- ) -> MetaResult {
- unimplemented!()
- }
- }
-
fn check_encode_region_write_requests(skip_wal: bool) {
let planned_batch = PlannedRegionBatch {
region_id: RegionId::new(1024, 1),
@@ -579,21 +445,22 @@ mod tests {
)
.unwrap(),
};
- let resolved_batch = ResolvedRegionBatch {
- planned: planned_batch,
- datanode: Peer {
- id: 1,
- addr: "node-1".to_string(),
- },
- };
- let writes = encode_region_write_requests(vec![resolved_batch], skip_wal).unwrap();
+ let writes = encode_region_write_requests(
+ vec![planned_batch],
+ skip_wal,
+ &HashMap::from([(1, Some(7))]),
+ )
+ .unwrap();
assert_eq!(1, writes.len());
- assert_eq!(1, writes[0].datanode.id);
- let Some(region_request::Body::BulkInsert(request)) = &writes[0].request.body else {
+ let Some(region_request::Body::BulkInsert(request)) = &writes[0].body else {
panic!("expected bulk insert request");
};
assert_eq!(RegionId::new(1024, 1).as_u64(), request.region_id);
assert_eq!(skip_wal, request.skip_wal);
+ assert_eq!(
+ Some(PartitionExprVersion { value: 7 }),
+ request.partition_expr_version
+ );
}
}
diff --git a/src/servers/src/batcher/logical_table/test_util.rs b/src/servers/src/batcher/logical_table/test_util.rs
index 1cc81f9b615..491bd6ff2b3 100644
--- a/src/servers/src/batcher/logical_table/test_util.rs
+++ b/src/servers/src/batcher/logical_table/test_util.rs
@@ -15,10 +15,9 @@
use std::collections::HashMap;
use std::sync::Arc;
-use api::region::RegionResponse;
use api::v1::flow::{DirtyWindowRequests, FlowRequest, FlowResponse};
use api::v1::meta::Peer;
-use api::v1::region::{InsertRequests, RegionRequest};
+use api::v1::region::InsertRequests;
use arrow::array::{StringArray, TimestampMillisecondArray};
use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema};
use arrow::record_batch::RecordBatch;
@@ -27,12 +26,9 @@ use common_meta::error::Result as MetaResult;
use common_meta::node_manager::{
DatanodeManager, DatanodeRef, Flownode, FlownodeManager, FlownodeRef,
};
-use snafu::ResultExt;
use tokio::sync::mpsc;
use crate::batcher::logical_table::batch_convert::RecordBatchWithTsIdx;
-use crate::batcher::logical_table::region_write::PhysicalFlushNodeRequester;
-use crate::error;
pub(in crate::batcher::logical_table) fn mock_tag_batch(
tag_name: &str,
@@ -149,14 +145,3 @@ impl FlownodeManager for FlowNotificationMockNodeManager {
self.flownode.clone()
}
}
-
-#[async_trait]
-impl PhysicalFlushNodeRequester for ConcurrentMockNodeManager {
- async fn handle(&self, peer: &Peer, request: RegionRequest) -> error::Result {
- let datanode = self.datanode(peer).await;
- datanode
- .handle(request)
- .await
- .context(error::CommonMetaSnafu)
- }
-}
diff --git a/src/standalone/src/datanode_manager.rs b/src/standalone/src/datanode_manager.rs
index 13a17ea2b43..a27a94b3c8d 100644
--- a/src/standalone/src/datanode_manager.rs
+++ b/src/standalone/src/datanode_manager.rs
@@ -77,7 +77,17 @@ impl RegionInvoker {
#[async_trait]
impl Datanode for RegionInvoker {
- async fn handle(&self, request: RegionRequest) -> MetaResult {
+ async fn handle_request(&self, request: Arc) -> MetaResult {
+ // Local execution consumes the request; standalone does not retain it for retries.
+ let request = Arc::try_unwrap(request)
+ .map_err(|_| {
+ InvalidRegionRequestSnafu {
+ reason: "local region request has multiple owners",
+ }
+ .build()
+ })
+ .map_err(BoxedError::new)
+ .context(meta_error::ExternalSnafu)?;
let span = request
.header
.as_ref()
diff --git a/src/standalone/src/options.rs b/src/standalone/src/options.rs
index 079cae4bca6..05d929cc7da 100644
--- a/src/standalone/src/options.rs
+++ b/src/standalone/src/options.rs
@@ -25,7 +25,7 @@ use flow::FlowConfig;
use frontend::frontend::FrontendOptions;
use frontend::service_config::{
InfluxdbOptions, JaegerOptions, MysqlOptions, OpentsdbOptions, OtlpOptions,
- PendingRowsBatcherOptions, PostgresOptions, PromStoreOptions,
+ PendingRowsBatcherOptions, PostgresOptions, PromStoreOptions, WriteRetryOptions,
};
use mito2::config::MitoConfig;
use pipeline::PipelineOptions;
@@ -168,6 +168,11 @@ impl StandaloneOptions {
otlp: cloned_opts.otlp,
prom_store: cloned_opts.prom_store,
meta_client: None,
+ // Local region calls consume requests directly and do not use remote write retries.
+ write_retry: WriteRetryOptions {
+ max_retries: 0,
+ ..Default::default()
+ },
logging: cloned_opts.logging,
user_provider: cloned_opts.user_provider,
query: cloned_opts.query,
@@ -217,6 +222,18 @@ mod tests {
use crate::options::*;
+ #[test]
+ fn test_frontend_write_retries_disabled() {
+ assert!(FrontendOptions::default().write_retry.max_retries > 0);
+ assert_eq!(
+ StandaloneOptions::default()
+ .frontend_options()
+ .write_retry
+ .max_retries,
+ 0
+ );
+ }
+
#[test]
fn test_logical_batcher_config_forwarding() {
let opts: StandaloneOptions = toml::from_str("[pending_rows_batcher]\nprotocols = ['influxdb']\n[pending_rows_batcher.logical_table]\nprotocols = ['otlp', 'prom']\npending_rows_flush_interval = '10ms'").unwrap();
diff --git a/tests-integration/src/otlp.rs b/tests-integration/src/otlp.rs
index 3e447bb0194..2349d580c5f 100644
--- a/tests-integration/src/otlp.rs
+++ b/tests-integration/src/otlp.rs
@@ -134,9 +134,9 @@ mod test {
#[async_trait::async_trait]
impl Datanode for GatedTraceDatanode {
- async fn handle(
+ async fn handle_request(
&self,
- request: RegionRequest,
+ request: Arc,
) -> common_meta::error::Result {
let gate = {
let mut gate = self.gate.lock().unwrap();
diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs
index 64f9f4c6af1..0729e6ea96b 100644
--- a/tests-integration/src/test_util.rs
+++ b/tests-integration/src/test_util.rs
@@ -892,6 +892,7 @@ pub fn build_test_prom_server(
LogicalTablePendingRowsBatcher::try_new(
frontend_ref.partition_manager().clone(),
frontend_ref.node_manager().clone(),
+ frontend_ref.inserter().write_dispatcher().clone(),
frontend_ref.catalog_manager().clone(),
frontend_ref.table_flownode_set_cache().clone(),
true,