fix(mito2): keep processing bulk requests after one fails to encode (#9638)

* fix(mito2): keep processing bulk requests after one fails to encode

When `push_bulk` fails to build the WAL entry of a bulk request, it answers
that request with the error. `process_bulk_requests` then returned from the
loop, so the remaining requests of the drained batch were dropped without an
answer and their callers only saw a receive error. Skip the failed request
and continue with the rest of the batch instead.

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* test(mito2): keep bulk install barriers independent per target

The bulk install test barrier was a single process-wide slot, so tests
that arm it while running in the same process replaced each other's
barrier. Keep one barrier per armed target instead.

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

---------

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

src/mito2/src/engine/basic_test.rs

@@ -16,16 +16,22 @@
 
 use std::collections::HashMap;
 use std::sync::atomic::Ordering;
+use std::time::Duration;
 
 use api::v1::helper::row;
 use api::v1::value::ValueData;
-use api::v1::{Rows, SemanticType};
+use api::v1::{ArrowIpc, Rows, SemanticType};
 use common_base::readable_size::ReadableSize;
 use common_error::ext::ErrorExt;
 use common_error::status_code::StatusCode;
-use common_recordbatch::RecordBatches;
+use common_recordbatch::{DfRecordBatch, RecordBatches};
+use common_test_util::flight::encode_to_flight_data;
 use common_time::Timestamp;
 use common_wal::options::WAL_OPTIONS_KEY;
+use datatypes::arrow::array::{
+    ArrayRef, DictionaryArray, Float64Array, StringArray, TimestampMillisecondArray,
+};
+use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit, UInt32Type};
 use datatypes::prelude::ConcreteDataType;
 use datatypes::schema::ColumnSchema;
 use futures::TryStreamExt;
@@ -35,13 +41,14 @@ use rstest::rstest;
 use rstest_reuse::{self, apply};
 use store_api::metadata::ColumnMetadata;
 use store_api::region_request::{
-    PathType, RegionCreateRequest, RegionFlushRequest, RegionOpenRequest, RegionPutRequest,
-    RegionRequirements,
+    PathType, RegionBulkInsertsRequest, RegionCreateRequest, RegionFlushRequest, RegionOpenRequest,
+    RegionPutRequest, RegionRequirements,
 };
 use store_api::storage::RegionId;
 
 use super::*;
 use crate::region::version::VersionControlData;
