fix(mito2): distinguish staging rejection from write backpressure

Signed-off-by: WenyXu <wenymedia@gmail.com>

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::<Error>().unwrap(),
         Error::RejectWrite { .. }

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::<Error>().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::<Error>().unwrap(),
+            Error::RejectWriteInStaging { .. }
+        );
+    }
 }
 
 #[tokio::test]

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);
         }
     }
 

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<S> RegionWorkerLoop<S> {
                         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<S> RegionWorkerLoop<S> {
             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<S> RegionWorkerLoop<S> {
                     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<S> RegionWorkerLoop<S> {
                 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();