From 372f33540a03b2a400611dfd2a27a591365d41e0 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 14 Oct 2026 14:02:28 +0800 Subject: [PATCH] feat(frontend): backport write retries to release/5e02efa9 (#9569) * feat(client): support region request reuse and route refresh hints (#9540) * feat(client): support shared region requests without deep cloning Signed-off-by: WenyXu * feat: propagate route refresh retry hints Signed-off-by: WenyXu * style: preserve local error variant imports in retry hints Signed-off-by: WenyXu * fix(mito2): distinguish staging rejection from write backpressure Signed-off-by: WenyXu * fix(meta): preserve legacy retry hints in instruction replies Signed-off-by: WenyXu * test(client): inline retry hint compatibility cases Signed-off-by: WenyXu * test(datanode): align batch-open retry hint assertion Signed-off-by: WenyXu * style: use imported types in region request tests Signed-off-by: WenyXu --------- Signed-off-by: WenyXu (cherry picked from commit 684c8ff2917c10dcf8103664f5eb2b01ef6720e9) Signed-off-by: WenyXu * refactor(operator): centralize write partitioning and dispatch (#9541) * refactor(operator): centralize write partitioning and dispatch Signed-off-by: WenyXu * fix(operator): preserve write dispatch behavior during refactor Signed-off-by: WenyXu * style(operator): import write dispatch error paths Signed-off-by: WenyXu * style(operator): align dispatch imports and test utilities Signed-off-by: WenyXu * style(operator): use imported types in partition tests Signed-off-by: WenyXu --------- Signed-off-by: WenyXu (cherry picked from commit a2d28e6060df4c1e40de23df995d0f04f0adb193) Signed-off-by: WenyXu * feat(operator): retry failed writes with refreshed partitions (#9542) * feat(operator): retry failed writes with refreshed partitions Signed-off-by: WenyXu * style(operator): simplify retry test type references Signed-off-by: WenyXu * refactor(operator): clarify write request rebuilding names Signed-off-by: WenyXu * refactor(operator): encapsulate write retry attempts Signed-off-by: WenyXu * style(operator): group inserter configuration methods Signed-off-by: WenyXu * refactor(operator): share routed peers and separate retry decisions Signed-off-by: WenyXu * refactor(operator): avoid temporary write routing allocations Signed-off-by: WenyXu --------- Signed-off-by: WenyXu (cherry picked from commit d86741d0ad92b98cff2cc0cff101734fd02e270a) Signed-off-by: WenyXu * 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 * 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 * 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 * fix(operator): stop detached write retries when callers cancel Signed-off-by: WenyXu * 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 --------- Signed-off-by: WenyXu (cherry picked from commit 863eebdecbb9c11e7e3249b6effb43a0c53568f3) Signed-off-by: WenyXu * style(mito2): format index error imports Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- Cargo.lock | 7 + Cargo.toml | 1 + config/config.md | 5 + config/frontend.example.toml | 12 + src/client/Cargo.toml | 1 + src/client/src/client.rs | 22 +- src/client/src/error.rs | 26 +- src/client/src/region.rs | 699 +++++-- src/client/src/region/codec.rs | 63 + src/common/error/src/ext.rs | 18 +- src/common/error/tests/ext.rs | 9 + src/common/grpc/src/flight.rs | 29 + src/common/meta/src/ddl/drop_table.rs | 4 +- src/common/meta/src/ddl/undrop_table.rs | 60 +- src/common/meta/src/ddl/utils.rs | 74 +- src/common/meta/src/error.rs | 50 +- src/common/meta/src/instruction.rs | 34 +- src/common/meta/src/node_manager.rs | 12 +- src/common/meta/src/peer.rs | 3 + .../src/reconciliation/reconcile_catalog.rs | 3 +- .../src/reconciliation/reconcile_database.rs | 3 +- .../reconcile_logical_tables.rs | 3 +- .../src/reconciliation/reconcile_table.rs | 3 +- src/common/meta/src/rpc/router.rs | 33 +- src/common/meta/src/test_util.rs | 3 +- src/common/procedure/src/error.rs | 21 +- src/datanode/src/error.rs | 30 +- src/datanode/src/region_server.rs | 2 +- src/frontend/AGENTS.md | 6 +- src/frontend/Cargo.toml | 1 + src/frontend/src/error.rs | 10 +- src/frontend/src/frontend.rs | 72 +- src/frontend/src/instance/builder.rs | 143 +- src/frontend/src/server.rs | 1 + src/frontend/src/service_config.rs | 1 + src/metric-engine/src/engine/bulk_insert.rs | 24 +- src/metric-engine/src/error.rs | 20 + src/mito2/src/engine/flush_test.rs | 2 + src/mito2/src/engine/staging_test.rs | 72 +- src/mito2/src/error.rs | 73 +- src/mito2/src/sst/index.rs | 4 +- src/mito2/src/worker/handle_write.rs | 14 +- src/operator/AGENTS.md | 3 + src/operator/Cargo.toml | 5 + src/operator/src/bulk_insert.rs | 260 +-- src/operator/src/delete.rs | 77 +- src/operator/src/insert.rs | 97 +- src/operator/src/lib.rs | 1 + src/operator/src/req_convert.rs | 2 + .../src/req_convert/common/partitioner.rs | 1081 ++++++++++- src/operator/src/tests.rs | 5 +- src/operator/src/tests/partition_manager.rs | 6 +- src/operator/src/write_dispatcher.rs | 615 +++++++ src/operator/src/write_dispatcher/options.rs | 106 ++ src/operator/src/write_dispatcher/tests.rs | 1606 +++++++++++++++++ src/servers/AGENTS.md | 3 + src/servers/benches/flush_batch_physical.rs | 32 +- src/servers/src/pending_rows_batcher.rs | 686 +++---- src/standalone/src/datanode_manager.rs | 12 +- src/standalone/src/options.rs | 19 +- tests-integration/src/test_util.rs | 1 + 61 files changed, 5316 insertions(+), 974 deletions(-) create mode 100644 src/client/src/region/codec.rs create mode 100644 src/operator/src/write_dispatcher.rs create mode 100644 src/operator/src/write_dispatcher/options.rs create mode 100644 src/operator/src/write_dispatcher/tests.rs diff --git a/Cargo.lock b/Cargo.lock index 5c0519c849c..3600aec93ae 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2097,6 +2097,7 @@ dependencies = [ "tokio", "tokio-stream", "tonic 0.14.2", + "tonic-prost", "tracing", ] @@ -5673,6 +5674,7 @@ dependencies = [ "store-api", "strfmt", "table", + "temp-env", "tokio", "tokio-util", "toml 0.8.23", @@ -10003,7 +10005,9 @@ dependencies = [ "arrow-ipc 58.3.0", "async-stream", "async-trait", + "backon", "bytes", + "cache", "catalog", "chrono", "client", @@ -10034,12 +10038,14 @@ dependencies = [ "futures", "futures-util", "humantime", + "humantime-serde", "itertools 0.14.0", "jsonb", "lazy_static", "meta-client", "meter-core", "meter-macros", + "mito-codec", "moka", "object-store", "partition", @@ -10060,6 +10066,7 @@ dependencies = [ "substrait 1.3.0-alpha.1", "table", "tokio", + "tokio-stream", "tokio-util", "tonic 0.14.2", "tracing", diff --git a/Cargo.toml b/Cargo.toml index a2606954d63..c073767826b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -262,6 +262,7 @@ url = "2.3" # Pin below 0.14.5 to avoid the `connection_timeout_future` "async fn resumed after # completion" panic (greptimedb#8265 / grpc-rust#2522). Remove once tonic ships a fix. tonic = { version = "0.14, <0.14.5", features = ["tls-aws-lc", "gzip", "zstd"] } +tonic-prost = "0.14" tower = "0.5" tower-http = "0.6" tracing = "0.1" diff --git a/config/config.md b/config/config.md index a8d12c3557e..5879685e616 100644 --- a/config/config.md +++ b/config/config.md @@ -361,6 +361,11 @@ | `query.experimental_spill_path` | String | Unset | Spill directory path. Ignored unless mode is "custom". | | `query.experimental_spill_max_temp_directory_size` | String | `1GiB` | Maximum total size of spill directory (default: "1GiB").
Ignored unless mode is "custom". | | `query.experimental_spill_compression` | String | `uncompressed` | Compression for spilled data files: "uncompressed" (default), "lz4_frame", "zstd".
Ignored unless mode is "custom". | +| `write_retry` | -- | -- | Retry retryable distributed write failures.
Backoff grows exponentially, with jitter to spread concurrent retries. | +| `write_retry.max_retries` | Integer | `3` | Maximum retries after the initial write attempt. Set to 0 to disable retries. | +| `write_retry.initial_backoff` | String | `1500ms` | Initial backoff before jitter is applied. Must be positive. | +| `write_retry.max_backoff` | String | `10s` | Maximum backoff before jitter is applied. Must be at least initial_backoff. | +| `write_retry.backoff_multiplier` | Float | `2.0` | Backoff multiplier. Must be finite and at least 1. | | `datanode` | -- | -- | Datanode options. | | `datanode.client` | -- | -- | Datanode client options. | | `datanode.client.connect_timeout` | String | `10s` | -- | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 7c6862916e4..44cc5d2ddaa 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -337,6 +337,18 @@ memory_pool_size = "50%" #+ experimental_spill_compression = "uncompressed" # --- End spill-to-disk --- +## Retry retryable distributed write failures. +## Backoff grows exponentially, with jitter to spread concurrent retries. +[write_retry] +## Maximum retries after the initial write attempt. Set to 0 to disable retries. +max_retries = 3 +## Initial backoff before jitter is applied. Must be positive. +initial_backoff = "1500ms" +## Maximum backoff before jitter is applied. Must be at least initial_backoff. +max_backoff = "10s" +## Backoff multiplier. Must be finite and at least 1. +backoff_multiplier = 2.0 + ## Datanode options. [datanode] ## Datanode client options. diff --git a/src/client/Cargo.toml b/src/client/Cargo.toml index 5b479b1fbf2..4e8402890ed 100644 --- a/src/client/Cargo.toml +++ b/src/client/Cargo.toml @@ -45,6 +45,7 @@ substrait.workspace = true tokio.workspace = true tokio-stream = { workspace = true, features = ["net"] } tonic.workspace = true +tonic-prost.workspace = true [dev-dependencies] common-grpc-expr.workspace = true diff --git a/src/client/src/client.rs b/src/client/src/client.rs index fc40c98a723..ed59f6701e7 100644 --- a/src/client/src/client.rs +++ b/src/client/src/client.rs @@ -20,7 +20,6 @@ use api::v1::HealthCheckRequest; use api::v1::flow::flow_client::FlowClient as PbFlowClient; use api::v1::health_check_client::HealthCheckClient; use api::v1::prometheus_gateway_client::PrometheusGatewayClient; -use api::v1::region::region_client::RegionClient as PbRegionClient; use arrow_flight::flight_service_client::FlightServiceClient; use common_grpc::channel_manager::{ ChannelConfig, ChannelManager, ClientTlsOption, load_client_tls_config, @@ -482,23 +481,12 @@ impl Client { Ok(FlightClient { addr, client }) } - pub(crate) fn raw_region_client(&self) -> Result<(String, PbRegionClient)> { + /// Creates the region gRPC transport with the configured message limits. + pub(crate) fn raw_region_client(&self) -> Result<(String, tonic::client::Grpc)> { let (addr, channel) = self.find_channel()?; - let client = PbRegionClient::new(channel) - .max_decoding_message_size( - self.inner - .control_channel_manager - .config() - .max_recv_message_size - .as_bytes() as usize, - ) - .max_encoding_message_size( - self.inner - .control_channel_manager - .config() - .max_send_message_size - .as_bytes() as usize, - ); + let client = tonic::client::Grpc::new(channel) + .max_decoding_message_size(self.max_grpc_recv_message_size()) + .max_encoding_message_size(self.max_grpc_send_message_size()); Ok((addr, client)) } diff --git a/src/client/src/error.rs b/src/client/src/error.rs index 168ba422715..fdf0df9582f 100644 --- a/src/client/src/error.rs +++ b/src/client/src/error.rs @@ -18,7 +18,7 @@ use common_error::define_from_tonic_status; use common_error::ext::{BoxedError, ErrorExt, RetryHint}; use common_error::status_code::StatusCode; use common_macro::stack_trace_debug; -use snafu::{Location, Snafu}; +use snafu::{Location, Snafu, location}; use tonic::Code; use tonic::metadata::errors::InvalidMetadataValue; @@ -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(), @@ -205,6 +218,15 @@ impl ErrorExt for Error { define_from_tonic_status!(Error, Tonic); +impl From for common_meta::error::Error { + fn from(error: Error) -> Self { + Self::External { + source: BoxedError::new(error), + location: location!(), + } + } +} + impl Error { /// Returns the gRPC status code if this error is caused by a gRPC request failure. pub fn tonic_code(&self) -> Option { diff --git a/src/client/src/region.rs b/src/client/src/region.rs index 4a4d243f36a..24876c73fcf 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -12,6 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod codec; + +use std::borrow::Borrow; use std::sync::Arc; use std::time::Duration; @@ -49,10 +52,12 @@ use crate::error::{ MissingFieldSnafu, Result, ServerSnafu, }; use crate::flight::{FlightMessageKind, FlightMessageReader, decode_flight_data}; +use crate::region::codec::RegionCodec; use crate::{Client, metrics}; const FLIGHT_DO_GET_TIMEOUT: Duration = Duration::from_secs(10); +/// Sends region operations through a configured gRPC client. #[derive(Debug)] pub struct RegionRequester { client: Client, @@ -62,20 +67,8 @@ pub struct RegionRequester { #[async_trait] impl Datanode for RegionRequester { - async fn handle(&self, request: RegionRequest) -> MetaResult { - self.handle_inner(request).await.map_err(|err| { - if err.should_retry() { - meta_error::Error::RetryLater { - source: BoxedError::new(err), - clean_poisons: false, - } - } else { - meta_error::Error::External { - source: BoxedError::new(err), - location: location!(), - } - } - }) + async fn handle_request(&self, request: Arc) -> MetaResult { + self.handle_inner(request).await.map_err(Into::into) } async fn handle_query(&self, request: QueryRequest) -> MetaResult { @@ -144,11 +137,49 @@ impl RegionRequester { recordbatches_from_flight_message_stream(addr, flight_message_stream).await } - async fn handle_inner(&self, request: RegionRequest) -> Result { + /// Sends an owned or shared request without deep-cloning its protobuf payload. + /// + /// Accepts both [`RegionRequest`] and [`Arc`]. Callers can retain a + /// shared request for another attempt; this method does not retry or determine + /// whether replaying the request is safe. + pub async fn handle(&self, request: R) -> Result + where + R: Borrow + Send + Sync + 'static, + { + self.handle_inner(request).await + } + + /// Sends a remote dynamic filter update to the region service. + pub async fn handle_remote_dyn_filter_update( + &self, + query_id: impl Into, + update: RemoteDynFilterUpdate, + ) -> Result { + self.handle_inner(build_remote_dyn_filter_update_request(query_id, update)) + .await + } + + /// Unregisters a remote dynamic filter in the region service. + pub async fn handle_remote_dyn_filter_unregister( + &self, + query_id: impl Into, + unregister: RemoteDynFilterUnregister, + ) -> Result { + self.handle_inner(build_remote_dyn_filter_unregister_request( + query_id, unregister, + )) + .await + } + + async fn handle_inner(&self, request: R) -> Result + where + R: Borrow + Send + Sync + 'static, + { let request_body = request + .borrow() .body .as_ref() - .with_context(|| MissingFieldSnafu { field: "body" })?; + .context(MissingFieldSnafu { field: "body" })?; let is_insert = matches!( request_body, region_request::Body::Inserts(_) | region_request::Body::BulkInsert(_) @@ -168,48 +199,42 @@ impl RegionRequester { } } - let response = client - .handle(request) - .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!(), - } - })? - .into_inner(); + let response = async { + client + .ready() + .await + .map_err(|e| tonic::Status::unknown(format!("Service was not ready: {e}")))?; + let mut request = tonic::Request::new(request); + request.extensions_mut().insert(tonic::GrpcMethod::new( + "greptime.v1.region.Region", + "Handle", + )); + client + .unary( + request, + tonic::codegen::http::uri::PathAndQuery::from_static( + "/greptime.v1.region.Region/Handle", + ), + RegionCodec::default(), + ) + .await + } + .await + .map_err(|status| map_region_error(status, addr))? + .into_inner(); check_response_header(&response.header)?; Ok(RegionResponse::from_region_response(response)) } +} - pub async fn handle(&self, request: RegionRequest) -> Result { - self.handle_inner(request).await - } - - pub async fn handle_remote_dyn_filter_update( - &self, - query_id: impl Into, - update: RemoteDynFilterUpdate, - ) -> Result { - self.handle_inner(build_remote_dyn_filter_update_request(query_id, update)) - .await - } - - pub async fn handle_remote_dyn_filter_unregister( - &self, - query_id: impl Into, - unregister: RemoteDynFilterUnregister, - ) -> Result { - self.handle_inner(build_remote_dyn_filter_unregister_request( - query_id, unregister, - )) - .await +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!(), } } @@ -418,13 +443,23 @@ pub fn check_response_header(header: &Option) -> Result<()> { } #[cfg(test)] -mod test { - use api::v1::Status as PbStatus; +#[allow(deprecated)] +mod tests { + use std::collections::VecDeque; + use std::sync::Mutex; + use api::v1::region::region_server::{Region, RegionServer}; use api::v1::region::{ - BulkInsertRequest, RegionResponse as PbRegionResponse, RemoteDynFilterUnregister, - RemoteDynFilterUpdate, region_request, remote_dyn_filter_request, + BulkInsertRequest, InsertRequest, InsertRequests, RegionResponse as PbRegionResponse, + RemoteDynFilterUnregister, RemoteDynFilterUpdate, region_request, + remote_dyn_filter_request, }; + use api::v1::value::ValueData; + use api::v1::{ + ColumnDataType, ColumnSchema as PbColumnSchema, Row, Rows, Status as PbStatus, Value, + }; + use common_error::ext::RetryHint; + use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_recordbatch::adapter::RecordBatchMetrics; use datatypes::arrow::array::Int32Array; use datatypes::prelude::{ConcreteDataType, VectorRef}; @@ -432,74 +467,526 @@ mod test { use datatypes::vectors::Int32Vector; use futures_util::stream; use tokio::net::TcpListener; + use tokio::task::JoinHandle; use tokio_stream::wrappers::TcpListenerStream; use tonic::codec::CompressionEncoding; use tonic::{Request, Response, Status}; use super::*; - use crate::Error::{self, IllegalDatabaseResponse, Server}; + use crate::Error; - #[derive(Clone)] - struct CompressionRecordingRegionService { - zstd_headers: Arc>>, + #[derive(Clone, Default)] + struct RecordingRegionService { + requests: Arc>>, + responses: Arc>>>, } #[tonic::async_trait] - impl Region for CompressionRecordingRegionService { + impl Region for RecordingRegionService { async fn handle( &self, request: Request, ) -> std::result::Result, Status> { let metadata = request.metadata(); - let sends_zstd = metadata - .get("grpc-encoding") - .and_then(|v| v.to_str().ok()) - .is_some_and(|value| value == "zstd"); + let sends_zstd = metadata.get("grpc-encoding").is_some_and(|v| v == "zstd"); let accepts_zstd = metadata .get("grpc-accept-encoding") - .and_then(|value| value.to_str().ok()) - .is_some_and(|value| value.split(',').any(|encoding| encoding == "zstd")); - self.zstd_headers + .and_then(|v| v.to_str().ok()) + .is_some_and(|v| v.split(',').any(|encoding| encoding == "zstd")); + self.requests .lock() .unwrap() - .push((sends_zstd, accepts_zstd)); + .push((request.into_inner(), sends_zstd, accepts_zstd)); + self.responses + .lock() + .unwrap() + .pop_front() + .unwrap_or_else(|| Ok(build_response(StatusCode::Success))) + .map(Response::new) + } + } - Ok(Response::new(PbRegionResponse { - header: Some(ResponseHeader { - status: Some(PbStatus { - status_code: StatusCode::Success as u32, - ..Default::default() - }), + struct RegionTestServer { + addr: String, + service: RecordingRegionService, + task: JoinHandle<()>, + } + + impl RegionTestServer { + async fn start() -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap().to_string(); + let service = RecordingRegionService::default(); + let server_service = service.clone(); + let task = tokio::spawn(async move { + tonic::transport::Server::builder() + .add_service( + RegionServer::new(server_service) + .accept_compressed(CompressionEncoding::Zstd) + .send_compressed(CompressionEncoding::Zstd), + ) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + .unwrap(); + }); + Self { + addr, + service, + task, + } + } + + fn create_requester( + &self, + config: ChannelConfig, + send: bool, + accept: bool, + ) -> RegionRequester { + RegionRequester::new( + Client::with_manager_and_urls( + ChannelManager::with_config(config, None), + [self.addr.clone()], + ), + send, + accept, + ) + } + } + + impl Drop for RegionTestServer { + fn drop(&mut self) { + self.task.abort(); + } + } + + fn build_response(code: StatusCode) -> PbRegionResponse { + PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus { + status_code: code as u32, + ..Default::default() }), + }), + affected_rows: 1, + ..Default::default() + } + } + + fn build_insert_request() -> RegionRequest { + RegionRequest { + header: Some(RegionRequestHeader { + dbname: "shared-request-test".to_string(), ..Default::default() - })) + }), + body: Some(region_request::Body::Inserts(InsertRequests { + requests: vec![InsertRequest { + region_id: 42, + rows: Some(Rows { + schema: vec![ + PbColumnSchema { + column_name: "message".to_string(), + datatype: ColumnDataType::String as i32, + ..Default::default() + }, + PbColumnSchema { + column_name: "payload".to_string(), + datatype: ColumnDataType::Binary as i32, + ..Default::default() + }, + ], + rows: vec![Row { + values: vec![ + Value { + value_data: Some(ValueData::StringValue( + "hello 世界".to_string(), + )), + }, + Value { + value_data: Some(ValueData::BinaryValue(vec![0, 1, 127, 255])), + }, + ], + }], + }), + ..Default::default() + }], + })), } } #[tokio::test] - async fn test_inserts_and_bulk_insert_use_transport_compression() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let zstd_headers = Arc::new(std::sync::Mutex::new(Vec::new())); - let service = CompressionRecordingRegionService { - zstd_headers: zstd_headers.clone(), - }; - let server = tokio::spawn(async move { - tonic::transport::Server::builder() - .add_service( - RegionServer::new(service) - .accept_compressed(CompressionEncoding::Zstd) - .send_compressed(CompressionEncoding::Zstd), - ) - .serve_with_incoming(TcpListenerStream::new(listener)) + async fn test_owned_and_shared_requests_preserve_payload_and_compression() { + let server = RegionTestServer::start().await; + let shared = Arc::new(build_insert_request()); + for (send, accept) in [(false, false), (true, false), (false, true), (true, true)] { + let requester = server.create_requester(ChannelConfig::default(), send, accept); + let channel = + tonic::transport::Endpoint::from_shared(format!("http://{}", server.addr)) + .unwrap() + .connect() + .await + .unwrap(); + let mut baseline = api::v1::region::region_client::RegionClient::new(channel); + if send { + baseline = baseline.send_compressed(CompressionEncoding::Zstd); + } + if accept { + baseline = baseline.accept_compressed(CompressionEncoding::Zstd); + } + assert_eq!( + baseline + .handle(build_insert_request()) + .await + .unwrap() + .into_inner() + .affected_rows, + 1 + ); + assert_eq!( + requester + .handle(build_insert_request()) + .await + .unwrap() + .affected_rows, + 1 + ); + assert_eq!( + requester + .handle(Arc::clone(&shared)) + .await + .unwrap() + .affected_rows, + 1 + ); + } + let requests = server.service.requests.lock().unwrap(); + assert_eq!(requests.len(), 12); + for ((request, sends, accepts), (send, accept)) in requests.iter().zip( + [(false, false), (true, false), (false, true), (true, true)] + .into_iter() + .flat_map(|settings| [settings, settings, settings]), + ) { + assert_eq!(request, shared.as_ref()); + assert_eq!((*sends, *accepts), (send, accept)); + } + } + + #[tokio::test] + async fn test_shared_request_can_be_resent_after_transport_error() { + let server = RegionTestServer::start().await; + let mut status = Status::unavailable("migrating"); + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT, + RetryHint::Retryable.as_str().parse().unwrap(), + ); + server + .service + .responses + .lock() + .unwrap() + .push_back(Err(status)); + let requester = server.create_requester(ChannelConfig::default(), false, false); + let shared = Arc::new(build_insert_request()); + let error = requester.handle(Arc::clone(&shared)).await.unwrap_err(); + assert_eq!(error.retry_hint(), RetryHint::Retryable); + assert!( + matches!(error, Error::RegionServer { addr, code: tonic::Code::Unavailable, .. } if addr == server.addr) + ); + requester.handle(Arc::clone(&shared)).await.unwrap(); + let requests = server.service.requests.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert_eq!(&requests[0].0, shared.as_ref()); + assert_eq!(&requests[1].0, shared.as_ref()); + } + + #[test] + fn test_retry_hint_wire_compatibility() { + for (name, wire_hint, expected_hint) in [ + ("legacy_missing", None, RetryHint::NonRetryable), + ("legacy_retryable", Some("retryable"), RetryHint::Retryable), + ( + "legacy_non_retryable", + Some("non_retryable"), + RetryHint::NonRetryable, + ), + ( + "route_refresh", + Some("retry_after_route_refresh"), + RetryHint::RetryAfterRouteRefresh, + ), + ( + "unknown_future_hint", + Some("future_retry_hint"), + RetryHint::NonRetryable, + ), + ] { + let mut status = Status::failed_precondition("region unavailable"); + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_CODE, + (StatusCode::RegionNotReady as u32) + .to_string() + .parse() + .unwrap(), + ); + if let Some(hint) = wire_hint { + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT, + hint.parse().unwrap(), + ); + } + let error: Error = status.into(); + assert_eq!(error.retry_hint(), expected_hint, "{name}"); + assert_eq!(error.status_code(), StatusCode::RegionNotReady, "{name}"); + assert_eq!( + error.tonic_code(), + Some(tonic::Code::FailedPrecondition), + "{name}" + ); + } + } + + #[tokio::test] + async fn test_transport_error_parity_and_datanode_mapping() { + let server = RegionTestServer::start().await; + let requester = server.create_requester(ChannelConfig::default(), false, false); + for hint in [ + RetryHint::Retryable, + RetryHint::RetryAfterRouteRefresh, + RetryHint::NonRetryable, + ] { + let mut status = Status::failed_precondition("region unavailable"); + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT, + hint.as_str().parse().unwrap(), + ); + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_CODE, + (StatusCode::RegionNotFound as u32) + .to_string() + .parse() + .unwrap(), + ); + server + .service + .responses + .lock() + .unwrap() + .extend((0..4).map(|_| Err(status.clone()))); + let mut baseline = api::v1::region::region_client::RegionClient::connect(format!( + "http://{}", + server.addr + )) + .await + .unwrap(); + let baseline: Error = baseline + .handle(build_insert_request()) .await - .unwrap(); + .unwrap_err() + .into(); + assert_eq!(baseline.status_code(), StatusCode::RegionNotFound); + assert_eq!(baseline.retry_hint(), hint); + assert_eq!(baseline.tonic_code(), Some(tonic::Code::FailedPrecondition)); + let owned = requester.handle(build_insert_request()).await.unwrap_err(); + let shared = requester + .handle(Arc::new(build_insert_request())) + .await + .unwrap_err(); + for error in [owned, shared] { + assert_eq!(error.status_code(), baseline.status_code()); + assert_eq!(error.retry_hint(), baseline.retry_hint()); + assert_eq!(error.tonic_code(), baseline.tonic_code()); + } + 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); + 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)) }); - let client = Client::with_urls([addr.to_string()]); - let requester = RegionRequester::new(client.clone(), true, true); - let send_only_requester = RegionRequester::new(client.clone(), true, false); - let accept_only_requester = RegionRequester::new(client.clone(), false, true); - let disabled_requester = RegionRequester::new(client, false, false); + 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, + .. + } = source.as_any().downcast_ref::().unwrap() + else { + panic!("expected a region server error"); + }; + assert_eq!(*actual_code, code); + assert_eq!(source.retry_hint(), original.retry_hint()); + assert_eq!(source.to_string(), original.to_string()); + } + } + } + } + + #[tokio::test] + async fn test_shared_request_preserves_response_header_errors() { + let server = RegionTestServer::start().await; + server.service.responses.lock().unwrap().extend([ + Ok(PbRegionResponse::default()), + Ok(build_response(StatusCode::RegionNotFound)), + ]); + let requester = server.create_requester(ChannelConfig::default(), false, false); + let shared = Arc::new(build_insert_request()); + assert!(matches!( + requester.handle(Arc::clone(&shared)).await.unwrap_err(), + Error::IllegalDatabaseResponse { .. } + )); + assert!(matches!( + requester.handle(shared).await.unwrap_err(), + Error::Server { + code: StatusCode::RegionNotFound, + .. + } + )); + } + + #[tokio::test] + async fn test_shared_request_respects_configured_message_limits() { + let server = RegionTestServer::start().await; + let shared = Arc::new(build_insert_request()); + for send_limit in [true, false] { + let mut config = ChannelConfig::default(); + if send_limit { + config.max_send_message_size = "1B".parse().unwrap(); + } else { + config.max_recv_message_size = "1B".parse().unwrap(); + } + let mut baseline = api::v1::region::region_client::RegionClient::connect(format!( + "http://{}", + server.addr + )) + .await + .unwrap() + .max_encoding_message_size(config.max_send_message_size.as_bytes() as usize) + .max_decoding_message_size(config.max_recv_message_size.as_bytes() as usize); + let baseline: Error = baseline + .handle(build_insert_request()) + .await + .unwrap_err() + .into(); + if !send_limit { + assert_eq!( + baseline.tonic_code(), + Some(tonic::Code::OutOfRange), + "{baseline:?}" + ); + } + let requester = server.create_requester(config, false, false); + let error = requester.handle(Arc::clone(&shared)).await.unwrap_err(); + assert_eq!( + error.tonic_code(), + baseline.tonic_code(), + "shared: {error:?}; generated: {baseline:?}" + ); + assert_eq!( + error.status_code(), + baseline.status_code(), + "shared: {error:?}; generated: {baseline:?}" + ); + let Error::RegionServer { source, .. } = &error else { + panic!("expected a region server error: {error:?}"); + }; + assert_eq!( + source.retry_hint(), + baseline.retry_hint(), + "shared: {error:?}; generated: {baseline:?}" + ); + assert_eq!( + server.service.requests.lock().unwrap().len(), + if send_limit { 0 } else { 2 } + ); + } + let requester = server.create_requester(ChannelConfig::default(), false, false); + requester.handle(shared).await.unwrap(); + } + + #[tokio::test] + async fn test_inserts_and_bulk_insert_use_transport_compression() { + let server = RegionTestServer::start().await; + let requester = server.create_requester(ChannelConfig::default(), true, true); + let send_only_requester = server.create_requester(ChannelConfig::default(), true, false); + let accept_only_requester = server.create_requester(ChannelConfig::default(), false, true); + let disabled_requester = server.create_requester(ChannelConfig::default(), false, false); let inserts = || RegionRequest { body: Some(region_request::Body::Inserts(Default::default())), ..Default::default() @@ -535,9 +1022,15 @@ mod test { (false, false), (false, false) ], - *zstd_headers.lock().unwrap() + server + .service + .requests + .lock() + .unwrap() + .iter() + .map(|(_, send, accept)| (*send, *accept)) + .collect::>() ); - server.abort(); } #[test] @@ -599,13 +1092,13 @@ mod test { let result = check_response_header(&None); assert!(matches!( result.unwrap_err(), - IllegalDatabaseResponse { .. } + Error::IllegalDatabaseResponse { .. } )); let result = check_response_header(&Some(ResponseHeader { status: None })); assert!(matches!( result.unwrap_err(), - IllegalDatabaseResponse { .. } + Error::IllegalDatabaseResponse { .. } )); let result = check_response_header(&Some(ResponseHeader { @@ -624,7 +1117,7 @@ mod test { })); assert!(matches!( result.unwrap_err(), - IllegalDatabaseResponse { .. } + Error::IllegalDatabaseResponse { .. } )); let result = check_response_header(&Some(ResponseHeader { @@ -633,7 +1126,7 @@ mod test { err_msg: "blabla".to_string(), }), })); - let Server { code, msg, .. } = result.unwrap_err() else { + let Error::Server { code, msg, .. } = result.unwrap_err() else { unreachable!() }; assert_eq!(code, StatusCode::Internal); diff --git a/src/client/src/region/codec.rs b/src/client/src/region/codec.rs new file mode 100644 index 00000000000..9476385df36 --- /dev/null +++ b/src/client/src/region/codec.rs @@ -0,0 +1,63 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::borrow::Borrow; +use std::marker::PhantomData; + +use api::v1::region::{RegionRequest, RegionResponse}; +use prost::Message; +use tonic::Status; +use tonic::codec::{Codec, EncodeBuf, Encoder}; +use tonic_prost::ProstDecoder; + +/// Encodes owned or shared region requests using the existing protobuf wire format. +pub(crate) struct RegionCodec(PhantomData); + +impl Default for RegionCodec { + fn default() -> Self { + Self(PhantomData) + } +} + +impl Codec for RegionCodec +where + R: Borrow + Send + 'static, +{ + type Encode = R; + type Decode = RegionResponse; + type Encoder = RegionEncoder; + type Decoder = ProstDecoder; + + fn encoder(&mut self) -> Self::Encoder { + RegionEncoder(PhantomData) + } + + fn decoder(&mut self) -> Self::Decoder { + ProstDecoder::default() + } +} + +/// Borrows the protobuf payload while consuming its owned or shared container. +pub(crate) struct RegionEncoder(PhantomData); + +impl> Encoder for RegionEncoder { + type Item = R; + type Error = Status; + + fn encode(&mut self, item: R, buf: &mut EncodeBuf<'_>) -> Result<(), Status> { + item.borrow() + .encode(buf) + .map_err(|error| Status::internal(error.to_string())) + } +} diff --git a/src/common/error/src/ext.rs b/src/common/error/src/ext.rs index 59895551156..bd5c12f7a3f 100644 --- a/src/common/error/src/ext.rs +++ b/src/common/error/src/ext.rs @@ -32,6 +32,10 @@ use crate::status_code::StatusCode; pub enum RetryHint { /// The operation may succeed if retried later. Retryable, + /// The operation may succeed after refreshing routing and partition information. + /// The request may need rebuilding when its partition assignment changes. + /// This hint does not imply that a failed operation had no side effects. + RetryAfterRouteRefresh, /// Retrying the same operation is not expected to help. /// /// This is the default for errors that do not explicitly opt in to retry. @@ -39,16 +43,23 @@ pub enum RetryHint { } const RETRY_HINT_RETRYABLE: &str = "retryable"; +const RETRY_HINT_AFTER_ROUTE_REFRESH: &str = "retry_after_route_refresh"; const RETRY_HINT_NON_RETRYABLE: &str = "non_retryable"; impl RetryHint { + /// Returns whether another attempt may succeed, possibly after refreshing routes. pub fn is_retryable(self) -> bool { - matches!(self, RetryHint::Retryable) + matches!( + self, + RetryHint::Retryable | RetryHint::RetryAfterRouteRefresh + ) } + /// Returns the retry hint encoded in gRPC error metadata. pub fn as_str(self) -> &'static str { match self { RetryHint::Retryable => RETRY_HINT_RETRYABLE, + RetryHint::RetryAfterRouteRefresh => RETRY_HINT_AFTER_ROUTE_REFRESH, RetryHint::NonRetryable => RETRY_HINT_NON_RETRYABLE, } } @@ -76,6 +87,7 @@ impl FromStr for RetryHint { fn from_str(s: &str) -> Result { match s { RETRY_HINT_RETRYABLE => Ok(RetryHint::Retryable), + RETRY_HINT_AFTER_ROUTE_REFRESH => Ok(RetryHint::RetryAfterRouteRefresh), RETRY_HINT_NON_RETRYABLE => Ok(RetryHint::NonRetryable), _ => Err(()), } @@ -119,7 +131,9 @@ pub trait ErrorExt: StackError { /// /// Implementations should return [`RetryHint::Retryable`] only when retrying the /// same operation may succeed without changing the request. The default is - /// [`RetryHint::NonRetryable`] to avoid accidental retry loops. + /// [`RetryHint::NonRetryable`] to avoid accidental retry loops. Use + /// [`RetryHint::RetryAfterRouteRefresh`] when routing or partition information + /// must be refreshed before another attempt. fn retry_hint(&self) -> RetryHint { RetryHint::NonRetryable } diff --git a/src/common/error/tests/ext.rs b/src/common/error/tests/ext.rs index 8be8afa4748..c34a47c400f 100644 --- a/src/common/error/tests/ext.rs +++ b/src/common/error/tests/ext.rs @@ -164,6 +164,15 @@ fn test_transparent_flag() { #[test] fn test_retry_hint_helpers() { assert!(RetryHint::Retryable.is_retryable()); + assert!(RetryHint::RetryAfterRouteRefresh.is_retryable()); + assert_eq!( + RetryHint::RetryAfterRouteRefresh.as_str(), + "retry_after_route_refresh" + ); + assert_eq!( + RetryHint::from_str("retry_after_route_refresh").unwrap(), + RetryHint::RetryAfterRouteRefresh + ); assert!(!RetryHint::NonRetryable.is_retryable()); assert_eq!(RetryHint::Retryable.as_str(), "retryable"); diff --git a/src/common/grpc/src/flight.rs b/src/common/grpc/src/flight.rs index 1245a86efb1..43357ff85e6 100644 --- a/src/common/grpc/src/flight.rs +++ b/src/common/grpc/src/flight.rs @@ -37,6 +37,35 @@ use vec1::{Vec1, vec1}; use crate::error; use crate::error::{DecodeFlightDataSnafu, InvalidFlightDataSnafu, Result}; +/// Encodes a batch into the schema, record-batch header and payload used by bulk inserts. +/// +/// Uses the default Flight compression and rejects dictionary batches, whose +/// additional Flight messages cannot be represented by this three-part format. +pub fn record_batch_to_ipc( + record_batch: DfRecordBatch, +) -> Result<(ProstBytes, ProstBytes, ProstBytes)> { + let mut encoder = FlightEncoder::default(); + let schema = encoder.encode_schema(record_batch.schema().as_ref()); + let mut iter = encoder + .encode(FlightMessage::RecordBatch(record_batch)) + .into_iter(); + let flight_data = iter.next().context(InvalidFlightDataSnafu { + reason: "Failed to encode empty flight data", + })?; + if iter.next().is_some() { + return error::NotSupportedSnafu { + feat: "bulk insert RecordBatch with dictionary arrays", + } + .fail(); + } + + Ok(( + schema.data_header, + flight_data.data_header, + flight_data.data_body, + )) +} + /// Flight metadata key used to carry flow query extensions as JSON pairs. pub const FLOW_EXTENSIONS_METADATA_KEY: &str = "x-greptime-flow-extensions"; /// Flight metadata key used to carry query snapshot read upper bounds as JSON. diff --git a/src/common/meta/src/ddl/drop_table.rs b/src/common/meta/src/ddl/drop_table.rs index b73bd95e1bc..622a27d804b 100644 --- a/src/common/meta/src/ddl/drop_table.rs +++ b/src/common/meta/src/ddl/drop_table.rs @@ -19,6 +19,8 @@ use std::collections::HashMap; use async_trait::async_trait; use common_error::ext::BoxedError; +#[cfg(feature = "enterprise")] +use common_error::ext::ErrorExt; use common_event_recorder::Event; use common_procedure::error::{ExternalSnafu, FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ @@ -53,7 +55,7 @@ use crate::rpc::router::{RegionRoute, operating_leader_region_roles}; #[cfg(feature = "enterprise")] fn ensure_retry_later(err: error::Error) -> error::Error { - if err.is_retry_later() { + if err.retry_hint().is_retryable() { err } else { error::Error::retry_later(err) diff --git a/src/common/meta/src/ddl/undrop_table.rs b/src/common/meta/src/ddl/undrop_table.rs index 73a361ffd0e..efb1d4824d7 100644 --- a/src/common/meta/src/ddl/undrop_table.rs +++ b/src/common/meta/src/ddl/undrop_table.rs @@ -17,8 +17,8 @@ use api::v1::region::{ use async_trait::async_trait; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, - Result as ProcedureResult, Status, + Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, + Procedure, Result as ProcedureResult, Status, }; use common_telemetry::tracing_context::TracingContext; use common_telemetry::warn; @@ -156,6 +156,7 @@ impl UndropTableProcedure { } .await; if let Err(err) = restore_result { + // Keep this classification aligned with map_restore_error(). let should_cleanup_opened_regions = !err.is_retry_later(); let err = self.map_restore_metadata_error(err); if should_cleanup_opened_regions @@ -274,6 +275,19 @@ impl UndropTableProcedure { Ok(()) } + fn map_restore_error(err: error::Error) -> ProcedureError { + // Restore retries keep opened regions, while other failures close them. Preserve + // this policy even when the underlying metadata error carries a retry hint: + // retrying after cleanup would restore metadata without reopening the regions. + if err.is_retry_later() { + map_to_procedure_error(err) + } else if err.need_clean_poisons() { + ProcedureError::external_and_clean_poisons(err) + } else { + ProcedureError::external(err) + } + } + fn map_restore_metadata_error(&self, err: error::Error) -> error::Error { match err { error::Error::TombstoneTargetAlreadyExists { .. } => { @@ -413,7 +427,12 @@ impl Procedure for UndropTableProcedure { async fn execute(&mut self, _: &ProcedureContext) -> ProcedureResult { match self.data.state { UndropTableState::Prepare => self.on_prepare().await, - UndropTableState::RestoreMetadata => self.on_restore_metadata().await, + UndropTableState::RestoreMetadata => { + return self + .on_restore_metadata() + .await + .map_err(Self::map_restore_error); + } UndropTableState::OpenRegions => self.on_open_regions().await, UndropTableState::InvalidateTableCache => self.on_broadcast().await, } @@ -606,7 +625,8 @@ mod tests { use api::v1::region::region_request; use async_trait::async_trait; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; - use common_error::ext::ErrorExt; + use common_error::ext::{BoxedError, ErrorExt}; + use common_error::mock::MockError; use common_error::status_code::StatusCode; use store_api::storage::RegionId; use table::table_name::TableName; @@ -635,6 +655,38 @@ mod tests { } } + #[test] + fn test_restore_metadata_retry_policy() { + let source = error::Error::retry_later(MockError::new(StatusCode::Internal)); + let err = Err::<(), _>(BoxedError::new(source)) + .context(error::ExternalSnafu) + .unwrap_err(); + assert!(err.retry_hint().is_retryable()); + assert!(!err.is_retry_later()); + let err = UndropTableProcedure::map_restore_error(err); + assert!(!err.is_retry_later()); + assert!(!err.need_clean_poisons()); + + for retryable in [false, true] { + for clean_poisons in [false, true] { + let source = BoxedError::new(MockError::new(StatusCode::Internal)); + let err = if retryable { + error::Error::RetryLater { + source, + clean_poisons, + } + } else { + Err::<(), _>(source) + .context(error::AbortProcedureSnafu { clean_poisons }) + .unwrap_err() + }; + let err = UndropTableProcedure::map_restore_error(err); + assert_eq!(err.is_retry_later(), retryable); + assert_eq!(err.need_clean_poisons(), clean_poisons); + } + } + } + #[test] fn test_map_restore_metadata_error_without_table_name() { let context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); diff --git a/src/common/meta/src/ddl/utils.rs b/src/common/meta/src/ddl/utils.rs index ebfbd6fd522..a8058d89b94 100644 --- a/src/common/meta/src/ddl/utils.rs +++ b/src/common/meta/src/ddl/utils.rs @@ -28,7 +28,7 @@ use api::v1::region::{ region_request, }; use common_catalog::consts::{METRIC_ENGINE, MITO_ENGINE}; -use common_error::ext::BoxedError; +use common_error::ext::{BoxedError, ErrorExt}; use common_procedure::error::Error as ProcedureError; use common_telemetry::tracing_context::TracingContext; use common_telemetry::{error, info, warn}; @@ -61,7 +61,7 @@ use crate::wal_provider::RegionWalOptions; pub fn add_peer_context_if_needed(datanode: Peer) -> impl FnOnce(Error) -> Error { move |err| { error!(err; "Failed to operate datanode, peer: {}", datanode); - if !err.is_retry_later() { + if !err.retry_hint().is_retryable() { return Err::<(), BoxedError>(BoxedError::new(err)) .context(OperateDatanodeSnafu { peer: datanode }) .unwrap_err(); @@ -84,7 +84,7 @@ pub(crate) fn is_metric_engine_logical_table( /// This function determines whether the error should be retried and if poison cleanup is needed, /// then maps it to the appropriate procedure error variant. pub fn map_to_procedure_error(e: Error) -> ProcedureError { - match (e.is_retry_later(), e.need_clean_poisons()) { + match (e.retry_hint().is_retryable(), e.need_clean_poisons()) { (true, true) => ProcedureError::retry_later_and_clean_poisons(e), (true, false) => ProcedureError::retry_later(e), (false, true) => ProcedureError::external_and_clean_poisons(e), @@ -256,7 +256,7 @@ pub fn handle_multiple_results(results: Vec>) -> MultipleRes match result { Ok(value) => ok_results.push(value), Err(err) => { - if err.is_retry_later() { + if err.retry_hint().is_retryable() { retryable_results.push(err); } else { non_retryable_results.push(err); @@ -475,8 +475,74 @@ pub fn extract_column_metadatas( #[cfg(test)] mod tests { + use common_error::mock::MockError; + use common_error::status_code::StatusCode; + use super::*; + fn external_error(retryable: bool) -> Error { + let source = MockError::new(StatusCode::Internal); + let source = if retryable { + BoxedError::new(Error::retry_later(source)) + } else { + BoxedError::new(source) + }; + Err::<(), _>(source) + .context(error::ExternalSnafu) + .unwrap_err() + } + + #[test] + fn test_map_to_procedure_error() { + for retryable in [false, true] { + for clean_poisons in [false, true] { + let err = if retryable { + Error::RetryLater { + source: BoxedError::new(external_error(false)), + clean_poisons, + } + } else { + Err::<(), _>(BoxedError::new(external_error(true))) + .context(error::AbortProcedureSnafu { clean_poisons }) + .unwrap_err() + }; + let err = map_to_procedure_error(err); + assert_eq!(err.is_retry_later(), retryable); + assert_eq!(err.need_clean_poisons(), clean_poisons); + } + let err = map_to_procedure_error(external_error(retryable)); + assert_eq!(err.is_retry_later(), retryable); + assert!(!err.need_clean_poisons()); + } + } + + #[test] + fn test_handle_multiple_results_retry_hint() { + assert!(matches!( + handle_multiple_results(vec![Ok(()), Err(external_error(true))]), + MultipleResults::PartialRetryable(_) + )); + assert!(matches!( + handle_multiple_results::<()>(vec![ + Err(external_error(true)), + Err(external_error(true)) + ]), + MultipleResults::AllRetryable(_) + )); + assert!(matches!( + handle_multiple_results(vec![ + Ok(()), + Err(external_error(true)), + Err(external_error(false)) + ]), + MultipleResults::PartialNonRetryable(_) + )); + assert!(matches!( + handle_multiple_results::<()>(vec![Err(external_error(false))]), + MultipleResults::AllNonRetryable(_) + )); + } + #[test] fn test_get_catalog_and_schema() { let test_catalog = "my_catalog"; diff --git a/src/common/meta/src/error.rs b/src/common/meta/src/error.rs index 3bd1475cb40..de3831589e7 100644 --- a/src/common/meta/src/error.rs +++ b/src/common/meta/src/error.rs @@ -1330,8 +1330,11 @@ impl ErrorExt for Error { use Error::*; match self { - RetryLater { .. } - | GetLatestCacheRetryExceeded { .. } + RetryLater { source, .. } => match source.retry_hint() { + RetryHint::RetryAfterRouteRefresh => RetryHint::RetryAfterRouteRefresh, + RetryHint::Retryable | RetryHint::NonRetryable => RetryHint::Retryable, + }, + GetLatestCacheRetryExceeded { .. } | NoLeader { .. } | ElectionNoLeader { .. } | ElectionLeaderLeaseExpired { .. } @@ -1354,15 +1357,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(), @@ -1447,6 +1450,7 @@ mod retry_hint_tests { use std::sync::Arc; use common_error::mock::MockError; + use snafu::ResultExt; use super::*; @@ -1464,6 +1468,42 @@ mod retry_hint_tests { assert_eq!(err.retry_hint(), RetryHint::Retryable); } + #[test] + fn test_leader_retry_hints() { + for (error, expected) in [ + ( + NoLeaderSnafu { table_id: 1_u32 }.build(), + RetryHint::Retryable, + ), + (ElectionNoLeaderSnafu.build(), RetryHint::Retryable), + ( + ElectionLeaderLeaseExpiredSnafu.build(), + RetryHint::Retryable, + ), + ( + ElectionLeaderLeaseChangedSnafu.build(), + RetryHint::Retryable, + ), + ] { + assert_eq!(error.retry_hint(), expected, "{error:?}"); + } + } + + #[test] + fn test_procedure_error_retry_hints() { + let source = common_procedure::Error::retry_later(MockError::new(StatusCode::Internal)); + let err = Err::<(), _>(source).context(PutPoisonSnafu).unwrap_err(); + assert_eq!(err.retry_hint(), RetryHint::Retryable); + + let source = common_procedure::Error::retry_later(MockError::new(StatusCode::Internal)); + let err = Err::<(), _>(source) + .context(ProcedureStateReceiverSnafu { + procedure_id: ProcedureId::random(), + }) + .unwrap_err(); + assert_eq!(err.retry_hint(), RetryHint::NonRetryable); + } + #[test] fn test_get_cache_forwards_retry_hint() { let source = Arc::new(Error::retry_later(MockError::new(StatusCode::Internal))); diff --git a/src/common/meta/src/instruction.rs b/src/common/meta/src/instruction.rs index ba4f2e32ab7..54c966bae35 100644 --- a/src/common/meta/src/instruction.rs +++ b/src/common/meta/src/instruction.rs @@ -46,7 +46,7 @@ pub struct InstructionError { pub message: String, /// Retry hint serialized as a string. #[serde( - serialize_with = "RetryHint::serialize_as_str", + serialize_with = "InstructionError::serialize_retry_hint", deserialize_with = "RetryHint::deserialize_from_str" )] pub retry_hint: RetryHint, @@ -105,6 +105,18 @@ impl InstructionError { retry_hint: error.retry_hint(), } } + + /// Keeps instruction replies readable by peers that only know the legacy hints. + fn serialize_retry_hint(hint: &RetryHint, serializer: S) -> Result + where + S: Serializer, + { + let hint = match hint { + RetryHint::RetryAfterRouteRefresh => RetryHint::Retryable, + hint => *hint, + }; + RetryHint::serialize_as_str(&hint, serializer) + } } impl Display for InstructionError { @@ -1259,6 +1271,26 @@ mod tests { ); } + #[test] + fn test_instruction_error_route_refresh_compatibility() { + let error = InstructionError::new( + StatusCode::RegionNotFound, + "region not found", + RetryHint::RetryAfterRouteRefresh, + ); + let serialized = serde_json::to_string(&error).unwrap(); + assert_eq!( + r#"{"code":4005,"message":"region not found","retry_hint":"retryable"}"#, + serialized + ); + let decoded: InstructionError = serde_json::from_str(&serialized).unwrap(); + assert_eq!(decoded.code, error.code); + assert_eq!(decoded.message, error.message); + assert_eq!(decoded.retry_hint, RetryHint::Retryable); + assert_eq!(error.retry_hint, RetryHint::RetryAfterRouteRefresh); + assert_eq!(error.retry_hint.as_str(), "retry_after_route_refresh"); + } + #[test] fn test_instruction_error_from_error() { let error = MockError::new(StatusCode::RegionNotFound); diff --git a/src/common/meta/src/node_manager.rs b/src/common/meta/src/node_manager.rs index 6cc6c3e1539..d18e2cc8471 100644 --- a/src/common/meta/src/node_manager.rs +++ b/src/common/meta/src/node_manager.rs @@ -28,12 +28,22 @@ use crate::peer::Peer; #[async_trait::async_trait] pub trait Datanode: Send + Sync { /// Handles DML, and DDL requests. - async fn handle(&self, request: RegionRequest) -> Result; + async fn handle_request(&self, request: Arc) -> Result; /// Handles query requests async fn handle_query(&self, request: QueryRequest) -> Result; } +impl dyn Datanode + '_ { + /// Handles an owned or shared request without cloning its payload. + pub async fn handle(&self, request: R) -> Result + where + R: Into> + Send, + { + self.handle_request(request.into()).await + } +} + pub type DatanodeRef = Arc; /// The trait for handling requests to flownode diff --git a/src/common/meta/src/peer.rs b/src/common/meta/src/peer.rs index 924c293bbe0..18ffd92f8bb 100644 --- a/src/common/meta/src/peer.rs +++ b/src/common/meta/src/peer.rs @@ -22,6 +22,9 @@ use crate::cluster::NodeInfo; use crate::error::Error; use crate::{DatanodeId, FlownodeId}; +/// Shared reference to a [`Peer`]. +pub type PeerRef = Arc; + /// [`PeerResolver`] provides methods to look up peer information by node ID. /// /// This trait allows resolving both datanode and flownode peers, regardless of their current activity status. diff --git a/src/common/meta/src/reconciliation/reconcile_catalog.rs b/src/common/meta/src/reconciliation/reconcile_catalog.rs index d3c2e526dd3..d154f5aa27f 100644 --- a/src/common/meta/src/reconciliation/reconcile_catalog.rs +++ b/src/common/meta/src/reconciliation/reconcile_catalog.rs @@ -16,6 +16,7 @@ use std::any::Any; use std::fmt::Debug; use std::time::Instant; +use common_error::ext::ErrorExt; use common_procedure::error::FromJsonSnafu; use common_procedure::{ Context as ProcedureContext, Error as ProcedureError, LockKey, Procedure, @@ -188,7 +189,7 @@ impl Procedure for ReconcileCatalogProcedure { Ok(status) } Err(e) => { - if e.is_retry_later() { + if e.retry_hint().is_retryable() { metrics::METRIC_META_RECONCILIATION_PROCEDURE_ERROR .with_label_values(&[procedure_name, step, metrics::ERROR_TYPE_RETRYABLE]) .inc(); diff --git a/src/common/meta/src/reconciliation/reconcile_database.rs b/src/common/meta/src/reconciliation/reconcile_database.rs index 7d4262877c1..c5c23c28023 100644 --- a/src/common/meta/src/reconciliation/reconcile_database.rs +++ b/src/common/meta/src/reconciliation/reconcile_database.rs @@ -23,6 +23,7 @@ use std::fmt::Debug; use std::time::Instant; use async_trait::async_trait; +use common_error::ext::ErrorExt; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ Context as ProcedureContext, Error as ProcedureError, LockKey, Procedure, @@ -228,7 +229,7 @@ impl Procedure for ReconcileDatabaseProcedure { Ok(status) } Err(e) => { - if e.is_retry_later() { + if e.retry_hint().is_retryable() { metrics::METRIC_META_RECONCILIATION_PROCEDURE_ERROR .with_label_values(&[procedure_name, step, metrics::ERROR_TYPE_RETRYABLE]) .inc(); diff --git a/src/common/meta/src/reconciliation/reconcile_logical_tables.rs b/src/common/meta/src/reconciliation/reconcile_logical_tables.rs index 8c16e5a450f..5c5c654f6df 100644 --- a/src/common/meta/src/reconciliation/reconcile_logical_tables.rs +++ b/src/common/meta/src/reconciliation/reconcile_logical_tables.rs @@ -22,6 +22,7 @@ use std::any::Any; use std::fmt::Debug; use async_trait::async_trait; +use common_error::ext::ErrorExt; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, @@ -207,7 +208,7 @@ impl Procedure for ReconcileLogicalTablesProcedure { Ok(status) } Err(e) => { - if e.is_retry_later() { + if e.retry_hint().is_retryable() { metrics::METRIC_META_RECONCILIATION_PROCEDURE_ERROR .with_label_values(&[procedure_name, step, metrics::ERROR_TYPE_RETRYABLE]) .inc(); diff --git a/src/common/meta/src/reconciliation/reconcile_table.rs b/src/common/meta/src/reconciliation/reconcile_table.rs index d48999b895f..a14fe13a3db 100644 --- a/src/common/meta/src/reconciliation/reconcile_table.rs +++ b/src/common/meta/src/reconciliation/reconcile_table.rs @@ -20,6 +20,7 @@ pub(crate) mod update_table_info; use std::fmt::Debug; +use common_error::ext::ErrorExt; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, @@ -318,7 +319,7 @@ impl Procedure for ReconcileTableProcedure { Ok(status) } Err(e) => { - if e.is_retry_later() { + if e.retry_hint().is_retryable() { metrics::METRIC_META_RECONCILIATION_PROCEDURE_ERROR .with_label_values(&[procedure_name, step, metrics::ERROR_TYPE_RETRYABLE]) .inc(); diff --git a/src/common/meta/src/rpc/router.rs b/src/common/meta/src/rpc/router.rs index 519b12e081f..b43b9cc0b4a 100644 --- a/src/common/meta/src/rpc/router.rs +++ b/src/common/meta/src/rpc/router.rs @@ -128,15 +128,25 @@ pub fn convert_to_region_leader_map(region_routes: &[RegionRoute]) -> HashMap>() } +/// Returns the leader of the first route matching the region number. pub fn find_region_leader( region_routes: &[RegionRoute], region_number: RegionNumber, ) -> Option { + find_region_leader_ref(region_routes, region_number).cloned() +} + +/// Borrows the leader of the first route matching the region number. +/// +/// Returns `None` if no route matches or the first matching route has no leader. +pub fn find_region_leader_ref( + region_routes: &[RegionRoute], + region_number: RegionNumber, +) -> Option<&Peer> { region_routes .iter() .find(|x| x.region.id.region_number() == region_number) .and_then(|r| r.leader_peer.as_ref()) - .cloned() } /// Returns the region numbers of the leader regions on the target datanode. @@ -616,6 +626,27 @@ mod tests { } } + #[test] + fn test_find_region_leader() { + let mut routes = vec![ + new_test_region_route(RegionId::new(1, 1)), + new_test_region_route(RegionId::new(1, 1)), + ]; + routes[1].leader_peer = Some(Peer::new(2, "a2")); + + assert_eq!( + find_region_leader_ref(&routes, 1), + routes[0].leader_peer.as_ref() + ); + assert_eq!(find_region_leader(&routes, 1), routes[0].leader_peer); + assert_eq!(find_region_leader_ref(&routes, 2), None); + assert_eq!(find_region_leader(&routes, 2), None); + + routes[0].leader_peer = None; + assert_eq!(find_region_leader_ref(&routes, 1), None); + assert_eq!(find_region_leader(&routes, 1), None); + } + #[test] fn test_leader_is_downgraded() { let mut region_route = RegionRoute { diff --git a/src/common/meta/src/test_util.rs b/src/common/meta/src/test_util.rs index 1009bc1c99a..0fbc49db913 100644 --- a/src/common/meta/src/test_util.rs +++ b/src/common/meta/src/test_util.rs @@ -111,7 +111,8 @@ struct MockNode { #[async_trait::async_trait] impl Datanode for MockNode { - async fn handle(&self, request: RegionRequest) -> Result { + async fn handle_request(&self, request: Arc) -> Result { + let request = Arc::try_unwrap(request).expect("mock request must be exclusively owned"); self.handler.handle(&self.peer, request).await } diff --git a/src/common/procedure/src/error.rs b/src/common/procedure/src/error.rs index 67fab54b156..b3f4a83d5f4 100644 --- a/src/common/procedure/src/error.rs +++ b/src/common/procedure/src/error.rs @@ -306,7 +306,8 @@ impl ErrorExt for Error { | Error::DeletePoison { source, .. } | Error::GetPoison { source, .. } | Error::CheckStatus { source, .. } => source.retry_hint(), - Error::ProcedureExec { source, .. } => source.retry_hint(), + // The watcher reports this only after the child has failed or been poisoned. + Error::ProcedureExec { .. } => RetryHint::NonRetryable, Error::StartRemoveOutdatedMetaTask { source, .. } | Error::StopRemoveOutdatedMetaTask { source, .. } => source.retry_hint(), Error::DeleteState { error, .. } => retry_hint_from_opendal_error(error), @@ -379,8 +380,11 @@ mod tests { use std::sync::Arc; use common_error::mock::MockError; + use tokio::sync::watch; use super::*; + use crate::procedure::ProcedureState; + use crate::watcher; #[test] fn test_retry_later_hint_is_retryable() { @@ -397,6 +401,21 @@ mod tests { assert_eq!(err.retry_hint(), RetryHint::Retryable); } + #[tokio::test] + async fn test_finished_procedure_hint_is_non_retryable() { + let source = Arc::new(Error::retry_later(MockError::new(StatusCode::Internal))); + for state in [ + ProcedureState::failed(source.clone()), + ProcedureState::poisoned(Default::default(), source), + ] { + let (sender, mut receiver) = watch::channel(ProcedureState::Running); + sender.send(state).unwrap(); + let err = watcher::wait(&mut receiver).await.unwrap_err(); + + assert_eq!(err.retry_hint(), RetryHint::NonRetryable); + } + } + #[test] fn test_retry_exceeded_hint_is_non_retryable() { let source = Arc::new(Error::retry_later(MockError::new(StatusCode::Internal))); diff --git a/src/datanode/src/error.rs b/src/datanode/src/error.rs index 6c176547edd..96ccd7ef354 100644 --- a/src/datanode/src/error.rs +++ b/src/datanode/src/error.rs @@ -524,8 +524,8 @@ impl ErrorExt for Error { use Error::*; match self { + RegionNotReady { .. } | RegionNotFound { .. } => RetryHint::RetryAfterRouteRefresh, RegionBusy { .. } - | RegionNotReady { .. } | ConcurrentQueryLimiterClosed { .. } | ConcurrentQueryLimiterTimeout { .. } => RetryHint::Retryable, NewPlanDecoder { source, .. } | ExecuteLogicalPlan { source, .. } => { @@ -565,14 +565,28 @@ mod tests { use super::*; #[test] - fn test_region_state_hints_are_retryable() { + fn test_region_state_retry_hints() { let region_id = RegionId::new(1024, 1); - - let err = RegionBusySnafu { region_id }.build(); - assert_eq!(err.retry_hint(), RetryHint::Retryable); - - let err = RegionNotReadySnafu { region_id }.build(); - assert_eq!(err.retry_hint(), RetryHint::Retryable); + for (error, status, hint) in [ + ( + RegionBusySnafu { region_id }.build(), + StatusCode::RegionBusy, + RetryHint::Retryable, + ), + ( + RegionNotReadySnafu { region_id }.build(), + StatusCode::RegionNotReady, + RetryHint::RetryAfterRouteRefresh, + ), + ( + RegionNotFoundSnafu { region_id }.build(), + StatusCode::RegionNotFound, + RetryHint::RetryAfterRouteRefresh, + ), + ] { + assert_eq!(error.status_code(), status); + assert_eq!(error.retry_hint(), hint); + } } #[test] diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 10bcf32b13d..f3e5e1a0f60 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -2527,7 +2527,7 @@ mod tests { .unwrap_err(); assert_matches!(&err, error::Error::HandleBatchOpenRequest { .. }); assert_eq!(err.status_code(), StatusCode::RegionNotFound); - assert_eq!(err.retry_hint(), RetryHint::NonRetryable); + assert_eq!(err.retry_hint(), RetryHint::RetryAfterRouteRefresh); } struct CurrentEngineTest { diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md index b5c0871fc11..c1bda73cba2 100644 --- a/src/frontend/AGENTS.md +++ b/src/frontend/AGENTS.md @@ -42,7 +42,8 @@ remote datanodes via `operator`/`client`. - **Insert** (`instance/grpc.rs`): `handle_inserts` / `handle_row_inserts` → `check_permission` → `operator`'s `Inserter` (schema validation, optional auto-create, partition routing) → local `RegionServer` (standalone) or RPC to - datanodes (distributed). + datanodes (distributed). `FrontendBuilder` configures the operator write dispatcher + with frontend retry options for distributed inserts, deletes, and bulk writes. ## Public surface @@ -72,7 +73,8 @@ cargo nextest run -p frontend `servers`; permissions, planning, and routing live here. - Standalone vs distributed diverge in datanode access (the `standalone` crate's local `RegionServer` adapter vs `NodeClients` RPC), MetaClient usage, and - whether heartbeat matters. In standalone, the cache invalidator is a no-op. + whether heartbeat matters. The builder defaults local invalidation to its + layered cache registry; callers may supply a custom invalidator. ## Maintenance contract diff --git a/src/frontend/Cargo.toml b/src/frontend/Cargo.toml index 7ad24a466e8..15714839c48 100644 --- a/src/frontend/Cargo.toml +++ b/src/frontend/Cargo.toml @@ -97,5 +97,6 @@ hyper-util = { workspace = true, features = ["tokio"] } reqwest.workspace = true serde_json.workspace = true strfmt = "0.2" +temp-env = "0.3" tower.workspace = true uuid.workspace = true diff --git a/src/frontend/src/error.rs b/src/frontend/src/error.rs index c0f9e2fe6bd..556852caa0c 100644 --- a/src/frontend/src/error.rs +++ b/src/frontend/src/error.rs @@ -314,6 +314,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)] @@ -382,7 +389,8 @@ impl ErrorExt for Error { | Error::IllegalAuthConfig { .. } | Error::ColumnNoneDefaultValue { .. } | Error::IncompleteGrpcRequest { .. } - | Error::InvalidTlsConfig { .. } => StatusCode::InvalidArguments, + | Error::InvalidTlsConfig { .. } + | Error::InvalidWriteRetryConfig { .. } => StatusCode::InvalidArguments, Error::NotSupported { .. } => StatusCode::Unsupported, diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index bb15abd7bbf..2499cd19d1d 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -36,9 +36,10 @@ use crate::heartbeat::HeartbeatTask; use crate::instance::Instance; use crate::service_config::{ InfluxdbOptions, JaegerOptions, MysqlOptions, OpentsdbOptions, OtlpOptions, PostgresOptions, - PromStoreOptions, + PromStoreOptions, WriteRetryOptions, }; +/// Configures frontend protocols, query execution, and distributed write retries. #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] #[serde(default)] pub struct FrontendOptions { @@ -71,6 +72,8 @@ pub struct FrontendOptions { pub meta_client: Option, pub logging: LoggingOptions, pub datanode: DatanodeClientOptions, + /// Retry policy for retryable distributed write failures. + pub write_retry: WriteRetryOptions, pub user_provider: Option, pub tracing: TracingOptions, pub query: QueryOptions, @@ -106,6 +109,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(), @@ -218,6 +222,72 @@ mod tests { type GrpcStream = Pin> + Send + Sync + 'static>>; + #[test] + fn test_write_retry_config() { + let defaults: FrontendOptions = toml::from_str("").unwrap(); + assert_eq!(defaults.write_retry, WriteRetryOptions::default()); + let options: FrontendOptions = toml::from_str( + r#" +[write_retry] +max_retries = 5 +initial_backoff = "250ms" +max_backoff = "5s" +backoff_multiplier = 3.0 +"#, + ) + .unwrap(); + assert_eq!( + options.write_retry, + WriteRetryOptions { + max_retries: 5, + initial_backoff: Duration::from_millis(250), + max_backoff: Duration::from_secs(5), + backoff_multiplier: 3.0, + } + ); + options.write_retry.validate().unwrap(); + let serialized = toml::to_string(&options).unwrap(); + let parsed: FrontendOptions = toml::from_str(&serialized).unwrap(); + assert_eq!(options.write_retry, parsed.write_retry); + let disabled: FrontendOptions = toml::from_str("[write_retry]\nmax_retries = 0").unwrap(); + assert_eq!(disabled.write_retry.max_retries, 0); + assert_eq!( + disabled.write_retry.initial_backoff, + defaults.write_retry.initial_backoff + ); + } + + #[test] + fn test_write_retry_config_from_env() { + temp_env::with_vars( + [ + ("FRONTEND_RETRY_TEST__WRITE_RETRY__MAX_RETRIES", Some("4")), + ( + "FRONTEND_RETRY_TEST__WRITE_RETRY__INITIAL_BACKOFF", + Some("500ms"), + ), + ("FRONTEND_RETRY_TEST__WRITE_RETRY__MAX_BACKOFF", Some("8s")), + ( + "FRONTEND_RETRY_TEST__WRITE_RETRY__BACKOFF_MULTIPLIER", + Some("1.5"), + ), + ], + || { + let options = + FrontendOptions::load_layered_options(None, "FRONTEND_RETRY_TEST").unwrap(); + assert_eq!( + options.write_retry, + WriteRetryOptions { + max_retries: 4, + initial_backoff: Duration::from_millis(500), + max_backoff: Duration::from_secs(8), + backoff_multiplier: 1.5, + } + ); + }, + ); + } + #[test] fn test_toml() { let opts = FrontendOptions::default(); diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index e49b75bbf89..62512a47e63 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -24,8 +24,8 @@ use catalog::system_schema::semantic_graph::EntityGraphProviderRef; use common_base::Plugins; use common_datasource::object_store::LocalFileAccess; use common_event_recorder::{EventRecorderImpl, EventRecorderRef}; -use common_meta::cache::{LayeredCacheRegistryRef, TableRouteCacheRef}; -use common_meta::cache_invalidator::{CacheInvalidatorRef, DummyCacheInvalidator}; +use common_meta::cache::{LayeredCacheRegistryRef, 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; @@ -44,6 +44,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; @@ -64,7 +65,7 @@ pub struct FrontendBuilder { options: FrontendOptions, kv_backend: KvBackendRef, layered_cache_registry: LayeredCacheRegistryRef, - local_cache_invalidator: Option, + local_cache_invalidator: CacheInvalidatorRef, catalog_manager: CatalogManagerRef, node_manager: NodeManagerRef, plugins: Option, @@ -74,6 +75,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, @@ -87,8 +91,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, @@ -134,9 +138,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 } } @@ -181,7 +186,11 @@ impl FrontendBuilder { } } + /// Validates configuration and assembles the frontend [`Instance`]. pub async fn try_build(self) -> Result { + if let Err(reason) = self.options.write_retry.validate() { + return error::InvalidWriteRetryConfigSnafu { reason }.fail(); + } let kv_backend = self.kv_backend; let node_manager = self.node_manager; let plugins = self.plugins.unwrap_or_default(); @@ -204,9 +213,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::() { @@ -222,18 +246,24 @@ impl FrontendBuilder { name: TABLE_FLOWNODE_SET_CACHE_NAME, })?; - let inserter = Arc::new(Inserter::new( - self.catalog_manager.clone(), - partition_manager.clone(), - node_manager.clone(), - table_flownode_cache, - self.options.auto_create_table, - )); - let deleter = Arc::new(Deleter::new( - self.catalog_manager.clone(), - partition_manager.clone(), - node_manager.clone(), - )); + let inserter = Arc::new( + Inserter::new( + self.catalog_manager.clone(), + partition_manager.clone(), + node_manager.clone(), + table_flownode_cache, + self.options.auto_create_table, + ) + .with_write_dispatcher(write_dispatcher.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(), @@ -373,3 +403,76 @@ impl FrontendBuilder { }) } } + +#[cfg(test)] +mod tests { + use api::v1::meta::Role; + use common_meta::cache_invalidator::{Context, DummyCacheInvalidator}; + use common_meta::instruction::CacheIdent; + use common_meta::key::table_route::{TableRouteKey, TableRouteValue}; + use common_meta::key::{MetadataKey, MetadataValue}; + use common_meta::rpc::router::{Region, RegionRoute}; + use common_meta::rpc::store::PutRequest; + use meta_client::client::MetaClientBuilder; + use store_api::storage::RegionId; + + use super::*; + + #[tokio::test] + async fn test_local_cache_invalidator() { + for custom_invalidator in [false, true] { + let mut builder = FrontendBuilder::new_test( + &FrontendOptions::default(), + Arc::new(MetaClientBuilder::new(0, Role::Frontend).build()), + ); + if custom_invalidator { + builder = builder.with_local_cache_invalidator(Arc::new(DummyCacheInvalidator)); + } + let backend = builder.kv_backend.clone(); + let routes: TableRouteCacheRef = builder.layered_cache_registry.get().unwrap(); + let partitions: PartitionInfoCacheRef = builder.layered_cache_registry.get().unwrap(); + let instance = builder.try_build().await.unwrap(); + let table_id = 1024; + for region_number in [1, 2] { + let route = TableRouteValue::physical(vec![RegionRoute { + region: Region { + id: RegionId::new(table_id, region_number), + ..Default::default() + }, + ..Default::default() + }]); + backend + .put( + PutRequest::new() + .with_key(TableRouteKey::new(table_id).to_bytes()) + .with_value(route.try_as_raw_value().unwrap()), + ) + .await + .unwrap(); + if region_number == 1 { + routes.get(table_id).await.unwrap().unwrap(); + partitions.get(table_id).await.unwrap().unwrap(); + } + } + assert!(routes.contains_key(&table_id)); + assert!(partitions.contains_key(&table_id)); + + instance + .cache_invalidator() + .invalidate(&Context::default(), &[CacheIdent::TableId(table_id)]) + .await + .unwrap(); + + assert_eq!(custom_invalidator, routes.contains_key(&table_id)); + assert_eq!(custom_invalidator, partitions.contains_key(&table_id)); + let route = routes.get(table_id).await.unwrap().unwrap(); + let region = route.as_physical_table_route_ref().unwrap().region_routes[0] + .region + .id; + assert_eq!( + if custom_invalidator { 1 } else { 2 }, + region.region_number() + ); + } + } +} diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index e112e5ae2a9..8cb33a1e28a 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -132,6 +132,7 @@ where PendingRowsBatcher::try_new( self.instance.partition_manager().clone(), self.instance.node_manager().clone(), + self.instance.inserter().write_dispatcher().clone(), self.instance.catalog_manager().clone(), self.instance.table_flownode_set_cache().clone(), opts.prom_store.with_metric_engine, diff --git a/src/frontend/src/service_config.rs b/src/frontend/src/service_config.rs index 03ec12e2c9f..4dc152d0df9 100644 --- a/src/frontend/src/service_config.rs +++ b/src/frontend/src/service_config.rs @@ -24,6 +24,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 postgres::PostgresOptions; pub use prom_store::PromStoreOptions; diff --git a/src/metric-engine/src/engine/bulk_insert.rs b/src/metric-engine/src/engine/bulk_insert.rs index 8472815b145..87b6679c61c 100644 --- a/src/metric-engine/src/engine/bulk_insert.rs +++ b/src/metric-engine/src/engine/bulk_insert.rs @@ -348,6 +348,19 @@ mod tests { #[tokio::test] async fn test_bulk_insert_physical_region_passthrough() { + check_bulk_insert_physical_region_passthrough(false).await; + } + + #[tokio::test] + async fn test_physical_bulk_partition_columns_preserve_stored_rows() { + let stripped = check_bulk_insert_physical_region_passthrough(false).await; + let retained = check_bulk_insert_physical_region_passthrough(true).await; + assert_eq!(retained, stripped); + } + + async fn check_bulk_insert_physical_region_passthrough( + retain_partition_columns: bool, + ) -> String { // Use flat format so that BulkMemtable is used (supports write_bulk). let mito_config = MitoConfig { default_flat_format: true, @@ -376,7 +389,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, @@ -385,6 +402,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); let response = env .metric() @@ -401,6 +422,7 @@ mod tests { .unwrap(); let batches = RecordBatches::try_collect(stream).await.unwrap(); assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 6); + batches.pretty_print().unwrap() } #[tokio::test] diff --git a/src/metric-engine/src/error.rs b/src/metric-engine/src/error.rs index 0803223ba33..f9d61d6af31 100644 --- a/src/metric-engine/src/error.rs +++ b/src/metric-engine/src/error.rs @@ -486,6 +486,9 @@ impl ErrorExt for Error { use Error::*; match self { + PhysicalRegionNotFound { .. } | LogicalRegionNotFound { .. } => { + RetryHint::RetryAfterRouteRefresh + } CreateMitoRegion { source, .. } | OpenMitoRegion { source, .. } | CleanUpMitoRegion { source, .. } @@ -510,3 +513,20 @@ impl ErrorExt for Error { } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_missing_region_retry_hints() { + let region_id = RegionId::new(1024, 1); + for error in [ + PhysicalRegionNotFoundSnafu { region_id }.build(), + LogicalRegionNotFoundSnafu { region_id }.build(), + ] { + assert_eq!(error.status_code(), StatusCode::RegionNotFound); + assert_eq!(error.retry_hint(), RetryHint::RetryAfterRouteRefresh); + } + } +} diff --git a/src/mito2/src/engine/flush_test.rs b/src/mito2/src/engine/flush_test.rs index d7155e10520..8d45f852592 100644 --- a/src/mito2/src/engine/flush_test.rs +++ b/src/mito2/src/engine/flush_test.rs @@ -24,6 +24,7 @@ use api::v1::Rows; use async_trait::async_trait; use common_base::Plugins; use common_base::readable_size::ReadableSize; +use common_error::ext::{ErrorExt, RetryHint}; use common_recordbatch::RecordBatches; use common_time::util::current_time_millis; use common_wal::options::{KafkaWalOptions, WAL_OPTIONS_KEY, WalOptions}; @@ -968,6 +969,7 @@ async fn test_region_write_buffer_rejects_only_full_region_queue() { .expect("the stalled queue should be rejected at the region hard limit") .unwrap() .unwrap_err(); + assert_eq!(err.retry_hint(), RetryHint::Retryable); assert_matches!( err.into_inner().as_any().downcast_ref::().unwrap(), Error::RejectWrite { .. } diff --git a/src/mito2/src/engine/staging_test.rs b/src/mito2/src/engine/staging_test.rs index bd823708509..6493c5e2ad2 100644 --- a/src/mito2/src/engine/staging_test.rs +++ b/src/mito2/src/engine/staging_test.rs @@ -18,11 +18,14 @@ use std::sync::Arc; use std::time::Duration; use std::{assert_matches, fs}; -use api::v1::Rows; -use common_error::ext::ErrorExt; +use api::v1::{ArrowIpc, Rows}; +use common_error::ext::{ErrorExt, RetryHint}; use common_error::status_code::StatusCode; use common_function::utils::partition_expr_version; +use common_grpc::flight; use common_recordbatch::RecordBatches; +use datatypes::arrow::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray}; +use datatypes::arrow::record_batch::RecordBatch; use datatypes::value::Value; use object_store::Buffer; use object_store::layers::mock::{ @@ -34,8 +37,9 @@ use store_api::region_engine::{ RegionEngine, RemapManifestsRequest, SetRegionRoleStateResponse, SettableRegionRoleState, }; use store_api::region_request::{ - ApplyStagingManifestRequest, EnterStagingRequest, RegionAlterRequest, RegionFlushRequest, - RegionPutRequest, RegionRequest, RegionTruncateRequest, StagingPartitionDirective, + ApplyStagingManifestRequest, EnterStagingRequest, RegionAlterRequest, RegionBulkInsertsRequest, + RegionFlushRequest, RegionPutRequest, RegionRequest, RegionTruncateRequest, + StagingPartitionDirective, }; use store_api::storage::{RegionId, ScanRequest}; @@ -254,7 +258,7 @@ fn default_partition_expr() -> String { } #[tokio::test] -async fn test_staging_reject_all_writes_rejects_put() { +async fn test_staging_reject_all_writes_rejects_put_and_bulk() { let mut env = TestEnv::new().await; let engine = env.create_engine(MitoConfig::default()).await; @@ -281,23 +285,49 @@ async fn test_staging_reject_all_writes_rejects_put() { schema: column_schemas, rows: build_rows(0, 2), }; - let err = engine - .handle_request( + let payload = RecordBatch::try_from_iter([ + ( + "tag_0", + Arc::new(StringArray::from(vec!["0", "1"])) as ArrayRef, + ), + ( + "field_0", + Arc::new(Float64Array::from(vec![0.0, 1.0])) as ArrayRef, + ), + ( + "ts", + Arc::new(TimestampMillisecondArray::from(vec![0, 1000])) as ArrayRef, + ), + ]) + .unwrap(); + let (schema, data_header, data) = flight::record_batch_to_ipc(payload.clone()).unwrap(); + for request in [ + RegionRequest::Put(RegionPutRequest { + skip_wal: false, + rows, + hint: None, + partition_expr_version: None, + }), + RegionRequest::BulkInserts(RegionBulkInsertsRequest { region_id, - RegionRequest::Put(RegionPutRequest { - skip_wal: false, - rows, - hint: None, - partition_expr_version: None, - }), - ) - .await - .unwrap_err(); - assert_eq!(err.status_code(), StatusCode::StorageUnavailable); - assert_matches!( - err.into_inner().as_any().downcast_ref::().unwrap(), - Error::RejectWrite { .. } - ); + payload, + raw_data: ArrowIpc { + schema, + data_header, + payload: data, + }, + partition_expr_version: None, + aligned_schema_version: None, + }), + ] { + let err = engine.handle_request(region_id, request).await.unwrap_err(); + assert_eq!(err.status_code(), StatusCode::StorageUnavailable); + assert_eq!(err.retry_hint(), RetryHint::RetryAfterRouteRefresh); + assert_matches!( + err.into_inner().as_any().downcast_ref::().unwrap(), + Error::RejectWriteInStaging { .. } + ); + } } #[tokio::test] diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index 21416bf3652..57811a23b6d 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -633,6 +633,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Region {} rejects writes while in staging", region_id))] + RejectWriteInStaging { + region_id: RegionId, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to compact region {}", region_id))] CompactRegion { region_id: RegionId, @@ -1541,7 +1548,7 @@ impl ErrorExt for Error { RegionClosed { .. } => StatusCode::Cancelled, StaleCompactionExecution { .. } => StatusCode::Cancelled, RegionTruncated { .. } => StatusCode::Cancelled, - RejectWrite { .. } => StatusCode::StorageUnavailable, + RejectWrite { .. } | RejectWriteInStaging { .. } => StatusCode::StorageUnavailable, CompactRegion { source, .. } => source.status_code(), EditRegion { source, .. } => source.status_code(), CompatReader { .. } => StatusCode::Unexpected, @@ -1641,6 +1648,13 @@ impl ErrorExt for Error { use Error::*; match self { + PartitionExprVersionMismatch { .. } + | RegionNotFound { .. } + | RejectWriteInStaging { .. } + | RegionState { + state: RegionRoleState::Follower, + .. + } => RetryHint::RetryAfterRouteRefresh, ReadParquet { .. } | WriteParquet { .. } | RejectWrite { .. } @@ -1721,6 +1735,63 @@ mod tests { use snafu::IntoError; use super::*; + use crate::region::RegionLeaderState; + + #[test] + fn test_partition_version_mismatch_retry_hint() { + let error = PartitionExprVersionMismatchSnafu { + region_id: RegionId::new(1, 1), + request_version: 1u64, + expected_version: 2u64, + } + .build(); + assert_eq!(error.status_code(), StatusCode::InvalidArguments); + assert_eq!(error.retry_hint(), RetryHint::RetryAfterRouteRefresh); + } + + #[test] + fn test_region_state_retry_hint() { + let region_id = RegionId::new(1, 1); + let writable = RegionRoleState::Leader(RegionLeaderState::Writable); + let follower_error = RegionStateSnafu { + region_id, + state: RegionRoleState::Follower, + expect: writable, + } + .build(); + assert_eq!( + follower_error.retry_hint(), + RetryHint::RetryAfterRouteRefresh + ); + assert_eq!(follower_error.status_code(), StatusCode::RegionNotReady); + let leader_error = RegionStateSnafu { + region_id, + state: writable, + expect: RegionRoleState::Follower, + } + .build(); + assert_eq!(leader_error.retry_hint(), RetryHint::Retryable); + for (error, status, hint) in [ + ( + RejectWriteSnafu { region_id }.build(), + StatusCode::StorageUnavailable, + RetryHint::Retryable, + ), + ( + RejectWriteInStagingSnafu { region_id }.build(), + StatusCode::StorageUnavailable, + RetryHint::RetryAfterRouteRefresh, + ), + ( + RegionNotFoundSnafu { region_id }.build(), + StatusCode::RegionNotFound, + RetryHint::RetryAfterRouteRefresh, + ), + ] { + assert_eq!(error.status_code(), status); + assert_eq!(error.retry_hint(), hint); + } + } #[test] fn test_manifest_update_persistence() { diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 7a3f2179614..05b82f3d346 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -54,8 +54,8 @@ use crate::cache::{CacheManagerRef, CacheStrategy}; use crate::config::VectorIndexConfig; use crate::config::{BloomFilterConfig, FulltextIndexConfig, InvertedIndexConfig}; use crate::error::{ - BuildIndexAsyncSnafu, DecodeSnafu, Error, InvalidRecordBatchSnafu, JoinSnafu, RegionClosedSnafu, - RegionDroppedSnafu, RegionTruncatedSnafu, Result, + BuildIndexAsyncSnafu, DecodeSnafu, Error, InvalidRecordBatchSnafu, JoinSnafu, + RegionClosedSnafu, RegionDroppedSnafu, RegionTruncatedSnafu, Result, }; use crate::metrics::{ INDEX_ARTIFACT_CLEANUP_FAILURE_TOTAL, INDEX_CREATE_MEMORY_USAGE, INDEX_PUBLICATION_STALE_TOTAL, diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index ff5ae650229..856fb75978d 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -26,7 +26,7 @@ use store_api::storage::RegionId; use crate::error::{ InvalidRequestSnafu, PartitionExprVersionMismatchSnafu, RegionNotFoundSnafu, RegionStateSnafu, - RejectWriteSnafu, Result, + RejectWriteInStagingSnafu, RejectWriteSnafu, Result, }; use crate::metrics; use crate::metrics::{ @@ -339,7 +339,7 @@ impl RegionWorkerLoop { if region.reject_all_writes_in_staging() { sender_req .sender - .send(RejectWriteSnafu { region_id }.fail()); + .send(RejectWriteInStagingSnafu { region_id }.fail()); continue; } @@ -402,7 +402,7 @@ impl RegionWorkerLoop { if region.reject_all_writes_in_staging() { sender_req .sender - .send(RejectWriteSnafu { region_id }.fail()); + .send(RejectWriteInStagingSnafu { region_id }.fail()); continue; } let expected_version = region.expected_partition_expr_version(); @@ -484,7 +484,9 @@ impl RegionWorkerLoop { RegionRoleState::Leader(RegionLeaderState::Writable) | RegionRoleState::Leader(RegionLeaderState::Staging) => { if region.reject_all_writes_in_staging() { - bulk_req.sender.send(RejectWriteSnafu { region_id }.fail()); + bulk_req + .sender + .send(RejectWriteInStagingSnafu { region_id }.fail()); continue; } let region_ctx = RegionWriteCtx::new( @@ -531,7 +533,9 @@ impl RegionWorkerLoop { continue; }; if region.reject_all_writes_in_staging() { - bulk_req.sender.send(RejectWriteSnafu { region_id }.fail()); + bulk_req + .sender + .send(RejectWriteInStagingSnafu { region_id }.fail()); continue; } let expected_version = region.expected_partition_expr_version(); diff --git a/src/operator/AGENTS.md b/src/operator/AGENTS.md index 761d8e5080a..9da3c4dfed8 100644 --- a/src/operator/AGENTS.md +++ b/src/operator/AGENTS.md @@ -14,6 +14,7 @@ insert/delete routing. | Statements | `src/operator/src/statement.rs`, `src/operator/src/statement/` | `StatementExecutor` and per-statement handlers | | Inserts | `src/operator/src/insert.rs`, `src/operator/src/bulk_insert.rs` | Insert, auto-create/alter, routing, dispatch | | Deletes | `src/operator/src/delete.rs` | Delete conversion and dispatch | +| Write dispatch | `src/operator/src/write_dispatcher.rs` | Shared insert/delete/bulk dispatch, bounded retries, route refresh and write request rebuilding | | Request conversion | `src/operator/src/req_convert.rs`, `src/operator/src/req_convert/` | Table/row/column requests to region requests | | Region requests | `src/operator/src/region_req_factory.rs` | Region-level DDL/DML request construction | | Procedures | `src/operator/src/procedure.rs` | Metasrv procedure administration | @@ -38,3 +39,5 @@ cargo sqlness bare -t ``` Check standalone and distributed routing when changing region dispatch. +Write retries operate on prepared region requests in `write_dispatcher`; keep +admission and Flow side effects outside retries, and preserve successful regions. diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml index c9258df2c0f..1df87836b87 100644 --- a/src/operator/Cargo.toml +++ b/src/operator/Cargo.toml @@ -25,6 +25,7 @@ arrow.workspace = true arrow-ipc.workspace = true async-stream.workspace = true async-trait.workspace = true +backon.workspace = true bytes.workspace = true catalog.workspace = true chrono.workspace = true @@ -55,12 +56,14 @@ file-engine.workspace = true futures.workspace = true futures-util.workspace = true humantime.workspace = true +humantime-serde.workspace = true itertools.workspace = true jsonb.workspace = true lazy_static.workspace = true 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 @@ -85,7 +88,9 @@ tonic.workspace = true tracing.workspace = true [dev-dependencies] +cache.workspace = true catalog = { workspace = true, features = ["testing"] } common-meta = { workspace = true, features = ["testing"] } common-test-util.workspace = true path-slash = "0.2" +tokio-stream = { workspace = true, features = ["net"] } diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index 22fae236196..dc570bf8761 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -14,13 +14,13 @@ use std::collections::HashSet; -use ahash::{HashMap, HashMapExt}; use api::v1::flow::{DirtyWindowRequest, DirtyWindowRequests}; use api::v1::region::{ BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request, }; use api::v1::{ArrowIpc, PartitionExprVersion}; use arrow::array::Array; +use arrow::compute; use arrow::record_batch::RecordBatch; use bytes::Bytes; use common_base::AffectedRows; @@ -28,12 +28,16 @@ use common_grpc::FlightData; use common_grpc::flight::{FlightEncoder, FlightMessage}; use common_telemetry::error; use common_telemetry::tracing_context::TracingContext; +use futures::future; 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; +use crate::write_dispatcher::WriteAttempt; use crate::{error, metrics}; impl Inserter { @@ -91,12 +95,6 @@ impl Inserter { .get(®ion_number) .copied() .unwrap_or_default(); - let datanode = self - .partition_manager - .find_region_leader(region_id) - .await - .context(error::FindRegionLeaderSnafu)?; - let request = RegionRequest { header: Some(RegionRequestHeader { tracing_context: TracingContext::from_current_span().to_w3c(), @@ -115,12 +113,16 @@ impl Inserter { })), }; - let _datanode_handle_timer = metrics::HANDLE_BULK_INSERT_ELAPSED - .with_label_values(&["datanode_handle"]) - .start_timer(); - let datanode = self.node_manager.datanode(&datanode).await; - let result = datanode - .handle(request) + let writes = self.write_dispatcher.prepare(vec![request]).await?; + let result = self + .write_dispatcher + .execute( + writes, + Some( + metrics::HANDLE_BULK_INSERT_ELAPSED.with_label_values(&["datanode_handle"]), + ), + None, + ) .await .context(error::RequestRegionSnafu) .map(|r| r.affected_rows); @@ -132,122 +134,142 @@ impl Inserter { return result; } - let mut mask_per_datanode = HashMap::with_capacity(region_masks.len()); - for (region_number, mask) in region_masks { - let region_id = RegionId::new(table_id, region_number); - let datanode = self - .partition_manager - .find_region_leader(region_id) - .await - .context(error::FindRegionLeaderSnafu)?; - mask_per_datanode - .entry(datanode) - .or_insert_with(Vec::new) - .push((region_id, mask)); - } + let regions = region_masks + .keys() + .map(|region| RegionId::new(table_id, *region)) + .collect::>(); + let mut peers = self + .write_dispatcher + .resolve_peers(®ions) + .await + .context(error::FindRegionLeaderSnafu)?; let wait_all_datanode_timer = metrics::HANDLE_BULK_INSERT_ELAPSED .with_label_values(&["wait_all_datanode"]) .start_timer(); - let mut handles = Vec::with_capacity(mask_per_datanode.len()); + // 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 (peer, masks) in mask_per_datanode { - for (region_id, mask) in masks { - if mask.select_none() { - continue; - } - let partition_expr_version = partition_versions - .get(®ion_id.region_number()) - .copied() - .unwrap_or_default(); - let rb = record_batch.clone(); - let schema_bytes = schema_bytes.clone(); - let node_manager = self.node_manager.clone(); - let peer = peer.clone(); - let raw_header_and_data = if mask.select_all() { - Some(( - raw_flight_data.data_header.clone(), - raw_flight_data.data_body.clone(), - )) - } else { - None - }; - let handle: common_runtime::JoinHandle> = - common_runtime::spawn_global(async move { - let (header, payload) = if mask.select_all() { - // SAFETY: raw data must be present, we can avoid re-encoding. - raw_header_and_data.unwrap() - } else { - let filter_timer = metrics::HANDLE_BULK_INSERT_ELAPSED - .with_label_values(&["filter"]) - .start_timer(); - let batch = arrow::compute::filter_record_batch(&rb, mask.array()) - .context(error::ComputeArrowSnafu)?; - filter_timer.observe_duration(); - metrics::BULK_REQUEST_ROWS - .with_label_values(&["rows_per_region"]) - .observe(batch.num_rows() as f64); - - let encode_timer = metrics::HANDLE_BULK_INSERT_ELAPSED - .with_label_values(&["encode"]) - .start_timer(); - let mut iter = FlightEncoder::default() - .encode(FlightMessage::RecordBatch(batch)) - .into_iter(); - let Some(flight_data) = iter.next() else { - // Safety: `iter` on a type of `Vec1`, which is guaranteed to have - // at least one element. - unreachable!() - }; - ensure!( - iter.next().is_none(), - error::NotSupportedSnafu { - feat: "bulk insert RecordBatch with dictionary arrays", - } - ); - encode_timer.observe_duration(); - (flight_data.data_header, flight_data.data_body) - }; - let _datanode_handle_timer = metrics::HANDLE_BULK_INSERT_ELAPSED - .with_label_values(&["datanode_handle"]) - .start_timer(); - let request = RegionRequest { - header: Some(RegionRequestHeader { - tracing_context: TracingContext::from_current_span().to_w3c(), - ..Default::default() - }), - body: Some(region_request::Body::BulkInsert(BulkInsertRequest { - region_id: region_id.as_u64(), - partition_expr_version: partition_expr_version - .map(|value| PartitionExprVersion { value }), - aligned_schema_version: None, - body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { - schema: schema_bytes, - data_header: header, - payload, - })), - })), - }; - - let datanode = node_manager.datanode(&peer).await; - datanode - .handle(request) - .await - .context(error::RequestRegionSnafu) - }); - handles.push(handle); + for (region_number, mask) in region_masks { + let region_id = RegionId::new(table_id, region_number); + if mask.select_none() { + continue; } + let partition_expr_version = partition_versions + .get(®ion_id.region_number()) + .copied() + .unwrap_or_default(); + let peer = peers + .remove(®ion_id) + .context(partition::error::FindLeaderSnafu { + region_id, + table_id, + }) + .context(error::FindRegionLeaderSnafu)?; + let dispatcher = self.write_dispatcher.clone(); + let rb = record_batch.clone(); + let schema_bytes = schema_bytes.clone(); + let raw_header_and_data = if mask.select_all() { + Some(( + raw_flight_data.data_header.clone(), + raw_flight_data.data_body.clone(), + )) + } else { + None + }; + let handle: common_runtime::JoinHandle> = + common_runtime::spawn_global(async move { + let (header, payload) = if mask.select_all() { + // SAFETY: raw data must be present, we can avoid re-encoding. + raw_header_and_data.unwrap() + } else { + let filter_timer = metrics::HANDLE_BULK_INSERT_ELAPSED + .with_label_values(&["filter"]) + .start_timer(); + let batch = compute::filter_record_batch(&rb, mask.array()) + .context(error::ComputeArrowSnafu)?; + filter_timer.observe_duration(); + metrics::BULK_REQUEST_ROWS + .with_label_values(&["rows_per_region"]) + .observe(batch.num_rows() as f64); + + let encode_timer = metrics::HANDLE_BULK_INSERT_ELAPSED + .with_label_values(&["encode"]) + .start_timer(); + let mut iter = FlightEncoder::default() + .encode(FlightMessage::RecordBatch(batch)) + .into_iter(); + let Some(flight_data) = iter.next() else { + // Safety: `iter` on a type of `Vec1`, which is guaranteed to have + // at least one element. + unreachable!() + }; + ensure!( + iter.next().is_none(), + error::NotSupportedSnafu { + feat: "bulk insert RecordBatch with dictionary arrays", + } + ); + encode_timer.observe_duration(); + (flight_data.data_header, flight_data.data_body) + }; + let request = RegionRequest { + header: Some(RegionRequestHeader { + tracing_context: TracingContext::from_current_span().to_w3c(), + ..Default::default() + }), + body: Some(region_request::Body::BulkInsert(BulkInsertRequest { + region_id: region_id.as_u64(), + partition_expr_version: partition_expr_version + .map(|value| PartitionExprVersion { value }), + aligned_schema_version: None, + body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { + schema: schema_bytes, + data_header: header, + payload, + })), + })), + }; + + Ok(dispatcher + .send( + peer, + request, + Some( + &metrics::HANDLE_BULK_INSERT_ELAPSED + .with_label_values(&["datanode_handle"]), + ), + ) + .await) + }); + handles.push(handle); } - let region_responses = futures::future::try_join_all(handles) - .await - .context(error::JoinTaskSnafu)?; + let dispatcher = self.write_dispatcher.clone(); + let rows_inserted = common_runtime::spawn_global(async move { + let results = future::try_join_all(handles) + .await + .context(error::JoinTaskSnafu)?; + let attempts = results.into_iter().collect::>>()?; + dispatcher + .finish( + attempts, + Some( + &metrics::HANDLE_BULK_INSERT_ELAPSED + .with_label_values(&["datanode_handle"]), + ), + Some(&cancellation), + ) + .await + .context(error::RequestRegionSnafu) + .map(|response| response.affected_rows) + }) + .await + .context(error::JoinTaskSnafu)??; wait_all_datanode_timer.observe_duration(); - let mut rows_inserted: usize = 0; - for res in region_responses { - rows_inserted += res?.affected_rows; - } crate::metrics::DIST_INGEST_ROW_COUNT .with_label_values(&[db_name.as_str()]) .inc_by(rows_inserted as u64); diff --git a/src/operator/src/delete.rs b/src/operator/src/delete.rs index 9641e0e27d4..9bf14471696 100644 --- a/src/operator/src/delete.rs +++ b/src/operator/src/delete.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use std::sync::Arc; use std::{iter, mem}; @@ -20,45 +20,55 @@ use api::v1::region::{DeleteRequests as RegionDeleteRequests, RegionRequestHeade use api::v1::{DeleteRequests, RowDeleteRequests}; use catalog::CatalogManagerRef; use common_meta::node_manager::{AffectedRows, NodeManagerRef}; -use common_meta::peer::Peer; use common_query::Output; use common_telemetry::tracing_context::TracingContext; -use futures_util::future; use partition::manager::PartitionRuleManagerRef; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt, ensure}; use table::TableRef; use table::requests::DeleteRequest as TableDeleteRequest; +use tokio_util::sync::CancellationToken; use crate::error::{ - CatalogSnafu, FindRegionLeaderSnafu, InvalidDeleteRequestSnafu, JoinTaskSnafu, - MissingTimeIndexColumnSnafu, RequestDeletesSnafu, Result, TableNotFoundSnafu, + CatalogSnafu, InvalidDeleteRequestSnafu, JoinTaskSnafu, MissingTimeIndexColumnSnafu, + RequestDeletesSnafu, Result, TableNotFoundSnafu, }; use crate::region_req_factory::RegionRequestFactory; use crate::req_convert::common::preprocess_row_delete_requests; use crate::req_convert::delete::{ColumnToRow, RowToRegion, TableToRegion}; +use crate::write_dispatcher::{WriteDispatcher, WriteDispatcherRef}; pub struct Deleter { catalog_manager: CatalogManagerRef, partition_manager: PartitionRuleManagerRef, - node_manager: NodeManagerRef, + write_dispatcher: WriteDispatcherRef, } pub type DeleterRef = Arc; impl Deleter { + /// Creates a row deleter with the default write dispatcher. pub fn new( catalog_manager: CatalogManagerRef, partition_manager: PartitionRuleManagerRef, node_manager: NodeManagerRef, ) -> Self { Self { + write_dispatcher: Arc::new(WriteDispatcher::new( + node_manager, + partition_manager.clone(), + )), catalog_manager, partition_manager, - node_manager, } } + /// Installs the dispatcher shared with the other write paths. + pub fn with_write_dispatcher(mut self, dispatcher: WriteDispatcherRef) -> Self { + self.write_dispatcher = dispatcher; + self + } + pub async fn handle_column_deletes( &self, requests: DeleteRequests, @@ -129,52 +139,25 @@ impl Deleter { ..Default::default() }); - let tasks = self - .group_requests_by_peer(requests) - .await? - .into_iter() - .map(|(peer, deletes)| { - let request = request_factory.build_delete(deletes); - let node_manager = self.node_manager.clone(); - common_runtime::spawn_global(async move { - node_manager - .datanode(&peer) - .await - .handle(request) - .await - .context(RequestDeletesSnafu) - }) - }); - let results = future::try_join_all(tasks).await.context(JoinTaskSnafu)?; - - let affected_rows = results - .into_iter() - .map(|resp| resp.map(|r| r.affected_rows)) - .sum::>()?; + let dispatcher = self.write_dispatcher.clone(); + let request = request_factory.build_delete(requests); + let writes = dispatcher.prepare(vec![request]).await?; + // 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); Ok(affected_rows) } - async fn group_requests_by_peer( - &self, - requests: RegionDeleteRequests, - ) -> Result> { - let mut deletes: HashMap = HashMap::new(); - - for req in requests.requests { - let peer = self - .partition_manager - .find_region_leader(req.region_id.into()) - .await - .context(FindRegionLeaderSnafu)?; - deletes.entry(peer).or_default().requests.push(req); - } - - Ok(deletes) - } - async fn trim_columns( &self, mut requests: RowDeleteRequests, diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index a4a69ef8ba9..a27a383b231 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -38,7 +38,7 @@ use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME; use common_grpc_expr::util::ColumnExpr; use common_meta::cache::TableFlownodeSetCacheRef; use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME; -use common_meta::node_manager::{AffectedRows, NodeManagerRef}; +use common_meta::node_manager::NodeManagerRef; use common_meta::peer::Peer; use common_meta::rpc::ddl::TriggerReason; use common_query::Output; @@ -47,7 +47,6 @@ use common_query::prelude::{greptime_timestamp, greptime_value}; use common_telemetry::tracing_context::TracingContext; use common_telemetry::{error, info, warn}; use datatypes::schema::SkippingIndexOptions; -use futures_util::future; use meter_macros::write_meter; use partition::manager::PartitionRuleManagerRef; use session::context::QueryContextRef; @@ -72,10 +71,11 @@ use table::requests::{ VALID_TABLE_OPTION_KEYS, is_semantic_option_key, validate_semantic_option, }; use table::table_reference::TableReference; +use tokio_util::sync::CancellationToken; use crate::error::{ - CatalogSnafu, ColumnOptionsSnafu, CreatePartitionRulesSnafu, FindRegionLeaderSnafu, - InvalidInsertRequestSnafu, JoinTaskSnafu, RequestInsertsSnafu, Result, TableNotFoundSnafu, + CatalogSnafu, ColumnOptionsSnafu, CreatePartitionRulesSnafu, InvalidInsertRequestSnafu, + JoinTaskSnafu, RequestInsertsSnafu, Result, TableNotFoundSnafu, }; use crate::expr_helper; use crate::region_req_factory::RegionRequestFactory; @@ -84,11 +84,13 @@ use crate::req_convert::insert::{ ColumnToRow, RowToRegion, StatementToRegion, TableToRegion, fill_reqs_with_impure_default, }; use crate::statement::StatementExecutor; +use crate::write_dispatcher::{WriteDispatcher, WriteDispatcherRef}; pub struct Inserter { catalog_manager: CatalogManagerRef, pub(crate) partition_manager: PartitionRuleManagerRef, pub(crate) node_manager: NodeManagerRef, + pub(crate) write_dispatcher: WriteDispatcherRef, pub(crate) table_flownode_set_cache: TableFlownodeSetCacheRef, /// Server-side upper bound for auto table creation on write. /// When `false`, missing tables are never auto-created regardless of the @@ -150,6 +152,7 @@ pub struct InstantAndNormalInsertRequests { } impl Inserter { + /// Creates a row inserter with the default write dispatcher. pub fn new( catalog_manager: CatalogManagerRef, partition_manager: PartitionRuleManagerRef, @@ -158,6 +161,10 @@ impl Inserter { auto_create_table: bool, ) -> Self { Self { + write_dispatcher: Arc::new(WriteDispatcher::new( + node_manager.clone(), + partition_manager.clone(), + )), catalog_manager, partition_manager, node_manager, @@ -166,6 +173,12 @@ impl Inserter { } } + /// Installs the shared dispatcher used by direct writes. + pub fn with_write_dispatcher(mut self, dispatcher: WriteDispatcherRef) -> Self { + self.write_dispatcher = dispatcher; + self + } + pub async fn handle_column_inserts( &self, requests: InsertRequests, @@ -430,30 +443,19 @@ impl Inserter { .await?; flow_mirror_task.detach(self.node_manager.clone())?; - // Write requests to datanode and wait for response - let write_tasks = self - .group_requests_by_peer(normal_requests) - .await? - .into_iter() - .map(|(peer, inserts)| { - let node_manager = self.node_manager.clone(); - let request = request_factory.build_insert(inserts); - common_runtime::spawn_global(async move { - node_manager - .datanode(&peer) - .await - .handle(request) - .await - .context(RequestInsertsSnafu) - }) - }); - let results = future::try_join_all(write_tasks) - .await - .context(JoinTaskSnafu)?; - let affected_rows = results - .into_iter() - .map(|resp| resp.map(|r| r.affected_rows)) - .sum::>()?; + let dispatcher = self.write_dispatcher.clone(); + let request = request_factory.build_insert(normal_requests); + let writes = dispatcher.prepare(vec![request]).await?; + // 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); @@ -463,40 +465,6 @@ impl Inserter { )) } - async fn group_requests_by_peer( - &self, - requests: RegionInsertRequests, - ) -> Result> { - // group by region ids first to reduce repeatedly call `find_region_leader` - // TODO(discord9): determine if a addition clone is worth it - let mut requests_per_region: HashMap = HashMap::new(); - for req in requests.requests { - let region_id = RegionId::from_u64(req.region_id); - requests_per_region - .entry(region_id) - .or_default() - .requests - .push(req); - } - - let mut inserts: HashMap = HashMap::new(); - - for (region_id, reqs) in requests_per_region { - let peer = self - .partition_manager - .find_region_leader(region_id) - .await - .context(FindRegionLeaderSnafu)?; - inserts - .entry(peer) - .or_default() - .requests - .extend(reqs.requests); - } - - Ok(inserts) - } - /// Returns `Some(reason)` if the config or request hint disables automatic /// table creation. Exempt private system tables are handled by /// [`Self::is_auto_create_exempt_private_table`]. @@ -1219,6 +1187,11 @@ impl Inserter { } } + /// Returns the dispatcher shared with batching execution paths. + pub fn write_dispatcher(&self) -> &WriteDispatcherRef { + &self.write_dispatcher + } + pub fn node_manager(&self) -> &NodeManagerRef { &self.node_manager } diff --git a/src/operator/src/lib.rs b/src/operator/src/lib.rs index 5e723faeb5c..80b820ff2aa 100644 --- a/src/operator/src/lib.rs +++ b/src/operator/src/lib.rs @@ -28,3 +28,4 @@ pub mod table; #[cfg(test)] pub(crate) mod tests; pub mod utils; +pub mod write_dispatcher; diff --git a/src/operator/src/req_convert.rs b/src/operator/src/req_convert.rs index 456ff0b328d..82681e09968 100644 --- a/src/operator/src/req_convert.rs +++ b/src/operator/src/req_convert.rs @@ -15,3 +15,5 @@ pub(crate) mod common; pub mod delete; pub mod insert; + +pub(crate) use common::partitioner::TablePartitioner; diff --git a/src/operator/src/req_convert/common/partitioner.rs b/src/operator/src/req_convert/common/partitioner.rs index 3353e2fcda7..d83cf1d207b 100644 --- a/src/operator/src/req_convert/common/partitioner.rs +++ b/src/operator/src/req_convert/common/partitioner.rs @@ -12,72 +12,1083 @@ // See the License for the specific language governing permissions and // limitations under the License. -use api::v1::region::{DeleteRequest, InsertRequest}; -use api::v1::{PartitionExprVersion, Rows}; +use std::collections::HashMap; +use std::sync::Arc; + +use api::v1::region::region_request::Body; +use api::v1::region::{ + BulkInsertRequest, DeleteRequest, DeleteRequests, InsertRequest, InsertRequests, RegionRequest, + bulk_insert_request, +}; +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 datatypes::value::Value; +use mito_codec::row_converter::{CompositeValues, PrimaryKeyCodec, SparsePrimaryKeyCodec}; +use partition::PartitionRuleRef; use partition::manager::PartitionRuleManager; -use snafu::ResultExt; -use store_api::storage::RegionId; +use partition::partition::RegionMask; +use partition::splitter::RowSplitter; +use snafu::{OptionExt, ResultExt, ensure}; +use store_api::storage::{RegionId, RegionNumber, TableId, consts}; use table::metadata::TableInfo; -use crate::error::{Result, SplitDeleteSnafu, SplitInsertSnafu}; +use crate::error::{self, Result, SplitDeleteSnafu, SplitInsertSnafu}; +/// Loads table routing rules and converts rows to region write requests. pub struct Partitioner<'a> { partition_manager: &'a PartitionRuleManager, } +/// A single snapshot of a table's partition rule and expression versions. +pub(crate) struct TablePartitioner { + table_id: TableId, + rule: PartitionRuleRef, + versions: HashMap>, +} + impl<'a> Partitioner<'a> { + /// Creates a partitioner backed by the shared routing caches. pub fn new(partition_manager: &'a PartitionRuleManager) -> Self { Self { partition_manager } } + /// Routes insert rows using the current table partitions. pub async fn partition_insert_requests( &self, table_info: &TableInfo, rows: Rows, skip_wal: bool, ) -> Result> { - let table_id = table_info.table_id(); - let requests = self - .partition_manager - .split_rows(table_info, rows) + TablePartitioner::try_new(self.partition_manager, table_info) .await .context(SplitInsertSnafu)? - .into_iter() - .map( - |(region_number, (rows, partition_expr_version))| InsertRequest { - skip_wal, - region_id: RegionId::new(table_id, region_number).into(), - rows: Some(rows), - partition_expr_version: partition_expr_version - .map(|value| PartitionExprVersion { value }), - }, - ) - .collect(); - Ok(requests) + .partition_insert_requests(rows, skip_wal) } + /// Routes delete rows using the current table partitions. pub async fn partition_delete_requests( &self, table_info: &TableInfo, rows: Rows, ) -> Result> { - let table_id = table_info.table_id(); - - let requests = self - .partition_manager - .split_rows(table_info, rows) + TablePartitioner::try_new(self.partition_manager, table_info) .await .context(SplitDeleteSnafu)? + .partition_delete_requests(rows) + } +} + +impl TablePartitioner { + /// Reads a table's partition rule and expression versions into one routing snapshot. + pub async fn try_new( + partition_manager: &PartitionRuleManager, + table_info: &TableInfo, + ) -> partition::error::Result { + let (rule, versions) = partition_manager + .find_table_partition_rule(table_info) + .await?; + Ok(Self { + table_id: table_info.table_id(), + rule, + versions: versions + .into_iter() + .map(|(region, version)| { + (region, version.map(|value| PartitionExprVersion { value })) + }) + .collect(), + }) + } + + /// Rebuilds write requests using this partition snapshot after a route refresh. + /// + /// Takes ownership so row payloads can move into their new regions. Requests whose + /// region and partition version still match retain their original payloads. + /// The table metadata comes from the same refresh that created this partitioner. + pub(crate) fn rebuild_write_requests( + &self, + request: RegionRequest, + table_info: &TableInfo, + ) -> Result> { + let RegionRequest { header, body } = request; + let bodies = match body.context(error::UnexpectedSnafu { + violated: "Cannot rebuild a request without a body", + })? { + Body::Inserts(inserts) => self.rebuild_insert_requests(inserts)?, + Body::Deletes(deletes) => self.rebuild_delete_requests(deletes)?, + Body::BulkInsert(bulk) => self.rebuild_bulk_insert_request(bulk, table_info)?, + _ => { + return error::UnexpectedSnafu { + violated: "Only writes can be rebuilt", + } + .fail(); + } + }; + Ok(bodies .into_iter() - .map( - |(region_number, (rows, partition_expr_version))| DeleteRequest { - region_id: RegionId::new(table_id, region_number).into(), + .map(|body| RegionRequest { + header: header.clone(), + body: Some(body), + }) + .collect()) + } + + /// Moves insert rows into requests for this snapshot's destination regions. + pub fn partition_insert_requests( + &self, + rows: Rows, + skip_wal: bool, + ) -> Result> { + Ok(RowSplitter::new(self.rule.clone()) + .split(rows) + .context(SplitInsertSnafu)? + .into_iter() + .map(|(region, rows)| InsertRequest { + region_id: RegionId::new(self.table_id, region).as_u64(), + rows: Some(rows), + partition_expr_version: self.version(region), + skip_wal, + }) + .collect()) + } + + /// Moves delete rows into requests for this snapshot's destination regions. + pub fn partition_delete_requests(&self, rows: Rows) -> Result> { + Ok(RowSplitter::new(self.rule.clone()) + .split(rows) + .context(SplitDeleteSnafu)? + .into_iter() + .map(|(region, rows)| DeleteRequest { + region_id: RegionId::new(self.table_id, region).as_u64(), + rows: Some(rows), + partition_expr_version: self.version(region), + }) + .collect()) + } + + /// Splits an Arrow batch, reusing the original IPC when one region selects all rows. + #[cfg(test)] + pub fn partition_bulk_insert_requests( + &self, + batch: RecordBatch, + original_ipc: Option, + ) -> Result> { + let masks = self + .rule + .split_record_batch(&batch) + .context(SplitInsertSnafu)?; + self.build_bulk_insert_requests(batch, masks, original_ipc) + } + + fn build_bulk_insert_requests( + &self, + batch: RecordBatch, + masks: HashMap, + mut original_ipc: Option, + ) -> Result> { + let mut requests = Vec::with_capacity(masks.len()); + for (region, mask) in masks { + if mask.select_none() { + continue; + } + let original = if mask.select_all() { + original_ipc.take() + } else { + None + }; + let ipc = if let Some(ipc) = original { + ipc + } else { + let selected = if mask.select_all() { + batch.clone() + } else { + compute::filter_record_batch(&batch, mask.array()) + .context(error::ComputeArrowSnafu)? + }; + let (schema, data_header, payload) = flight::record_batch_to_ipc(selected) + .map_err(BoxedError::new) + .context(error::ExternalSnafu)?; + ArrowIpc { + schema, + data_header, + payload, + } + }; + requests.push(BulkInsertRequest { + region_id: RegionId::new(self.table_id, region).as_u64(), + partition_expr_version: self.version(region), + // The destination datanode validates against its current schema. + aligned_schema_version: None, + body: Some(bulk_insert_request::Body::ArrowIpc(ipc)), + }); + } + Ok(requests) + } + + fn rebuild_insert_requests(&self, inserts: InsertRequests) -> Result> { + let mut bodies = Vec::new(); + for insert in inserts.requests { + let region_id = RegionId::from_u64(insert.region_id); + if self.is_partition_expr_unchanged(region_id, insert.partition_expr_version) { + bodies.push(Body::Inserts(InsertRequests { + requests: vec![insert], + })); + continue; + } + let rows = insert.rows.context(error::UnexpectedSnafu { + violated: "Cannot rebuild an insert without rows", + })?; + for request in self.partition_insert_requests(rows, insert.skip_wal)? { + bodies.push(Body::Inserts(InsertRequests { + requests: vec![request], + })); + } + } + Ok(bodies) + } + + fn rebuild_delete_requests(&self, deletes: DeleteRequests) -> Result> { + let mut bodies = Vec::new(); + for delete in deletes.requests { + let region_id = RegionId::from_u64(delete.region_id); + if self.is_partition_expr_unchanged(region_id, delete.partition_expr_version) { + bodies.push(Body::Deletes(DeleteRequests { + requests: vec![delete], + })); + continue; + } + let rows = delete.rows.context(error::UnexpectedSnafu { + violated: "Cannot rebuild a delete without rows", + })?; + for request in self.partition_delete_requests(rows)? { + bodies.push(Body::Deletes(DeleteRequests { + requests: vec![request], + })); + } + } + Ok(bodies) + } + + fn rebuild_bulk_insert_request( + &self, + bulk: BulkInsertRequest, + table_info: &TableInfo, + ) -> Result> { + let region_id = RegionId::from_u64(bulk.region_id); + if self.is_partition_expr_unchanged(region_id, bulk.partition_expr_version) { + return Ok(vec![Body::BulkInsert(bulk)]); + } + let bulk_insert_request::Body::ArrowIpc(ipc) = + bulk.body.context(error::UnexpectedSnafu { + violated: "Cannot rebuild a bulk insert without Arrow IPC data", + })?; + let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema) + .and_then(|mut decoder| decoder.try_decode_record_batch(&ipc.data_header, &ipc.payload)) + .map_err(BoxedError::new) + .context(error::ExternalSnafu)?; + 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 + .build_bulk_insert_requests(batch, masks, Some(ipc))? + .into_iter() + .map(Body::BulkInsert) + .collect()) + } + + /// Restores omitted physical-table tags for routing without changing the write payload. + fn restore_partition_tags( + &self, + batch: &RecordBatch, + table_info: &TableInfo, + ) -> Result> { + let meta = &table_info.meta; + if !table_info.is_physical_table() || meta.column_ids.len() != meta.schema.num_columns() { + return Ok(None); + } + let missing = self + .rule + .partition_columns() + .iter() + .filter(|name| batch.column_by_name(name).is_none()) + .filter_map(|name| { + let index = meta.schema.column_index_by_name(name)?; + (meta.primary_key_indices.contains(&index) + && meta.schema.column_schemas()[index].data_type.is_string()) + .then_some((name, meta.column_ids[index])) + }) + .collect::>(); + if missing.is_empty() { + return Ok(None); + } + let Some(primary_key) = batch.column_by_name(consts::PRIMARY_KEY_COLUMN_NAME) else { + return Ok(None); + }; + let primary_key = + primary_key + .as_any() + .downcast_ref::() + .context(error::UnexpectedSnafu { + violated: "Sparse primary key column must be Binary", + })?; + let mut builders = missing + .iter() + .map(|_| StringBuilder::new()) + .collect::>(); + let codec = SparsePrimaryKeyCodec::schemaless(); + for row in 0..batch.num_rows() { + ensure!( + !primary_key.is_null(row), + error::UnexpectedSnafu { + violated: "Sparse primary key must not be null", + } + ); + let key = primary_key.value(row); + codec + .decode_ids(key) + .map_err(BoxedError::new) + .context(error::ExternalSnafu)?; + let CompositeValues::Sparse(values) = codec + .decode(key) + .map_err(BoxedError::new) + .context(error::ExternalSnafu)? + else { + return error::UnexpectedSnafu { + violated: "Sparse primary key codec returned dense values", + } + .fail(); + }; + for ((_, column_id), builder) in missing.iter().zip(&mut builders) { + let label = match values.get(column_id) { + Some(Value::String(value)) => Some(value.as_utf8()), + _ => None, + }; + 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 + /// sent to the new peer. A removed region or changed partition expression version + /// requires rebuilding the request using this snapshot's partition rule. + fn is_partition_expr_unchanged( + &self, + region_id: RegionId, + version: Option, + ) -> bool { + region_id.table_id() == self.table_id + && self.versions.get(®ion_id.region_number()) == Some(&version) + } + + fn version(&self, region: RegionNumber) -> Option { + self.versions.get(®ion).copied().flatten() + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use api::v1::region::{ + AlignedSchemaVersion, BulkInsertRequest, DeleteRequest, InsertRequest, RegionRequestHeader, + }; + use api::v1::value::ValueData; + use api::v1::{ArrowIpc, ColumnDataType, ColumnSchema, PartitionExprVersion, Row, Rows, Value}; + use common_error::ext::{ErrorExt, RetryHint}; + use common_error::status_code::StatusCode; + use common_grpc::flight; + use common_meta::cache::{self, TableInfoCache}; + use common_meta::key::TableMetadataManager; + use common_meta::key::table_route::TableRouteValue; + use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::kv_backend::test_util::MockKvBackendBuilder; + use common_meta::rpc::router::{Region, RegionRoute}; + use common_meta::rpc::store::RangeResponse; + 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::{PrimaryKeyCodec, SparsePrimaryKeyCodec, sparse}; + 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::tests::{ + create_partition_rule_manager, new_test_table_info, 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 = new_test_table_info(1, "physical_tags", [1, 2].into_iter()); + table_info.meta.schema = Arc::new(schema::Schema::new(vec![ + schema::ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + schema::ColumnSchema::new("zone", ConcreteDataType::string_datatype(), true), + ])); + table_info.meta.primary_key_indices = vec![0, 1]; + table_info.meta.partition_key_indices = vec![1]; + table_info.meta.column_ids = vec![10, 20]; + table_info.meta.options.extra_options.insert( + metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY.to_string(), + String::new(), + ); + table_info + } + + fn build_sparse_batch() -> RecordBatch { + let codec = SparsePrimaryKeyCodec::schemaless(); + let keys = [Some("east"), Some("west"), Some(""), None, None] + .into_iter() + .enumerate() + .map(|(index, zone)| { + let mut key = Vec::new(); + let values = [ + ( + sparse::RESERVED_COLUMN_ID_TABLE_ID, + ValueRef::UInt32(42 + index as u32), + ), + ( + sparse::RESERVED_COLUMN_ID_TSID, + ValueRef::UInt64(index as u64), + ), + (10, ValueRef::String("node")), + ]; + codec.encode_value_refs(&values, &mut key).unwrap(); + if let Some(zone) = zone { + codec + .encode_raw_tag_value([(20, zone.as_bytes())].into_iter(), &mut key) + .unwrap(); + } + if index == 4 { + key.extend_from_slice(&20_u32.to_be_bytes()); + key.push(0); + } + key + }) + .collect::>(); + RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("host", DataType::Utf8, false), + Field::new(consts::PRIMARY_KEY_COLUMN_NAME, DataType::Binary, false), + ])), + vec![ + Arc::new(StringArray::from(vec!["node"; keys.len()])), + Arc::new(BinaryArray::from_iter_values(keys.iter())), + ], + ) + .unwrap() + } + + #[tokio::test] + async fn test_partition_tags_use_physical_table_metadata() { + let backend = Arc::new(MemoryKvBackend::new()); + let table_route_cache = test_new_table_route_cache(backend.clone()); + let partition_info_cache = test_new_partition_info_cache(table_route_cache.clone()); + let manager = + PartitionRuleManager::new(backend.clone(), table_route_cache, partition_info_cache); + let table_info = build_sparse_table_info(); + let routes = [ + expr::col("zone").lt(DataValue::String("m".into())), + expr::col("zone").gt_eq(DataValue::String("m".into())), + ] + .into_iter() + .enumerate() + .map(|(index, expr)| RegionRoute { + region: Region { + id: RegionId::new(1, index as u32 + 1), + partition_expr: expr.as_json_str().unwrap(), + ..Default::default() + }, + ..Default::default() + }) + .collect(); + TableMetadataManager::new(backend) + .create_table_metadata( + table_info.clone(), + TableRouteValue::physical(routes), + Default::default(), + ) + .await + .unwrap(); + let batch = build_sparse_batch(); + let partitioner = TablePartitioner::try_new(&manager, &table_info) + .await + .unwrap(); + let restored = partitioner + .restore_partition_tags(&batch, &table_info) + .unwrap() + .unwrap(); + let zones = restored + .column_by_name("zone") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + zones.iter().collect::>(), + vec![Some("east"), Some("west"), Some(""), None, None] + ); + + for (physical, tag, valid_ids) in [ + (false, true, true), + (true, false, true), + (true, true, false), + ] { + let mut info = table_info.clone(); + if !physical { + info.meta + .options + .extra_options + .remove(metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY); + } + if !tag { + info.meta.primary_key_indices = vec![0]; + } + if !valid_ids { + info.meta.column_ids.clear(); + } + let partitioner = TablePartitioner::try_new(&manager, &info).await.unwrap(); + assert!( + partitioner + .restore_partition_tags(&batch, &info) + .unwrap() + .is_none() + ); + assert!( + partitioner + .partition_bulk_insert_requests(batch.clone(), None) + .is_err() + ); + } + } + + #[test] + fn test_rebuild_sparse_bulk_with_changed_partition_column() { + let partitioner = build_sparse_partitioner(); + let table_info = build_sparse_table_info(); + let batch = build_sparse_batch(); + let restored = partitioner + .restore_partition_tags(&batch, &table_info) + .unwrap() + .unwrap(); + let zones = restored + .column_by_name("zone") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + zones.iter().collect::>(), + vec![Some("east"), Some("west"), Some(""), None, None] + ); + let (schema, data_header, payload) = flight::record_batch_to_ipc(batch.clone()).unwrap(); + let rebuilt = partitioner + .rebuild_write_requests( + build_request(Body::BulkInsert(BulkInsertRequest { + region_id: RegionId::new(1, 1).as_u64(), + partition_expr_version: Some(PartitionExprVersion { value: 1 }), + 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_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() { + let table_info = new_test_table_info(1, "table_1", [1].into_iter()); + for (backend_unavailable, expected_status, expected_hint) in [ + (false, StatusCode::TableNotFound, RetryHint::NonRetryable), + (true, StatusCode::Internal, RetryHint::Retryable), + ] { + let backend = Arc::new( + MockKvBackendBuilder::default() + .range_fn(Arc::new(move |_| { + if backend_unavailable { + common_meta::error::NoLeaderSnafu { table_id: 1u32 }.fail() + } else { + Ok(RangeResponse { + kvs: Vec::new(), + more: false, + }) + } + })) + .build() + .unwrap(), + ); + let routes = Arc::new(cache::new_table_route_cache( + "partitioner_test_routes".to_string(), + Cache::new(16), + backend.clone(), + )); + let partitions = Arc::new(partition::cache::new_partition_info_cache( + "partitioner_test_partitions".to_string(), + Cache::new(16), + routes.clone(), + )); + let manager = PartitionRuleManager::new(backend, routes, partitions); + let partitioner = Partitioner::new(&manager); + let insert_error = partitioner + .partition_insert_requests(&table_info, Rows::default(), false) + .await + .unwrap_err(); + let delete_error = partitioner + .partition_delete_requests(&table_info, Rows::default()) + .await + .unwrap_err(); + + assert!(matches!(insert_error, Error::SplitInsert { .. })); + assert!(matches!(delete_error, Error::SplitDelete { .. })); + for error in [insert_error, delete_error] { + assert_eq!(error.status_code(), expected_status); + assert_eq!(error.retry_hint(), expected_hint); + } + } + } + + async fn create_fixture() -> (PartitionRuleManagerRef, TableInfoCache) { + let backend = Arc::new(MemoryKvBackend::new()); + let manager = create_partition_rule_manager(backend.clone()).await; + TableMetadataManager::new(backend.clone()) + .create_table_metadata( + new_test_table_info(42, "logical_reroute", [1, 2, 3].into_iter()), + TableRouteValue::logical(1), + Default::default(), + ) + .await + .unwrap(); + let cache = cache::new_table_info_cache( + "reroute_test_table_info".to_string(), + Cache::new(16), + backend, + ); + (manager, cache) + } + + fn build_request(body: Body) -> RegionRequest { + RegionRequest { + header: Some(RegionRequestHeader { + dbname: "greptime.public".to_string(), + ..Default::default() + }), + body: Some(body), + } + } + + fn build_rows(values: &[i32]) -> Rows { + Rows { + schema: vec![ColumnSchema { + column_name: "a".to_string(), + datatype: ColumnDataType::Int32 as i32, + ..Default::default() + }], + rows: values + .iter() + .map(|value| Row { + values: vec![Value { + value_data: Some(ValueData::I32Value(*value)), + }], + }) + .collect(), + } + } + + fn build_row_request( + rows: Rows, + insert: bool, + region: u32, + version: Option, + ) -> RegionRequest { + let region_id = RegionId::new(1, region).as_u64(); + let partition_expr_version = version.map(|value| PartitionExprVersion { value }); + let body = if insert { + Body::Inserts(InsertRequests { + requests: vec![InsertRequest { + region_id, + partition_expr_version, rows: Some(rows), - partition_expr_version: partition_expr_version - .map(|value| PartitionExprVersion { value }), - }, - ) - .collect(); - Ok(requests) + skip_wal: true, + }], + }) + } else { + Body::Deletes(DeleteRequests { + requests: vec![DeleteRequest { + region_id, + partition_expr_version, + rows: Some(rows), + }], + }) + }; + build_request(body) + } + + fn row_request_parts(request: &RegionRequest) -> (u64, Option, &Rows) { + match request.body.as_ref().unwrap() { + Body::Inserts(inserts) => { + assert_eq!(inserts.requests.len(), 1); + let insert = &inserts.requests[0]; + assert!(insert.skip_wal); + ( + insert.region_id, + insert.partition_expr_version.map(|v| v.value), + insert.rows.as_ref().unwrap(), + ) + } + Body::Deletes(deletes) => { + assert_eq!(deletes.requests.len(), 1); + let delete = &deletes.requests[0]; + ( + delete.region_id, + delete.partition_expr_version.map(|v| v.value), + delete.rows.as_ref().unwrap(), + ) + } + _ => panic!("expected row write"), + } + } + + #[tokio::test] + async fn test_rebuild_rows_updates_partitions_without_losing_rows() { + let (manager, cache) = create_fixture().await; + let table_info = cache.get(1).await.unwrap().unwrap(); + let (_, versions) = manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + for insert in [true, false] { + let request = build_row_request(build_rows(&[1, 20, 70]), insert, 99, Some(0)); + let header = request.header.clone(); + let rerouted = rebuild_with_cached_table(&manager, &cache, request) + .await + .unwrap(); + assert_eq!(rerouted.len(), 3); + let mut actual = Vec::new(); + for request in rerouted { + assert_eq!(request.header, header); + let (id, version, rows) = row_request_parts(&request); + let region_id = RegionId::from_u64(id); + assert_eq!(region_id.table_id(), 1); + assert_eq!(version, versions[®ion_id.region_number()]); + assert_eq!(rows.schema, build_rows(&[]).schema); + assert_eq!(rows.rows.len(), 1); + actual.push((region_id.region_number(), rows.rows[0].clone())); + } + actual.sort_by_key(|(region, _)| *region); + assert_eq!( + actual, + vec![ + (1, build_rows(&[70]).rows.remove(0)), + (2, build_rows(&[20]).rows.remove(0)), + (3, build_rows(&[1]).rows.remove(0)) + ] + ); + } + } + + #[tokio::test] + async fn test_unchanged_row_routes_preserve_allocations() { + let (manager, cache) = create_fixture().await; + let table_info = cache.get(1).await.unwrap().unwrap(); + let (_, versions) = manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + for insert in [true, false] { + let request = build_row_request(build_rows(&[1]), insert, 3, versions[&3]); + let original_rows = row_request_parts(&request).2; + let row_ptr = original_rows.rows.as_ptr(); + let value_ptr = original_rows.rows[0].values.as_ptr(); + let rerouted = rebuild_with_cached_table(&manager, &cache, request) + .await + .unwrap(); + assert_eq!(rerouted.len(), 1); + let rows = row_request_parts(&rerouted[0]).2; + assert_eq!(rows.rows.as_ptr(), row_ptr); + assert_eq!(rows.rows[0].values.as_ptr(), value_ptr); + assert_eq!(rows, &build_rows(&[1])); + } + } + + fn build_bulk_request(region: u32, version: Option, values: Vec) -> RegionRequest { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(values))], + ) + .unwrap(); + let (schema, data_header, payload) = flight::record_batch_to_ipc(batch).unwrap(); + build_request(Body::BulkInsert(BulkInsertRequest { + region_id: RegionId::new(1, region).as_u64(), + partition_expr_version: version.map(|value| PartitionExprVersion { value }), + aligned_schema_version: Some(AlignedSchemaVersion { + schema_version: 123, + }), + body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { + schema, + data_header, + payload, + })), + })) + } + + #[tokio::test] + async fn test_rebuild_bulk_reencodes_only_changed_partitions() { + let (manager, cache) = create_fixture().await; + let table_info = cache.get(1).await.unwrap().unwrap(); + let (_, versions) = manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + let request = build_bulk_request(99, Some(0), vec![1, 20, 70]); + let header = request.header.clone(); + let rerouted = rebuild_with_cached_table(&manager, &cache, request) + .await + .unwrap(); + assert_eq!(rerouted.len(), 3); + let mut actual = Vec::new(); + for request in rerouted { + assert_eq!(request.header, header); + let Some(Body::BulkInsert(bulk)) = request.body else { + panic!("expected bulk insert") + }; + let region = RegionId::from_u64(bulk.region_id); + assert_eq!(region.table_id(), 1); + assert_eq!(bulk.aligned_schema_version, None); + assert_eq!( + bulk.partition_expr_version.map(|v| v.value), + versions[®ion.region_number()] + ); + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = bulk.body else { + panic!("expected Arrow IPC") + }; + let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema) + .unwrap() + .try_decode_record_batch(&ipc.data_header, &ipc.payload) + .unwrap(); + assert_eq!(batch.num_rows(), 1); + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + actual.push((region.region_number(), values.value(0))); + } + actual.sort_unstable(); + assert_eq!(actual, [(1, 70), (2, 20), (3, 1)]); + + // A leader-only change retains the original IPC buffers and schema hint. + let request = build_bulk_request(3, versions[&3], vec![1]); + let Some(Body::BulkInsert(bulk)) = request.body.as_ref() else { + unreachable!() + }; + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &bulk.body else { + unreachable!() + }; + let pointers = ( + ipc.schema.as_ptr(), + ipc.data_header.as_ptr(), + ipc.payload.as_ptr(), + ); + let rerouted = rebuild_with_cached_table(&manager, &cache, request) + .await + .unwrap(); + assert_eq!(rerouted.len(), 1); + let Some(Body::BulkInsert(bulk)) = rerouted[0].body.as_ref() else { + unreachable!() + }; + assert_eq!( + bulk.aligned_schema_version, + Some(AlignedSchemaVersion { + schema_version: 123 + }) + ); + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &bulk.body else { + unreachable!() + }; + assert_eq!( + ( + ipc.schema.as_ptr(), + ipc.data_header.as_ptr(), + ipc.payload.as_ptr() + ), + pointers + ); + } + + #[tokio::test] + async fn test_rebuild_logical_table_preserves_logical_region_ids() { + let (manager, cache) = create_fixture().await; + for insert in [true, false] { + let mut request = build_row_request(build_rows(&[1, 20, 70]), insert, 99, Some(0)); + match request.body.as_mut().unwrap() { + Body::Inserts(inserts) => { + inserts.requests[0].region_id = RegionId::new(42, 99).as_u64() + } + Body::Deletes(deletes) => { + deletes.requests[0].region_id = RegionId::new(42, 99).as_u64() + } + _ => unreachable!(), + } + let rerouted = rebuild_with_cached_table(&manager, &cache, request) + .await + .unwrap(); + assert_eq!(rerouted.len(), 3); + let mut regions = rerouted + .iter() + .map(|request| { + let region = RegionId::from_u64(row_request_parts(request).0); + assert_eq!(region.table_id(), 42); + region.region_number() + }) + .collect::>(); + regions.sort_unstable(); + assert_eq!(regions, [1, 2, 3]); + } + } + + async fn rebuild_with_cached_table( + manager: &PartitionRuleManager, + cache: &TableInfoCache, + request: RegionRequest, + ) -> Result> { + let region_id = match request.body.as_ref().unwrap() { + Body::Inserts(inserts) => inserts.requests[0].region_id, + Body::Deletes(deletes) => deletes.requests[0].region_id, + Body::BulkInsert(bulk) => bulk.region_id, + _ => unreachable!(), + }; + let table_id = RegionId::from_u64(region_id).table_id(); + let table_info = cache.get(table_id).await.unwrap().unwrap(); + let partitioner = TablePartitioner::try_new(manager, &table_info) + .await + .context(error::InvalidPartitionSnafu)?; + partitioner.rebuild_write_requests(request, &table_info) } } diff --git a/src/operator/src/tests.rs b/src/operator/src/tests.rs index 9263fa4cca7..f1cd4f11b07 100644 --- a/src/operator/src/tests.rs +++ b/src/operator/src/tests.rs @@ -16,4 +16,7 @@ mod kv_backend; mod partition_manager; pub(crate) use kv_backend::prepare_mocked_backend; -pub(crate) use partition_manager::{create_partition_rule_manager, new_test_table_info}; +pub(crate) use partition_manager::{ + create_partition_rule_manager, new_test_table_info, test_new_partition_info_cache, + test_new_table_route_cache, +}; diff --git a/src/operator/src/tests/partition_manager.rs b/src/operator/src/tests/partition_manager.rs index ecf87aedac9..7fdcb96d93f 100644 --- a/src/operator/src/tests/partition_manager.rs +++ b/src/operator/src/tests/partition_manager.rs @@ -107,7 +107,7 @@ fn new_test_region_wal_options(regions: Vec) -> RegionWalOptions { HashMap::default() } -fn test_new_table_route_cache(kv_backend: KvBackendRef) -> TableRouteCacheRef { +pub(crate) fn test_new_table_route_cache(kv_backend: KvBackendRef) -> TableRouteCacheRef { let cache = CacheBuilder::new(128).build(); Arc::new(new_table_route_cache( "table_route_cache".to_string(), @@ -116,7 +116,9 @@ fn test_new_table_route_cache(kv_backend: KvBackendRef) -> TableRouteCacheRef { )) } -fn test_new_partition_info_cache(table_route_cache: TableRouteCacheRef) -> PartitionInfoCacheRef { +pub(crate) fn test_new_partition_info_cache( + table_route_cache: TableRouteCacheRef, +) -> PartitionInfoCacheRef { let cache = CacheBuilder::new(128).build(); Arc::new(new_partition_info_cache( "partition_info_cache".to_string(), diff --git a/src/operator/src/write_dispatcher.rs b/src/operator/src/write_dispatcher.rs new file mode 100644 index 00000000000..2c65a79b3ce --- /dev/null +++ b/src/operator/src/write_dispatcher.rs @@ -0,0 +1,615 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +mod options; +#[cfg(test)] +mod tests; + +use std::collections::hash_map::Entry; +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use api::region::RegionResponse; +use api::v1::region::region_request::Body; +use api::v1::region::{DeleteRequests, InsertRequests, RegionRequest}; +use backon::{BackoffBuilder, ExponentialBuilder}; +use common_error::ext::{BoxedError, ErrorExt, RetryHint}; +use common_meta::cache::TableInfoCacheRef; +use common_meta::cache_invalidator::{CacheInvalidatorRef, Context}; +use common_meta::error::{Error, ExternalSnafu, Result, TableInfoNotFoundSnafu, UnexpectedSnafu}; +use common_meta::instruction::CacheIdent; +use common_meta::node_manager::NodeManagerRef; +use common_meta::peer::{Peer, PeerRef}; +use common_meta::rpc::router; +use futures::future; +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}; +use crate::req_convert::TablePartitioner; + +/// Dispatches already partitioned writes, retaining only failed payloads for retries. +pub struct WriteDispatcher { + node_manager: NodeManagerRef, + partition_manager: PartitionRuleManagerRef, + retry: Option, +} + +/// Shared reference to a [`WriteDispatcher`]. +pub type WriteDispatcherRef = Arc; + +/// The result of one write RPC attempt. +pub(crate) enum WriteAttempt { + Success(RegionResponse), + Failed { + error: Error, + /// Retained when retries are enabled and the original client error allows retry. + /// The remaining retry budget is checked separately. + request: Option, + }, +} + +impl WriteDispatcher { + /// Creates an owned-request dispatcher without retries, including for standalone. + pub fn new(node_manager: NodeManagerRef, partition_manager: PartitionRuleManagerRef) -> Self { + Self { + node_manager, + partition_manager, + retry: None, + } + } + + /// Enables shared remote sends and retries using frontend configuration and caches. + pub fn with_retry( + mut self, + options: WriteRetryOptions, + table_info_cache: TableInfoCacheRef, + cache_invalidator: CacheInvalidatorRef, + ) -> Self { + if options.max_retries > 0 { + self.retry = Some(WriteRetry { + node_manager: self.node_manager.clone(), + partition_manager: self.partition_manager.clone(), + options, + table_info_cache, + cache_invalidator, + }); + } else { + self.retry = None; + } + self + } + + /// Writes planned requests without repeating admission, default values, or Flow side effects. + pub async fn dispatch(&self, requests: Vec) -> Result { + self.dispatch_inner(requests, None).await + } + + /// Records each datanode send attempt, including retries, in the supplied histogram. + pub async fn dispatch_with_metrics( + &self, + requests: Vec, + histogram: Histogram, + ) -> Result { + self.dispatch_inner(requests, Some(histogram)).await + } + + /// Resolves routes before writes are detached from the caller. + pub(crate) async fn prepare( + &self, + requests: Vec, + ) -> error::Result> { + let requests = requests + .into_iter() + .map(|mut request| { + let header = request.header.take(); + (header, split_request(request)) + }) + .collect::>(); + let regions = requests + .iter() + .flat_map(|(_, requests)| requests) + .map(write_region_id) + .collect::>>() + .context(RequestRegionSnafu)?; + let peers = self + .resolve_peers(®ions) + .await + .context(FindRegionLeaderSnafu)?; + let mut writes = Vec::new(); + // Preserve each original envelope and peer batching when retries are disabled. + for (header, requests) in requests { + let mut by_peer: HashMap = HashMap::new(); + for request in requests { + let peer = get_peer( + &peers, + write_region_id(&request).context(RequestRegionSnafu)?, + ) + .context(RequestRegionSnafu)?; + match by_peer.entry(peer) { + Entry::Vacant(entry) => { + entry.insert(request); + } + Entry::Occupied(mut entry) => match (&mut entry.get_mut().body, request.body) { + (Some(Body::Inserts(existing)), Some(Body::Inserts(next))) => { + existing.requests.extend(next.requests) + } + (Some(Body::Deletes(existing)), Some(Body::Deletes(next))) => { + existing.requests.extend(next.requests) + } + _ => { + return UnexpectedSnafu { + err_msg: "Cannot combine different write kinds", + } + .fail() + .context(RequestRegionSnafu); + } + }, + } + } + for (peer, mut request) in by_peer { + request.header = header.clone(); + writes.push((peer, request)); + } + } + Ok(writes) + } + + /// Sends requests using routes already resolved by the caller. + /// Cancellation stops future retries while allowing in-flight requests to complete. + pub(crate) async fn execute( + &self, + writes: Vec<(PeerRef, RegionRequest)>, + histogram: Option, + cancellation: Option<&CancellationToken>, + ) -> Result { + let histogram = histogram.as_ref(); + let writes = if self.retry.is_some() { + writes + .into_iter() + .flat_map(|(peer, request)| { + split_request(request) + .into_iter() + .map(move |request| (Arc::clone(&peer), request)) + }) + .collect() + } else { + writes + }; + let attempts = future::join_all( + writes + .into_iter() + .map(|(peer, request)| self.send(peer, request, histogram)), + ) + .await; + self.finish(attempts, histogram, cancellation).await + } + + /// Sends the first attempt immediately so bulk encoding can overlap other sends. + pub(crate) async fn send( + &self, + peer: PeerRef, + request: RegionRequest, + histogram: Option<&Histogram>, + ) -> WriteAttempt { + if let Some(retry) = &self.retry { + return retry + .send( + RoutedWrite { + peer, + request: Arc::new(request), + }, + histogram, + ) + .await; + } + let _timer = histogram.map(Histogram::start_timer); + let datanode = self.node_manager.datanode(&peer).await; + match datanode.handle(request).await { + Ok(response) => WriteAttempt::Success(response), + Err(error) => WriteAttempt::Failed { + error, + request: None, + }, + } + } + + /// Collects first attempts and retries their failures together, refreshing routes once per round. + pub(crate) async fn finish( + &self, + attempts: Vec, + histogram: Option<&Histogram>, + cancellation: Option<&CancellationToken>, + ) -> Result { + if let Some(retry) = &self.retry { + if retry.should_retry(&attempts, 0) { + return retry.retry(attempts, histogram, cancellation).await; + } + let mut response = RegionResponse::new(0); + let failures = merge_attempts(attempts, &mut response)?; + return match failures.into_iter().next() { + Some((_, error)) => Err(error), + None => Ok(response), + }; + } + let mut response = RegionResponse::new(0); + for attempt in attempts { + match attempt { + WriteAttempt::Success(result) => merge_response(&mut response, result), + WriteAttempt::Failed { error, .. } => return Err(error), + } + } + Ok(response) + } + + /// Resolves only the requested peers, loading each table route once. + pub(crate) async fn resolve_peers( + &self, + regions: &[RegionId], + ) -> partition::error::Result> { + resolve_peers(&self.partition_manager, regions).await + } + + async fn dispatch_inner( + &self, + requests: Vec, + histogram: Option, + ) -> Result { + let writes = self + .prepare(requests) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + self.execute(writes, histogram, None).await + } +} + +/// Sends shared requests and coordinates retries with refreshed routes. +struct WriteRetry { + node_manager: NodeManagerRef, + partition_manager: PartitionRuleManagerRef, + options: WriteRetryOptions, + table_info_cache: TableInfoCacheRef, + cache_invalidator: CacheInvalidatorRef, +} + +/// A single-region write request paired with its destination peer. +pub(crate) struct RoutedWrite { + peer: PeerRef, + request: Arc, +} + +impl WriteRetry { + async fn send(&self, write: RoutedWrite, histogram: Option<&Histogram>) -> WriteAttempt { + let _timer = histogram.map(Histogram::start_timer); + 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, request } + } + } + } + + fn should_retry(&self, attempts: &[WriteAttempt], retries: usize) -> bool { + if retries >= self.options.max_retries { + return false; + } + let mut retryable = false; + for attempt in attempts { + match attempt { + WriteAttempt::Success(_) => {} + WriteAttempt::Failed { + request: Some(_), .. + } => retryable = true, + WriteAttempt::Failed { request: None, .. } => return false, + } + } + retryable + } + + async fn retry( + &self, + mut attempts: Vec, + histogram: Option<&Histogram>, + cancellation: Option<&CancellationToken>, + ) -> Result { + let mut response = RegionResponse::new(0); + let mut backoff = ExponentialBuilder::default() + .with_min_delay(self.options.initial_backoff) + .with_max_delay(self.options.max_backoff) + .with_factor(self.options.backoff_multiplier) + .with_max_times(self.options.max_retries) + .with_jitter() + .build(); + let mut retries = 0; + while self.should_retry(&attempts, retries) { + let Some(delay) = backoff.next() else { + break; + }; + 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)?; + let mut first_error = None; + let mut refresh = Vec::new(); + let mut pending = Vec::new(); + for (write, error) in failures { + if error.retry_hint() == RetryHint::RetryAfterRouteRefresh { + refresh.push(write); + } else { + pending.push(write); + } + 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; + retries += 1; + } + match merge_attempts(attempts, &mut response)?.into_iter().next() { + Some((_, error)) => Err(error), + None => Ok(response), + } + } + + async fn reroute_failed_writes(&self, writes: Vec) -> Result> { + let regions = writes + .iter() + .map(|write| write_region_id(&write.request)) + .collect::>>()?; + self.invalidate_routes(®ions).await?; + let mut partitioners = HashMap::new(); + for table_id in regions + .iter() + .map(|region| region.table_id()) + .collect::>() + { + let table_info = self + .table_info_cache + .get(table_id) + .await? + .with_context(|| TableInfoNotFoundSnafu { + table: table_id.to_string(), + })?; + let partitioner = TablePartitioner::try_new(&self.partition_manager, &table_info) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + partitioners.insert(table_id, (table_info, partitioner)); + } + let mut rebuilt_requests = Vec::new(); + for write in writes { + let request = Arc::try_unwrap(write.request).map_err(|_| { + UnexpectedSnafu { + err_msg: "Write request is still shared after its RPC completed", + } + .build() + })?; + let table_id = write_region_id(&request)?.table_id(); + let (table_info, partitioner) = + partitioners.get(&table_id).context(UnexpectedSnafu { + err_msg: "Refreshed write has no partition snapshot", + })?; + let requests = partitioner + .rebuild_write_requests(request, table_info) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + rebuilt_requests.extend(requests); + } + self.prepare(rebuilt_requests).await + } + + async fn prepare(&self, requests: Vec) -> Result> { + let regions = requests + .iter() + .map(write_region_id) + .collect::>>()?; + let peers = resolve_peers(&self.partition_manager, ®ions) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + requests + .into_iter() + .zip(regions) + .map(|(request, region)| { + Ok(RoutedWrite { + peer: get_peer(&peers, region)?, + request: Arc::new(request), + }) + }) + .collect() + } + + async fn invalidate_routes(&self, regions: &[RegionId]) -> Result<()> { + let mut table_ids = regions + .iter() + .map(|region| region.table_id()) + .collect::>(); + let mut physical_ids = HashSet::new(); + for table_id in &table_ids { + let (physical_id, _) = self + .partition_manager + .find_physical_table_route_with_id(*table_id) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + physical_ids.insert(physical_id); + } + table_ids.extend(physical_ids); + let caches = table_ids + .into_iter() + .map(CacheIdent::TableId) + .collect::>(); + self.cache_invalidator + .invalidate(&Context::default(), &caches) + .await + } +} + +async fn resolve_peers( + partition_manager: &PartitionRuleManager, + regions: &[RegionId], +) -> partition::error::Result> { + let table_ids = regions + .iter() + .map(|region| region.table_id()) + .collect::>(); + let mut physical_tables = HashMap::new(); + let mut routes = HashMap::new(); + for table_id in table_ids { + let (physical_id, route) = partition_manager + .find_physical_table_route_with_id(table_id) + .await?; + physical_tables.insert(table_id, physical_id); + routes.entry(physical_id).or_insert(route); + } + let mut shared_peers: HashMap<&Peer, PeerRef> = HashMap::new(); + let mut peers = HashMap::with_capacity(regions.len()); + for region in regions { + if peers.contains_key(region) { + continue; + } + let peer = physical_tables + .get(®ion.table_id()) + .and_then(|physical_id| routes.get(physical_id)) + .and_then(|route| { + router::find_region_leader_ref(&route.region_routes, region.region_number()) + }) + .context(partition::error::FindLeaderSnafu { + region_id: *region, + table_id: region.table_id(), + })?; + let peer = shared_peers + .entry(peer) + .or_insert_with(|| Arc::new(peer.clone())); + peers.insert(*region, Arc::clone(peer)); + } + Ok(peers) +} + +fn get_peer(peers: &HashMap, region: RegionId) -> Result { + peers.get(®ion).cloned().context(UnexpectedSnafu { + err_msg: "Planned write has no resolved peer", + }) +} + +fn merge_attempts( + attempts: Vec, + response: &mut RegionResponse, +) -> Result> { + let mut failures = Vec::new(); + let mut terminal_error = None; + for attempt in attempts { + match attempt { + WriteAttempt::Success(result) => merge_response(response, result), + WriteAttempt::Failed { + error, + request: Some(write), + } => failures.push((write, error)), + WriteAttempt::Failed { + error, + request: None, + } => { + terminal_error.get_or_insert(error); + } + } + } + // A terminal error takes precedence over retained retryable failures. + match terminal_error { + Some(error) => Err(error), + None => Ok(failures), + } +} + +fn merge_response(response: &mut RegionResponse, result: RegionResponse) { + response.affected_rows += result.affected_rows; + response.extensions.extend(result.extensions); + response.metadata = result.metadata; +} + +fn split_request(request: RegionRequest) -> Vec { + let RegionRequest { header, body } = request; + let bodies: Vec = match body { + Some(Body::Inserts(inserts)) => inserts + .requests + .into_iter() + .map(|request| { + Body::Inserts(InsertRequests { + requests: vec![request], + }) + }) + .collect(), + Some(Body::Deletes(deletes)) => deletes + .requests + .into_iter() + .map(|request| { + Body::Deletes(DeleteRequests { + requests: vec![request], + }) + }) + .collect(), + body => return vec![RegionRequest { header, body }], + }; + bodies + .into_iter() + .map(|body| RegionRequest { + header: header.clone(), + body: Some(body), + }) + .collect() +} + +fn write_region_id(request: &RegionRequest) -> Result { + let region_id = match &request.body { + Some(Body::Inserts(inserts)) => match inserts.requests.as_slice() { + [request] => Some(request.region_id), + _ => None, + }, + Some(Body::Deletes(deletes)) => match deletes.requests.as_slice() { + [request] => Some(request.region_id), + _ => None, + }, + Some(Body::BulkInsert(bulk)) => Some(bulk.region_id), + _ => None, + }; + region_id.map(RegionId::from_u64).ok_or_else(|| { + UnexpectedSnafu { + err_msg: "Expected a single-region write request", + } + .build() + }) +} diff --git a/src/operator/src/write_dispatcher/options.rs b/src/operator/src/write_dispatcher/options.rs new file mode 100644 index 00000000000..20c8ba535f3 --- /dev/null +++ b/src/operator/src/write_dispatcher/options.rs @@ -0,0 +1,106 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::time::Duration; + +use serde::{Deserialize, Serialize}; + +/// Exponential backoff for retryable distributed writes. +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] +#[serde(default)] +pub struct WriteRetryOptions { + /// Maximum additional attempts after the original write. Zero disables retries. + pub max_retries: usize, + /// Initial backoff before jitter is applied. + #[serde(with = "humantime_serde")] + pub initial_backoff: Duration, + /// Maximum backoff before jitter is applied. + #[serde(with = "humantime_serde")] + pub max_backoff: Duration, + /// Multiplier applied to the backoff after each retry. + pub backoff_multiplier: f32, +} + +impl Default for WriteRetryOptions { + fn default() -> Self { + Self { + max_retries: 3, + initial_backoff: Duration::from_millis(1500), + max_backoff: Duration::from_secs(10), + backoff_multiplier: 2.0, + } + } +} + +impl WriteRetryOptions { + /// Checks that the configured backoff is positive, bounded, and nondecreasing. + pub fn validate(&self) -> Result<(), &'static str> { + if self.initial_backoff.is_zero() { + return Err("write_retry.initial_backoff must be positive"); + } + if self.max_backoff < self.initial_backoff { + return Err("write_retry.max_backoff must be at least initial_backoff"); + } + if !self.backoff_multiplier.is_finite() || self.backoff_multiplier < 1.0 { + return Err("write_retry.backoff_multiplier must be finite and at least 1"); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_validate_write_retry_options() { + for options in [ + WriteRetryOptions::default(), + WriteRetryOptions { + max_retries: 0, + ..Default::default() + }, + WriteRetryOptions { + backoff_multiplier: 1.0, + ..Default::default() + }, + ] { + options.validate().unwrap(); + } + for options in [ + WriteRetryOptions { + initial_backoff: Duration::ZERO, + ..Default::default() + }, + WriteRetryOptions { + max_backoff: Duration::from_secs(1), + ..Default::default() + }, + WriteRetryOptions { + backoff_multiplier: 0.5, + ..Default::default() + }, + WriteRetryOptions { + backoff_multiplier: f32::NAN, + ..Default::default() + }, + WriteRetryOptions { + backoff_multiplier: f32::INFINITY, + ..Default::default() + }, + ] { + assert!(options.validate().is_err(), "{options:?}"); + } + } +} diff --git a/src/operator/src/write_dispatcher/tests.rs b/src/operator/src/write_dispatcher/tests.rs new file mode 100644 index 00000000000..694cebf5a1c --- /dev/null +++ b/src/operator/src/write_dispatcher/tests.rs @@ -0,0 +1,1606 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::{HashMap, VecDeque}; +use std::future::Future; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use std::task::Poll; +use std::time::{Duration, Instant}; + +use api::v1::region::region_server::{Region, RegionServer}; +use api::v1::region::{ + AlignedSchemaVersion, BulkInsertRequest, DeleteRequest, InsertRequest, RegionRequestHeader, + RegionResponse as PbRegionResponse, bulk_insert_request, +}; +use api::v1::value::ValueData; +use api::v1::{ + ArrowIpc, ColumnDataType, ColumnSchema, PartitionExprVersion, ResponseHeader, Row, Rows, + Status as PbStatus, Value, helper, +}; +use client::client_manager::NodeClients; +use common_error::ext::{ErrorExt, RetryHint}; +use common_error::status_code::StatusCode; +use common_grpc::channel_manager::ChannelConfig; +use common_grpc::flight::{self, FlightDecoder}; +use common_meta::cache::{ + self as meta_cache, LayeredCacheRegistryBuilder, TableInfoCacheRef, TableRouteCacheRef, +}; +use common_meta::cache_invalidator::{CacheInvalidator, CacheInvalidatorRef, Context}; +use common_meta::ddl::test_util::datanode_handler::DatanodeWatcher; +use common_meta::error::{Result, UnexpectedSnafu}; +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::{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::{PrimaryKeyCodec, SparsePrimaryKeyCodec, sparse}; +use partition::cache::PartitionInfoCacheRef; +use partition::expr::{self, Operand, PartitionExpr, RestrictedOp}; +use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef}; +use store_api::metric_engine_consts; +use store_api::storage::{RegionId, consts}; +use tokio::net::TcpListener; +use tokio::sync::{Barrier, Notify, Semaphore}; +use tokio::task::JoinHandle; +use tokio_stream::wrappers::TcpListenerStream; +use tonic::{Code, Request, Response, Status}; + +use super::*; +use crate::error::Error; +use crate::tests::{create_partition_rule_manager, new_test_table_info, prepare_mocked_backend}; + +#[derive(Clone, Default)] +struct RecordingRegionService { + requests: Arc>>, + rejections: Arc>>>, + release: Arc>>>, +} + +#[async_trait::async_trait] +impl Region for RecordingRegionService { + async fn handle( + &self, + request: Request, + ) -> std::result::Result, Status> { + let request = request.into_inner(); + let region_id = write_region_id(&request).unwrap().as_u64(); + self.requests.lock().unwrap().push(request); + 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) + { + return Err(status); + } + Ok(Response::new(PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus { + status_code: StatusCode::Success as u32, + ..Default::default() + }), + }), + affected_rows: 1, + ..Default::default() + })) + } +} + +#[derive(Default)] +struct RecordingCacheInvalidator { + calls: Mutex>>, + inner: Option, + route_update: Mutex>, +} + +#[async_trait::async_trait] +impl CacheInvalidator for RecordingCacheInvalidator { + async fn invalidate(&self, ctx: &Context, caches: &[CacheIdent]) -> Result<()> { + self.calls.lock().unwrap().push(caches.to_vec()); + let route_update = self.route_update.lock().unwrap().take(); + if let Some((backend, updated)) = route_update { + let manager = TableRouteManager::new(backend.clone()); + let storage = manager.table_route_storage(); + let current = storage.get_with_raw_bytes(1).await?.unwrap(); + let (txn, _) = storage.build_update_txn(1, ¤t, &updated)?; + assert!(backend.txn(txn).await?.succeeded); + } + if let Some(inner) = &self.inner { + inner.invalidate(ctx, caches).await + } else { + UnexpectedSnafu { + err_msg: "stop after recording route refresh", + } + .fail() + } + } + + fn invalidate_all(&self) -> Result<()> { + // This recorder owns no cache entries. + Ok(()) + } +} + +struct RegionTestServer { + peer: Peer, + service: Arc, + invalidator: Arc, + task: JoinHandle<()>, +} + +impl RegionTestServer { + async fn start() -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let peer = Peer::new(1, listener.local_addr().unwrap().to_string()); + let service = Arc::new(RecordingRegionService::default()); + let server_service = service.clone(); + let task = tokio::spawn(async move { + tonic::transport::Server::builder() + .add_service(RegionServer::new(server_service.as_ref().clone())) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + .unwrap(); + }); + Self { + peer, + service, + invalidator: Arc::default(), + task, + } + } + + async fn create_dispatcher(&self, max_retries: usize) -> WriteDispatcher { + let backend = prepare_mocked_backend().await; + let partition_manager = create_partition_rule_manager(backend.clone()).await; + update_table_leader(&backend, &self.peer).await; + let registry = cache::build_fundamental_cache_registry(backend); + let table_info_cache = registry.get::().unwrap(); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + WriteDispatcher::new(clients, partition_manager).with_retry( + build_retry_options(max_retries), + table_info_cache, + self.invalidator.clone(), + ) + } + + fn fail_region(&self, region_number: u32, hints: impl IntoIterator) { + let region_id = RegionId::new(1, region_number).as_u64(); + self.service.rejections.lock().unwrap().insert( + 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(), + ); + } + + fn attempt_counts(&self) -> HashMap { + let mut counts = HashMap::new(); + for request in self.service.requests.lock().unwrap().iter() { + *counts + .entry(write_region_id(request).unwrap().as_u64()) + .or_default() += 1; + } + counts + } +} + +impl Drop for RegionTestServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +fn build_retry_options(max_retries: usize) -> WriteRetryOptions { + WriteRetryOptions { + max_retries, + initial_backoff: Duration::from_millis(1), + max_backoff: Duration::from_millis(2), + ..Default::default() + } +} + +fn create_routing_caches( + backend: KvBackendRef, +) -> ( + PartitionRuleManagerRef, + TableInfoCacheRef, + CacheInvalidatorRef, +) { + let registry = Arc::new( + cache::with_default_composite_cache_registry( + LayeredCacheRegistryBuilder::default() + .add_cache_registry(cache::build_fundamental_cache_registry(backend.clone())), + ) + .unwrap() + .build(), + ); + let partition_manager = Arc::new(PartitionRuleManager::new( + backend.clone(), + registry.get::().unwrap(), + registry.get::().unwrap(), + )); + let table_info_cache = registry.get::().unwrap(); + (partition_manager, table_info_cache, registry) +} + +async fn update_table_leader(backend: &KvBackendRef, peer: &Peer) { + let manager = TableRouteManager::new(backend.clone()); + let storage = manager.table_route_storage(); + let current = storage.get_with_raw_bytes(1).await.unwrap().unwrap(); + let mut routes = current.region_routes().unwrap().clone(); + for route in &mut routes { + route.leader_peer = Some(peer.clone()); + } + let updated = current.update(routes).unwrap(); + let (txn, _) = storage.build_update_txn(1, ¤t, &updated).unwrap(); + assert!(backend.txn(txn).await.unwrap().succeeded); +} + +fn build_insert_request(regions: &[u32]) -> RegionRequest { + RegionRequest { + body: Some(Body::Inserts(InsertRequests { + requests: regions + .iter() + .map(|region| InsertRequest { + region_id: RegionId::new(1, *region).as_u64(), + ..Default::default() + }) + .collect(), + })), + ..Default::default() + } +} + +#[tokio::test] +async fn test_write_retry_disabled_and_enabled_ablation() { + common_telemetry::init_default_ut_logging(); + for run in 0..3 { + for max_retries in [0, 1] { + let server = RegionTestServer::start().await; + server.fail_region(1, [RetryHint::Retryable]); + let dispatcher = server.create_dispatcher(max_retries).await; + let result = dispatcher.dispatch(vec![build_insert_request(&[1])]).await; + let attempts = server.attempt_counts()[&RegionId::new(1, 1).as_u64()]; + common_telemetry::info!( + "write_retry_ablation run={run} max_retries={max_retries} success={} attempts={attempts}", + result.is_ok() + ); + assert_eq!(result.is_ok(), max_retries > 0); + assert_eq!(attempts, max_retries + 1); + assert!(server.invalidator.calls.lock().unwrap().is_empty()); + } + } +} + +#[tokio::test] +async fn test_write_retry_does_not_replay_successful_regions() { + let server = RegionTestServer::start().await; + server.fail_region(2, [RetryHint::Retryable]); + let dispatcher = server.create_dispatcher(2).await; + let response = dispatcher + .dispatch(vec![build_insert_request(&[1, 2])]) + .await + .unwrap(); + assert_eq!(response.affected_rows, 2); + assert_eq!( + server.attempt_counts(), + HashMap::from([ + (RegionId::new(1, 1).as_u64(), 1), + (RegionId::new(1, 2).as_u64(), 2), + ]) + ); + assert!(server.invalidator.calls.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn test_write_retry_attempt_budget_and_nonretryable_errors() { + for (hint, expected_attempts) in [(RetryHint::Retryable, 3), (RetryHint::NonRetryable, 1)] { + let server = RegionTestServer::start().await; + server.fail_region(1, [hint; 4]); + let dispatcher = server.create_dispatcher(2).await; + let error = dispatcher + .dispatch(vec![build_insert_request(&[1])]) + .await + .unwrap_err(); + assert_eq!(error.retry_hint(), hint); + assert_eq!( + server.attempt_counts()[&RegionId::new(1, 1).as_u64()], + expected_attempts + ); + assert!(server.invalidator.calls.lock().unwrap().is_empty()); + } +} + +#[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] { + for regions in [[1, 2], [2, 1]] { + let server = RegionTestServer::start().await; + server.fail_region(1, [hint]); + server.fail_region(2, [RetryHint::NonRetryable]); + let dispatcher = server.create_dispatcher(2).await; + let error = dispatcher + .dispatch(vec![build_insert_request(®ions)]) + .await + .unwrap_err(); + let terminal_region = RegionId::new(1, 2).as_u64(); + assert!( + format!("{error:?}") + .contains(&format!("region {terminal_region} rejected the write")), + "{error:?}" + ); + assert_eq!( + server.attempt_counts(), + HashMap::from([(RegionId::new(1, 1).as_u64(), 1), (terminal_region, 1),]) + ); + assert!(server.invalidator.calls.lock().unwrap().is_empty()); + } + } +} + +#[tokio::test] +async fn test_should_retry_requires_a_retained_request_and_budget() { + let backend = prepare_mocked_backend().await; + let (partition_manager, table_info_cache, cache_invalidator) = create_routing_caches(backend); + let retry = WriteRetry { + node_manager: Arc::new(NodeClients::default()), + partition_manager, + options: build_retry_options(2), + table_info_cache, + cache_invalidator, + }; + let failure = |request| { + let error: client::Error = Status::unavailable("write rejected").into(); + WriteAttempt::Failed { + error: error.into(), + request, + } + }; + let retryable = || { + failure(Some(RoutedWrite { + peer: Arc::new(Peer::new(1, "peer")), + request: Arc::new(build_insert_request(&[1])), + })) + }; + for (name, attempts, retries, expected) in [ + ("empty", vec![], 0, false), + ( + "all succeeded", + vec![WriteAttempt::Success(RegionResponse::new(1))], + 0, + false, + ), + ("retryable", vec![retryable()], 0, true), + ("last retry", vec![retryable()], 1, true), + ("budget exhausted", vec![retryable()], 2, false), + ("terminal first", vec![failure(None), retryable()], 0, false), + ("terminal last", vec![retryable(), failure(None)], 0, false), + ] { + assert_eq!(retry.should_retry(&attempts, retries), expected, "{name}"); + } +} + +#[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; + for region in [1, 2] { + server.fail_region(region, [RetryHint::RetryAfterRouteRefresh]); + } + let dispatcher = server.create_dispatcher(2).await; + let error = dispatcher + .dispatch(vec![build_insert_request(&[1, 2])]) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("stop after recording route refresh") + ); + let calls = server.invalidator.calls.lock().unwrap(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0], vec![CacheIdent::TableId(1)]); + assert_eq!(server.service.requests.lock().unwrap().len(), 2); +} + +#[tokio::test] +async fn test_finish_retries_initial_failures_in_one_round() { + let server = RegionTestServer::start().await; + for region in [2, 3] { + server.fail_region(region, [RetryHint::RetryAfterRouteRefresh]); + } + let backend = prepare_mocked_backend().await; + create_partition_rule_manager(backend.clone()).await; + update_table_leader(&backend, &server.peer).await; + let (partition_manager, table_info_cache, registry) = create_routing_caches(backend); + let table_info = table_info_cache.get(1).await.unwrap().unwrap(); + let (_, versions) = partition_manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + let invalidator = Arc::new(RecordingCacheInvalidator { + inner: Some(registry), + ..Default::default() + }); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + let dispatcher = WriteDispatcher::new(clients, partition_manager).with_retry( + build_retry_options(2), + table_info_cache, + invalidator.clone(), + ); + let mut attempts = Vec::new(); + for (region, value) in [(1, 101), (2, 11), (3, 1)] { + let request = + build_partitioned_request(WriteKind::BulkInsert, region, versions[®ion], &[value]); + attempts.push( + dispatcher + .send(Arc::new(server.peer.clone()), request, None) + .await, + ); + } + assert!(invalidator.calls.lock().unwrap().is_empty()); + assert_eq!( + server.attempt_counts(), + HashMap::from([ + (RegionId::new(1, 1).as_u64(), 1), + (RegionId::new(1, 2).as_u64(), 1), + (RegionId::new(1, 3).as_u64(), 1), + ]) + ); + let response = dispatcher.finish(attempts, None, None).await.unwrap(); + assert_eq!(response.affected_rows, 3); + assert_eq!( + server.attempt_counts(), + HashMap::from([ + (RegionId::new(1, 1).as_u64(), 1), + (RegionId::new(1, 2).as_u64(), 2), + (RegionId::new(1, 3).as_u64(), 2), + ]) + ); + assert_eq!( + *invalidator.calls.lock().unwrap(), + vec![vec![CacheIdent::TableId(1)]] + ); +} + +#[tokio::test] +async fn test_write_retry_refreshes_stale_route_and_sends_only_failed_region_to_new_leader() { + let old_server = RegionTestServer::start().await; + let new_server = RegionTestServer::start().await; + old_server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]); + let backend = prepare_mocked_backend().await; + // Reuse the operator fixture to populate table metadata and partition rules. + create_partition_rule_manager(backend.clone()).await; + update_table_leader(&backend, &old_server.peer).await; + + let (partition_manager, table_info_cache, registry) = create_routing_caches(backend.clone()); + let failed_region = RegionId::new(1, 2); + assert_eq!( + partition_manager + .find_region_leader(failed_region) + .await + .unwrap(), + old_server.peer + ); + let table_info = table_info_cache.get(1).await.unwrap().unwrap(); + let (_, versions) = partition_manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + + // Metadata has migrated while the frontend still holds the old route. + update_table_leader(&backend, &new_server.peer).await; + assert_eq!( + partition_manager + .find_region_leader(failed_region) + .await + .unwrap(), + old_server.peer + ); + + let invalidator = Arc::new(RecordingCacheInvalidator { + calls: Mutex::default(), + inner: Some(registry), + ..Default::default() + }); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + let dispatcher = WriteDispatcher::new(clients, partition_manager).with_retry( + build_retry_options(2), + table_info_cache, + invalidator.clone(), + ); + let mut request = build_insert_request(&[1, 2]); + if let Some(Body::Inserts(inserts)) = &mut request.body { + for insert in &mut inserts.requests { + insert.partition_expr_version = versions + [&RegionId::from_u64(insert.region_id).region_number()] + .map(|value| PartitionExprVersion { value }); + } + } + let response = dispatcher.dispatch(vec![request]).await.unwrap(); + assert_eq!(response.affected_rows, 2); + assert_eq!( + old_server.attempt_counts(), + HashMap::from([ + (RegionId::new(1, 1).as_u64(), 1), + (failed_region.as_u64(), 1), + ]) + ); + assert_eq!( + new_server.attempt_counts(), + HashMap::from([(failed_region.as_u64(), 1)]) + ); + assert_eq!( + *invalidator.calls.lock().unwrap(), + vec![vec![CacheIdent::TableId(1)]] + ); +} + +#[derive(Clone, Copy, Debug)] +enum WriteKind { + Insert, + Delete, + BulkInsert, +} + +fn build_partitioned_request( + kind: WriteKind, + region: u32, + version: Option, + values: &[i32], +) -> RegionRequest { + let region_id = RegionId::new(1, region).as_u64(); + let partition_expr_version = version.map(|value| PartitionExprVersion { value }); + let rows = Rows { + schema: vec![ColumnSchema { + column_name: "a".to_string(), + datatype: ColumnDataType::Int32 as i32, + ..Default::default() + }], + rows: values + .iter() + .map(|value| Row { + values: vec![Value { + value_data: Some(ValueData::I32Value(*value)), + }], + }) + .collect(), + }; + let body = match kind { + WriteKind::Insert => Body::Inserts(InsertRequests { + requests: vec![InsertRequest { + region_id, + partition_expr_version, + rows: Some(rows), + skip_wal: true, + }], + }), + WriteKind::Delete => Body::Deletes(DeleteRequests { + requests: vec![DeleteRequest { + region_id, + partition_expr_version, + rows: Some(rows), + }], + }), + WriteKind::BulkInsert => { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(values.to_vec()))], + ) + .unwrap(); + let (schema, data_header, payload) = flight::record_batch_to_ipc(batch).unwrap(); + Body::BulkInsert(BulkInsertRequest { + region_id, + partition_expr_version, + aligned_schema_version: Some(AlignedSchemaVersion { + schema_version: 123, + }), + body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { + schema, + data_header, + payload, + })), + }) + } + }; + RegionRequest { + header: None, + body: Some(body), + } +} + +fn request_parts(request: &RegionRequest) -> (u32, Option, Vec) { + let region = write_region_id(request).unwrap().region_number(); + let (version, rows) = match request.body.as_ref().unwrap() { + Body::Inserts(inserts) => { + assert!(inserts.requests[0].skip_wal); + ( + inserts.requests[0].partition_expr_version, + inserts.requests[0].rows.as_ref().unwrap(), + ) + } + Body::Deletes(deletes) => ( + deletes.requests[0].partition_expr_version, + deletes.requests[0].rows.as_ref().unwrap(), + ), + Body::BulkInsert(bulk) => { + assert_eq!(bulk.aligned_schema_version, None); + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &bulk.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(); + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + return ( + region, + bulk.partition_expr_version.map(|v| v.value), + values.values().to_vec(), + ); + } + _ => panic!("expected write request"), + }; + ( + region, + version.map(|v| v.value), + rows.rows + .iter() + .map(|row| { + let Some(ValueData::I32Value(value)) = row.values[0].value_data else { + panic!("expected int32") + }; + value + }) + .collect(), + ) +} + +#[tokio::test] +async fn test_write_retry_rebuilds_only_failed_rows_after_route_refresh() { + for kind in [WriteKind::Insert, WriteKind::Delete, WriteKind::BulkInsert] { + let old_server = RegionTestServer::start().await; + let new_server = RegionTestServer::start().await; + old_server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]); + let backend = prepare_mocked_backend().await; + create_partition_rule_manager(backend.clone()).await; + update_table_leader(&backend, &old_server.peer).await; + let (partition_manager, table_info_cache, registry) = + create_routing_caches(backend.clone()); + let table_info = table_info_cache.get(1).await.unwrap().unwrap(); + let (_, old_versions) = partition_manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + let manager = TableRouteManager::new(backend.clone()); + let current = manager + .table_route_storage() + .get_with_raw_bytes(1) + .await + .unwrap() + .unwrap(); + let old_routes = current.region_routes().unwrap(); + let replaced = old_routes + .iter() + .find(|route| route.region.id.region_number() == 2) + .unwrap(); + let mut routes = old_routes + .iter() + .filter(|route| route.region.id.region_number() != 2) + .cloned() + .collect::>(); + for (region, lower, upper) in [(4, 10, 30), (5, 30, 50)] { + let mut route = replaced.clone(); + route.region.id = RegionId::new(1, region); + route.leader_peer = Some(new_server.peer.clone()); + route.region.partition_expr = PartitionExpr::new( + Operand::Expr(PartitionExpr::new( + Operand::Column("a".to_string()), + RestrictedOp::GtEq, + Operand::Value(datatypes::value::Value::Int32(lower)), + )), + RestrictedOp::And, + Operand::Expr(PartitionExpr::new( + Operand::Column("a".to_string()), + RestrictedOp::Lt, + Operand::Value(datatypes::value::Value::Int32(upper)), + )), + ) + .as_json_str() + .unwrap(); + routes.push(route); + } + let invalidator = Arc::new(RecordingCacheInvalidator { + calls: Mutex::default(), + inner: Some(registry), + route_update: Mutex::new(Some((backend, current.update(routes).unwrap()))), + }); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + 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]), + build_partitioned_request(kind, 2, old_versions[&2], &[20, 40]), + ]) + .await + .unwrap(); + assert_eq!(response.affected_rows, 3, "{kind:?}"); + assert_eq!( + old_server.attempt_counts(), + HashMap::from([ + (RegionId::new(1, 1).as_u64(), 1), + (RegionId::new(1, 2).as_u64(), 1), + ]), + "{kind:?}" + ); + let (_, versions) = partition_manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + let mut writes = new_server + .service + .requests + .lock() + .unwrap() + .iter() + .map(request_parts) + .collect::>(); + writes.sort_by_key(|write| write.0); + assert_eq!( + writes, + vec![(4, versions[&4], vec![20]), (5, versions[&5], vec![40])], + "{kind:?}" + ); + assert_eq!( + *invalidator.calls.lock().unwrap(), + vec![vec![CacheIdent::TableId(1)]] + ); + } +} + +fn build_sparse_bulk_request( + region: u32, + version: Option, + tags: &[(&str, &str)], +) -> (RegionRequest, RecordBatch) { + let codec = SparsePrimaryKeyCodec::schemaless(); + let keys = tags + .iter() + .enumerate() + .map(|(index, (host, zone))| { + let mut key = Vec::new(); + codec + .encode_value_refs( + &[ + (sparse::RESERVED_COLUMN_ID_TABLE_ID, ValueRef::UInt32(42)), + ( + sparse::RESERVED_COLUMN_ID_TSID, + ValueRef::UInt64(index as u64), + ), + (10, ValueRef::String(host)), + (20, ValueRef::String(zone)), + ], + &mut key, + ) + .unwrap(); + key + }) + .collect::>(); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new(consts::PRIMARY_KEY_COLUMN_NAME, DataType::Binary, false), + Field::new("host", DataType::Utf8, false), + ])), + vec![ + Arc::new(BinaryArray::from_iter_values(keys)), + Arc::new(StringArray::from_iter_values( + tags.iter().map(|(host, _)| *host), + )), + ], + ) + .unwrap(); + let (schema, data_header, payload) = flight::record_batch_to_ipc(batch.clone()).unwrap(); + let request = RegionRequest { + body: Some(Body::BulkInsert(BulkInsertRequest { + region_id: RegionId::new(1, region).as_u64(), + partition_expr_version: version.map(|value| PartitionExprVersion { value }), + body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { + schema, + data_header, + payload, + })), + ..Default::default() + })), + ..Default::default() + }; + (request, batch) +} + +#[tokio::test] +async fn test_write_retry_restores_changed_partition_tags_for_failed_regions() { + let old_server = RegionTestServer::start().await; + let new_server = RegionTestServer::start().await; + old_server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]); + let backend: KvBackendRef = Arc::new(MemoryKvBackend::new()); + let metadata = TableMetadataManager::new(backend.clone()); + let mut table_info = 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 = prepare_mocked_backend().await; + let partition_manager = create_partition_rule_manager(backend.clone()).await; + TableMetadataManager::new(backend) + .create_table_metadata( + new_test_table_info(42, "dispatcher_logical", [1, 2, 3].into_iter()), + TableRouteValue::logical(1), + Default::default(), + ) + .await + .unwrap(); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + let dispatcher = WriteDispatcher::new(clients, partition_manager.clone()); + let regions = [ + RegionId::new(1, 1), + RegionId::new(1, 2), + RegionId::new(42, 1), + RegionId::new(42, 2), + RegionId::new(42, 1), + ]; + let peers = dispatcher.resolve_peers(®ions).await.unwrap(); + assert_eq!(peers.len(), 4); + for region in regions { + assert_eq!( + peers[®ion].as_ref(), + &partition_manager.find_region_leader(region).await.unwrap() + ); + } + assert!(Arc::ptr_eq( + &peers[&RegionId::new(1, 1)], + &peers[&RegionId::new(42, 1)], + )); + let missing = RegionId::new(42, 999); + let expected = partition_manager + .find_region_leader(missing) + .await + .unwrap_err(); + let actual = dispatcher.resolve_peers(&[missing]).await.unwrap_err(); + assert_eq!(actual.status_code(), expected.status_code()); + assert_eq!(actual.retry_hint(), expected.retry_hint()); + assert!(dispatcher.resolve_peers(&[]).await.unwrap().is_empty()); +} + +#[tokio::test] +async fn test_resolve_peers_shares_only_identical_peers() { + let backend = prepare_mocked_backend().await; + create_partition_rule_manager(backend.clone()).await; + let manager = TableRouteManager::new(backend.clone()); + let storage = manager.table_route_storage(); + let current = storage.get_with_raw_bytes(1).await.unwrap().unwrap(); + let mut routes = current.region_routes().unwrap().clone(); + let peer = Peer::new(7, "same-address"); + let other = Peer::new(7, "other-address"); + for route in &mut routes { + route.leader_peer = Some(if route.region.id.region_number() == 3 { + other.clone() + } else { + peer.clone() + }); + } + let updated = current.update(routes).unwrap(); + let (txn, _) = storage.build_update_txn(1, ¤t, &updated).unwrap(); + assert!(backend.txn(txn).await.unwrap().succeeded); + let (partition_manager, _, _) = create_routing_caches(backend); + let dispatcher = WriteDispatcher::new(Arc::new(NodeClients::default()), partition_manager); + let regions = [1, 2, 3].map(|region| RegionId::new(1, region)); + let peers = dispatcher.resolve_peers(®ions).await.unwrap(); + assert_eq!(peers.len(), 3); + assert_eq!(peers[®ions[0]].as_ref(), &peer); + assert_eq!(peers[®ions[1]].as_ref(), &peer); + assert_eq!(peers[®ions[2]].as_ref(), &other); + assert!(Arc::ptr_eq(&peers[®ions[0]], &peers[®ions[1]])); + assert!(!Arc::ptr_eq(&peers[®ions[0]], &peers[®ions[2]])); +} + +#[tokio::test] +async fn test_dispatch_preserves_write_payloads_and_envelopes() { + let backend = prepare_mocked_backend().await; + let partition_manager = create_partition_rule_manager(backend.clone()).await; + TableMetadataManager::new(backend) + .create_table_metadata( + new_test_table_info(42, "dispatcher_logical", [1, 2, 3].into_iter()), + TableRouteValue::logical(1), + Default::default(), + ) + .await + .unwrap(); + let (sender, mut received) = tokio::sync::mpsc::channel(8); + let watcher = DatanodeWatcher::new(sender).with_handler(|_, _| Ok(RegionResponse::new(1))); + let dispatcher = WriteDispatcher::new( + Arc::new(MockDatanodeManager::new(watcher)), + partition_manager, + ); + // Physical and logical regions share a peer but retain distinct region IDs. + let regions = [RegionId::new(1, 1), RegionId::new(42, 1)]; + let rows = Rows { + schema: vec![helper::tag_column_schema("a", ColumnDataType::Int32)], + rows: vec![Row { + values: vec![ValueData::I32Value(7).into()], + }], + }; + let version = Some(PartitionExprVersion { value: 17 }); + let bodies = [ + Body::Inserts(InsertRequests { + requests: regions + .iter() + .map(|region| InsertRequest { + region_id: region.as_u64(), + rows: Some(rows.clone()), + partition_expr_version: version, + skip_wal: true, + }) + .collect(), + }), + Body::Deletes(DeleteRequests { + requests: regions + .iter() + .map(|region| DeleteRequest { + region_id: region.as_u64(), + rows: Some(rows.clone()), + partition_expr_version: version, + }) + .collect(), + }), + ]; + for body in bodies { + let requests = ["first", "second"].map(|dbname| RegionRequest { + header: Some(RegionRequestHeader { + dbname: dbname.to_string(), + ..Default::default() + }), + body: Some(body.clone()), + }); + let writes = dispatcher.prepare(requests.to_vec()).await.unwrap(); + assert!(received.try_recv().is_err()); + let response = dispatcher.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)); + } + assert!(received.try_recv().is_err()); + } + let requests = regions.map(|region| RegionRequest { + header: Some(RegionRequestHeader { + dbname: "bulk".to_string(), + ..Default::default() + }), + body: Some(Body::BulkInsert(BulkInsertRequest { + region_id: region.as_u64(), + partition_expr_version: version, + aligned_schema_version: None, + body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { + schema: "schema".into(), + data_header: "header".into(), + payload: "payload".into(), + })), + })), + }); + let response = dispatcher.dispatch(requests.to_vec()).await.unwrap(); + assert_eq!(response.affected_rows, 2); + for request in requests { + assert_eq!(received.try_recv().unwrap(), (Peer::new(1, ""), request)); + } + assert!(received.try_recv().is_err()); + + let error = dispatcher + .prepare(vec![build_insert_request(&[1, 999])]) + .await + .unwrap_err(); + assert!(matches!(error, Error::FindRegionLeader { .. })); + assert!(received.try_recv().is_err()); +} + +#[derive(Clone)] +struct GatedDatanodeHandler { + entered: Arc, + rejected: Arc, + completed: Arc, + release: Arc, +} + +#[async_trait::async_trait] +impl MockDatanodeHandler for GatedDatanodeHandler { + async fn handle(&self, peer: &Peer, _request: RegionRequest) -> Result { + self.entered.wait().await; + if peer.id == 1 { + self.rejected.store(true, Ordering::SeqCst); + return UnexpectedSnafu { + err_msg: "first region rejected", + } + .fail(); + } + self.release.notified().await; + self.completed.store(true, Ordering::SeqCst); + Ok(RegionResponse::new(1)) + } + + async fn handle_query( + &self, + _peer: &Peer, + _request: QueryRequest, + ) -> Result { + unimplemented!() + } +} + +#[tokio::test] +async fn test_dispatch_sends_concurrently_and_waits_for_siblings_after_error() { + let backend = prepare_mocked_backend().await; + let partition_manager = create_partition_rule_manager(backend).await; + let handler = GatedDatanodeHandler { + entered: Arc::new(Barrier::new(2)), + rejected: Arc::default(), + completed: Arc::default(), + release: Arc::default(), + }; + let dispatcher = WriteDispatcher::new( + Arc::new(MockDatanodeManager::new(handler.clone())), + partition_manager, + ); + let dispatch = dispatcher.dispatch(vec![build_insert_request(&[1, 2])]); + tokio::pin!(dispatch); + tokio::time::timeout( + Duration::from_secs(1), + futures::future::poll_fn(|cx| { + assert!( + dispatch.as_mut().poll(cx).is_pending(), + "dispatch returned while a sibling remained blocked" + ); + if handler.rejected.load(Ordering::SeqCst) { + Poll::Ready(()) + } else { + Poll::Pending + } + }), + ) + .await + .expect("both region sends must reach the barrier concurrently"); + assert!(!handler.completed.load(Ordering::SeqCst)); + handler.release.notify_one(); + let error = tokio::time::timeout(Duration::from_secs(1), dispatch) + .await + .unwrap() + .unwrap_err(); + assert!(error.to_string().contains("first region rejected")); + assert!(handler.completed.load(Ordering::SeqCst)); +} + +#[tokio::test] +async fn test_resolve_peers_cache_lookup_ablation() { + common_telemetry::init_default_ut_logging(); + let backend = prepare_mocked_backend().await; + create_partition_rule_manager(backend.clone()).await; + TableMetadataManager::new(backend.clone()) + .create_table_metadata( + new_test_table_info(42, "resolution_ablation", [1, 2, 3].into_iter()), + TableRouteValue::logical(1), + Default::default(), + ) + .await + .unwrap(); + // A dedicated cache label isolates counters from other concurrently running tests. + let cache_name = "write_dispatcher_resolution_ablation_table_route"; + let routes = Arc::new(meta_cache::new_table_route_cache( + cache_name.to_string(), + moka::future::Cache::new(16), + backend.clone(), + )); + let partitions = Arc::new(partition::cache::new_partition_info_cache( + "write_dispatcher_resolution_ablation_partitions".to_string(), + moka::future::Cache::new(16), + routes.clone(), + )); + let partition_manager = Arc::new(PartitionRuleManager::new(backend, routes, partitions)); + let dispatcher = + WriteDispatcher::new(Arc::new(NodeClients::default()), partition_manager.clone()); + let regions = (0..256) + .map(|index| RegionId::new(if index % 2 == 0 { 1 } else { 42 }, index % 3 + 1)) + .collect::>(); + let expected = dispatcher.resolve_peers(®ions).await.unwrap(); + let counter = common_meta::metrics::CACHE_CONTAINER_CACHE_GET.with_label_values(&[cache_name]); + for run in 0..6 { + // Alternate order; compare the same warm cache, routes and deterministic workload. + let modes = if run % 2 == 0 { + ["delete", "insert", "batched"] + } else { + ["batched", "insert", "delete"] + }; + for mode in modes { + let before = counter.get(); + let started = Instant::now(); + let (elapsed, peers_match) = if mode == "batched" { + let peers = dispatcher.resolve_peers(®ions).await.unwrap(); + (started.elapsed(), peers == expected) + } else { + let mut peers = HashMap::new(); + for region in ®ions { + if mode == "insert" && peers.contains_key(region) { + continue; + } + peers.insert( + *region, + partition_manager.find_region_leader(*region).await.unwrap(), + ); + } + let elapsed = started.elapsed(); + let peers_match = peers.len() == expected.len() + && peers.iter().all(|(region, peer)| { + expected + .get(region) + .is_some_and(|expected| expected.as_ref() == peer) + }); + (elapsed, peers_match) + }; + let gets = counter.get() - before; + assert!(peers_match, "{mode}"); + let expected_gets = match mode { + "delete" => 384, + "insert" => 9, + "batched" => 3, + _ => unreachable!(), + }; + assert_eq!(gets, expected_gets); + common_telemetry::info!( + "route_resolution_ablation run={run} mode={mode} regions={} cache_gets={gets} elapsed_ns={}", + regions.len(), + elapsed.as_nanos() + ); + } + } +} + +#[tokio::test] +async fn test_write_retry_rebuilds_logical_rows_in_mixed_table_envelope() { + for kind in [WriteKind::Insert, WriteKind::Delete] { + let server = RegionTestServer::start().await; + let backend = prepare_mocked_backend().await; + create_partition_rule_manager(backend.clone()).await; + TableMetadataManager::new(backend.clone()) + .create_table_metadata( + new_test_table_info(42, "retry_logical", [1, 2, 3].into_iter()), + TableRouteValue::logical(1), + Default::default(), + ) + .await + .unwrap(); + update_table_leader(&backend, &server.peer).await; + let (partition_manager, table_info_cache, registry) = create_routing_caches(backend); + let table_info = table_info_cache.get(1).await.unwrap().unwrap(); + let (_, versions) = partition_manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + let invalidator = Arc::new(RecordingCacheInvalidator { + inner: Some(registry), + ..Default::default() + }); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + let dispatcher = WriteDispatcher::new(clients, partition_manager).with_retry( + build_retry_options(2), + table_info_cache, + invalidator.clone(), + ); + 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(); + 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(), + ..Default::default() + }); + let mut physical = build_partitioned_request(kind, 1, versions[&1], &[70]); + physical.header = header.clone(); + // The stale logical request must move to region 3 while keeping table 42. + let mut logical = build_partitioned_request(kind, 2, Some(0), &[1]); + let mut rebuilt = build_partitioned_request(kind, 3, versions[&3], &[1]); + for request in [&mut logical, &mut rebuilt] { + request.header = header.clone(); + let region_id = match request.body.as_mut().unwrap() { + Body::Inserts(inserts) => &mut inserts.requests[0].region_id, + Body::Deletes(deletes) => &mut deletes.requests[0].region_id, + _ => unreachable!(), + }; + *region_id = RegionId::new(42, RegionId::from_u64(*region_id).region_number()).as_u64(); + } + let mut envelope = physical.clone(); + match ( + envelope.body.as_mut().unwrap(), + logical.body.clone().unwrap(), + ) { + (Body::Inserts(physical), Body::Inserts(logical)) => { + physical.requests.extend(logical.requests); + } + (Body::Deletes(physical), Body::Deletes(logical)) => { + physical.requests.extend(logical.requests); + } + _ => unreachable!(), + } + let response = dispatcher.dispatch(vec![envelope]).await.unwrap(); + assert_eq!(response.affected_rows, 2, "{kind:?}"); + assert_eq!( + server.attempt_counts(), + HashMap::from([ + (physical_id, 1), + (stale_logical_id, 1), + (rebuilt_logical_id, 1) + ]), + "{kind:?}" + ); + let mut requests = server.service.requests.lock().unwrap().clone(); + requests.sort_by_key(|request| write_region_id(request).unwrap().as_u64()); + assert_eq!(requests, vec![physical, logical, rebuilt], "{kind:?}"); + let calls = invalidator.calls.lock().unwrap(); + assert_eq!(calls.len(), 1, "{kind:?}"); + assert_eq!(calls[0].len(), 2, "{kind:?}"); + assert!(calls[0].contains(&CacheIdent::TableId(1))); + assert!(calls[0].contains(&CacheIdent::TableId(42))); + } +} + +#[tokio::test] +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 table_info_cache = retry.table_info_cache.clone(); + let cache_invalidator = retry.cache_invalidator.clone(); + 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!( + dispatcher + .dispatch(vec![build_insert_request(&[1])]) + .await + .is_err() + ); + assert_eq!(server.attempt_counts()[&RegionId::new(1, 1).as_u64()], 1); +} diff --git a/src/servers/AGENTS.md b/src/servers/AGENTS.md index 2a446033f62..50e7b029bf9 100644 --- a/src/servers/AGENTS.md +++ b/src/servers/AGENTS.md @@ -35,3 +35,6 @@ cargo nextest run -p servers ``` Keep protocol translation here and permissions/database behavior in frontend. +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 `pending_rows_batcher.rs`. diff --git a/src/servers/benches/flush_batch_physical.rs b/src/servers/benches/flush_batch_physical.rs index 45a3f37c780..eaf36cc7786 100644 --- a/src/servers/benches/flush_batch_physical.rs +++ b/src/servers/benches/flush_batch_physical.rs @@ -15,8 +15,6 @@ use std::collections::HashMap; use std::sync::Arc; -use api::region::RegionResponse; -use api::v1::meta::Peer; use api::v1::region::RegionRequest; use arrow::array::{Float64Array, StringArray, TimestampMillisecondArray}; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema, TimeUnit}; @@ -28,12 +26,11 @@ use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema as DtColumnSchema, Schema as DtSchema}; use partition::error::Result as PartitionResult; use partition::partition::{PartitionRule, PartitionRuleRef, RegionMask}; -use servers::error::{self, Result}; +use servers::error; use servers::pending_rows_batcher::{ - PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider, + PhysicalFlushCatalogProvider, PhysicalFlushPartitionProvider, PhysicalFlushWriteDispatcher, PhysicalTableMetadata, RecordBatchWithTsIdx, TableBatch, flush_batch_physical, }; -use store_api::storage::RegionId; use table::test_util::table_info::test_table_info; use tokio::runtime::Runtime; @@ -95,24 +92,17 @@ impl PhysicalFlushPartitionProvider for BenchPartitionProvider { async fn find_table_partition_rule( &self, _table_info: &table::metadata::TableInfo, - ) -> PartitionResult { - Ok(Arc::new(SingleRegionPartitionRule)) - } - - async fn find_region_leader(&self, _region_id: RegionId) -> Result { - Ok(Peer { - id: 1, - addr: "bench-node".to_string(), - }) + ) -> PartitionResult<(PartitionRuleRef, HashMap>)> { + Ok((Arc::new(SingleRegionPartitionRule), HashMap::new())) } } -struct BenchNodeRequester; +struct BenchWriteDispatcher; #[async_trait] -impl PhysicalFlushNodeRequester for BenchNodeRequester { - async fn handle(&self, _peer: &Peer, _request: RegionRequest) -> error::Result { - Ok(RegionResponse::new(0)) +impl PhysicalFlushWriteDispatcher for BenchWriteDispatcher { + async fn dispatch(&self, _requests: Vec) -> error::Result { + Ok(0) } } @@ -224,7 +214,7 @@ fn bench_flush_batch_physical(c: &mut Criterion) { table: make_physical_table_metadata(num_tags), }; let partition_provider = BenchPartitionProvider; - let node_requester = BenchNodeRequester; + let write_dispatcher = BenchWriteDispatcher; let mut group = c.benchmark_group("flush_batch_physical"); @@ -244,7 +234,7 @@ fn bench_flush_batch_physical(c: &mut Criterion) { "phy", &ctx, &partition_provider, - &node_requester, + &write_dispatcher, &catalog_provider, ) .await @@ -271,7 +261,7 @@ fn bench_flush_batch_physical(c: &mut Criterion) { "phy", &ctx, &partition_provider, - &node_requester, + &write_dispatcher, &catalog_provider, ) .await diff --git a/src/servers/src/pending_rows_batcher.rs b/src/servers/src/pending_rows_batcher.rs index 03eb90dce28..7063f9773a1 100644 --- a/src/servers/src/pending_rows_batcher.rs +++ b/src/servers/src/pending_rows_batcher.rs @@ -18,19 +18,18 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use api::v1::flow::{DirtyWindowRequest, DirtyWindowRequests}; -use api::v1::meta::Peer; use api::v1::region::{ BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request, }; -use api::v1::{ArrowIpc, ColumnSchema, RowInsertRequests, Rows}; +use api::v1::{ArrowIpc, ColumnSchema, PartitionExprVersion, RowInsertRequests, Rows}; use arrow::array::Array; use arrow::compute::{concat_batches, filter_record_batch}; use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema, TimeUnit}; use arrow::record_batch::RecordBatch; use async_trait::async_trait; -use bytes::Bytes; use catalog::CatalogManagerRef; -use common_grpc::flight::{FlightEncoder, FlightMessage}; +use common_grpc::error::Error as GrpcError; +use common_grpc::flight; use common_meta::cache::TableFlownodeSetCacheRef; use common_meta::node_manager::NodeManagerRef; use common_query::prelude::{GREPTIME_PHYSICAL_TABLE, greptime_timestamp, greptime_value}; @@ -40,6 +39,7 @@ use dashmap::DashMap; use dashmap::mapref::entry::Entry; use futures::StreamExt; use metric_engine::batch_modifier::{TagColumnInfo, modify_batch_sparse}; +use operator::write_dispatcher::WriteDispatcherRef; use partition::manager::PartitionRuleManagerRef; use partition::partition::PartitionRuleRef; use session::context::QueryContextRef; @@ -80,7 +80,6 @@ pub fn pending_rows_batch_sync_enabled() -> bool { .unwrap_or(true) } const WORKER_IDLE_TIMEOUT_MULTIPLIER: u32 = 3; -const PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT: usize = 3; const MAX_CONCURRENT_FLOW_NOTIFICATIONS: usize = 8; #[async_trait] pub trait PendingRowsSchemaAlterer: Send + Sync { @@ -126,23 +125,21 @@ pub trait PhysicalFlushCatalogProvider: Send + Sync { ) -> catalog::error::Result>; } +/// Loads the physical partition rule and its expression versions. #[async_trait] pub trait PhysicalFlushPartitionProvider: Send + Sync { + /// Loads a routing snapshot for the physical table. async fn find_table_partition_rule( &self, table_info: &TableInfo, - ) -> partition::error::Result; - - async fn find_region_leader(&self, region_id: RegionId) -> Result; + ) -> partition::error::Result<(PartitionRuleRef, HashMap>)>; } +/// Dispatches the region requests produced by a physical flush. #[async_trait] -pub trait PhysicalFlushNodeRequester: Send + Sync { - async fn handle( - &self, - peer: &Peer, - request: RegionRequest, - ) -> Result; +pub trait PhysicalFlushWriteDispatcher: Send + Sync { + /// Dispatches one physical flush together, sharing routing refreshes across its writes. + async fn dispatch(&self, requests: Vec) -> Result; } #[derive(Clone)] @@ -185,36 +182,30 @@ impl PhysicalFlushPartitionProvider for PartitionManagerPhysicalFlushAdapter { async fn find_table_partition_rule( &self, table_info: &TableInfo, - ) -> partition::error::Result { + ) -> partition::error::Result<(PartitionRuleRef, HashMap>)> { self.partition_manager .find_table_partition_rule(table_info) .await - .map(|(rule, _)| rule) - } - - async fn find_region_leader(&self, region_id: RegionId) -> Result { - let peer = self.partition_manager.find_region_leader(region_id).await?; - Ok(peer) } } +/// Sends a physical flush through the shared write dispatcher. #[derive(Clone)] -struct NodeManagerPhysicalFlushAdapter { - node_manager: NodeManagerRef, +struct WriteDispatcherPhysicalFlushAdapter { + write_dispatcher: WriteDispatcherRef, } #[async_trait] -impl PhysicalFlushNodeRequester for NodeManagerPhysicalFlushAdapter { - async fn handle( - &self, - peer: &Peer, - request: RegionRequest, - ) -> error::Result { - let datanode = self.node_manager.datanode(peer).await; - datanode - .handle(request) +impl PhysicalFlushWriteDispatcher for WriteDispatcherPhysicalFlushAdapter { + async fn dispatch(&self, requests: Vec) -> Result { + self.write_dispatcher + .dispatch_with_metrics( + requests, + PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED.with_label_values(&["flush_write_region"]), + ) .await .context(error::CommonMetaSnafu) + .map(|response| response.affected_rows) } } @@ -347,7 +338,7 @@ pub struct PendingRowsBatcher { flush_interval: Duration, max_batch_rows: usize, partition_manager: PartitionRuleManagerRef, - node_manager: NodeManagerRef, + write_dispatcher: WriteDispatcherRef, catalog_manager: CatalogManagerRef, flow_notification_tx: mpsc::Sender, flush_semaphore: Arc, @@ -364,6 +355,7 @@ impl PendingRowsBatcher { 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, @@ -404,7 +396,7 @@ impl PendingRowsBatcher { flush_interval, max_batch_rows, partition_manager, - node_manager, + write_dispatcher, catalog_manager, flow_notification_tx, prom_store_with_metric_engine, @@ -876,7 +868,7 @@ impl PendingRowsBatcher { 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(), self.flush_interval, @@ -932,7 +924,7 @@ fn start_worker( mut rx: mpsc::Receiver, shutdown: broadcast::Sender<()>, partition_manager: PartitionRuleManagerRef, - node_manager: NodeManagerRef, + write_dispatcher: WriteDispatcherRef, catalog_manager: CatalogManagerRef, flow_notification_tx: mpsc::Sender, flush_interval: Duration, @@ -982,7 +974,7 @@ fn start_worker( spawn_flush( flush, partition_manager.clone(), - node_manager.clone(), + write_dispatcher.clone(), catalog_manager.clone(), flow_notification_tx.clone(), flush_semaphore.clone(), @@ -994,7 +986,7 @@ fn start_worker( flush_batch_with_managers( flush, partition_manager.clone(), - node_manager.clone(), + write_dispatcher.clone(), catalog_manager.clone(), flow_notification_tx.clone(), ).await; @@ -1031,7 +1023,7 @@ fn start_worker( spawn_flush( flush, partition_manager.clone(), - node_manager.clone(), + write_dispatcher.clone(), catalog_manager.clone(), flow_notification_tx.clone(), flush_semaphore.clone(), @@ -1043,7 +1035,7 @@ fn start_worker( flush_batch_with_managers( flush, partition_manager.clone(), - node_manager.clone(), + write_dispatcher.clone(), catalog_manager.clone(), flow_notification_tx.clone(), ).await; @@ -1104,7 +1096,7 @@ fn drain_batch(batch: &mut Option) -> Option { async fn spawn_flush( flush: FlushBatch, partition_manager: PartitionRuleManagerRef, - node_manager: NodeManagerRef, + write_dispatcher: WriteDispatcherRef, catalog_manager: CatalogManagerRef, flow_notification_tx: mpsc::Sender, semaphore: Arc, @@ -1116,7 +1108,7 @@ async fn spawn_flush( flush_batch_with_managers( flush, partition_manager, - node_manager, + write_dispatcher, catalog_manager, flow_notification_tx, ) @@ -1128,7 +1120,7 @@ async fn spawn_flush( flush_batch_with_managers( flush, partition_manager, - node_manager, + write_dispatcher, catalog_manager, flow_notification_tx, ) @@ -1137,11 +1129,6 @@ async fn spawn_flush( } } -struct FlushRegionWrite { - datanode: Peer, - request: RegionRequest, -} - struct PlannedRegionBatch { region_id: RegionId, batch: RecordBatch, @@ -1154,15 +1141,6 @@ impl PlannedRegionBatch { } } -struct ResolvedRegionBatch { - planned: PlannedRegionBatch, - datanode: Peer, -} - -fn should_dispatch_concurrently(region_write_count: usize) -> bool { - region_write_count > 1 -} - /// Classifies columns in a logical-table batch for sparse primary-key conversion. /// /// Returns: @@ -1267,72 +1245,20 @@ fn columns_taxonomy( Ok((tag_columns, essential_column_indices)) } -fn strip_partition_columns_from_batch(batch: RecordBatch) -> Result { - ensure!( - batch.num_columns() >= PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT, - error::InternalSnafu { - err_msg: format!( - "Expected at least {} columns in physical batch, got {}", - PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT, - batch.num_columns() - ), - } - ); - let essential_indices: Vec = (0..PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT).collect(); - batch.project(&essential_indices).context(error::ArrowSnafu) -} - -async fn flush_region_writes_concurrently( - node_manager: &(impl PhysicalFlushNodeRequester + ?Sized), - writes: Vec, -) -> Result { - let mut affected_rows = 0; - if !should_dispatch_concurrently(writes.len()) { - for write in writes { - let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED - .with_label_values(&["flush_write_region"]) - .start_timer(); - affected_rows += node_manager - .handle(&write.datanode, write.request) - .await? - .affected_rows; - } - return Ok(affected_rows); - } - - let write_futures = writes.into_iter().map(|write| async move { - let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED - .with_label_values(&["flush_write_region"]) - .start_timer(); - - let response = node_manager.handle(&write.datanode, write.request).await?; - Ok::<_, Error>(response.affected_rows) - }); - - // todo(hl): should be bounded. - let affected_rows = futures::future::try_join_all(write_futures) - .await? - .into_iter() - .sum(); - Ok(affected_rows) -} - async fn flush_batch_with_managers( flush: FlushBatch, partition_manager: PartitionRuleManagerRef, - node_manager: NodeManagerRef, + write_dispatcher: WriteDispatcherRef, catalog_manager: CatalogManagerRef, flow_notification_tx: mpsc::Sender, ) { let partition_provider = PartitionManagerPhysicalFlushAdapter { partition_manager }; - let node_requester = NodeManagerPhysicalFlushAdapter { - node_manager: node_manager.clone(), - }; + let dispatcher = WriteDispatcherPhysicalFlushAdapter { write_dispatcher }; let catalog_provider = CatalogManagerPhysicalFlushAdapter { catalog_manager }; flush_batch( flush, &partition_provider, - &node_requester, + &dispatcher, &catalog_provider, flow_notification_tx, ) @@ -1342,7 +1268,7 @@ async fn flush_batch_with_managers( async fn flush_batch( flush: FlushBatch, partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized), - node_manager: &(impl PhysicalFlushNodeRequester + ?Sized), + write_dispatcher: &(impl PhysicalFlushWriteDispatcher + ?Sized), catalog_manager: &(impl PhysicalFlushCatalogProvider + ?Sized), flow_notification_tx: mpsc::Sender, ) { @@ -1366,7 +1292,7 @@ async fn flush_batch( &physical_table_name, &ctx, partition_manager, - node_manager, + write_dispatcher, catalog_manager, ) .await; @@ -1553,7 +1479,7 @@ pub async fn flush_batch_physical( physical_table_name: &str, ctx: &QueryContextRef, partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized), - node_manager: &(impl PhysicalFlushNodeRequester + ?Sized), + write_dispatcher: &(impl PhysicalFlushWriteDispatcher + ?Sized), catalog_manager: &(impl PhysicalFlushCatalogProvider + ?Sized), ) -> Result { // 1. Resolve the physical table and get column ID mapping @@ -1588,7 +1514,7 @@ pub async fn flush_batch_physical( })?; // 2. Get the physical table's partition rule (one lookup instead of N) - let partition_rule = { + let (partition_rule, partition_versions) = { let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED .with_label_values(&["flush_physical_fetch_partition_rule"]) .start_timer(); @@ -1609,16 +1535,16 @@ pub async fn flush_batch_physical( // 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( - combined_batch, - physical_table_id, - partition_rule.as_ref(), - partition_columns, - )?; + let planned_batches = + plan_region_batches(combined_batch, physical_table_id, partition_rule.as_ref())?; - let resolved_batches = resolve_region_targets(planned_batches, partition_manager).await?; - let region_writes = encode_region_write_requests(resolved_batches)?; - flush_region_writes_concurrently(node_manager, region_writes).await + let mut region_writes = encode_region_write_requests(planned_batches, &partition_versions)?; + for request in &mut region_writes { + if let Some(header) = &mut request.header { + header.dbname = ctx.get_db_string(); + } + } + write_dispatcher.dispatch(region_writes).await } /// Transforms logical table batches into physical format (sparse primary key encoding). @@ -1710,16 +1636,6 @@ fn split_combined_batch_by_region( Ok(map) } -fn prepare_physical_region_routing_batch( - combined_batch: RecordBatch, - partition_columns: &[String], -) -> Result { - if partition_columns.is_empty() { - return Ok(combined_batch); - } - strip_partition_columns_from_batch(combined_batch) -} - fn plan_region_batch( stripped_batch: &RecordBatch, physical_table_id: TableId, @@ -1754,15 +1670,15 @@ fn plan_region_batches( combined_batch: RecordBatch, physical_table_id: TableId, partition_rule: &dyn partition::partition::PartitionRule, - partition_columns: &[String], ) -> Result> { let region_masks = split_combined_batch_by_region(&combined_batch, partition_rule)?; - let stripped_batch = prepare_physical_region_routing_batch(combined_batch, partition_columns)?; + // Retain partition columns in the RPC payload so routing can be rebuilt after repartition. + // Mito's sparse bulk conversion projects the stored fields and encoded primary key. let mut planned_batches = Vec::new(); for (region_number, mask) in region_masks { if let Some(planned_batch) = - plan_region_batch(&stripped_batch, physical_table_id, region_number, &mask)? + plan_region_batch(&combined_batch, physical_table_id, region_number, &mask)? { planned_batches.push(planned_batch); } @@ -1771,38 +1687,18 @@ fn plan_region_batches( Ok(planned_batches) } -async fn resolve_region_targets( - planned_batches: Vec, - partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized), -) -> Result> { - let mut resolved_batches = Vec::with_capacity(planned_batches.len()); - for planned in planned_batches { - let datanode = { - let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED - .with_label_values(&["flush_physical_resolve_region_leader"]) - .start_timer(); - partition_manager - .find_region_leader(planned.region_id) - .await? - }; - - resolved_batches.push(ResolvedRegionBatch { planned, datanode }); - } - - Ok(resolved_batches) -} - fn encode_region_write_requests( - resolved_batches: Vec, -) -> Result> { - let mut region_writes = Vec::with_capacity(resolved_batches.len()); - for resolved in resolved_batches { - let region_id = resolved.planned.region_id; + planned_batches: Vec, + partition_versions: &HashMap>, +) -> Result> { + let mut region_writes = Vec::with_capacity(planned_batches.len()); + for planned in planned_batches { + let region_id = planned.region_id; let (schema_bytes, data_header, payload) = { let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED .with_label_values(&["flush_physical_encode_ipc"]) .start_timer(); - record_batch_to_ipc(resolved.planned.batch)? + flight::record_batch_to_ipc(planned.batch).map_err(map_ipc_error)? }; let request = RegionRequest { @@ -1812,7 +1708,11 @@ fn encode_region_write_requests( }), body: Some(region_request::Body::BulkInsert(BulkInsertRequest { 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, @@ -1824,10 +1724,7 @@ fn encode_region_write_requests( })), }; - region_writes.push(FlushRegionWrite { - datanode: resolved.datanode, - request, - }); + region_writes.push(request); } Ok(region_writes) @@ -1844,28 +1741,14 @@ fn notify_waiters(waiters: Vec, result: Result<()>) { } } -fn record_batch_to_ipc(record_batch: RecordBatch) -> Result<(Bytes, Bytes, Bytes)> { - let mut encoder = FlightEncoder::default(); - let schema = encoder.encode_schema(record_batch.schema().as_ref()); - let mut iter = encoder - .encode(FlightMessage::RecordBatch(record_batch)) - .into_iter(); - let Some(flight_data) = iter.next() else { - return Err(Error::Internal { - err_msg: "Failed to encode empty flight data".to_string(), - }); - }; - if iter.next().is_some() { - return Err(Error::NotSupported { - feat: "bulk insert RecordBatch with dictionary arrays".to_string(), - }); +fn map_ipc_error(error: GrpcError) -> Error { + match error { + GrpcError::NotSupported { feat } => Error::NotSupported { feat }, + GrpcError::InvalidFlightData { reason, .. } => Error::Internal { err_msg: reason }, + error => Error::Internal { + err_msg: error.to_string(), + }, } - - Ok(( - schema.data_header, - flight_data.data_header, - flight_data.data_body, - )) } #[cfg(test)] @@ -1876,21 +1759,21 @@ mod tests { use std::sync::{Arc, Mutex}; use std::time::Duration; - use api::region::RegionResponse; use api::v1::flow::{DirtyWindowRequests, FlowRequest, FlowResponse}; use api::v1::meta::Peer; - use api::v1::region::{InsertRequests, RegionRequest, region_request}; + use api::v1::region::{InsertRequests, RegionRequest, bulk_insert_request, region_request}; use api::v1::value::ValueData; use api::v1::{ - ColumnDataType, ColumnSchema, Row, RowInsertRequest, RowInsertRequests, Rows, SemanticType, - Value, + ColumnDataType, ColumnSchema, PartitionExprVersion, Row, RowInsertRequest, + RowInsertRequests, Rows, SemanticType, Value, }; - use arrow::array::{BinaryArray, BooleanArray, StringArray, TimestampMillisecondArray}; + 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 catalog::error::Result as CatalogResult; use catalog::memory::MemoryCatalogManager; + use common_grpc::flight::FlightDecoder; use common_meta::cache::{ TableFlownodeSetCacheRef, new_table_flownode_set_cache, new_table_route_cache, }; @@ -1899,16 +1782,13 @@ mod tests { use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::kv_backend::{KvBackend, TxnService}; use common_meta::node_manager::{ - Datanode, DatanodeManager, DatanodeRef, Flownode, FlownodeManager, FlownodeRef, - NodeManagerRef, + DatanodeManager, DatanodeRef, Flownode, FlownodeManager, FlownodeRef, NodeManagerRef, }; use common_meta::rpc::store::{ BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, BatchPutRequest, BatchPutResponse, DeleteRangeRequest, DeleteRangeResponse, PutRequest, PutResponse, RangeRequest, RangeResponse, }; - use common_query::request::QueryRequest; - use common_recordbatch::SendableRecordBatchStream; use dashmap::DashMap; use datatypes::schema::{ColumnSchema as DtColumnSchema, Schema as DtSchema}; use moka::future::CacheBuilder; @@ -1917,23 +1797,20 @@ mod tests { use partition::manager::PartitionRuleManager; use partition::partition::{PartitionRule, PartitionRuleRef, RegionMask}; use smallvec::SmallVec; - use snafu::ResultExt; use store_api::storage::RegionId; use table::metadata::TableId; use table::test_util::table_info::test_table_info; use tokio::sync::{Notify, Semaphore, broadcast, mpsc, oneshot}; - use tokio::time::{advance, sleep}; + use tokio::time::advance; use super::{ - BatchKey, Error, FlushBatch, FlushRegionWrite, FlushWaiter, PendingBatch, - PendingRowsBatcher, PendingWorker, PhysicalFlushCatalogProvider, - PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider, PhysicalTableMetadata, - PlannedRegionBatch, RecordBatchWithTsIdx, ResolvedRegionBatch, TableBatch, WorkerCommand, + BatchKey, Error, FlushBatch, FlushWaiter, PendingBatch, PendingRowsBatcher, PendingWorker, + PhysicalFlushCatalogProvider, PhysicalFlushPartitionProvider, PhysicalFlushWriteDispatcher, + PhysicalTableMetadata, PlannedRegionBatch, RecordBatchWithTsIdx, TableBatch, WorkerCommand, columns_taxonomy, drain_batch, encode_region_write_requests, extract_timestamps, - flush_batch, flush_batch_physical, flush_region_writes_concurrently, greptime_timestamp, + flush_batch, flush_batch_physical, greptime_timestamp, notify_flow_dirty_windows_after_flush, plan_region_batches, remove_worker_if_same_channel, - should_close_worker_on_idle_timeout, should_dispatch_concurrently, - start_flow_notification_worker, start_worker, strip_partition_columns_from_batch, + should_close_worker_on_idle_timeout, start_flow_notification_worker, start_worker, transform_logical_batches_to_physical, try_enqueue_flow_notification, }; use crate::error; @@ -2261,45 +2138,8 @@ mod tests { } } - struct TwoRegionPartitionRule { - partition_columns: Vec, - } - - impl PartitionRule for TwoRegionPartitionRule { - fn as_any(&self) -> &dyn std::any::Any { - self - } - - fn partition_columns(&self) -> &[String] { - &self.partition_columns - } - - fn find_region( - &self, - _values: &[datatypes::prelude::Value], - ) -> partition::error::Result { - unimplemented!() - } - - fn split_record_batch( - &self, - _record_batch: &RecordBatch, - ) -> partition::error::Result> - { - Ok(HashMap::from([ - (1, RegionMask::new(BooleanArray::from(vec![true, false]), 1)), - (2, RegionMask::new(BooleanArray::from(vec![false, true]), 1)), - ( - 3, - RegionMask::new(BooleanArray::from(vec![false, false]), 0), - ), - ])) - } - } - struct MockFlushPartitionProvider { partition_rule_calls: Arc, - region_leader_calls: Arc, } #[async_trait] @@ -2307,40 +2147,28 @@ mod tests { async fn find_table_partition_rule( &self, _table_info: &table::metadata::TableInfo, - ) -> PartitionResult { + ) -> PartitionResult<(PartitionRuleRef, HashMap>)> { self.partition_rule_calls.fetch_add(1, Ordering::SeqCst); - Ok(Arc::new(SingleRegionPartitionRule)) - } - - async fn find_region_leader(&self, _region_id: RegionId) -> error::Result { - self.region_leader_calls.fetch_add(1, Ordering::SeqCst); - Ok(Peer { - id: 1, - addr: "node-1".to_string(), - }) + Ok((Arc::new(SingleRegionPartitionRule), HashMap::new())) } } #[derive(Default)] - struct MockFlushNodeRequester { + struct MockFlushWriteDispatcher { writes: Arc, fail: bool, } #[async_trait] - impl PhysicalFlushNodeRequester for MockFlushNodeRequester { - async fn handle( - &self, - _peer: &Peer, - _request: RegionRequest, - ) -> error::Result { + impl PhysicalFlushWriteDispatcher for MockFlushWriteDispatcher { + async fn dispatch(&self, _requests: Vec) -> error::Result { self.writes.fetch_add(1, Ordering::SeqCst); if self.fail { return Err(Error::Internal { err_msg: "physical write failed".to_string(), }); } - Ok(RegionResponse::new(0)) + Ok(0) } } @@ -2427,44 +2255,6 @@ mod tests { assert_eq!("cpu", pending_batch.tables[&43].table_name); } - #[derive(Clone)] - struct ConcurrentMockDatanode { - delay: Duration, - inflight: Arc, - max_inflight: Arc, - } - - #[async_trait] - impl Datanode for ConcurrentMockDatanode { - async fn handle(&self, _request: RegionRequest) -> MetaResult { - let now = self.inflight.fetch_add(1, Ordering::SeqCst) + 1; - loop { - let max = self.max_inflight.load(Ordering::SeqCst); - if now <= max { - break; - } - if self - .max_inflight - .compare_exchange(max, now, Ordering::SeqCst, Ordering::SeqCst) - .is_ok() - { - break; - } - } - - sleep(self.delay).await; - self.inflight.fetch_sub(1, Ordering::SeqCst); - Ok(RegionResponse::new(0)) - } - - async fn handle_query( - &self, - _request: QueryRequest, - ) -> MetaResult { - unimplemented!() - } - } - #[derive(Clone)] struct ConcurrentMockNodeManager { datanodes: Arc>, @@ -2804,9 +2594,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, }, @@ -2865,9 +2654,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, }, @@ -2886,21 +2674,6 @@ mod tests { ); } - #[async_trait] - impl PhysicalFlushNodeRequester for ConcurrentMockNodeManager { - async fn handle( - &self, - peer: &Peer, - request: RegionRequest, - ) -> error::Result { - let datanode = self.datanode(peer).await; - datanode - .handle(request) - .await - .context(error::CommonMetaSnafu) - } - } - #[test] fn test_remove_worker_if_same_channel_removes_matching_entry() { let workers = DashMap::new(); @@ -3054,8 +2827,11 @@ mod tests { workers.clone(), worker_rx, shutdown.clone(), - partition_manager, - node_manager, + partition_manager.clone(), + Arc::new(operator::write_dispatcher::WriteDispatcher::new( + node_manager, + partition_manager, + )), catalog_manager, flow_notification_tx, flush_interval, @@ -3110,118 +2886,6 @@ mod tests { ); } - #[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_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![ @@ -3553,8 +3217,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 node = MockFlushWriteDispatcher::default(); let ctx = session::context::QueryContext::arc(); flush_batch_physical( @@ -3563,7 +3226,6 @@ mod tests { &ctx, &MockFlushPartitionProvider { partition_rule_calls: partition_calls.clone(), - region_leader_calls: leader_calls.clone(), }, &node, &MockFlushCatalogProvider { @@ -3574,23 +3236,18 @@ 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)); } #[derive(Default)] - struct AffectedRowsFlushNodeRequester { + struct AffectedRowsFlushWriteDispatcher { affected_rows: usize, } #[async_trait] - impl PhysicalFlushNodeRequester for AffectedRowsFlushNodeRequester { - async fn handle( - &self, - _peer: &Peer, - _request: RegionRequest, - ) -> error::Result { - Ok(RegionResponse::new(self.affected_rows)) + impl PhysicalFlushWriteDispatcher for AffectedRowsFlushWriteDispatcher { + async fn dispatch(&self, _requests: Vec) -> error::Result { + Ok(self.affected_rows) } } @@ -3610,9 +3267,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)), }, @@ -3632,8 +3288,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 node = MockFlushWriteDispatcher::default(); let ctx = session::context::QueryContext::arc(); let err = flush_batch_physical( @@ -3642,7 +3297,6 @@ mod tests { &ctx, &MockFlushPartitionProvider { partition_rule_calls: partition_calls.clone(), - region_leader_calls: leader_calls.clone(), }, &node, &MockFlushCatalogProvider { table: None }, @@ -3655,7 +3309,6 @@ 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)); } @@ -3676,8 +3329,7 @@ mod tests { }, ]; let partition_calls = Arc::new(AtomicUsize::new(0)); - let leader_calls = Arc::new(AtomicUsize::new(0)); - let node = MockFlushNodeRequester::default(); + let node = MockFlushWriteDispatcher::default(); let ctx = session::context::QueryContext::arc(); let err = flush_batch_physical( @@ -3686,7 +3338,6 @@ mod tests { &ctx, &MockFlushPartitionProvider { partition_rule_calls: partition_calls.clone(), - region_leader_calls: leader_calls.clone(), }, &node, &MockFlushCatalogProvider { @@ -3698,15 +3349,12 @@ 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)); } - #[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), @@ -3714,22 +3362,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 = super::transform_logical_batches_to_physical( + &table_batches, + &HashMap::from([("host".to_string(), 1), ("zone".to_string(), 2)]), + &HashSet::from(["host"]), + ) + .unwrap(); + let combined_batch = super::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()); @@ -3737,14 +3401,88 @@ 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, + &HashMap::from([(1, Some(7)), (2, Some(8))]), + ) + .unwrap(); + for (index, write) in writes.into_iter().enumerate() { + let Some(region_request::Body::BulkInsert(request)) = write.body else { + panic!("expected bulk") + }; + assert_eq!( + request.partition_expr_version.unwrap().value, + 7 + index as u64 + ); + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = request.body else { + panic!("expected IPC") + }; + let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema) + .unwrap() + .try_decode_record_batch(&ipc.data_header, &ipc.payload) + .unwrap(); + // The partition-only Arrow projection must retain the complete sparse key, + // including the zone tag that a later partition rule can require. + assert!(batch.column_by_name("zone").is_none()); + assert_eq!(batch, expected.slice(index, 1)); + let tags = batch + .column_by_name("host") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(tags.len(), 1); + assert_eq!(tags.value(0), if index == 0 { "node-1" } else { "node-2" }); + } } #[test] fn test_encode_region_write_requests_builds_bulk_insert_requests() { + check_encode_region_write_requests(); + } + + struct TwoRegionPartitionRule { + partition_columns: Vec, + } + + impl PartitionRule for TwoRegionPartitionRule { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn partition_columns(&self) -> &[String] { + &self.partition_columns + } + + fn find_region( + &self, + _values: &[datatypes::prelude::Value], + ) -> partition::error::Result { + unimplemented!() + } + + fn split_record_batch( + &self, + _record_batch: &RecordBatch, + ) -> partition::error::Result> + { + Ok(HashMap::from([ + (1, RegionMask::new(BooleanArray::from(vec![true, false]), 1)), + (2, RegionMask::new(BooleanArray::from(vec![false, true]), 1)), + ( + 3, + RegionMask::new(BooleanArray::from(vec![false, false]), 0), + ), + ])) + } + } + + fn check_encode_region_write_requests() { let planned_batch = PlannedRegionBatch { region_id: RegionId::new(1024, 1), batch: RecordBatch::try_new( @@ -3765,20 +3503,18 @@ 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]).unwrap(); + let writes = + encode_region_write_requests(vec![planned_batch], &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!( + Some(PartitionExprVersion { value: 7 }), + request.partition_expr_version + ); } } diff --git a/src/standalone/src/datanode_manager.rs b/src/standalone/src/datanode_manager.rs index 13a17ea2b43..a27a94b3c8d 100644 --- a/src/standalone/src/datanode_manager.rs +++ b/src/standalone/src/datanode_manager.rs @@ -77,7 +77,17 @@ impl RegionInvoker { #[async_trait] impl Datanode for RegionInvoker { - async fn handle(&self, request: RegionRequest) -> MetaResult { + async fn handle_request(&self, request: Arc) -> MetaResult { + // Local execution consumes the request; standalone does not retain it for retries. + let request = Arc::try_unwrap(request) + .map_err(|_| { + InvalidRegionRequestSnafu { + reason: "local region request has multiple owners", + } + .build() + }) + .map_err(BoxedError::new) + .context(meta_error::ExternalSnafu)?; let span = request .header .as_ref() diff --git a/src/standalone/src/options.rs b/src/standalone/src/options.rs index 9f20d73952a..b0dbec1271e 100644 --- a/src/standalone/src/options.rs +++ b/src/standalone/src/options.rs @@ -25,7 +25,7 @@ use flow::FlowConfig; use frontend::frontend::FrontendOptions; use frontend::service_config::{ InfluxdbOptions, JaegerOptions, MysqlOptions, OpentsdbOptions, OtlpOptions, PostgresOptions, - PromStoreOptions, + PromStoreOptions, WriteRetryOptions, }; use mito2::config::MitoConfig; use pipeline::PipelineOptions; @@ -162,6 +162,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, @@ -210,6 +215,18 @@ mod tests { use super::*; + #[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_event_recorder_event_types_preserve_filter_semantics() { let all: StandaloneOptions = toml::from_str("").unwrap(); diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index 9d12848ce61..7f36c74cb7b 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -674,6 +674,7 @@ async fn setup_test_prom_app_with_frontend_inner( PendingRowsBatcher::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,