+use crate::region_write_ctx::test_hooks::arm_bulk_install_barrier;
 use crate::test_util::{
     CreateRequestBuilder, LogStoreFactory, TestEnv, build_delete_rows_for_key, build_rows,
     build_rows_for_key, delete_rows, delete_rows_schema, flush_region, kafka_log_store_factory,
@@ -1872,3 +1879,168 @@ async fn restart_without_region_close(
         .unwrap();
     engine
 }
+
+/// Builds a bulk insert request for a region created by [CreateRequestBuilder]
+/// with its default columns `(tag_0, field_0, ts)`.
+fn build_bulk_insert_request(region_id: RegionId, num_rows: usize) -> RegionBulkInsertsRequest {
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("tag_0", DataType::Utf8, true),
+        Field::new("field_0", DataType::Float64, true),
+        Field::new(
+            "ts",
+            DataType::Timestamp(TimeUnit::Millisecond, None),
+            false,
+        ),
+    ]));
+    let tag = Arc::new(StringArray::from_iter_values(
+        (0..num_rows).map(|v| format!("a{v}")),
+    )) as ArrayRef;
+    let field = Arc::new(Float64Array::from_iter_values(
+        (0..num_rows).map(|v| v as f64),
+    )) as ArrayRef;
+    let ts = Arc::new(TimestampMillisecondArray::from_iter_values(
+        (0..num_rows).map(|v| v as i64 * 1000),
+    )) as ArrayRef;
+    let payload = DfRecordBatch::try_new(schema, vec![tag, field, ts]).unwrap();
+    let (schema, record_batch) = encode_to_flight_data(payload.clone());
+
+    RegionBulkInsertsRequest {
+        skip_wal: false,
+        region_id,
+        payload,
+        raw_data: ArrowIpc {
+            schema: schema.data_header,
+            data_header: record_batch.data_header,
+            payload: record_batch.data_body,
+        },
+        partition_expr_version: None,
+        aligned_schema_version: None,
+    }
+}
+
+/// Builds a bulk insert request whose WAL entry can't be encoded. The batch
+/// misses `field_0`, so the worker fills it and clears `raw_data`, and its
+/// all-null dictionary `tag_0` column can't be encoded into a single IPC
+/// message.
+fn build_unencodable_bulk_insert_request(region_id: RegionId) -> RegionBulkInsertsRequest {
+    let schema = Arc::new(Schema::new(vec![
+        Field::new(
+            "tag_0",
+            DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)),
+            true,
+        ),
+        Field::new(
+            "ts",
+            DataType::Timestamp(TimeUnit::Millisecond, None),
+            false,
+        ),
+    ]));
+    let tag = Arc::new(DictionaryArray::<UInt32Type>::from_iter([None::<&str>])) as ArrayRef;
+    let ts = Arc::new(TimestampMillisecondArray::from_iter_values([0])) as ArrayRef;
+    let payload = DfRecordBatch::try_new(schema, vec![tag, ts]).unwrap();
+
+    RegionBulkInsertsRequest {
+        skip_wal: false,
+        region_id,
+        payload,
+        // The worker clears it after filling the missing column.
+        raw_data: ArrowIpc::default(),
+        partition_expr_version: None,
+        aligned_schema_version: None,
+    }
+}
+
+async fn scan_num_rows(engine: &MitoEngine, region_id: RegionId) -> usize {
+    let stream = engine
+        .scan_to_stream(region_id, ScanRequest::default())
+        .await
+        .unwrap();
+    let batches = RecordBatches::try_collect(stream).await.unwrap();
+    batches.iter().map(|batch| batch.num_rows()).sum()
+}
+
+#[tokio::test]
+async fn test_bulk_insert_encode_failure_keeps_worker_batch() {
+    common_telemetry::init_default_ut_logging();
+
+    let mut env = TestEnv::new().await;
+    // A single worker so requests of all regions share one batch.
+    let engine = env
+        .create_engine(MitoConfig {
+            num_workers: 1,
+            ..Default::default()
+        })
+        .await;
+
+    let blocker_region_id = RegionId::new(1, 1);
+    let region_id = RegionId::new(1, 2);
+    let other_region_id = RegionId::new(1, 3);
+    for id in [blocker_region_id, region_id, other_region_id] {
+        engine
+            .handle_request(
+                id,
+                RegionRequest::Create(CreateRequestBuilder::new().build()),
+            )
+            .await
+            .unwrap();
+    }
+
+    // Pauses the worker in a bulk write of the blocker region so the requests
+    // below queue up and reach the worker as one batch.
+    let mut barrier = arm_bulk_install_barrier(
+        blocker_region_id,
+        engine
+            .get_region(blocker_region_id)
+            .unwrap()
+            .version_control
+            .clone(),
+    );
+    let blocker_engine = engine.clone();
+    let blocker = tokio::spawn(async move {
+        blocker_engine
+            .handle_request(
+                blocker_region_id,
+                RegionRequest::BulkInserts(build_bulk_insert_request(blocker_region_id, 1)),
+            )
+            .await
+    });
+    tokio::time::timeout(Duration::from_secs(10), barrier.wait_until_reached())
+        .await
+        .expect("bulk write never reached the install barrier");
+
+    let mut receivers = Vec::new();
+    for (id, request) in [
+        (region_id, build_unencodable_bulk_insert_request(region_id)),
+        (region_id, build_bulk_insert_request(region_id, 3)),
+        (
+            other_region_id,
+            build_bulk_insert_request(other_region_id, 2),
+        ),
+    ] {
+        let metadata = engine.inner.get_metadata(id).ok();
+        let (request, receiver) = WorkerRequest::try_from_region_request(
+            id,
+            RegionRequest::BulkInserts(request),
+            metadata,
+        )
+        .unwrap();
+        engine
+            .inner
+            .workers
+            .submit_to_worker(id, request)
+            .await
+            .unwrap();
+        receivers.push(receiver);
+    }
+    barrier.release();
+    assert_eq!(1, blocker.await.unwrap().unwrap().affected_rows);
+
+    let [failed, same_region, other_region] = receivers.try_into().unwrap();
+    let err = failed.await.unwrap().unwrap_err();
+    assert_eq!(StatusCode::Unsupported, err.status_code(), "{err:?}");
+    // Requests after the failed one in the same batch are still written.
+    assert_eq!(3, same_region.await.unwrap().unwrap());
+    assert_eq!(2, other_region.await.unwrap().unwrap());
+    assert_eq!(3, scan_num_rows(&engine, region_id).await);
+    assert_eq!(2, scan_num_rows(&engine, other_region_id).await);
+}

