From 3c5f505bd4da8d6e658b5854c88a8bfde60add7e Mon Sep 17 00:00:00 2001 From: WenyXu Date: Mon, 12 Oct 2026 12:02:04 +0000 Subject: [PATCH] fix(mito2): distinguish staging rejection from write backpressure Signed-off-by: WenyXu --- src/mito2/src/engine/flush_test.rs | 2 + src/mito2/src/engine/staging_test.rs | 73 ++++++++++++++++++++-------- src/mito2/src/error.rs | 23 +++++++-- src/mito2/src/worker/handle_write.rs | 14 ++++-- 4 files changed, 82 insertions(+), 30 deletions(-) diff --git a/src/mito2/src/engine/flush_test.rs b/src/mito2/src/engine/flush_test.rs index b040ae5ec54..529b57b3578 100644 --- a/src/mito2/src/engine/flush_test.rs +++ b/src/mito2/src/engine/flush_test.rs @@ -24,6 +24,7 @@ use api::v1::Rows; use async_trait::async_trait; use common_base::Plugins; use common_base::readable_size::ReadableSize; +use common_error::ext::{ErrorExt, RetryHint}; use common_recordbatch::RecordBatches; use common_time::util::current_time_millis; use common_wal::options::{KafkaWalOptions, WAL_OPTIONS_KEY, WalOptions}; @@ -1003,6 +1004,7 @@ async fn test_region_write_buffer_rejects_only_full_region_queue() { .expect("the stalled queue should be rejected at the region hard limit") .unwrap() .unwrap_err(); + assert_eq!(err.retry_hint(), RetryHint::Retryable); assert_matches!( err.into_inner().as_any().downcast_ref::().unwrap(), Error::RejectWrite { .. } diff --git a/src/mito2/src/engine/staging_test.rs b/src/mito2/src/engine/staging_test.rs index bd823708509..1b2a7ffee13 100644 --- a/src/mito2/src/engine/staging_test.rs +++ b/src/mito2/src/engine/staging_test.rs @@ -18,11 +18,14 @@ use std::sync::Arc; use std::time::Duration; use std::{assert_matches, fs}; -use api::v1::Rows; -use common_error::ext::ErrorExt; +use api::v1::{ArrowIpc, Rows}; +use common_error::ext::{ErrorExt, RetryHint}; use common_error::status_code::StatusCode; use common_function::utils::partition_expr_version; +use common_grpc::flight; use common_recordbatch::RecordBatches; +use datatypes::arrow::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray}; +use datatypes::arrow::record_batch::RecordBatch; use datatypes::value::Value; use object_store::Buffer; use object_store::layers::mock::{ @@ -34,8 +37,9 @@ use store_api::region_engine::{ RegionEngine, RemapManifestsRequest, SetRegionRoleStateResponse, SettableRegionRoleState, }; use store_api::region_request::{ - ApplyStagingManifestRequest, EnterStagingRequest, RegionAlterRequest, RegionFlushRequest, - RegionPutRequest, RegionRequest, RegionTruncateRequest, StagingPartitionDirective, + ApplyStagingManifestRequest, EnterStagingRequest, RegionAlterRequest, RegionBulkInsertsRequest, + RegionFlushRequest, RegionPutRequest, RegionRequest, RegionTruncateRequest, + StagingPartitionDirective, }; use store_api::storage::{RegionId, ScanRequest}; @@ -254,7 +258,7 @@ fn default_partition_expr() -> String { } #[tokio::test] -async fn test_staging_reject_all_writes_rejects_put() { +async fn test_staging_reject_all_writes_rejects_put_and_bulk() { let mut env = TestEnv::new().await; let engine = env.create_engine(MitoConfig::default()).await; @@ -281,23 +285,50 @@ async fn test_staging_reject_all_writes_rejects_put() { schema: column_schemas, rows: build_rows(0, 2), }; - let err = engine - .handle_request( + let payload = RecordBatch::try_from_iter([ + ( + "tag_0", + Arc::new(StringArray::from(vec!["0", "1"])) as ArrayRef, + ), + ( + "field_0", + Arc::new(Float64Array::from(vec![0.0, 1.0])) as ArrayRef, + ), + ( + "ts", + Arc::new(TimestampMillisecondArray::from(vec![0, 1000])) as ArrayRef, + ), + ]) + .unwrap(); + let (schema, data_header, data) = flight::record_batch_to_ipc(payload.clone()).unwrap(); + for request in [ + RegionRequest::Put(RegionPutRequest { + skip_wal: false, + rows, + hint: None, + partition_expr_version: None, + }), + RegionRequest::BulkInserts(RegionBulkInsertsRequest { + skip_wal: false, region_id, - RegionRequest::Put(RegionPutRequest { - skip_wal: false, - rows, - hint: None, - partition_expr_version: None, - }), - ) - .await - .unwrap_err(); - assert_eq!(err.status_code(), StatusCode::StorageUnavailable); - assert_matches!( - err.into_inner().as_any().downcast_ref::().unwrap(), - Error::RejectWrite { .. } - ); + payload, + raw_data: ArrowIpc { + schema, + data_header, + payload: data, + }, + partition_expr_version: None, + aligned_schema_version: None, + }), + ] { + let err = engine.handle_request(region_id, request).await.unwrap_err(); + assert_eq!(err.status_code(), StatusCode::StorageUnavailable); + assert_eq!(err.retry_hint(), RetryHint::RetryAfterRouteRefresh); + assert_matches!( + err.into_inner().as_any().downcast_ref::().unwrap(), + Error::RejectWriteInStaging { .. } + ); + } } #[tokio::test] diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index ef4c847f758..e3fc49a1af9 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -641,6 +641,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Region {} rejects writes while in staging", region_id))] + RejectWriteInStaging { + region_id: RegionId, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to compact region {}", region_id))] CompactRegion { region_id: RegionId, @@ -1541,7 +1548,7 @@ impl ErrorExt for Error { RegionClosed { .. } => StatusCode::Cancelled, StaleCompactionExecution { .. } => StatusCode::Cancelled, RegionTruncated { .. } => StatusCode::Cancelled, - RejectWrite { .. } => StatusCode::StorageUnavailable, + RejectWrite { .. } | RejectWriteInStaging { .. } => StatusCode::StorageUnavailable, CompactRegion { source, .. } => source.status_code(), EditRegion { source, .. } => source.status_code(), CompatReader { .. } => StatusCode::Unexpected, @@ -1641,13 +1648,14 @@ impl ErrorExt for Error { match self { PartitionExprVersionMismatch { .. } | RegionNotFound { .. } - | RejectWrite { .. } + | RejectWriteInStaging { .. } | RegionState { state: RegionRoleState::Follower, .. } => RetryHint::RetryAfterRouteRefresh, ReadParquet { .. } | WriteParquet { .. } + | RejectWrite { .. } | Download { .. } | Upload { .. } | RegionState { .. } @@ -1761,18 +1769,25 @@ mod tests { } .build(); assert_eq!(leader_error.retry_hint(), RetryHint::Retryable); - for (error, status) in [ + for (error, status, hint) in [ ( RejectWriteSnafu { region_id }.build(), StatusCode::StorageUnavailable, + RetryHint::Retryable, + ), + ( + RejectWriteInStagingSnafu { region_id }.build(), + StatusCode::StorageUnavailable, + RetryHint::RetryAfterRouteRefresh, ), ( RegionNotFoundSnafu { region_id }.build(), StatusCode::RegionNotFound, + RetryHint::RetryAfterRouteRefresh, ), ] { assert_eq!(error.status_code(), status); - assert_eq!(error.retry_hint(), RetryHint::RetryAfterRouteRefresh); + assert_eq!(error.retry_hint(), hint); } } diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 41af569c649..1383a3e9a5d 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -26,7 +26,7 @@ use store_api::storage::RegionId; use crate::error::{ InvalidRequestSnafu, PartitionExprVersionMismatchSnafu, RegionNotFoundSnafu, RegionStateSnafu, - RejectWriteSnafu, Result, + RejectWriteInStagingSnafu, RejectWriteSnafu, Result, }; use crate::metrics; use crate::metrics::{ @@ -335,7 +335,7 @@ impl RegionWorkerLoop { if region.reject_all_writes_in_staging() { sender_req .sender - .send(RejectWriteSnafu { region_id }.fail()); + .send(RejectWriteInStagingSnafu { region_id }.fail()); continue; } @@ -399,7 +399,7 @@ impl RegionWorkerLoop { if region.reject_all_writes_in_staging() { sender_req .sender - .send(RejectWriteSnafu { region_id }.fail()); + .send(RejectWriteInStagingSnafu { region_id }.fail()); continue; } let expected_version = region.expected_partition_expr_version(); @@ -481,7 +481,9 @@ impl RegionWorkerLoop { RegionRoleState::Leader(RegionLeaderState::Writable) | RegionRoleState::Leader(RegionLeaderState::Staging) => { if region.reject_all_writes_in_staging() { - bulk_req.sender.send(RejectWriteSnafu { region_id }.fail()); + bulk_req + .sender + .send(RejectWriteInStagingSnafu { region_id }.fail()); continue; } let region_ctx = RegionWriteCtx::new( @@ -529,7 +531,9 @@ impl RegionWorkerLoop { continue; }; if region.reject_all_writes_in_staging() { - bulk_req.sender.send(RejectWriteSnafu { region_id }.fail()); + bulk_req + .sender + .send(RejectWriteInStagingSnafu { region_id }.fail()); continue; } let expected_version = region.expected_partition_expr_version();