diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index bd0e6827880..c54691a5a8f 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/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::::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); +} diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index f8d7eb9c840..8c306effd1d 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -470,13 +470,13 @@ pub(crate) mod test_hooks { release: watch::Sender, } - static ACTIVE_BARRIER: Mutex> = Mutex::new(None); + static ACTIVE_BARRIERS: Mutex> = Mutex::new(Vec::new()); static NEXT_BARRIER_ID: AtomicU64 = AtomicU64::new(1); - fn lock_active_barrier() -> std::sync::MutexGuard<'static, Option> { + fn lock_active_barriers() -> std::sync::MutexGuard<'static, Vec> { // 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; diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 721235e71b5..2b9551308f0 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -564,7 +564,7 @@ impl RegionWorkerLoop { // Collect requests by region. if !region_ctx.push_bulk(bulk_req.sender, bulk_req.request, None, bulk_req.skip_wal) { - return; + continue; } } }