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,