src/mito2/src/region_write_ctx.rs

@@ -470,13 +470,13 @@ pub(crate) mod test_hooks {
         release: watch::Sender<bool>,
     }
 
-    static ACTIVE_BARRIER: Mutex<Option<ActiveBarrier>> = Mutex::new(None);
+    static ACTIVE_BARRIERS: Mutex<Vec<ActiveBarrier>> = Mutex::new(Vec::new());
     static NEXT_BARRIER_ID: AtomicU64 = AtomicU64::new(1);
 
-    fn lock_active_barrier() -> std::sync::MutexGuard<'static, Option<ActiveBarrier>> {
+    fn lock_active_barriers() -> std::sync::MutexGuard<'static, Vec<ActiveBarrier>> {
         // Never let a poisoned mutex (e.g. a panic in another test while
         // holding the lock) hang or break unrelated tests.
-        ACTIVE_BARRIER
+        ACTIVE_BARRIERS
             .lock()
             .unwrap_or_else(|poisoned| poisoned.into_inner())
     }
@@ -513,8 +513,8 @@ pub(crate) mod test_hooks {
         }
     }
 
-    /// Arms the bulk-install barrier and returns the owning guard; any
-    /// previously armed barrier is replaced.
+    /// Arms a bulk-install barrier for the target and returns the owning
+    /// guard; barriers of concurrently running tests stay independent.
     pub(crate) fn arm_bulk_install_barrier(
         target_region_id: RegionId,
         target_version_control: VersionControlRef,
@@ -522,8 +522,7 @@ pub(crate) mod test_hooks {
         let (reached_tx, reached_rx) = watch::channel(false);
         let (release_tx, _release_rx) = watch::channel(false);
         let id = NEXT_BARRIER_ID.fetch_add(1, Ordering::Relaxed);
-        let mut active = lock_active_barrier();
-        *active = Some(ActiveBarrier {
+        lock_active_barriers().push(ActiveBarrier {
             id,
             target_region_id,
             target_version_control,
@@ -539,10 +538,7 @@ pub(crate) mod test_hooks {
     }
 
     fn disarm_barrier(id: u64) {
-        let mut active = lock_active_barrier();
-        if active.as_ref().is_some_and(|barrier| barrier.id == id) {
-            *active = None;
-        }
+        lock_active_barriers().retain(|barrier| barrier.id != id);
     }
 
     /// Pauses a bulk write before installing its parts until the barrier is
@@ -552,19 +548,14 @@ pub(crate) mod test_hooks {
         version_control: &VersionControlRef,
     ) {
         let (reached_tx, release_rx) = {
-            let active = lock_active_barrier();
-            match active.as_ref() {
-                Some(barrier)
-                    if barrier.target_region_id == region_id
-                        && std::sync::Arc::ptr_eq(
-                            &barrier.target_version_control,
-                            version_control,
-                        ) =>
-                {
-                    (barrier.reached.clone(), barrier.release.subscribe())
-                }
-                _ => return,
-            }
+            let active = lock_active_barriers();
+            let Some(barrier) = active.iter().find(|barrier| {
+                barrier.target_region_id == region_id
+                    && std::sync::Arc::ptr_eq(&barrier.target_version_control, version_control)
+            }) else {
+                return;
+            };
+            (barrier.reached.clone(), barrier.release.subscribe())
         };
         let _ = reached_tx.send(true);
         let mut release_rx = release_rx;

src/mito2/src/worker/handle_write.rs

@@ -564,7 +564,7 @@ impl<S> RegionWorkerLoop<S> {
 
             // Collect requests by region.
             if !region_ctx.push_bulk(bulk_req.sender, bulk_req.request, None, bulk_req.skip_wal) {
-                return;
+                continue;
             }
         }
     }