From f866ef2bd9853d990bcbb8623aa71713a82e84ff Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 21 Oct 2026 05:51:15 +0000 Subject: [PATCH] feat: return per-entry region write failures (#9728) * feat: return per-entry region write failures Signed-off-by: WenyXu * fix: preserve retry hints for single-entry region writes Signed-off-by: WenyXu * fix: preserve region write error compatibility Signed-off-by: WenyXu * refactor: simplify region partial result handling Signed-off-by: WenyXu * fix: align partial failure summaries with input order Signed-off-by: WenyXu * refactor: separate region write response handling Signed-off-by: WenyXu * test: fix partial result expectations and region map guard Signed-off-by: WenyXu * chore: limit lockfile changes to greptime proto Signed-off-by: WenyXu * fix: make partial region results an explicit client opt-in Signed-off-by: WenyXu * chore: pin greptime-proto to merged partial-result changes Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- Cargo.lock | 2 +- Cargo.toml | 2 +- src/api/src/region.rs | 29 +- src/client/src/error.rs | 56 +- src/client/src/lib.rs | 12 +- src/client/src/region.rs | 477 ++++++++- src/common/error/src/status_code.rs | 7 +- src/common/meta/src/error.rs | 13 + src/datanode/src/error.rs | 5 + src/datanode/src/region_server.rs | 1013 ++++++++++++++---- src/datanode/src/region_server/flow_write.rs | 1 + src/meta-srv/src/procedure/utils.rs | 2 + src/metric-engine/AGENTS.md | 1 + src/metric-engine/src/engine.rs | 17 +- src/metric-engine/src/engine/put.rs | 298 +++--- src/operator/src/write_dispatcher.rs | 5 +- src/operator/src/write_dispatcher/tests.rs | 19 +- src/servers/src/error.rs | 1 + src/servers/src/mysql/writer.rs | 4 +- src/servers/src/postgres/types/error.rs | 3 +- src/standalone/Cargo.toml | 2 + src/standalone/src/datanode_manager.rs | 69 +- src/store-api/src/region_request.rs | 74 +- tests/runner/src/client.rs | 1 + 24 files changed, 1647 insertions(+), 466 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 73bd260e9de..7a8da222c03 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6083,7 +6083,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=9533078d6fd4bdaddbca0750cb9bf7ad64b2e568#9533078d6fd4bdaddbca0750cb9bf7ad64b2e568" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=acb247415b8f4c2f0750783dab4e1f38c71f5860#acb247415b8f4c2f0750783dab4e1f38c71f5860" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", diff --git a/Cargo.toml b/Cargo.toml index 3cb4e77fb77..20582d8eaa4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -162,7 +162,7 @@ fs2 = "0.4" fst = "0.4.7" futures = "0.3" futures-util = "0.3" -greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "9533078d6fd4bdaddbca0750cb9bf7ad64b2e568" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "acb247415b8f4c2f0750783dab4e1f38c71f5860" } hdfs-native = "0.14.6" hex = "0.4" hostname = "0.4.0" diff --git a/src/api/src/region.rs b/src/api/src/region.rs index e7000cd7448..fb0d5b5e3f4 100644 --- a/src/api/src/region.rs +++ b/src/api/src/region.rs @@ -15,7 +15,19 @@ use std::collections::HashMap; use common_base::AffectedRows; -use greptime_proto::v1::region::RegionResponse as RegionResponseV1; +use greptime_proto::v1::region::{RegionRequestHeader, RegionResponse as RegionResponseV1}; + +/// Request extension advertising support for per-entry write failures. +pub const ACCEPT_PARTIAL_RESULT: &str = "accept_partial_result"; + +/// Advertises support for per-entry write failures in the request header. +pub fn set_accept_partial_result(header: &mut RegionRequestHeader) { + header + .query_context + .get_or_insert_with(Default::default) + .extensions + .insert(ACCEPT_PARTIAL_RESULT.to_string(), "true".to_string()); +} /// This result struct is derived from [RegionResponseV1] #[derive(Debug)] @@ -25,6 +37,21 @@ pub struct RegionResponse { pub metadata: Vec, } +/// Successful entries and failed-entry errors indexed by input order. +#[derive(Debug)] +pub struct BatchRegionResponse { + pub affected_rows: AffectedRows, + pub extensions: HashMap>, + pub failures: Vec<(usize, E)>, +} + +impl BatchRegionResponse { + /// Returns the failure with the smallest original request index. + pub fn first_failure(&self) -> Option<&(usize, E)> { + self.failures.iter().min_by_key(|(index, _)| *index) + } +} + impl RegionResponse { pub fn from_region_response(region_response: RegionResponseV1) -> Self { Self { diff --git a/src/client/src/error.rs b/src/client/src/error.rs index fdf0df9582f..d9fa8c96050 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, location}; +use snafu::{Location, OptionExt, Snafu, location}; use tonic::Code; use tonic::metadata::errors::InvalidMetadataValue; @@ -26,6 +26,15 @@ use tonic::metadata::errors::InvalidMetadataValue; #[snafu(visibility(pub))] #[stack_trace_debug] pub enum Error { + #[snafu(display( + "Region request has {} failed entries; {}", + response.failures.len(), + response.first_failure().map(|(index, error)| format!("entry {index}: {}", error.output_msg())).unwrap_or_default() + ))] + PartialRegion { + response: Box>, + }, + #[snafu(display("Illegal Flight messages, reason: {}", reason))] IllegalFlightMessages { reason: String, @@ -109,11 +118,12 @@ pub enum Error { location: Location, }, - // Server error carried in Tonic Status's metadata. + // Server error carried in a protobuf response status. #[snafu(display("{}", msg))] Server { code: StatusCode, msg: String, + retry_hint: RetryHint, #[snafu(implicit)] location: Location, }, @@ -163,6 +173,9 @@ pub type Result = std::result::Result; impl ErrorExt for Error { fn status_code(&self) -> StatusCode { match self { + Error::PartialRegion { response } => response + .first_failure() + .map_or(StatusCode::Internal, |(_, error)| error.status_code()), Error::IllegalFlightMessages { .. } | Error::MissingField { .. } | Error::IllegalDatabaseResponse { .. } => StatusCode::Internal, @@ -188,7 +201,8 @@ impl ErrorExt for Error { fn retry_hint(&self) -> RetryHint { match self { - Error::Tonic { retry_hint, .. } => *retry_hint, + Error::PartialRegion { .. } => RetryHint::Retryable, + Error::Tonic { retry_hint, .. } | Error::Server { retry_hint, .. } => *retry_hint, Error::FlightGet { source, .. } | Error::FlightStream { source, .. } | Error::FlowServer { source, .. } @@ -220,9 +234,22 @@ 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!(), + match error { + Error::PartialRegion { response } => Self::PartialRegion { + response: Box::new(api::region::BatchRegionResponse { + affected_rows: response.affected_rows, + extensions: response.extensions, + failures: response + .failures + .into_iter() + .map(|(index, error)| (index, error.into())) + .collect(), + }), + }, + error => Self::External { + source: BoxedError::new(error), + location: location!(), + }, } } } @@ -258,3 +285,20 @@ impl Error { ) } } + +impl Error { + /// Decodes an application error carried in a protobuf response. + pub(crate) fn try_from_response_status(status: &api::v1::Status) -> Result { + let code = StatusCode::from_u32(status.status_code).with_context(|| { + IllegalDatabaseResponseSnafu { + err_msg: format!("Unknown server status: {}", status.status_code), + } + })?; + Ok(ServerSnafu { + code, + msg: status.err_msg.clone(), + retry_hint: status.retry_hint.parse().unwrap_or(RetryHint::NonRetryable), + } + .build()) + } +} diff --git a/src/client/src/lib.rs b/src/client/src/lib.rs index 4b9d543e35f..208a3233ff7 100644 --- a/src/client/src/lib.rs +++ b/src/client/src/lib.rs @@ -35,7 +35,7 @@ use snafu::OptionExt; pub use self::client::{Client, ClientOptions}; pub use self::database::{Database, OutputMetrics, OutputWithMetrics}; pub use self::error::{Error, Result}; -use crate::error::{IllegalDatabaseResponseSnafu, ServerSnafu}; +use crate::error::IllegalDatabaseResponseSnafu; pub fn from_grpc_response(response: GreptimeResponse) -> Result { let header = response.header.context(IllegalDatabaseResponseSnafu { @@ -53,14 +53,6 @@ pub fn from_grpc_response(response: GreptimeResponse) -> Result { Response::AffectedRows(AffectedRows { value }) => Ok(value), } } else { - let status_code = - StatusCode::from_u32(status.status_code).context(IllegalDatabaseResponseSnafu { - err_msg: format!("invalid status: {:?}", status), - })?; - ServerSnafu { - code: status_code, - msg: status.err_msg, - } - .fail() + Err(Error::try_from_response_status(&status)?) } } diff --git a/src/client/src/region.rs b/src/client/src/region.rs index f9ebbb0d0bf..2e517eb30d0 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -18,8 +18,9 @@ use std::borrow::Borrow; use std::sync::Arc; use std::time::Duration; -use api::region::RegionResponse; +use api::region::{BatchRegionResponse, RegionResponse}; use api::v1::ResponseHeader; +use api::v1::region::RegionResponse as PbRegionResponse; use api::v1::region::{ RegionRequest, RegionRequestHeader, RemoteDynFilterRequest, RemoteDynFilterUnregister, RemoteDynFilterUpdate, region_request, remote_dyn_filter_request, @@ -49,7 +50,7 @@ use tonic::codec::CompressionEncoding; use crate::error::{ self, FlightGetSnafu, IllegalDatabaseResponseSnafu, IllegalFlightMessagesSnafu, - MissingFieldSnafu, Result, ServerSnafu, + MissingFieldSnafu, Result, }; use crate::flight::{FlightMessageReader, decode_flight_data}; use crate::region::codec::RegionCodec; @@ -141,7 +142,8 @@ impl RegionRequester { /// /// 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. + /// whether replaying the request is safe. Callers accepting partial results in the + /// request header must inspect failed entries before retrying. pub async fn handle(&self, request: R) -> Result where R: Borrow + Send + Sync + 'static, @@ -180,6 +182,11 @@ impl RegionRequester { .body .as_ref() .context(MissingFieldSnafu { field: "body" })?; + let request_count = match request_body { + region_request::Body::Inserts(requests) => Some(requests.requests.len()), + region_request::Body::Deletes(requests) => Some(requests.requests.len()), + _ => None, + }; let is_insert = matches!( request_body, region_request::Body::Inserts(_) | region_request::Body::BulkInsert(_) @@ -199,6 +206,13 @@ impl RegionRequester { } } + let accept_partial_result = request + .borrow() + .header + .as_ref() + .and_then(|header| header.query_context.as_ref()) + .and_then(|context| context.extensions.get(api::region::ACCEPT_PARTIAL_RESULT)) + .is_some_and(|value| value == "true"); let response = async { client .ready() @@ -223,9 +237,10 @@ impl RegionRequester { .map_err(|status| map_region_error(status, addr))? .into_inner(); - check_response_header(&response.header)?; - - Ok(RegionResponse::from_region_response(response)) + match request_count { + Some(count) if accept_partial_result => decode_write_response(response, count), + _ => decode_region_response(response), + } } } @@ -385,6 +400,7 @@ fn build_remote_dyn_filter_request( } } +/// Checks that the response contains a successful application status. pub fn check_response_header(header: &Option) -> Result<()> { let status = header .as_ref() @@ -396,35 +412,107 @@ pub fn check_response_header(header: &Option) -> Result<()> { if StatusCode::is_success(status.status_code) { Ok(()) } else { - let code = - StatusCode::from_u32(status.status_code).context(IllegalDatabaseResponseSnafu { - err_msg: format!("unknown server status: {:?}", status), - })?; - ServerSnafu { - code, - msg: status.err_msg.clone(), - } - .fail() + Err(error::Error::try_from_response_status(status)?) } } +/// Decodes an ordinary region response without per-entry write failures. +pub fn decode_region_response(response: PbRegionResponse) -> Result { + if !response.failures.is_empty() { + return IllegalDatabaseResponseSnafu { + err_msg: "Failure details are only valid for inserts/deletes", + } + .fail(); + } + check_response_header(&response.header)?; + Ok(RegionResponse::from_region_response(response)) +} + +/// Decodes an inserts/deletes response, validating failures against the original entry count. +/// A failed single-entry request returns its underlying error. +fn decode_write_response(response: PbRegionResponse, entry_count: usize) -> Result { + if response.failures.is_empty() { + if response + .header + .as_ref() + .and_then(|header| header.status.as_ref()) + .is_some_and(|status| status.status_code == StatusCode::PartialFailure as u32) + { + return IllegalDatabaseResponseSnafu { + err_msg: "Partial failure has no failed entries", + } + .fail(); + } + return decode_region_response(response); + } + let status = response + .header + .as_ref() + .and_then(|header| header.status.as_ref()) + .context(IllegalDatabaseResponseSnafu { + err_msg: "Missing region response status", + })?; + if status.status_code != StatusCode::PartialFailure as u32 { + return IllegalDatabaseResponseSnafu { + err_msg: "Failure details require a partial-failure status", + } + .fail(); + } + let mut seen = std::collections::HashSet::with_capacity(response.failures.len()); + let mut failures = Vec::with_capacity(response.failures.len()); + for error in response.failures { + let index = error.request_index as usize; + if index >= entry_count || !seen.insert(index) { + return IllegalDatabaseResponseSnafu { + err_msg: "Invalid or duplicate batch error index", + } + .fail(); + } + let status = error.status.context(IllegalDatabaseResponseSnafu { + err_msg: "Missing batch error status", + })?; + if StatusCode::is_success(status.status_code) { + return IllegalDatabaseResponseSnafu { + err_msg: "Batch error contains a success status", + } + .fail(); + } + failures.push((index, error::Error::try_from_response_status(&status)?)); + } + // A single failed entry has no successful peers to preserve. + if entry_count == 1 + && let Some((_, error)) = failures.pop() + { + return Err(error); + } + Err(error::Error::PartialRegion { + response: Box::new(BatchRegionResponse { + affected_rows: response.affected_rows as _, + extensions: response.extensions, + failures, + }), + }) +} + #[cfg(test)] #[allow(deprecated)] mod tests { use std::collections::VecDeque; use std::sync::Mutex; + use api::v1::ResponseHeader; use api::v1::region::region_server::{Region, RegionServer}; use api::v1::region::{ - BulkInsertRequest, InsertRequest, InsertRequests, RegionResponse as PbRegionResponse, - RemoteDynFilterUnregister, RemoteDynFilterUpdate, region_request, - remote_dyn_filter_request, + BulkInsertRequest, InsertRequest, InsertRequests, RegionRequestFailure, + 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_error::status_code::StatusCode; use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_recordbatch::adapter::RecordBatchMetrics; use datatypes::arrow::array::Int32Array; @@ -579,6 +667,20 @@ mod tests { } } + fn with_partial_result(mut request: RegionRequest) -> RegionRequest { + request + .header + .get_or_insert_with(Default::default) + .query_context + .get_or_insert_with(Default::default) + .extensions + .insert( + api::region::ACCEPT_PARTIAL_RESULT.to_string(), + "true".to_string(), + ); + request + } + #[tokio::test] async fn test_owned_and_shared_requests_preserve_payload_and_compression() { let server = RegionTestServer::start().await; @@ -636,6 +738,54 @@ mod tests { } } + #[tokio::test] + async fn test_partial_result_requires_explicit_acceptance() { + let server = RegionTestServer::start().await; + let requester = server.create_requester(ChannelConfig::default(), false, false); + for accepted in [false, true] { + let request = if accepted { + with_partial_result(build_insert_request()) + } else { + build_insert_request() + }; + let shared = Arc::new(request); + server + .service + .responses + .lock() + .unwrap() + .push_back(Ok(PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus { + status_code: StatusCode::PartialFailure as u32, + ..Default::default() + }), + }), + failures: vec![RegionRequestFailure { + request_index: 0, + status: Some(PbStatus { + status_code: StatusCode::RegionNotFound as u32, + retry_hint: RetryHint::RetryAfterRouteRefresh.as_str().to_string(), + ..Default::default() + }), + }], + ..Default::default() + })); + let error = requester.handle(Arc::clone(&shared)).await.unwrap_err(); + if accepted { + assert_eq!(error.status_code(), StatusCode::RegionNotFound); + assert_eq!(error.retry_hint(), RetryHint::RetryAfterRouteRefresh); + } else { + assert!(matches!(error, Error::IllegalDatabaseResponse { .. })); + } + assert_eq!( + &server.service.requests.lock().unwrap().last().unwrap().0, + shared.as_ref() + ); + assert_eq!(Arc::strong_count(&shared), 1); + } + } + #[tokio::test] async fn test_shared_request_can_be_resent_after_transport_error() { let server = RegionTestServer::start().await; @@ -762,7 +912,10 @@ mod tests { assert_eq!(error.tonic_code(), baseline.tonic_code()); } let datanode: &dyn Datanode = &requester; - let error = datanode.handle(build_insert_request()).await.unwrap_err(); + let error = datanode + .handle_request(Arc::new(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 { .. })); @@ -832,10 +985,16 @@ mod tests { body: Some(body), }; let shared = Arc::new(request); - let shared_error = datanode.handle(Arc::clone(&shared)).await.unwrap_err(); + let shared_error = datanode + .handle_request(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(); + let owned_error = datanode + .handle_request(Arc::new(request)) + .await + .unwrap_err(); for error in [owned_error, shared_error] { assert_eq!( error.retry_hint(), @@ -1053,8 +1212,281 @@ mod tests { .unwrap() } + #[tokio::test] + async fn test_partial_response_validation_and_compression() { + let server = RegionTestServer::start().await; + let requester = server.create_requester(ChannelConfig::default(), true, true); + let mut request = build_insert_request(); + if let Some(region_request::Body::Inserts(inserts)) = &mut request.body { + inserts.requests.push(inserts.requests[0].clone()); + } + let request = Arc::new(with_partial_result(request)); + let failure = RegionRequestFailure { + request_index: 0, + status: Some(PbStatus { + status_code: StatusCode::RegionNotFound as u32, + err_msg: "moved".to_string(), + retry_hint: "retry_after_route_refresh".to_string(), + }), + }; + for (failures, valid) in [ + (vec![], true), + (vec![failure.clone()], true), + (vec![failure.clone(), failure.clone()], false), + ( + vec![RegionRequestFailure { + request_index: u32::MAX, + ..failure.clone() + }], + false, + ), + ( + vec![RegionRequestFailure { + request_index: 0, + status: None, + }], + false, + ), + ( + vec![RegionRequestFailure { + request_index: 0, + status: Some(PbStatus::default()), + }], + false, + ), + ( + vec![RegionRequestFailure { + request_index: 0, + status: Some(PbStatus { + status_code: u32::MAX, + ..Default::default() + }), + }], + false, + ), + ] { + let has_error = !failures.is_empty(); + server + .service + .responses + .lock() + .unwrap() + .push_back(Ok(PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus { + status_code: if has_error { + StatusCode::PartialFailure + } else { + StatusCode::Success + } as u32, + ..Default::default() + }), + }), + affected_rows: 3, + extensions: std::collections::HashMap::from([("test".to_string(), vec![1, 2])]), + failures, + ..Default::default() + })); + let response = requester.handle_inner(Arc::clone(&request)).await; + if valid { + if has_error { + let error = response.unwrap_err(); + assert!(error.output_msg().contains("entry 0")); + assert!(error.output_msg().contains("moved")); + let error: common_meta::error::Error = error.into(); + assert!(error.output_msg().contains("moved")); + let common_meta::error::Error::PartialRegion { response } = error else { + panic!("Expected partial region failure"); + }; + assert_eq!(response.affected_rows, 3); + assert_eq!(response.extensions["test"], vec![1, 2]); + assert_eq!( + response.failures[0].1.retry_hint(), + RetryHint::RetryAfterRouteRefresh + ); + } else { + let response = response.unwrap(); + assert_eq!(response.affected_rows, 3); + assert_eq!(response.extensions["test"], vec![1, 2]); + } + } else { + assert!(matches!( + response.unwrap_err(), + Error::IllegalDatabaseResponse { .. } + )); + } + } + for (received, send, accept) in server.service.requests.lock().unwrap().iter() { + assert_eq!(received, request.as_ref()); + assert!(*send && *accept); + } + assert_eq!(Arc::strong_count(&request), 1); + } + + #[test] + fn test_partial_response_envelope() { + let failure = RegionRequestFailure { + request_index: 0, + status: Some(PbStatus { + status_code: StatusCode::RegionNotFound as u32, + retry_hint: RetryHint::RetryAfterRouteRefresh.as_str().to_string(), + ..Default::default() + }), + }; + for (status, failures, count) in [ + (Some(StatusCode::Success), vec![failure.clone()], Some(2)), + (Some(StatusCode::PartialFailure), vec![], Some(2)), + ( + Some(StatusCode::PartialFailure), + vec![failure.clone()], + None, + ), + (None, vec![failure.clone()], Some(2)), + ] { + let response = PbRegionResponse { + header: status.map(|code| ResponseHeader { + status: Some(PbStatus { + status_code: code as u32, + ..Default::default() + }), + }), + failures, + ..Default::default() + }; + assert!(matches!( + match count { + Some(count) => decode_write_response(response, count), + None => decode_region_response(response), + }, + Err(Error::IllegalDatabaseResponse { .. }) + )); + } + let response = PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus { + status_code: StatusCode::PartialFailure as u32, + ..Default::default() + }), + }), + failures: vec![failure], + ..Default::default() + }; + // A partial envelope must not be interpreted as success by the header-only decoder. + assert!(check_response_header(&response.header).is_err()); + let error = decode_write_response(response, 1).unwrap_err(); + assert_eq!(error.status_code(), StatusCode::RegionNotFound); + assert_eq!(error.retry_hint(), RetryHint::RetryAfterRouteRefresh); + } + + #[test] + fn test_partial_error_preserves_status_and_retry_hint() { + let expected_message = "Region request has 2 failed entries; entry 0: first input failed"; + for code in [StatusCode::StorageUnavailable, StatusCode::InvalidArguments] { + let response = PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus { + status_code: StatusCode::PartialFailure as u32, + ..Default::default() + }), + }), + failures: vec![ + RegionRequestFailure { + request_index: 1, + status: Some(PbStatus { + status_code: StatusCode::Internal as u32, + err_msg: "later input failed".to_string(), + ..Default::default() + }), + }, + RegionRequestFailure { + request_index: 0, + status: Some(PbStatus { + status_code: code as u32, + err_msg: "first input failed".to_string(), + retry_hint: RetryHint::NonRetryable.as_str().to_string(), + }), + }, + ], + ..Default::default() + }; + let error = decode_write_response(response, 2).unwrap_err(); + assert_eq!(code, error.status_code()); + assert_eq!(expected_message, error.to_string()); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + let error: common_meta::error::Error = error.into(); + assert_eq!(code, error.status_code()); + assert_eq!(expected_message, error.to_string()); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + let common_meta::error::Error::PartialRegion { response } = error else { + panic!("expected partial result") + }; + assert!( + response + .failures + .iter() + .all(|(_, error)| error.retry_hint() == RetryHint::NonRetryable) + ); + } + } + + #[tokio::test] + async fn test_region_rpc_errors() { + let server = RegionTestServer::start().await; + let requester = server.create_requester(ChannelConfig::default(), false, false); + let request = Arc::new(build_insert_request()); + let mut business_error = Status::unimplemented("unsupported write"); + business_error.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_CODE, + (StatusCode::Unsupported as u32) + .to_string() + .parse() + .unwrap(), + ); + for status in [ + Status::unimplemented("unknown method"), + business_error, + Status::unavailable("lost response"), + Status::deadline_exceeded("timeout"), + ] { + let code = status.code(); + server + .service + .responses + .lock() + .unwrap() + .push_back(Err(status)); + let error = requester + .handle_inner(Arc::clone(&request)) + .await + .unwrap_err(); + assert!(matches!(error, Error::RegionServer { code: actual, .. } if actual == code)); + } + assert_eq!(server.service.requests.lock().unwrap().len(), 4); + } + #[test] fn test_check_response_header() { + for (hint, expected) in [ + ("retryable", RetryHint::Retryable), + ( + "retry_after_route_refresh", + RetryHint::RetryAfterRouteRefresh, + ), + ("non_retryable", RetryHint::NonRetryable), + ("", RetryHint::NonRetryable), + ("future_hint", RetryHint::NonRetryable), + ] { + let error = check_response_header(&Some(ResponseHeader { + status: Some(PbStatus { + status_code: StatusCode::RegionNotFound as u32, + err_msg: "region moved".to_string(), + retry_hint: hint.to_string(), + }), + })) + .unwrap_err(); + assert_eq!(error.retry_hint(), expected, "{hint}"); + } + let result = check_response_header(&None); assert!(matches!( result.unwrap_err(), @@ -1071,6 +1503,7 @@ mod tests { status: Some(PbStatus { status_code: StatusCode::Success as u32, err_msg: String::default(), + ..Default::default() }), })); assert!(result.is_ok()); @@ -1079,6 +1512,7 @@ mod tests { status: Some(PbStatus { status_code: u32::MAX, err_msg: String::default(), + ..Default::default() }), })); assert!(matches!( @@ -1090,6 +1524,7 @@ mod tests { status: Some(PbStatus { status_code: StatusCode::Internal as u32, err_msg: "blabla".to_string(), + ..Default::default() }), })); let Error::Server { code, msg, .. } = result.unwrap_err() else { diff --git a/src/common/error/src/status_code.rs b/src/common/error/src/status_code.rs index b2c99c8e06b..592e40fec45 100644 --- a/src/common/error/src/status_code.rs +++ b/src/common/error/src/status_code.rs @@ -45,6 +45,8 @@ pub enum StatusCode { DeadlineExceeded = 1008, /// Service got suspended for various reason. For example, resources exceed limit. Suspended = 1009, + /// One or more entries failed; inspect per-entry results before retrying. + PartialFailure = 1010, // ====== End of common status code ================ // ====== Begin of SQL related status code ========= @@ -150,7 +152,8 @@ impl StatusCode { | StatusCode::RuntimeResourcesExhausted | StatusCode::External => true, - StatusCode::Success + StatusCode::PartialFailure + | StatusCode::Success | StatusCode::Unsupported | StatusCode::InvalidArguments | StatusCode::InvalidSyntax @@ -300,7 +303,7 @@ macro_rules! define_into_tonic_status { pub fn status_to_tonic_code(status_code: StatusCode) -> Code { match status_code { StatusCode::Success => Code::Ok, - StatusCode::Unknown | StatusCode::External => Code::Unknown, + StatusCode::Unknown | StatusCode::External | StatusCode::PartialFailure => Code::Unknown, StatusCode::Unsupported => Code::Unimplemented, StatusCode::Unexpected | StatusCode::IllegalState diff --git a/src/common/meta/src/error.rs b/src/common/meta/src/error.rs index 3de9e4a3a46..3b8f55e8a4c 100644 --- a/src/common/meta/src/error.rs +++ b/src/common/meta/src/error.rs @@ -41,6 +41,15 @@ pub use retry_hint::{retry_hint_from_postgres_error, retry_hint_from_postgres_po #[snafu(visibility(pub))] #[stack_trace_debug] pub enum Error { + #[snafu(display( + "Region request has {} failed entries; {}", + response.failures.len(), + response.first_failure().map(|(index, error)| format!("entry {index}: {}", error.output_msg())).unwrap_or_default() + ))] + PartialRegion { + response: Box>, + }, + #[snafu(display("Empty key is not allowed"))] EmptyKey { #[snafu(implicit)] @@ -1195,6 +1204,9 @@ impl ErrorExt for Error { fn status_code(&self) -> StatusCode { use Error::*; match self { + Error::PartialRegion { response } => response + .first_failure() + .map_or(StatusCode::Internal, |(_, error)| error.status_code()), IllegalServerState { .. } | EtcdTxnOpResponse { .. } | EtcdFailed { .. } @@ -1361,6 +1373,7 @@ impl ErrorExt for Error { use Error::*; match self { + Error::PartialRegion { .. } => RetryHint::Retryable, RetryLater { source, .. } => match source.retry_hint() { RetryHint::RetryAfterRouteRefresh => RetryHint::RetryAfterRouteRefresh, RetryHint::Retryable | RetryHint::NonRetryable => RetryHint::Retryable, diff --git a/src/datanode/src/error.rs b/src/datanode/src/error.rs index 3f61d4b3d76..aaf1ff4d8df 100644 --- a/src/datanode/src/error.rs +++ b/src/datanode/src/error.rs @@ -30,6 +30,9 @@ use tokio::time::error::Elapsed; #[snafu(visibility(pub))] #[stack_trace_debug] pub enum Error { + #[snafu(display("Failed to write a batch group"))] + BatchWrite { source: std::sync::Arc }, + #[snafu(display("Failed to execute async task"))] AsyncTaskExecute { #[snafu(implicit)] @@ -479,6 +482,7 @@ impl ErrorExt for Error { fn status_code(&self) -> StatusCode { use Error::*; match self { + Self::BatchWrite { source } => source.status_code(), NewPlanDecoder { source, .. } | ExecuteLogicalPlan { source, .. } => { source.status_code() } @@ -561,6 +565,7 @@ impl ErrorExt for Error { use Error::*; match self { + Self::BatchWrite { source } => source.retry_hint(), RegionNotReady { .. } | RegionNotFound { .. } => RetryHint::RetryAfterRouteRefresh, RegionBusy { .. } | ConcurrentQueryLimiterClosed { .. } diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 0b7d36e9f2c..15bf786c389 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -29,12 +29,12 @@ use std::sync::{Arc, RwLock}; use std::task::{Context, Poll}; use std::time::{Duration, Instant}; -use api::region::RegionResponse; +use api::region::{BatchRegionResponse, RegionResponse}; use api::v1::meta::TopicStat; use api::v1::region::sync_request::ManifestInfo; use api::v1::region::{ - ListMetadataRequest, RegionRequest as RegionRequestV1, RegionResponse as RegionResponseV1, - SyncRequest, region_request, + ListMetadataRequest, RegionRequest as RegionRequestV1, RegionRequestFailure, + RegionResponse as RegionResponseV1, SyncRequest, region_request, }; use api::v1::{ResponseHeader, Status}; use arrow_flight::{FlightData, Ticket}; @@ -58,7 +58,6 @@ use datafusion::datasource::TableProvider; use datafusion_common::tree_node::TreeNode; use datatypes::schema::SchemaRef; use either::Either; -use futures_util::future::try_join_all; use futures_util::stream::FuturesUnordered; use futures_util::{Stream, StreamExt}; use metric_engine::engine::MetricEngine; @@ -82,7 +81,7 @@ use servers::grpc::region_server::RegionServerHandler; use session::context::{ FLIGHT_METRICS_HEARTBEAT_INTERVAL, QueryContext, QueryContextBuilder, QueryContextRef, }; -use snafu::{OptionExt, ResultExt, ensure}; +use snafu::{IntoError, OptionExt, ResultExt, ensure}; use store_api::metric_engine_consts::{ FILE_ENGINE_NAME, LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, }; @@ -92,11 +91,11 @@ use store_api::region_engine::{ SyncRegionFromRequest, }; use store_api::region_request::{ - AffectedRows, BatchRegionDdlRequest, RegionCatchupRequest, RegionCloseRequest, - RegionOpenRequest, RegionRequest, + self as store_region_request, AffectedRows, BatchRegionDdlRequest, RegionCatchupRequest, + RegionCloseRequest, RegionOpenRequest, RegionRequest, }; use store_api::storage::RegionId; -use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc}; use tokio::time::{self, timeout}; use tonic::{Request, Response, Result as TonicResult}; @@ -540,7 +539,7 @@ impl RegionServer { let span = tracing_context.attach(info_span!("RegionServer::handle_batch_ddl_requests")); self.inner - .handle_batch_request(batch_request) + .handle_batch_ddl_requests(batch_request) .trace(span) .await } @@ -548,53 +547,50 @@ impl RegionServer { async fn handle_requests_in_parallel( &self, request: region_request::Body, - ) -> Result { - let requests = - RegionRequest::try_from_request_body(request).context(BuildRegionRequestsSnafu)?; - + ) -> Result> { + let requests = prepare_batch_write_requests(request); let requests = match self.try_handle_mito_batch_puts(requests).await? { Either::Left(response) => return Ok(response), Either::Right(requests) => requests, }; - // Try to optimize batch Put requests for metric engine - // Returns either Some(response) or None(requests_back) - match self.try_handle_metric_batch_puts(requests).await? { - Either::Left(response) => Ok(response), - Either::Right(requests) => { - // Fallback: original parallel processing - let tracing_context = TracingContext::from_current_span(); - let join_tasks = - requests - .into_iter() - .map(|(region_id, req): (RegionId, RegionRequest)| { - let self_to_move = self; - let span = tracing_context.attach(info_span!( - "RegionServer::handle_region_request", - region_id = region_id.to_string() - )); - async move { - self_to_move - .handle_request(region_id, req) - .trace(span) - .await - } - }); - - let results = try_join_all(join_tasks).await?; - let mut affected_rows = 0; - let mut extensions = HashMap::new(); - for result in results { - affected_rows += result.affected_rows; - extensions.extend(result.extensions); + let requests = match self.try_handle_metric_batch_puts(requests).await? { + Either::Left(response) => return Ok(response), + Either::Right(requests) => requests, + }; + let tracing_context = TracingContext::from_current_span(); + let results = futures::future::join_all(requests.into_iter().map(|request| { + let span = tracing_context.attach(info_span!( + "RegionServer::handle_region_request", + region_id = request.region_id.to_string() + )); + async move { + ( + request.index, + self.handle_request(request.region_id, request.request) + .trace(span) + .await, + ) + } + })) + .await; + let mut response = BatchRegionResponse { + affected_rows: 0, + extensions: HashMap::new(), + failures: Vec::new(), + }; + for (index, result) in results { + match result { + Ok(result) => { + response.affected_rows += result.affected_rows; + response.extensions.extend(result.extensions); + } + Err(error) => { + log_batch_write_error(&error, index); + response.failures.push((index, error)); } - - Ok(RegionResponse { - affected_rows, - extensions, - metadata: Vec::new(), - }) } } + Ok(response) } async fn handle_requests_in_serial( @@ -625,65 +621,35 @@ impl RegionServer { }) } - /// Attempts to optimize batch Put requests for metric engine. - /// - /// Returns Either::Left(response) if optimization succeeded, - /// or Either::Right(original_requests) to fall back to parallel processing. - /// - /// This avoids cloning requests when optimization cannot be applied. + /// Attempts metric batch writes, returning the requests when ordinary execution is needed. async fn try_handle_metric_batch_puts( &self, - requests: Vec<(RegionId, RegionRequest)>, - ) -> Result>> { - if requests.is_empty() { + requests: Vec, + ) -> Result, Vec>> { + let Some(first) = requests.first() else { return Ok(Either::Right(requests)); - } - - // Quick check: verify first request is Put and is metric engine - if !matches!(requests[0].1, RegionRequest::Put(_)) { - return Ok(Either::Right(requests)); - } - let first_region_id = requests[0].0; - let request_type = requests[0].1.request_type(); - - // SAFETY: If the first request belongs to metric engine, then ALL requests - // in this batch are guaranteed to belong to metric engine. This invariant - // is maintained by the request batching logic upstream. - let engine = match self - .inner - .get_engine(first_region_id, &RegionChange::None)? - { - CurrentEngine::Engine(e) => e, - _ => return Ok(Either::Right(requests)), }; - - if engine.name() != METRIC_ENGINE_NAME { + if !matches!(first.request, RegionRequest::Put(_)) { return Ok(Either::Right(requests)); } - // Check if ALL requests are Put (now we know it's worth checking) - let mut all_puts = true; - for (_, req) in &requests { - if !matches!(req, RegionRequest::Put(_)) { - all_puts = false; - break; - } - } + let first_region_id = first.region_id; + let request_type = first.request.request_type(); - if !all_puts { + // Upstream batching keeps requests from the same engine together. + // Ordinary execution reports engine lookup failures per entry. + let Ok(CurrentEngine::Engine(engine)) = + self.inner.get_engine(first_region_id, &RegionChange::None) + else { + return Ok(Either::Right(requests)); + }; + if engine.name() != METRIC_ENGINE_NAME + || !requests + .iter() + .all(|request| matches!(request.request, RegionRequest::Put(_))) + { return Ok(Either::Right(requests)); } - - // Now extract Put requests by consuming ownership (zero clone!) - let put_requests = requests.into_iter().map(|(region_id, req)| { - if let RegionRequest::Put(put) = req { - (region_id, put) - } else { - unreachable!("Already checked all are Put") - } - }); - - // Downcast to MetricEngine and call batch API let metric_engine = engine .as_any() .downcast_ref::() @@ -692,75 +658,115 @@ impl RegionServer { })? .clone(); + let indices: Vec<_> = requests + .iter() + .map(|request| (request.index, request.region_id)) + .collect(); + let put_requests = requests.into_iter().map(|request| match request.request { + RegionRequest::Put(put) => (request.region_id, put), + _ => unreachable!("Already checked all are Put"), + }); let tracing_context = TracingContext::from_current_span(); let batch_size = put_requests.len(); let span = tracing_context.attach(info_span!( "RegionServer::handle_metric_batch_puts", - batch_size = batch_size, + batch_size = batch_size )); - let result = common_runtime::spawn_ingest(async move { + let results = common_runtime::spawn_ingest(async move { metric_engine .put_regions_batch(put_requests) .trace(span) .await }) .await - .context(RuntimeJoinSnafu { request_type })? - .map_err(BoxedError::new) - .context(HandleRegionRequestSnafu { - region_id: first_region_id, - }) - .map(RegionResponse::new); - record_ingest_result(request_type, batch_size, &result); - result.map(Either::Left) + .context(RuntimeJoinSnafu { request_type })?; + let mut response = BatchRegionResponse { + affected_rows: 0, + extensions: HashMap::new(), + failures: Vec::new(), + }; + for group in results { + let region_id = indices[group.request_indices[0]].1; + let result = group + .result + .map_err(BoxedError::new) + .context(HandleRegionRequestSnafu { region_id }) + .map(RegionResponse::new); + record_ingest_result(request_type, group.request_indices.len(), &result); + match result { + Ok(result) => { + response.affected_rows += result.affected_rows; + } + Err(error) => { + log_batch_write_error(&error, indices[group.request_indices[0]].0); + let error = Arc::new(error); + response + .failures + .extend(group.request_indices.into_iter().map(|index| { + ( + indices[index].0, + error::BatchWriteSnafu.into_error(Arc::clone(&error)), + ) + })); + } + } + } + Ok(Either::Left(response)) } - /// Puts into ready mito regions together through - /// [`RegionEngine::handle_batch_put_requests`], so that each mito worker gets the puts of the - /// request to its regions at once and writes them in one batch. One ingest task awaits every - /// response at once and accounts each, for the regions it answers, as the per-region path - /// does; it hands the first error back at once and keeps accounting the rest, also when the - /// caller goes away. - /// - /// Returns Either::Right(original_requests) unless - /// [`mito_engine_for_batch_puts`](Self::mito_engine_for_batch_puts) accepts the requests. + /// Batches puts into ready mito regions and retains every entry's outcome. async fn try_handle_mito_batch_puts( &self, - requests: Vec<(RegionId, RegionRequest)>, - ) -> Result>> { + requests: Vec, + ) -> Result, Vec>> { let Some(engine) = self.mito_engine_for_batch_puts(&requests) else { return Ok(Either::Right(requests)); }; - - let request_type = requests[0].1.request_type(); + let request_type = requests[0].request.request_type(); + let mut indices: HashMap> = HashMap::new(); let puts = requests .into_iter() - .map(|(region_id, request)| match request { - RegionRequest::Put(put) => (region_id, put), - _ => unreachable!("Already checked all are Put"), + .map(|request| { + indices + .entry(request.region_id) + .or_default() + .push_back(request.index); + match request.request { + RegionRequest::Put(put) => (request.region_id, put), + _ => unreachable!("Already checked all are Put"), + } }) .collect(); let started = Instant::now(); let tracing_context = TracingContext::from_current_span(); - let (result_tx, result_rx) = oneshot::channel(); - let batch = common_runtime::spawn_ingest(async move { + common_runtime::spawn_ingest(async move { let mut pending = engine .handle_batch_put_requests(puts) .into_iter() .map(|(region_ids, response)| { + let request_indices = region_ids + .iter() + .filter_map(|region_id| { + indices + .get_mut(region_id) + .and_then(|indices| indices.pop_front()) + }) + .collect::>(); let span = tracing_context.attach(info_span!( "RegionServer::handle_region_request", region_id = region_ids[0].to_string() )); - async move { (region_ids, response.await) }.trace(span) + async move { (region_ids, request_indices, response.await) }.trace(span) }) .collect::>(); let histogram = crate::metrics::HANDLE_REGION_REQUEST_ELAPSED.with_label_values(&[request_type]); - let mut result_tx = Some(result_tx); - let mut affected_rows = 0; - let mut extensions = HashMap::new(); - while let Some((region_ids, result)) = pending.next().await { + let mut response = BatchRegionResponse { + affected_rows: 0, + extensions: HashMap::new(), + failures: Vec::new(), + }; + while let Some((region_ids, request_indices, result)) = pending.next().await { let elapsed = started.elapsed().as_secs_f64(); for _ in ®ion_ids { histogram.observe(elapsed); @@ -770,53 +776,43 @@ impl RegionServer { }); record_ingest_result(request_type, region_ids.len(), &result); match result { - Ok(response) => { - affected_rows += response.affected_rows; - extensions.extend(response.extensions); + Ok(result) => { + response.affected_rows += result.affected_rows; + response.extensions.extend(result.extensions); } - Err(err) => { - if let Some(result_tx) = result_tx.take() { - let _ = result_tx.send(Err(err)); - } + Err(error) => { + log_batch_write_error(&error, request_indices[0]); + let error = Arc::new(error); + response + .failures + .extend(request_indices.into_iter().map(|index| { + (index, error::BatchWriteSnafu.into_error(Arc::clone(&error))) + })); } } } - if let Some(result_tx) = result_tx { - let _ = result_tx.send(Ok(RegionResponse { - affected_rows, - extensions, - metadata: Vec::new(), - })); - } - }); - - match result_rx.await { - Ok(result) => result.map(Either::Left), - // The batch task answers on every path it completes, so it panicked or was cancelled. - Err(_) => { - batch.await.context(RuntimeJoinSnafu { request_type })?; - UnexpectedSnafu { - violated: "the batch put task completed without answering", - } - .fail() - } - } + response.failures.sort_unstable_by_key(|(index, _)| *index); + response + }) + .await + .context(RuntimeJoinSnafu { request_type }) + .map(Either::Left) } /// Returns the mito engine when there are at least two requests, every one a put into a ready /// mito region. fn mito_engine_for_batch_puts( &self, - requests: &[(RegionId, RegionRequest)], + requests: &[IndexedRegionRequest], ) -> Option { if requests.len() < 2 { return None; } - let first_region_id = requests[0].0; - let all_puts_to_ready_mito_regions = requests.iter().all(|(region_id, request)| { - matches!(request, RegionRequest::Put(_)) + let first_region_id = requests[0].region_id; + let all_puts_to_ready_mito_regions = requests.iter().all(|request| { + matches!(request.request, RegionRequest::Put(_)) && matches!( - self.inner.region_map.get(region_id).as_deref(), + self.inner.region_map.get(&request.region_id).as_deref(), Some(RegionEngineWithStatus::Ready(engine)) if engine.name() == MITO_ENGINE_NAME ) }); @@ -1062,22 +1058,45 @@ impl RegionServerHandler for RegionServer { { return self.handle_flow_write_request(request).await; } + let accept_partial_result = request + .header + .as_ref() + .and_then(|header| header.query_context.as_ref()) + .and_then(|context| context.extensions.get(api::region::ACCEPT_PARTIAL_RESULT)) + .is_some_and(|value| value == "true"); let body = request.body.context(servers_error::InvalidQuerySnafu { reason: "Expecting non-empty region request body.", })?; - self.handle(body).await + self.handle_inner(body, accept_partial_result).await } async fn handle(&self, request: region_request::Body) -> ServerResult { + self.handle_inner(request, false).await + } +} + +impl RegionServer { + async fn handle_inner( + &self, + request: region_request::Body, + accept_partial_result: bool, + ) -> ServerResult { + let request_count = match &request { + region_request::Body::Inserts(requests) => Some(requests.requests.len()), + region_request::Body::Deletes(requests) => Some(requests.requests.len()), + _ => None, + }; + if let Some(request_count) = request_count { + return self + .handle_write_request(request, request_count, accept_partial_result) + .await; + } let failed_requests_cnt = crate::metrics::REGION_SERVER_REQUEST_FAILURE_COUNT .with_label_values(&[request.as_ref()]); let response = match &request { region_request::Body::Creates(_) | region_request::Body::Drops(_) | region_request::Body::Alters(_) => self.handle_batch_ddl_requests(request).await, - region_request::Body::Inserts(_) | region_request::Body::Deletes(_) => { - self.handle_requests_in_parallel(request).await - } region_request::Body::Sync(sync_request) => { self.handle_sync_region_request(sync_request).await } @@ -1092,11 +1111,8 @@ impl RegionServerHandler for RegionServer { _ => self.handle_requests_in_serial(request).await, } .map_err(BoxedError::new) - .inspect_err(|_| { - failed_requests_cnt.inc(); - }) + .inspect_err(|_| failed_requests_cnt.inc()) .context(ExecuteGrpcRequestSnafu)?; - Ok(RegionResponseV1 { header: Some(ResponseHeader { status: Some(Status { @@ -1107,8 +1123,35 @@ impl RegionServerHandler for RegionServer { affected_rows: response.affected_rows as _, extensions: response.extensions, metadata: response.metadata, + failures: Vec::new(), }) } + + async fn handle_write_request( + &self, + request: region_request::Body, + entry_count: usize, + accept_partial_result: bool, + ) -> ServerResult { + let failed_requests_cnt = crate::metrics::REGION_SERVER_REQUEST_FAILURE_COUNT + .with_label_values(&[request.as_ref()]); + let mut result = self + .handle_requests_in_parallel(request) + .await + .map_err(BoxedError::new) + .inspect_err(|_| failed_requests_cnt.inc()) + .context(ExecuteGrpcRequestSnafu)?; + if !result.failures.is_empty() { + failed_requests_cnt.inc(); + // Old clients and single-entry writes retain tonic errors and retry metadata. + if (!accept_partial_result || entry_count == 1) + && let Some((_, error)) = result.failures.drain(..).min_by_key(|(index, _)| *index) + { + return Err(BoxedError::new(error)).context(ExecuteGrpcRequestSnafu); + } + } + Ok(encode_batch_write_response(result)) + } } #[async_trait] @@ -1166,6 +1209,97 @@ impl FlightCraft for RegionServer { } } +/// A decoded write retains its protobuf entry index even when preceding entries are empty. +struct IndexedRegionRequest { + index: usize, + region_id: RegionId, + request: RegionRequest, +} + +fn prepare_batch_write_requests(body: region_request::Body) -> Vec { + match body { + region_request::Body::Inserts(inserts) => inserts + .requests + .into_iter() + .enumerate() + .filter_map(|(index, request)| { + store_region_request::make_region_put(request).map(|(region_id, request)| { + IndexedRegionRequest { + index, + region_id, + request: RegionRequest::Put(request), + } + }) + }) + .collect(), + region_request::Body::Deletes(deletes) => deletes + .requests + .into_iter() + .enumerate() + .filter_map(|(index, request)| { + store_region_request::make_region_delete(request).map(|(region_id, request)| { + IndexedRegionRequest { + index, + region_id, + request: RegionRequest::Delete(request), + } + }) + }) + .collect(), + _ => unreachable!("Only inserts/deletes are prepared as batch writes"), + } +} + +fn encode_batch_write_response(result: BatchRegionResponse) -> RegionResponseV1 { + RegionResponseV1 { + header: Some(ResponseHeader { + status: Some(Status { + status_code: if result.failures.is_empty() { + StatusCode::Success + } else { + StatusCode::PartialFailure + } as u32, + err_msg: if let Some((index, error)) = result.first_failure() { + format!( + "{} region requests failed; entry {index}: {}", + result.failures.len(), + error.output_msg() + ) + } else { + String::new() + }, + ..Default::default() + }), + }), + metadata: Vec::new(), + affected_rows: result.affected_rows as _, + extensions: result.extensions, + failures: result + .failures + .into_iter() + .map(|(index, error)| RegionRequestFailure { + request_index: index as u32, + status: Some(Status { + status_code: error.status_code() as u32, + err_msg: error.output_msg(), + retry_hint: error.retry_hint().as_str().to_string(), + }), + }) + .collect(), + } +} + +fn log_batch_write_error(error: &error::Error, index: usize) { + if error.status_code().should_log_error() { + error!(error; "Failed to handle batch entry, request_index: {}", index); + } else { + debug!( + "Failed to handle batch entry, request_index: {}, err: {}", + index, error + ); + } +} + #[derive(Clone)] enum RegionEngineWithStatus { // An opening, or creating region. @@ -1679,7 +1813,7 @@ impl RegionServerInner { // Handle requests in batch. // // limitation: all create requests must be in the same engine. - pub async fn handle_batch_request( + pub async fn handle_batch_ddl_requests( &self, batch_request: BatchRegionDdlRequest, ) -> Result { @@ -2207,7 +2341,7 @@ mod tests { use std::sync::Arc; use api::v1::region::{InsertRequest, InsertRequests}; - use api::v1::{Rows, SemanticType}; + use api::v1::{PartitionExprVersion, Rows, SemanticType}; use common_error::ext::{ErrorExt, RetryHint}; use common_recordbatch::RecordBatches; use common_recordbatch::adapter::{RecordBatchMetrics, RegionWatermarkEntry}; @@ -2617,6 +2751,170 @@ mod tests { assert_eq!(err.status_code(), StatusCode::RegionNotReady); } + #[test] + fn test_batch_write_summary_uses_first_input_failure() { + let first = error::RegionNotFoundSnafu { + region_id: RegionId::new(1, 0), + } + .build(); + let expected = format!("2 region requests failed; entry 0: {}", first.output_msg()); + let later = error::RegionBusySnafu { + region_id: RegionId::new(1, 1), + } + .build(); + let response = encode_batch_write_response(BatchRegionResponse { + affected_rows: 0, + extensions: HashMap::new(), + failures: vec![(1, later), (0, first)], + }); + let status = response.header.unwrap().status.unwrap(); + assert_eq!(status.status_code, StatusCode::PartialFailure as u32); + assert_eq!(status.err_msg, expected); + assert_eq!(response.failures[0].request_index, 1); + assert_eq!(response.failures[1].request_index, 0); + } + + #[tokio::test] + async fn test_write_failure_response_compatibility() { + let server = mock_region_server(); + let missing = RegionId::new(1, 1); + let registering = RegionId::new(1, 2); + let (engine, _receiver) = MockRegionEngine::new(MITO_ENGINE_NAME); + server + .inner + .region_map + .insert(registering, RegionEngineWithStatus::Registering(engine)); + for (region_id, code) in [ + (missing, StatusCode::RegionNotFound), + (registering, StatusCode::RegionNotReady), + ] { + for delete in [false, true] { + for accept_partial in [false, true] { + // A preceding empty entry must still count toward the original envelope size. + for with_empty_entry in [false, true] { + let rows = Some(Rows { + schema: Vec::new(), + rows: Vec::new(), + }); + let body = if delete { + let request = api::v1::region::DeleteRequest { + region_id: region_id.as_u64(), + rows, + ..Default::default() + }; + region_request::Body::Deletes(api::v1::region::DeleteRequests { + requests: if with_empty_entry { + vec![Default::default(), request] + } else { + vec![request] + }, + }) + } else { + let request = InsertRequest { + region_id: region_id.as_u64(), + rows, + ..Default::default() + }; + region_request::Body::Inserts(InsertRequests { + requests: if with_empty_entry { + vec![Default::default(), request] + } else { + vec![request] + }, + }) + }; + let mut header = api::v1::region::RegionRequestHeader::default(); + if accept_partial { + header + .query_context + .get_or_insert_with(Default::default) + .extensions + .insert( + api::region::ACCEPT_PARTIAL_RESULT.to_string(), + "true".to_string(), + ); + } + let result = server + .handle_region_request(RegionRequestV1 { + header: Some(header), + body: Some(body), + }) + .await; + if accept_partial && with_empty_entry { + let response = result.unwrap(); + assert_eq!( + StatusCode::PartialFailure as u32, + response.header.unwrap().status.unwrap().status_code + ); + assert_eq!(1, response.failures.len()); + assert_eq!(1, response.failures[0].request_index); + let status = response.failures[0].status.as_ref().unwrap(); + assert_eq!(code as u32, status.status_code); + assert_eq!( + RetryHint::RetryAfterRouteRefresh.as_str(), + status.retry_hint + ); + } else { + let status: tonic::Status = result.unwrap_err().into(); + assert_eq!( + common_error::status_code::status_to_tonic_code(code), + status.code() + ); + assert_eq!( + (code as u32).to_string(), + status + .metadata() + .get(common_error::GREPTIME_DB_HEADER_ERROR_CODE) + .unwrap() + .to_str() + .unwrap() + ); + assert_eq!( + RetryHint::RetryAfterRouteRefresh.as_str(), + status + .metadata() + .get(common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT) + .unwrap() + .to_str() + .unwrap() + ); + } + } + } + } + } + } + + #[tokio::test] + async fn test_legacy_batch_returns_first_input_error() { + let server = mock_region_server(); + let first = RegionId::new(1, 1); + let second = RegionId::new(1, 2); + let (engine, _receiver) = MockRegionEngine::with_custom_apply_fn("mock", |engine| { + engine.handle_request_delay = Some(Duration::from_millis(10)); + engine.handle_request_mock_fn = Some(Box::new(|region_id, _| { + error::RegionBusySnafu { region_id }.fail() + })); + }); + server.register_test_region(first, engine); + let requests = [first, second] + .map(|region_id| InsertRequest { + region_id: region_id.as_u64(), + rows: Some(Rows::default()), + ..Default::default() + }) + .to_vec(); + let error = server + .handle_region_request(RegionRequestV1 { + header: Some(Default::default()), + body: Some(region_request::Body::Inserts(InsertRequests { requests })), + }) + .await + .unwrap_err(); + assert_eq!(StatusCode::RegionBusy, error.status_code()); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + } + #[tokio::test] async fn test_region_request_failed() { common_telemetry::init_default_ut_logging(); @@ -2642,8 +2940,7 @@ mod tests { .await .unwrap_err(); - let status = mock_region_server.inner.region_map.get(®ion_id); - assert!(status.is_none()); + assert!(!mock_region_server.inner.region_map.contains_key(®ion_id)); mock_region_server .inner @@ -2662,8 +2959,79 @@ mod tests { .await .unwrap_err(); - let status = mock_region_server.inner.region_map.get(®ion_id); - assert!(status.is_some()); + assert!(mock_region_server.inner.region_map.contains_key(®ion_id)); + + let insert = api::v1::region::InsertRequest { + region_id: region_id.as_u64(), + rows: Some(Rows { + schema: Vec::new(), + rows: Vec::new(), + }), + ..Default::default() + }; + let success_region = RegionId::new(1, 2); + let (success_engine, _receiver) = + MockRegionEngine::with_custom_apply_fn("mock", |engine| { + engine.handle_request_delay = Some(std::time::Duration::from_millis(10)); + engine.handle_request_mock_fn = Some(Box::new(|_, _| Ok(5))); + }); + mock_region_server.register_test_region(success_region, success_engine); + let successful_insert = api::v1::region::InsertRequest { + region_id: success_region.as_u64(), + ..insert.clone() + }; + let response = mock_region_server + .handle_region_request(RegionRequestV1 { + header: Some(api::v1::region::RegionRequestHeader { + query_context: Some(api::v1::QueryContext { + extensions: HashMap::from([( + api::region::ACCEPT_PARTIAL_RESULT.to_string(), + "true".to_string(), + )]), + ..Default::default() + }), + ..Default::default() + }), + body: Some(region_request::Body::Inserts(InsertRequests { + requests: vec![ + Default::default(), + insert.clone(), + successful_insert, + insert, + ], + })), + }) + .await + .unwrap(); + assert_eq!(response.affected_rows, 5); + assert_eq!( + response + .header + .as_ref() + .unwrap() + .status + .as_ref() + .unwrap() + .status_code, + StatusCode::PartialFailure as u32 + ); + assert_eq!( + response + .failures + .iter() + .map(|error| error.request_index) + .collect::>(), + vec![1, 3] + ); + assert!( + response + .failures + .iter() + .all(|error| error.status.as_ref().is_some_and(|status| { + status.status_code == StatusCode::Unexpected as u32 + && status.retry_hint == RetryHint::NonRetryable.as_str() + })) + ); } /// A region that fails to close is reported in the result of the stop, @@ -3231,6 +3599,114 @@ mod tests { } } + #[tokio::test] + async fn test_metric_batch_failure_indices() { + let mut env = TestEnv::with_prefix("region-server-metric-results").await; + let mito = env.create_engine(MitoConfig::default()).await; + let metric = Arc::new( + MetricEngine::try_new(mito, metric_engine::config::EngineConfig::default()).unwrap(), + ); + let mut server = mock_region_server(); + server.register_engine(metric.clone()); + let mut inserts = Vec::new(); + for table in [10, 20] { + let physical = RegionId::new(table, 0); + let logical = RegionId::new(table + 1, 0); + let mut create = CreateRequestBuilder::new().build(); + create.engine = METRIC_ENGINE_NAME.to_string(); + create.column_metadatas = vec![ + store_api::metadata::ColumnMetadata { + column_id: 0, + semantic_type: SemanticType::Timestamp, + column_schema: ColumnSchema::new( + "greptime_timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + }, + store_api::metadata::ColumnMetadata { + column_id: 1, + semantic_type: SemanticType::Field, + column_schema: ColumnSchema::new( + "greptime_value", + ConcreteDataType::float64_datatype(), + false, + ), + }, + ]; + create.primary_key.clear(); + create.table_dir = format!("metric-results/{table}/"); + create.options = HashMap::from([( + store_api::metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY.to_string(), + String::new(), + )]); + server + .handle_request(physical, RegionRequest::Create(create.clone())) + .await + .unwrap(); + create.table_dir = format!("metric-results/{}/", table + 1); + create.options = HashMap::from([( + store_api::metric_engine_consts::LOGICAL_TABLE_METADATA_KEY.to_string(), + physical.as_u64().to_string(), + )]); + server + .handle_request(logical, RegionRequest::Create(create.clone())) + .await + .unwrap(); + inserts.push(InsertRequest { + region_id: logical.as_u64(), + rows: Some(Rows { + schema: rows_schema(&create), + rows: vec![api::v1::Row { + values: vec![ + api::v1::Value { + value_data: Some( + api::v1::value::ValueData::TimestampMillisecondValue(1), + ), + }, + api::v1::Value { + value_data: Some(api::v1::value::ValueData::F64Value(1.0)), + }, + ], + }], + }), + ..Default::default() + }); + } + // Both entries of the second physical group fail at write_data. + metric + .set_region_role(RegionId::new(20, 0), RegionRole::Follower) + .unwrap(); + let response = server + .handle_requests_in_parallel(region_request::Body::Inserts(InsertRequests { + requests: vec![ + Default::default(), + inserts[0].clone(), + inserts[1].clone(), + inserts[1].clone(), + ], + })) + .await + .unwrap(); + assert_eq!(1, response.affected_rows); + let mut failures = response.failures; + failures.sort_unstable_by_key(|(index, _)| *index); + assert_eq!( + vec![2, 3], + failures.iter().map(|(index, _)| *index).collect::>() + ); + assert!( + failures + .iter() + .all(|(_, error)| error.status_code() == StatusCode::RegionNotReady) + ); + assert!( + failures + .iter() + .all(|(_, error)| error.retry_hint() == RetryHint::RetryAfterRouteRefresh) + ); + } + #[tokio::test] async fn test_mito_puts_of_one_request_share_a_wal_append() { // One worker that takes one message at a time, so that only puts that @@ -3293,49 +3769,126 @@ mod tests { } #[tokio::test] - async fn test_mito_batch_puts_fail_without_waiting_for_other_regions() { - let mut env = TestEnv::with_prefix("region-server-mito-batch-puts-error").await; - let config = MitoConfig { - num_workers: 2, - ..Default::default() - }; - let (region_server, log_store, wal_options) = - region_server_on_object_store_wal(&mut env, config).await; - // Regions 1 and 2 of table 1 belong to different workers. The put to region 1 waits for a - // WAL object whose create is held; Mito does not know region 2, so its put fails at once. - let held = RegionId::new(1, 1); - let unknown = RegionId::new(1, 2); - let held_insert = create_region(®ion_server, held, &wal_options).await; - let unknown_insert = InsertRequest { - region_id: unknown.as_u64(), - ..held_insert.clone() - }; - let mito = region_server - .inner - .region_map - .get(&held) - .unwrap() - .clone() - .into_engine(); - region_server - .inner - .region_map - .insert(unknown, RegionEngineWithStatus::Ready(mito)); - log_store.hold_creates(); + async fn test_mito_batch_failure_indices_for_same_region() { + for failed_index in [0, 1] { + let mut env = TestEnv::with_prefix("region-server-mito-same-region-failure").await; + let (server, _log_store, wal_options) = + region_server_on_object_store_wal(&mut env, MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + let insert = create_region(&server, region_id, &wal_options).await; + let mut requests = vec![insert.clone(), insert]; + requests[failed_index].partition_expr_version = Some(PartitionExprVersion { value: 1 }); + let response = server + .handle_requests_in_parallel(region_request::Body::Inserts(InsertRequests { + requests, + })) + .await + .unwrap(); + assert_eq!(response.affected_rows, 3); + assert_eq!(response.failures.len(), 1); + assert_eq!(response.failures[0].0, failed_index); + assert_eq!( + response.failures[0].1.status_code(), + StatusCode::InvalidArguments + ); + let engine = server + .inner + .region_map + .get(®ion_id) + .unwrap() + .clone() + .into_engine(); + let mito = engine.as_any().downcast_ref::().unwrap(); + let stream = mito + .scan_to_stream(region_id, Default::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 3 + ); + } + } - let requests = vec![held_insert, unknown_insert]; - let err = tokio::time::timeout( - Duration::from_secs(10), - region_server.handle_requests_in_parallel(region_request::Body::Inserts( - InsertRequests { requests }, - )), - ) - .await - .unwrap() - .unwrap_err(); + #[tokio::test] + async fn test_mito_batch_puts_retain_successful_regions() { + for accept_partial_result in [false, true] { + let mut env = TestEnv::with_prefix("region-server-mito-batch-puts-error").await; + let config = MitoConfig { + num_workers: 2, + ..Default::default() + }; + let (region_server, log_store, wal_options) = + region_server_on_object_store_wal(&mut env, config).await; + // Regions 1 and 2 of table 1 belong to different workers. The put to region 1 waits for a + // WAL object whose create is held; Mito does not know region 2, so its put fails at once. + let held = RegionId::new(1, 1); + let unknown = RegionId::new(1, 2); + let held_insert = create_region(®ion_server, held, &wal_options).await; + let unknown_insert = InsertRequest { + region_id: unknown.as_u64(), + ..held_insert.clone() + }; + let mito = region_server + .inner + .region_map + .get(&held) + .unwrap() + .clone() + .into_engine(); + region_server + .inner + .region_map + .insert(unknown, RegionEngineWithStatus::Ready(mito.clone())); + log_store.hold_creates(); - assert_eq!(StatusCode::RegionNotFound, err.status_code()); - log_store.release_creates(); + let requests = vec![Default::default(), held_insert, unknown_insert]; + let mut header = api::v1::region::RegionRequestHeader::default(); + if accept_partial_result { + api::region::set_accept_partial_result(&mut header); + } + let write = region_server.handle_region_request(RegionRequestV1 { + header: Some(header), + body: Some(region_request::Body::Inserts(InsertRequests { requests })), + }); + tokio::pin!(write); + assert!( + tokio::time::timeout(Duration::from_millis(100), &mut write) + .await + .is_err() + ); + log_store.release_creates(); + let result = tokio::time::timeout(Duration::from_secs(10), write) + .await + .unwrap(); + if accept_partial_result { + let response = result.unwrap(); + assert_eq!(3, response.affected_rows); + assert_eq!( + StatusCode::PartialFailure as u32, + response.header.unwrap().status.unwrap().status_code + ); + assert_eq!(1, response.failures.len()); + assert_eq!(2, response.failures[0].request_index); + assert_eq!( + StatusCode::RegionNotFound as u32, + response.failures[0].status.as_ref().unwrap().status_code + ); + } else { + assert_eq!( + StatusCode::RegionNotFound, + result.unwrap_err().status_code() + ); + } + let mito = mito.as_any().downcast_ref::().unwrap(); + let stream = mito.scan_to_stream(held, Default::default()).await.unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 3 + ); + } } #[tokio::test] @@ -3356,13 +3909,19 @@ mod tests { // One put takes the per-region path and two take the batch path. for requests in [vec![insert.clone()], vec![insert.clone(), insert]] { - let err = region_server + let count = requests.len(); + let response = region_server .handle_requests_in_parallel(region_request::Body::Inserts(InsertRequests { requests, })) .await - .unwrap_err(); - assert!(format!("{err:?}").contains("is stopped"), "{err:?}"); + .unwrap(); + assert_eq!(0, response.affected_rows); + assert_eq!(count, response.failures.len()); + for (index, (request_index, error)) in response.failures.iter().enumerate() { + assert_eq!(index, *request_index); + assert!(format!("{error:?}").contains("is stopped"), "{error:?}"); + } } } @@ -3433,7 +3992,11 @@ mod tests { skip_wal: false, partition_expr_version: None, }; - (region_id, RegionRequest::Put(request)) + IndexedRegionRequest { + index: 0, + region_id, + request: RegionRequest::Put(request), + } }; let delete = |region_id| { let request = RegionDeleteRequest { @@ -3441,7 +4004,11 @@ mod tests { hint: None, partition_expr_version: None, }; - (region_id, RegionRequest::Delete(request)) + IndexedRegionRequest { + index: 0, + region_id, + request: RegionRequest::Delete(request), + } }; for requests in [ diff --git a/src/datanode/src/region_server/flow_write.rs b/src/datanode/src/region_server/flow_write.rs index adb9d88568a..6bcb112bc78 100644 --- a/src/datanode/src/region_server/flow_write.rs +++ b/src/datanode/src/region_server/flow_write.rs @@ -129,6 +129,7 @@ fn response(reply: Reply) -> ServerResult { affected_rows: 0, extensions, metadata: Vec::new(), + failures: Vec::new(), }) } diff --git a/src/meta-srv/src/procedure/utils.rs b/src/meta-srv/src/procedure/utils.rs index ff88d1b38ba..9375e1a4392 100644 --- a/src/meta-srv/src/procedure/utils.rs +++ b/src/meta-srv/src/procedure/utils.rs @@ -356,11 +356,13 @@ pub mod mock { status: Some(PbStatus { status_code: 0, err_msg: String::default(), + ..Default::default() }), }), affected_rows: 0, extensions: Default::default(), metadata: Vec::new(), + failures: Vec::new(), }) } } diff --git a/src/metric-engine/AGENTS.md b/src/metric-engine/AGENTS.md index d1049478978..9c3820a0cc3 100644 --- a/src/metric-engine/AGENTS.md +++ b/src/metric-engine/AGENTS.md @@ -51,6 +51,7 @@ Mito2 data region. Queries against a physical region pass straight through. ## Public surface - Entry: `MetricEngine` in `src/metric-engine/src/engine.rs`, built via `try_new(mito, config)`. +- `MetricEngine::put_regions_batch` returns each physical group's outcome with its input indices; a failed group does not stop other groups. - Trait: `impl RegionEngine for MetricEngine` (name = `"metric"`). - Depends on `mito2`, `store-api`, and `mito-codec` (sparse primary key codec). diff --git a/src/metric-engine/src/engine.rs b/src/metric-engine/src/engine.rs index ec180839526..d01ca4650e6 100644 --- a/src/metric-engine/src/engine.rs +++ b/src/metric-engine/src/engine.rs @@ -488,6 +488,17 @@ impl RegionEngine for MetricEngine { } } +/// Outcome of a group of input requests. +/// +/// Requests rejected before physical-region grouping produce singleton failures. +#[derive(Debug)] +pub struct PutGroupResult { + /// Positions in the input batch, including separate entries for the same region. + pub request_indices: Vec, + /// Aggregate rows on success; no other physical group's outcome is implied. + pub result: Result, +} + impl MetricEngine { pub fn try_new(mito: MitoEngine, mut config: EngineConfig) -> Result { let metadata_region = MetadataRegion::new(mito.clone()); @@ -520,13 +531,11 @@ impl MetricEngine { self.inner.mito.clone() } - /// Batch put operation for multiple logical regions. - /// Requests are grouped by physical region; a failure can leave earlier - /// physical-region groups committed. + /// Writes all physical-region groups and retains each group's outcome. pub async fn put_regions_batch( &self, requests: impl ExactSizeIterator, - ) -> Result { + ) -> Vec { self.inner.put_regions_batch(requests).await } diff --git a/src/metric-engine/src/engine/put.rs b/src/metric-engine/src/engine/put.rs index b709d99d847..27ee3fba56c 100644 --- a/src/metric-engine/src/engine/put.rs +++ b/src/metric-engine/src/engine/put.rs @@ -29,7 +29,7 @@ use store_api::region_request::{ }; use store_api::storage::{RegionId, TableId}; -use crate::engine::MetricEngineInner; +use crate::engine::{MetricEngineInner, PutGroupResult}; use crate::error::{ ColumnNotFoundSnafu, CreateDefaultSnafu, ForbiddenPhysicalWriteSnafu, InvalidRequestSnafu, LogicalRegionNotFoundSnafu, PhysicalRegionNotFoundSnafu, Result, UnexpectedRequestSnafu, @@ -39,6 +39,12 @@ use crate::metrics::{FORBIDDEN_OPERATION_COUNT, MITO_OPERATION_ELAPSED}; use crate::row_modifier::{RowsIter, TableIdInput}; use crate::utils::to_data_region_id; +#[derive(Default)] +struct PutGroup { + request_indices: Vec, + requests: Vec<(RegionId, RegionPutRequest)>, +} + impl MetricEngineInner { /// Dispatch region put request pub async fn put_region( @@ -61,67 +67,63 @@ impl MetricEngineInner { } } - /// Batch write multiple logical regions to the same physical region. - /// - /// Dispatch region put requests in batch. - /// - /// Requests may span multiple physical regions. We group them by physical - /// region and write sequentially. This method fails fast on validation or - /// preparation errors within a group and stops at the first failure. - /// Writes in earlier physical-region groups are not rolled back if a later - /// group fails. + /// Retains every group's outcome instead of returning on the first failure. pub async fn put_regions_batch( &self, - requests: impl ExactSizeIterator, - ) -> Result { + mut requests: impl ExactSizeIterator, + ) -> Vec { let len = requests.len(); - if len == 0 { - return Ok(0); + return Vec::new(); } - let _timer = MITO_OPERATION_ELAPSED .with_label_values(&["put_batch"]) .start_timer(); - - // Fast path: single request, no batching overhead if len == 1 { - let (region_id, req) = requests.into_iter().next().unwrap(); - let is_putting_physical_region = - self.state.read().unwrap().exist_physical_region(region_id); - if is_putting_physical_region { - FORBIDDEN_OPERATION_COUNT.inc(); - return ForbiddenPhysicalWriteSnafu.fail(); + let (region_id, request) = requests.next().unwrap(); + return vec![PutGroupResult { + request_indices: vec![0], + result: self.put_region(region_id, request).await, + }]; + } + let mut groups = HashMap::new(); + let mut results = Vec::new(); + for (index, (region_id, request)) in requests.enumerate() { + if let Err(error) = self.group_put_request(&mut groups, index, region_id, request) { + results.push(PutGroupResult { + request_indices: vec![index], + result: Err(error), + }); } - - return self.put_logical_region(region_id, req).await; } - - let mut requests_per_physical: HashMap> = - HashMap::new(); - for (region_id, request) in requests { - let is_putting_physical_region = - self.state.read().unwrap().exist_physical_region(region_id); - if is_putting_physical_region { - FORBIDDEN_OPERATION_COUNT.inc(); - return ForbiddenPhysicalWriteSnafu.fail(); - } - let physical_region_id = self.find_physical_region_id(region_id)?; - requests_per_physical - .entry(physical_region_id) - .or_default() - .push((region_id, request)); + for (physical_region_id, group) in groups { + let result = self + .put_regions_batch_single_physical(physical_region_id, group.requests) + .await; + results.push(PutGroupResult { + request_indices: group.request_indices, + result, + }); } + results + } - let mut total_affected_rows: AffectedRows = 0; - for (physical_region_id, requests) in requests_per_physical { - let affected_rows = self - .put_regions_batch_single_physical(physical_region_id, requests) - .await?; - total_affected_rows += affected_rows; + fn group_put_request( + &self, + groups: &mut HashMap, + index: usize, + region_id: RegionId, + request: RegionPutRequest, + ) -> Result<()> { + if self.state.read().unwrap().exist_physical_region(region_id) { + FORBIDDEN_OPERATION_COUNT.inc(); + return ForbiddenPhysicalWriteSnafu.fail(); } - - Ok(total_affected_rows) + let physical_region_id = self.find_physical_region_id(region_id)?; + let group = groups.entry(physical_region_id).or_default(); + group.request_indices.push(index); + group.requests.push((region_id, request)); + Ok(()) } /// Write a batch of requests that all belong to the same physical region. @@ -1001,7 +1003,9 @@ mod tests { }, ) }); - let affected_rows = engine.inner.put_regions_batch(requests).await.unwrap(); + let results = engine.put_regions_batch(requests).await; + assert_eq!(results.len(), 1); + let affected_rows = results.into_iter().next().unwrap().result.unwrap(); assert_eq!(affected_rows, 6); assert_eq!( scan_timestamp_values(&engine, logical_region_id).await, @@ -1376,12 +1380,13 @@ mod tests { for ((_, request), skip_wal) in mixed_requests.iter_mut().zip(policies) { request.skip_wal = skip_wal; } - let err = env + let results = env .metric() - .inner .put_regions_batch(mixed_requests.into_iter()) - .await - .unwrap_err(); + .await; + assert_eq!(results.len(), 1); + assert_eq!(results[0].request_indices, vec![0, 1]); + let err = results.into_iter().next().unwrap().result.unwrap_err(); assert!(err.to_string().contains("inconsistent WAL policy in batch")); for logical_region_id in [logical_region_1, logical_region_2] { assert!( @@ -1392,12 +1397,12 @@ mod tests { } } - let affected_rows = env + let results = env .metric() - .inner .put_regions_batch(build_requests().into_iter()) - .await - .unwrap(); + .await; + assert_eq!(results.len(), 1); + let affected_rows = results.into_iter().next().unwrap().result.unwrap(); assert_eq!(affected_rows, 5); let request = ScanRequest::default(); @@ -1736,11 +1741,9 @@ mod tests { ]; // Batch write - let affected_rows = engine - .inner - .put_regions_batch(requests.into_iter()) - .await - .unwrap(); + let results = engine.put_regions_batch(requests.into_iter()).await; + assert_eq!(results.len(), 1); + let affected_rows = results.into_iter().next().unwrap().result.unwrap(); assert_eq!(affected_rows, 10); // Verify physical region contains data from all logical regions @@ -1760,72 +1763,70 @@ mod tests { async fn test_batch_write_with_partial_failure() { let env = TestEnv::new().await; env.init_metric_region().await; - let engine = env.metric(); - let physical_region_id = env.default_physical_region_id(); - let logical_region_1 = env.default_logical_region_id(); - let logical_region_2 = RegionId::new(1024, 1); - let nonexistent_region = RegionId::new(9999, 9999); - - env.create_logical_region(physical_region_id, logical_region_2) + let physical = env.default_physical_region_id(); + let logical = env.default_logical_region_id(); + let other_physical = RegionId::new(4096, 0); + let other_logical = RegionId::new(4097, 0); + env.create_physical_region(other_physical, "batch_results", vec![]) .await; - - // Prepare batch with one invalid region - let schema = test_util::row_schema_with_tags(&["job"]); - let requests = vec![ - ( - logical_region_1, - RegionPutRequest { - skip_wal: false, - rows: Rows { - schema: schema.clone(), - rows: test_util::build_rows(1, 3), - }, - hint: None, - partition_expr_version: None, + env.create_logical_region(other_physical, other_logical) + .await; + let request = |index: i64, version| { + let mut rows = test_util::build_rows(1, 1); + rows[0].values[0].value_data = + Some(api::v1::value::ValueData::TimestampMillisecondValue(index)); + RegionPutRequest { + rows: Rows { + schema: test_util::row_schema_with_tags(&["job"]), + rows, }, - ), - ( - nonexistent_region, - RegionPutRequest { - skip_wal: false, - rows: Rows { - schema: schema.clone(), - rows: test_util::build_rows(1, 2), - }, - hint: None, - partition_expr_version: None, - }, - ), - ( - logical_region_2, - RegionPutRequest { - skip_wal: false, - rows: Rows { - schema: schema.clone(), - rows: test_util::build_rows(1, 5), - }, - hint: None, - partition_expr_version: None, - }, - ), - ]; - - // Batch write - let result = engine.inner.put_regions_batch(requests.into_iter()).await; - assert!(result.is_err()); - - // Invalid region is detected before any write, so the physical region remains empty. - // Fail-fast is per physical-region group; cross-group partial success is possible. - let request = ScanRequest::default(); - let stream = env + skip_wal: false, + hint: None, + partition_expr_version: version, + } + }; + let mut results = env .metric() - .scan_to_stream(physical_region_id, request) - .await - .unwrap(); - let batches = RecordBatches::try_collect(stream).await.unwrap(); - - assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 0); + .put_regions_batch( + vec![ + (logical, request(0, None)), + (RegionId::new(9999, 0), request(1, None)), + (other_logical, request(2, Some(1))), + (other_logical, request(3, None)), + (logical, request(4, None)), + (physical, request(5, None)), + ] + .into_iter(), + ) + .await; + results.sort_by_key(|result| result.request_indices[0]); + assert_eq!( + results + .iter() + .map(|result| result.request_indices.clone()) + .collect::>(), + vec![vec![0, 4], vec![1], vec![2, 3], vec![5]] + ); + assert_eq!(results[0].result.as_ref().unwrap(), &2); + assert!(results[1].result.is_err()); + assert!(results[2].result.is_err()); + assert!(matches!( + results[3].result, + Err(crate::error::Error::ForbiddenPhysicalWrite { .. }) + )); + for (region, expected) in [(physical, 2), (other_physical, 0)] { + let stream = env + .metric() + .scan_to_stream(region, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + expected + ); + } } #[tokio::test] @@ -1849,11 +1850,10 @@ mod tests { }, )]; - let err = engine - .inner - .put_regions_batch(requests.into_iter()) - .await - .unwrap_err(); + let results = engine.put_regions_batch(requests.into_iter()).await; + assert_eq!(results.len(), 1); + assert_eq!(results[0].request_indices, vec![0]); + let err = results.into_iter().next().unwrap().result.unwrap_err(); assert!(matches!( err, @@ -1897,15 +1897,15 @@ mod tests { ), ]; - let err = engine - .inner - .put_regions_batch(requests.into_iter()) - .await - .unwrap_err(); - + let mut results = engine.put_regions_batch(requests.into_iter()).await; + results.sort_by_key(|group| group.request_indices[0]); + assert_eq!(results.len(), 2); + assert_eq!(results[0].request_indices, vec![0]); + assert_eq!(results[0].result.as_ref().unwrap(), &1); + assert_eq!(results[1].request_indices, vec![1]); assert!(matches!( - err, - crate::error::Error::ForbiddenPhysicalWrite { .. } + results[1].result, + Err(crate::error::Error::ForbiddenPhysicalWrite { .. }) )); } @@ -1932,11 +1932,10 @@ mod tests { }, )]; - let affected_rows = engine - .inner - .put_regions_batch(requests.into_iter()) - .await - .unwrap(); + let results = engine.put_regions_batch(requests.into_iter()).await; + assert_eq!(results.len(), 1); + assert_eq!(results[0].request_indices, vec![0]); + let affected_rows = results.into_iter().next().unwrap().result.unwrap(); assert_eq!(affected_rows, 5); } @@ -1946,15 +1945,12 @@ mod tests { env.init_metric_region().await; let engine = env.metric(); - // Empty batch should return zero affected rows - let requests = vec![]; - let affected_rows = engine - .inner - .put_regions_batch(requests.into_iter()) - .await - .unwrap(); - - assert_eq!(affected_rows, 0); + assert!( + engine + .put_regions_batch(vec![].into_iter()) + .await + .is_empty() + ); } #[tokio::test] diff --git a/src/operator/src/write_dispatcher.rs b/src/operator/src/write_dispatcher.rs index 9a042cfb4e7..d8a6ac9905e 100644 --- a/src/operator/src/write_dispatcher.rs +++ b/src/operator/src/write_dispatcher.rs @@ -203,10 +203,13 @@ impl WriteDispatcher { pub(crate) async fn send( &self, peer: PeerRef, - request: RegionRequest, + mut request: RegionRequest, histogram: Option<&Histogram>, ) -> WriteAttempt { if let Some(retry) = &self.retry { + if matches!(&request.body, Some(Body::Inserts(_) | Body::Deletes(_))) { + api::region::set_accept_partial_result(request.header.get_or_insert_default()); + } return retry .send( RoutedWrite { diff --git a/src/operator/src/write_dispatcher/tests.rs b/src/operator/src/write_dispatcher/tests.rs index 0f9af3b7190..a249a8dfe8e 100644 --- a/src/operator/src/write_dispatcher/tests.rs +++ b/src/operator/src/write_dispatcher/tests.rs @@ -356,6 +356,15 @@ async fn test_write_retry_disabled_and_enabled_ablation() { ); assert_eq!(result.is_ok(), max_retries > 0); assert_eq!(attempts, max_retries + 1); + for request in server.service.requests.lock().unwrap().iter() { + let accepts_partial = request + .header + .as_ref() + .and_then(|header| header.query_context.as_ref()) + .and_then(|context| context.extensions.get(api::region::ACCEPT_PARTIAL_RESULT)) + .is_some_and(|value| value == "true"); + assert_eq!(accepts_partial, max_retries > 0); + } assert!(server.invalidator.calls.lock().unwrap().is_empty()); } } @@ -1709,11 +1718,11 @@ async fn check_write_retry_rebuilds_logical_rows_in_mixed_table_envelope(kind: W ); 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, rebuilt_upper, logical, rebuilt_lower], - "{kind:?}" - ); + let mut expected = vec![physical, rebuilt_upper, logical, rebuilt_lower]; + for request in &mut expected { + api::region::set_accept_partial_result(request.header.get_or_insert_default()); + } + assert_eq!(requests, expected, "{kind:?}"); let calls = invalidator.calls.lock().unwrap(); assert_eq!(calls.len(), 2, "{kind:?}"); assert_eq!(calls[0].len(), 2, "{kind:?}"); diff --git a/src/servers/src/error.rs b/src/servers/src/error.rs index 6e558c8234b..2e17179206e 100644 --- a/src/servers/src/error.rs +++ b/src/servers/src/error.rs @@ -1014,6 +1014,7 @@ pub fn status_code_to_http_status(status_code: &StatusCode) -> HttpStatusCode { StatusCode::Internal | StatusCode::Unexpected | StatusCode::IllegalState + | StatusCode::PartialFailure | StatusCode::Unknown | StatusCode::RuntimeResourcesExhausted | StatusCode::EngineExecuteQuery => HttpStatusCode::INTERNAL_SERVER_ERROR, diff --git a/src/servers/src/mysql/writer.rs b/src/servers/src/mysql/writer.rs index d9253c5bea7..049f925c650 100644 --- a/src/servers/src/mysql/writer.rs +++ b/src/servers/src/mysql/writer.rs @@ -512,7 +512,9 @@ pub fn create_mysql_column_def(schema: &SchemaRef) -> Result> { fn mysql_error_kind(status_code: &StatusCode) -> ErrorKind { match status_code { StatusCode::Success => ErrorKind::ER_YES, - StatusCode::Unknown | StatusCode::External => ErrorKind::ER_UNKNOWN_ERROR, + StatusCode::PartialFailure | StatusCode::Unknown | StatusCode::External => { + ErrorKind::ER_UNKNOWN_ERROR + } StatusCode::Unsupported => ErrorKind::ER_NOT_SUPPORTED_YET, StatusCode::Cancelled | StatusCode::DeadlineExceeded => ErrorKind::ER_QUERY_INTERRUPTED, StatusCode::RuntimeResourcesExhausted => ErrorKind::ER_OUT_OF_RESOURCES, diff --git a/src/servers/src/postgres/types/error.rs b/src/servers/src/postgres/types/error.rs index a76bb4362c4..00d1bbbb1c1 100644 --- a/src/servers/src/postgres/types/error.rs +++ b/src/servers/src/postgres/types/error.rs @@ -382,7 +382,8 @@ impl From for PgErrorCode { StatusCode::DeadlineExceeded => PgErrorCode::Ec57000, StatusCode::External => PgErrorCode::Ec58000, - StatusCode::Unknown + StatusCode::PartialFailure + | StatusCode::Unknown | StatusCode::Unexpected | StatusCode::IllegalState | StatusCode::Internal => PgErrorCode::EcXX000, diff --git a/src/standalone/Cargo.toml b/src/standalone/Cargo.toml index a36983203ee..d89366d2e7e 100644 --- a/src/standalone/Cargo.toml +++ b/src/standalone/Cargo.toml @@ -46,5 +46,7 @@ tokio.workspace = true url.workspace = true [dev-dependencies] +datanode = { workspace = true, features = ["testing"] } +mito2 = { workspace = true, features = ["test"] } temp-env.workspace = true toml.workspace = true diff --git a/src/standalone/src/datanode_manager.rs b/src/standalone/src/datanode_manager.rs index 14104eacc2f..6eb704ed03c 100644 --- a/src/standalone/src/datanode_manager.rs +++ b/src/standalone/src/datanode_manager.rs @@ -17,7 +17,7 @@ use std::sync::Arc; use api::region::RegionResponse; use api::v1::region::{RegionRequest, RegionResponse as RegionResponseV1}; use async_trait::async_trait; -use client::region::check_response_header; +use client::region; use common_error::ext::BoxedError; use common_meta::error::{self as meta_error, Result as MetaResult}; use common_meta::node_manager::{ @@ -96,10 +96,7 @@ impl Datanode for RegionInvoker { .await .map_err(BoxedError::new) .context(meta_error::ExternalSnafu)?; - check_response_header(&response.header) - .map_err(BoxedError::new) - .context(meta_error::ExternalSnafu)?; - Ok(RegionResponse::from_region_response(response)) + region::decode_region_response(response).map_err(Into::into) } async fn handle_query(&self, request: QueryRequest) -> MetaResult { @@ -121,3 +118,65 @@ impl Datanode for RegionInvoker { .context(meta_error::ExternalSnafu) } } + +#[cfg(test)] +mod tests { + use api::v1::Rows; + use api::v1::region::{InsertRequest, InsertRequests, region_request}; + use common_error::ext::ErrorExt; + use common_error::mock::MockError; + use common_error::status_code::StatusCode; + use datanode::tests::{MockRegionEngine, mock_region_server}; + use mito2::engine::MITO_ENGINE_NAME; + use mito2::test_util::CreateRequestBuilder; + use store_api::region_request::RegionRequest as StoreRegionRequest; + use store_api::storage::RegionId; + + use super::*; + + #[tokio::test] + async fn test_write_preserves_original_error() { + for code in [StatusCode::StorageUnavailable, StatusCode::InvalidArguments] { + let mut server = mock_region_server(); + let (engine, _receiver) = MockRegionEngine::with_mock_fn( + MITO_ENGINE_NAME, + Box::new(move |region_id, request| { + if matches!(request, StoreRegionRequest::Put(_)) { + Err(BoxedError::new(MockError::new(code))) + .context(datanode::error::HandleRegionRequestSnafu { region_id }) + } else { + Ok(0) + } + }), + ); + server.register_engine(engine); + let mut requests = Vec::new(); + for number in 0..2 { + let region_id = RegionId::new(1, number); + server + .handle_request( + region_id, + StoreRegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + requests.push(InsertRequest { + region_id: region_id.as_u64(), + rows: Some(Rows::default()), + ..Default::default() + }); + } + let datanode: DatanodeRef = RegionInvoker::arc(server); + let error = datanode + .handle(RegionRequest { + header: Some(Default::default()), + body: Some(region_request::Body::Inserts(InsertRequests { requests })), + }) + .await + .unwrap_err(); + assert_eq!(code, error.status_code()); + assert_eq!(MockError::new(code).retry_hint(), error.retry_hint()); + assert!(!matches!(error, meta_error::Error::PartialRegion { .. })); + } + } +} diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index ebb2f7c3684..f0650f8d7dc 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -213,47 +213,55 @@ impl RegionRequest { } } +/// Converts one insert, treating absent rows as a no-op like vector inserts. +pub fn make_region_put( + request: api::v1::region::InsertRequest, +) -> Option<(RegionId, RegionPutRequest)> { + request.rows.map(|rows| { + ( + request.region_id.into(), + RegionPutRequest { + rows, + hint: None, + skip_wal: request.skip_wal, + partition_expr_version: request.partition_expr_version.map(|version| version.value), + }, + ) + }) +} + fn make_region_puts(inserts: InsertRequests) -> Result> { - let requests = inserts + Ok(inserts .requests .into_iter() - .filter_map(|r| { - let region_id = r.region_id.into(); - r.rows.map(|rows| { - ( - region_id, - RegionRequest::Put(RegionPutRequest { - rows, - hint: None, - skip_wal: r.skip_wal, - partition_expr_version: r.partition_expr_version.map(|v| v.value), - }), - ) - }) - }) - .collect(); - Ok(requests) + .filter_map(make_region_put) + .map(|(region_id, request)| (region_id, RegionRequest::Put(request))) + .collect()) +} + +/// Converts one delete, treating absent rows as a no-op like vector deletes. +pub fn make_region_delete( + request: api::v1::region::DeleteRequest, +) -> Option<(RegionId, RegionDeleteRequest)> { + request.rows.map(|rows| { + ( + request.region_id.into(), + RegionDeleteRequest { + rows, + hint: None, + partition_expr_version: request.partition_expr_version.map(|version| version.value), + }, + ) + }) } fn make_region_deletes(deletes: DeleteRequests) -> Result> { - let requests = deletes + Ok(deletes .requests .into_iter() - .filter_map(|r| { - let region_id = r.region_id.into(); - r.rows.map(|rows| { - ( - region_id, - RegionRequest::Delete(RegionDeleteRequest { - rows, - hint: None, - partition_expr_version: r.partition_expr_version.map(|v| v.value), - }), - ) - }) - }) - .collect(); - Ok(requests) + .filter_map(make_region_delete) + .map(|(region_id, request)| (region_id, RegionRequest::Delete(request))) + .collect()) } fn parse_region_create(create: CreateRequest) -> Result<(RegionId, RegionCreateRequest)> { diff --git a/tests/runner/src/client.rs b/tests/runner/src/client.rs index a595c877940..03f7f7ed1f5 100644 --- a/tests/runner/src/client.rs +++ b/tests/runner/src/client.rs @@ -169,6 +169,7 @@ impl MultiProtocolClient { result = ServerSnafu { code: status_code, msg, + retry_hint: e.retry_hint(), } .fail(); }