mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-10-13 23:42:32 +00:00
feat(frontend): retry transient distributed write failures (#9543)
* refactor(operator): route shared requests through NodeManager Remove the concrete NodeClients requester API and use the existing node manager for owned and shared requests. Preserve the original write hint while the legacy error adapter remains in place. Signed-off-by: WenyXu <wenymedia@gmail.com> * fix(meta): preserve retry policy through error hints Propagate client errors through External and let DDL read retry hints while preserving cleanup and terminal-child policies. Region transport failures retain the legacy Retryable classification, which now also participates in bounded write retries. Signed-off-by: WenyXu <wenymedia@gmail.com> * fix(operator): restore sparse partition tags when rebuilding writes Decode retained primary keys when refreshed partition expressions need different tags. Keep the partitioner and multi-region sparse write regression tests with the fix. Signed-off-by: WenyXu <wenymedia@gmail.com> * fix(operator): stop detached write retries when callers cancel Signed-off-by: WenyXu <wenymedia@gmail.com> * feat(frontend): configure retries across distributed write paths Wire the shared dispatcher through frontend and logical batching, preserving partition columns and request headers. Keep standalone retries disabled and include configuration, batching, and storage regression coverage. Signed-off-by: WenyXu <wenymedia@gmail.com> --------- Signed-off-by: WenyXu <wenymedia@gmail.com>
This commit is contained in:
Generated
+1
@@ -10136,6 +10136,7 @@ dependencies = [
|
||||
"meta-client",
|
||||
"meter-core",
|
||||
"meter-macros",
|
||||
"mito-codec",
|
||||
"moka",
|
||||
"object-store",
|
||||
"partition",
|
||||
|
||||
@@ -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").<br/>Ignored unless mode is "custom". |
|
||||
| `query.experimental_spill_compression` | String | `uncompressed` | Compression for spilled data files: "uncompressed" (default), "lz4_frame", "zstd".<br/>Ignored unless mode is "custom". |
|
||||
| `write_retry` | -- | -- | Retry retryable distributed write failures.<br/>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`. |
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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 {
|
||||
|
||||
+17
-8
@@ -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<Error> 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!(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+110
-24
@@ -67,7 +67,7 @@ pub struct RegionRequester {
|
||||
|
||||
#[async_trait]
|
||||
impl Datanode for RegionRequester {
|
||||
async fn handle(&self, request: RegionRequest) -> MetaResult<RegionResponse> {
|
||||
async fn handle_request(&self, request: Arc<RegionRequest>) -> MetaResult<RegionResponse> {
|
||||
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<S>(
|
||||
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::<Error>().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:?}"
|
||||
);
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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<Status> {
|
||||
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(())));
|
||||
|
||||
@@ -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<T: Debug>(results: Vec<Result<T>>) -> 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";
|
||||
|
||||
@@ -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)));
|
||||
|
||||
@@ -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<RegionResponse>;
|
||||
async fn handle_request(&self, request: Arc<RegionRequest>) -> Result<RegionResponse>;
|
||||
|
||||
/// Handles query requests
|
||||
async fn handle_query(&self, request: QueryRequest) -> Result<SendableRecordBatchStream>;
|
||||
}
|
||||
|
||||
impl dyn Datanode + '_ {
|
||||
/// Handles an owned or shared request without cloning its payload.
|
||||
pub async fn handle<R>(&self, request: R) -> Result<RegionResponse>
|
||||
where
|
||||
R: Into<Arc<RegionRequest>> + Send,
|
||||
{
|
||||
self.handle_request(request.into()).await
|
||||
}
|
||||
}
|
||||
|
||||
pub type DatanodeRef = Arc<dyn Datanode>;
|
||||
|
||||
/// The trait for handling requests to flownode
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -111,7 +111,8 @@ struct MockNode<T> {
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl<T: MockDatanodeHandler> Datanode for MockNode<T> {
|
||||
async fn handle(&self, request: RegionRequest) -> Result<RegionResponse> {
|
||||
async fn handle_request(&self, request: Arc<RegionRequest>) -> Result<RegionResponse> {
|
||||
let request = Arc::try_unwrap(request).expect("mock request must be exclusively owned");
|
||||
self.handler.handle(&self.peer, request).await
|
||||
}
|
||||
|
||||
|
||||
@@ -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)));
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
|
||||
@@ -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<MetaClientOptions>,
|
||||
pub logging: LoggingOptions,
|
||||
pub datanode: DatanodeClientOptions,
|
||||
/// Retry policy for retryable distributed write failures.
|
||||
pub write_retry: WriteRetryOptions,
|
||||
pub user_provider: Option<String>,
|
||||
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<T> =
|
||||
Pin<Box<dyn Stream<Item = std::result::Result<T, Status>> + 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(
|
||||
|
||||
@@ -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<CacheInvalidatorRef>,
|
||||
local_cache_invalidator: CacheInvalidatorRef,
|
||||
catalog_manager: CatalogManagerRef,
|
||||
node_manager: NodeManagerRef,
|
||||
plugins: Option<Plugins>,
|
||||
@@ -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<Instance> {
|
||||
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::<RegionQueryHandlerFactoryRef>() {
|
||||
@@ -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()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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::<usize>(),
|
||||
if skip_wal && !close { 0 } else { 6 },
|
||||
);
|
||||
batches.pretty_print().unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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<api::v1::region::RegionRequest>,
|
||||
) -> common_meta::error::Result<api::region::RegionResponse> {
|
||||
let rows = match _request.body.as_ref() {
|
||||
Some(api::v1::region::region_request::Body::Inserts(requests)) => requests
|
||||
|
||||
@@ -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<Vec<RegionRequest>> {
|
||||
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<ArrowIpc>,
|
||||
original_ipc: Option<ArrowIpc>,
|
||||
skip_wal: bool,
|
||||
) -> Result<Vec<BulkInsertRequest>> {
|
||||
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<RegionNumber, RegionMask>,
|
||||
mut original_ipc: Option<ArrowIpc>,
|
||||
skip_wal: bool,
|
||||
) -> Result<Vec<BulkInsertRequest>> {
|
||||
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<Vec<Body>> {
|
||||
fn rebuild_bulk_insert_request(
|
||||
&self,
|
||||
bulk: BulkInsertRequest,
|
||||
table_info: &TableInfo,
|
||||
) -> Result<Vec<Body>> {
|
||||
let region_id = RegionId::from_u64(bulk.region_id);
|
||||
if self.is_partition_expr_unchanged(region_id, bulk.partition_expr_version) {
|
||||
return Ok(vec![Body::BulkInsert(bulk)]);
|
||||
@@ -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<Option<RecordBatch>> {
|
||||
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::<Vec<_>>();
|
||||
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::<BinaryArray>()
|
||||
.context(error::UnexpectedSnafu {
|
||||
violated: "Sparse primary key column must be Binary",
|
||||
})?;
|
||||
let mut builders = missing
|
||||
.iter()
|
||||
.map(|_| StringBuilder::new())
|
||||
.collect::<Vec<_>>();
|
||||
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::<Vec<_>>();
|
||||
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::<StringArray>()
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
zones.iter().collect::<Vec<_>>(),
|
||||
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::<StringArray>()
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
zones.iter().collect::<Vec<_>>(),
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -113,7 +113,7 @@ fn new_test_region_wal_options(regions: Vec<RegionNumber>) -> 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(),
|
||||
|
||||
@@ -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<NodeClients>,
|
||||
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<Histogram>,
|
||||
cancellation: Option<&CancellationToken>,
|
||||
) -> Result<RegionResponse> {
|
||||
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<WriteAttempt>,
|
||||
histogram: Option<&Histogram>,
|
||||
cancellation: Option<&CancellationToken>,
|
||||
) -> Result<RegionResponse> {
|
||||
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<NodeClients>,
|
||||
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<WriteAttempt>,
|
||||
histogram: Option<&Histogram>,
|
||||
cancellation: Option<&CancellationToken>,
|
||||
) -> Result<RegionResponse> {
|
||||
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);
|
||||
|
||||
@@ -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<Mutex<Vec<RegionRequest>>>,
|
||||
rejections: Arc<Mutex<HashMap<u64, VecDeque<RetryHint>>>>,
|
||||
rejections: Arc<Mutex<HashMap<u64, VecDeque<Status>>>>,
|
||||
release: Arc<Mutex<Option<Arc<Semaphore>>>>,
|
||||
}
|
||||
|
||||
#[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::<TableInfoCacheRef>().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<Item = RetryHint>) {
|
||||
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<u64>,
|
||||
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::<Vec<_>>();
|
||||
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!(
|
||||
|
||||
@@ -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**
|
||||
|
||||
@@ -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<PartitionRuleRef> {
|
||||
Ok(Arc::new(SingleRegionPartitionRule))
|
||||
}
|
||||
|
||||
async fn find_region_leader(&self, _region_id: RegionId) -> Result<Peer> {
|
||||
Ok(Peer {
|
||||
id: 1,
|
||||
addr: "bench-node".to_string(),
|
||||
})
|
||||
) -> PartitionResult<(PartitionRuleRef, HashMap<u32, Option<u64>>)> {
|
||||
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<RegionResponse> {
|
||||
Ok(RegionResponse::new(0))
|
||||
impl PhysicalFlushWriteDispatcher for BenchWriteDispatcher {
|
||||
async fn dispatch(&self, _requests: Vec<RegionRequest>) -> error::Result<usize> {
|
||||
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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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<TableBatch>,
|
||||
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<FlushWaiter>,
|
||||
/// Rows and completion waiters owned by one physical flush.
|
||||
pub(crate) struct Batch {
|
||||
pub(crate) table_batches: Vec<TableBatch>,
|
||||
pub(crate) total_row_count: usize,
|
||||
pub(crate) db_string: String,
|
||||
pub(crate) ctx: QueryContextRef,
|
||||
pub(crate) waiters: Vec<FlushWaiter>,
|
||||
}
|
||||
|
||||
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<usize> {
|
||||
// 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<AtomicUsize>,
|
||||
region_leader_calls: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -679,40 +665,28 @@ mod tests {
|
||||
async fn find_table_partition_rule(
|
||||
&self,
|
||||
_table_info: &table::metadata::TableInfo,
|
||||
) -> PartitionResult<PartitionRuleRef> {
|
||||
) -> PartitionResult<(PartitionRuleRef, HashMap<u32, Option<u64>>)> {
|
||||
self.partition_rule_calls.fetch_add(1, Ordering::SeqCst);
|
||||
Ok(Arc::new(SingleRegionPartitionRule))
|
||||
}
|
||||
|
||||
async fn find_region_leader(&self, _region_id: RegionId) -> error::Result<Peer> {
|
||||
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<AtomicUsize>,
|
||||
fail: bool,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl PhysicalFlushNodeRequester for MockFlushNodeRequester {
|
||||
async fn handle(
|
||||
&self,
|
||||
_peer: &Peer,
|
||||
_request: RegionRequest,
|
||||
) -> error::Result<RegionResponse> {
|
||||
impl PhysicalFlushWriteDispatcher for MockFlushWriteDispatcher {
|
||||
async fn dispatch(&self, _requests: Vec<RegionRequest>) -> error::Result<usize> {
|
||||
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<RegionResponse> {
|
||||
Ok(RegionResponse::new(self.affected_rows))
|
||||
impl PhysicalFlushWriteDispatcher for AffectedRowsFlushWriteDispatcher {
|
||||
async fn dispatch(&self, _requests: Vec<RegionRequest>) -> error::Result<usize> {
|
||||
Ok(self.affected_rows)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<RecordBatch> {
|
||||
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<usize> = (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![
|
||||
|
||||
@@ -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<WorkerCommand>,
|
||||
}
|
||||
|
||||
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<WorkerCommand>,
|
||||
workers: Arc<WorkerRegistry<BatchKey, WorkerCommand>>,
|
||||
mut rx: mpsc::Receiver<WorkerCommand>,
|
||||
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,
|
||||
|
||||
@@ -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<Option<PhysicalTableMetadata>>;
|
||||
}
|
||||
|
||||
/// 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<PartitionRuleRef>;
|
||||
|
||||
async fn find_region_leader(&self, region_id: RegionId) -> Result<Peer>;
|
||||
) -> partition::error::Result<(PartitionRuleRef, HashMap<u32, Option<u64>>)>;
|
||||
}
|
||||
|
||||
/// 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<api::region::RegionResponse>;
|
||||
pub trait PhysicalFlushWriteDispatcher: Send + Sync {
|
||||
/// Dispatches one physical flush together, sharing routing refreshes across its writes.
|
||||
async fn dispatch(&self, requests: Vec<RegionRequest>) -> Result<usize>;
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
@@ -115,125 +111,61 @@ impl PhysicalFlushPartitionProvider for PartitionManagerPhysicalFlushAdapter {
|
||||
async fn find_table_partition_rule(
|
||||
&self,
|
||||
table_info: &TableInfo,
|
||||
) -> partition::error::Result<PartitionRuleRef> {
|
||||
) -> partition::error::Result<(PartitionRuleRef, HashMap<u32, Option<u64>>)> {
|
||||
self.partition_manager
|
||||
.find_table_partition_rule(table_info)
|
||||
.await
|
||||
.map(|(rule, _)| rule)
|
||||
}
|
||||
|
||||
async fn find_region_leader(&self, region_id: RegionId) -> Result<Peer> {
|
||||
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<api::region::RegionResponse> {
|
||||
let datanode = self.node_manager.datanode(peer).await;
|
||||
datanode
|
||||
.handle(request)
|
||||
impl PhysicalFlushWriteDispatcher for WriteDispatcherPhysicalFlushAdapter {
|
||||
async fn dispatch(&self, requests: Vec<RegionRequest>) -> Result<usize> {
|
||||
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<FlushRegionWrite>,
|
||||
) -> Result<usize> {
|
||||
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<HashMap<u32, partition::partition::RegionMask>> {
|
||||
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<RecordBatch> {
|
||||
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<Vec<PlannedRegionBatch>> {
|
||||
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<PlannedRegionBatch>,
|
||||
partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized),
|
||||
) -> Result<Vec<ResolvedRegionBatch>> {
|
||||
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<ResolvedRegionBatch>,
|
||||
skip_wal: bool,
|
||||
) -> Result<Vec<FlushRegionWrite>> {
|
||||
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<u32, Option<u64>>,
|
||||
) -> Result<Vec<RegionRequest>> {
|
||||
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::<StringArray>()
|
||||
.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<AtomicUsize>,
|
||||
max_inflight: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Datanode for ConcurrentMockDatanode {
|
||||
async fn handle(&self, _request: RegionRequest) -> MetaResult<RegionResponse> {
|
||||
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<SendableRecordBatchStream> {
|
||||
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
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<RegionResponse> {
|
||||
let datanode = self.datanode(peer).await;
|
||||
datanode
|
||||
.handle(request)
|
||||
.await
|
||||
.context(error::CommonMetaSnafu)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -77,7 +77,17 @@ impl RegionInvoker {
|
||||
|
||||
#[async_trait]
|
||||
impl Datanode for RegionInvoker {
|
||||
async fn handle(&self, request: RegionRequest) -> MetaResult<RegionResponse> {
|
||||
async fn handle_request(&self, request: Arc<RegionRequest>) -> MetaResult<RegionResponse> {
|
||||
// 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()
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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<RegionRequest>,
|
||||
) -> common_meta::error::Result<RegionResponse> {
|
||||
let gate = {
|
||||
let mut gate = self.gate.lock().unwrap();
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user