From 8da533c35ffc901db22827c62ed2877f2e33e6c5 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 9 Nov 2026 16:38:59 +0000 Subject: [PATCH] fix(mito2): keep rows written during staging in the WAL (#10253) * fix(mito2): keep rows written during staging in the WAL A flush during staging committed its SSTs to the staging manifest, then dropped the memtable and released the WAL. If staging did not end successfully those rows were gone. A flush that committed after staging had ended lost its rows even though staging succeeded. Staging regions no longer flush. They are skipped by memory-pressure, periodic and region-limit flushes, and their writes are rejected instead of queued at twice the region limit or when the node's staging memtables reach half of the engine's flush limit. A manual or repartition flush on a staging region with data gets a retryable error; the migration pre-flush and the WAL prune flush reply success without flushing. Closing a staging region leaves staging first, so skip-WAL rows are flushed to the normal manifest. Signed-off-by: jeremyhi * fix(mito2): keep accounting and closing staging regions during staged edits Signed-off-by: jeremyhi * fix(mito2): track staged edits without relying on the partition info Signed-off-by: jeremyhi * test(mito2): synchronize staging close tests with worker admission Signed-off-by: jeremyhi * test(mito2): observe the manifest lock wait when closing staging regions Signed-off-by: jeremyhi * fix(mito2): detect staged edits through the staging partition info Signed-off-by: jeremyhi * docs(mito2): record that staging regions must not flush Signed-off-by: jeremyhi * fix(mito2): do not restore staging after a close abandons the edit Signed-off-by: jeremyhi * fix(mito2): reject region edits during a close-time flush Signed-off-by: jeremyhi * style(mito2): format the staging tests Signed-off-by: jeremyhi * fix(mito2): reject staging edits once staging is abandoned Signed-off-by: jeremyhi * fix(mito2): bind apply-staging edits to their partition expression Signed-off-by: jeremyhi * docs(mito2): note that staging regions pin their topic's prune point Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/mito2/AGENTS.md | 3 + src/mito2/src/engine/partition_filter_test.rs | 73 +- src/mito2/src/engine/staging_test.rs | 1105 +++++++++++++++-- src/mito2/src/flush.rs | 35 +- src/mito2/src/region.rs | 60 +- src/mito2/src/request.rs | 12 +- src/mito2/src/worker.rs | 9 + src/mito2/src/worker/handle_apply_staging.rs | 3 +- src/mito2/src/worker/handle_close.rs | 10 + src/mito2/src/worker/handle_flush.rs | 103 +- src/mito2/src/worker/handle_manifest.rs | 40 +- src/store-api/src/region_engine.rs | 4 +- 12 files changed, 1252 insertions(+), 205 deletions(-) diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index ddb7c7c69b8..7018afe5ef3 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -126,6 +126,9 @@ Tests live next to the code as `*_test.rs` (e.g. `src/mito2/src/engine/flush_tes - In flat batches, equal `__primary_key` dictionary keys imply the same series, but not the reverse: `wrap_pk_binary_to_dict` gives every row its own key. Compare key bytes when grouping rows by series. +- Staging regions never flush: a flush would move rows into the staging manifest and + release their WAL, losing them if staging does not end successfully. Every flush + trigger in `worker/handle_flush.rs` must skip them, and `handle_close` leaves staging first. ## Maintenance contract diff --git a/src/mito2/src/engine/partition_filter_test.rs b/src/mito2/src/engine/partition_filter_test.rs index 70a4b7c487f..9e02d59e46e 100644 --- a/src/mito2/src/engine/partition_filter_test.rs +++ b/src/mito2/src/engine/partition_filter_test.rs @@ -20,13 +20,15 @@ use datatypes::arrow::array::AsArray; use datatypes::value::Value; use partition::expr::{Operand, col}; use partition::function::IcebergHash; -use store_api::region_engine::RegionEngine; +use store_api::region_engine::{RegionEngine, RemapManifestsRequest}; use store_api::region_request::{ - EnterStagingRequest, RegionFlushRequest, RegionRequest, StagingPartitionDirective, + ApplyStagingManifestRequest, EnterStagingRequest, RegionFlushRequest, RegionRequest, + StagingPartitionDirective, }; use store_api::storage::{RegionId, ScanRequest}; use crate::config::MitoConfig; +use crate::engine::MitoEngine; use crate::test_util::{CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema}; /// Helper to create a partition expression for testing. @@ -39,6 +41,30 @@ fn range_expr_string(col_name: &str, start: f64, end: f64) -> String { .unwrap() } +/// Ends staging the way repartition does, importing the region's own files. +async fn apply_staging(engine: &MitoEngine, region_id: RegionId, partition_expr: String) { + let result = engine + .remap_manifests(RemapManifestsRequest { + region_id, + input_regions: vec![region_id], + region_mapping: [(region_id, vec![region_id])].into_iter().collect(), + new_partition_exprs: [(region_id, partition_expr.clone())].into_iter().collect(), + }) + .await + .unwrap(); + engine + .handle_request( + region_id, + RegionRequest::ApplyStagingManifest(ApplyStagingManifestRequest { + partition_expr, + central_region_id: region_id, + manifest_path: result.manifest_paths[®ion_id].clone(), + }), + ) + .await + .unwrap(); +} + #[tokio::test] async fn test_partition_filter_basic() { test_partition_filter_basic_with_format(false).await; @@ -51,7 +77,7 @@ async fn test_partition_filter_basic() { /// 1. Creates a region with partition expr `tag_0 >= "0"` (covers all data) /// 2. Writes data (tag_0 = "0"..="9") and flushes - SST has old partition expr /// 3. Enters staging mode with partition expr `tag_0 >= "5"` (narrower) -/// 4. Writes more data and flushes in staging mode +/// 4. Writes more data in staging mode /// 5. Exits staging mode - region now has new partition expr /// 6. Scans all data - old SST data not matching new expr should be filtered async fn test_partition_filter_basic_with_format(flat_format: bool) { @@ -118,16 +144,7 @@ async fn test_partition_filter_basic_with_format(flat_format: bool) { }; put_rows(&engine, region_id, rows_data).await; - // Flush in staging mode - engine - .handle_request( - region_id, - RegionRequest::Flush(RegionFlushRequest::default()), - ) - .await - .unwrap(); - - // Scan data in staging mode - should only see initial 5 rows (staging SST not visible) + // Scan data in staging mode - rows written in staging stay readable in the memtable let projection = Some(vec![1]); let request = ScanRequest { projection, @@ -138,20 +155,16 @@ async fn test_partition_filter_basic_with_format(flat_format: bool) { let batches = RecordBatches::try_collect(stream).await.unwrap(); let total_rows: usize = batches.iter().map(|rb| rb.num_rows()).sum(); assert_eq!( - total_rows, 5, - "Should see 5 rows before exiting staging mode" + total_rows, 11, + "Should see 11 rows before exiting staging mode" ); // Exit staging mode - use store_api::region_engine::SettableRegionRoleState; - engine - .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Leader) - .await - .unwrap(); + apply_staging(&engine, region_id, new_partition_expr).await; // Scan after exiting staging - the old SST (tag_0 = "0".."4") should have // rows filtered by partition expr (tag_0 >= "5"), which means none of them pass. - // But the staging SST (tag_0 = "5".."10") satisfies the partition expr. + // But the memtable rows (tag_0 = "5".."10") satisfy the partition expr. let projection = Some(vec![1]); let request = ScanRequest { projection, @@ -163,10 +176,10 @@ async fn test_partition_filter_basic_with_format(flat_format: bool) { let total_rows: usize = batches.iter().map(|rb| rb.num_rows()).sum(); // After exit: // - Old SST (rows 0-4): partition expr is "tag_0 >= 5", so these are filtered out - // - Staging SST (rows 5-10): These satisfy partition expr, so they're visible + // - Memtable (rows 5-10): These satisfy partition expr, so they're visible assert_eq!( total_rows, 6, - "Should see 6 rows after exiting staging (rows 5-10 from staging SST), flat_format: {}", + "Should see 6 rows after exiting staging (rows 5-10 from the memtable), flat_format: {}", flat_format ); } @@ -220,13 +233,16 @@ async fn test_partition_filter_iceberg_hash_with_format(flat_format: bool) { IcebergHash::evaluate_value(&Value::String(i.to_string().as_str().into())).unwrap() < Value::Int32(0) }; - let hash_expr = Operand::IcebergHash(IcebergHash::new("tag_0")).lt(Value::Int32(0)); + let hash_expr = Operand::IcebergHash(IcebergHash::new("tag_0")) + .lt(Value::Int32(0)) + .as_json_str() + .unwrap(); engine .handle_request( region_id, RegionRequest::EnterStaging(EnterStagingRequest { partition_directive: StagingPartitionDirective::UpdatePartitionExpr( - hash_expr.as_json_str().unwrap(), + hash_expr.clone(), ), }), ) @@ -244,12 +260,7 @@ async fn test_partition_filter_iceberg_hash_with_format(flat_format: bool) { }, ) .await; - flush().await; - use store_api::region_engine::SettableRegionRoleState; - engine - .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Leader) - .await - .unwrap(); + apply_staging(&engine, region_id, hash_expr).await; let request = ScanRequest { projection: Some(vec![0]), diff --git a/src/mito2/src/engine/staging_test.rs b/src/mito2/src/engine/staging_test.rs index 5594e5aa7bb..10351f918d9 100644 --- a/src/mito2/src/engine/staging_test.rs +++ b/src/mito2/src/engine/staging_test.rs @@ -14,16 +14,20 @@ //! Integration tests for staging state functionality. +use std::collections::HashMap; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::time::Duration; use std::{assert_matches, fs}; -use api::v1::{ArrowIpc, Rows}; +use api::v1::{ArrowIpc, ColumnSchema, Rows}; +use async_trait::async_trait; 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 common_time::util::current_time_millis; use datatypes::arrow::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::value::Value; @@ -33,28 +37,36 @@ use object_store::layers::mock::{ MockLayerBuilder, OpDelete, Result as MockResult, Write, Writer, }; use partition::expr::{PartitionExpr, col}; +use store_api::mito_engine_options::WRITE_BUFFER_SIZE_KEY; use store_api::region_engine::{ - RegionEngine, RemapManifestsRequest, SetRegionRoleStateResponse, SettableRegionRoleState, + RegionEngine, RegionRole, RemapManifestsRequest, SetRegionRoleStateResponse, + SettableRegionRoleState, }; use store_api::region_request::{ ApplyStagingManifestRequest, EnterStagingRequest, RegionAlterRequest, RegionBulkInsertsRequest, - RegionFlushRequest, RegionPutRequest, RegionRequest, RegionTruncateRequest, - StagingPartitionDirective, + RegionCloseRequest, RegionFlushReason, RegionFlushRequest, RegionPutRequest, RegionRequest, + RegionTruncateRequest, StagingPartitionDirective, }; -use store_api::storage::{RegionId, ScanRequest}; +use store_api::storage::{FileId, RegionId, ScanRequest}; use crate::config::MitoConfig; -use crate::engine::listener::NotifyEnterStagingResultListener; +use crate::engine::MitoEngine; +use crate::engine::listener::{EventListener, FlushListener, NotifyEnterStagingResultListener}; use crate::error::Error; +use crate::flush::WriteBufferManager; use crate::manifest::action::{ RegionChange, RegionEdit, RegionMetaAction, RegionMetaActionList, RegionPartitionExprChange, }; -use crate::region::{RegionLeaderState, RegionRoleState, parse_partition_expr}; -use crate::request::WorkerRequest; +use crate::region::{MitoRegionRef, RegionLeaderState, RegionRoleState}; +use crate::request::{RegionEditRequest, WorkerRequest}; use crate::sst::FormatType; +use crate::sst::file::FileMeta; use crate::test_util::{ - CheckpointTaskBlocker, CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema, + CheckpointTaskBlocker, CreateRequestBuilder, MockWriteBufferManager, TestEnv, build_rows, + flush_region, put_rows, reopen_region, rows_schema, }; +use crate::time_provider::mock::MockTimeProvider; +use crate::worker::MAX_INITIAL_CHECK_DELAY_SECS; fn range_expr(col_name: &str, start: i64, end: i64) -> PartitionExpr { col(col_name) @@ -436,13 +448,6 @@ async fn test_staging_write_partition_expr_version_with_format(flat_format: bool .await .unwrap(); assert_eq!(result.affected_rows, 3); - engine - .handle_request( - region_id, - RegionRequest::Flush(RegionFlushRequest::default()), - ) - .await - .unwrap(); let result = engine .remap_manifests(RemapManifestsRequest { @@ -551,9 +556,6 @@ async fn test_staging_manifest_directory_with_format(flat_format: bool) { let region_id = RegionId::new(1024, 0); let request = CreateRequestBuilder::new().build(); - // Get column schemas before consuming the request - let column_schemas = rows_schema(&request); - // Check manifest files after region creation (before staging mode) // Create region engine @@ -637,22 +639,6 @@ async fn test_staging_manifest_directory_with_format(flat_format: bool) { Error::StagingPartitionExprMismatch { .. } ); - // Put some data and flush in staging mode - let rows_data = Rows { - schema: column_schemas.clone(), - rows: build_rows(0, 3), - }; - put_rows(&engine, region_id, rows_data).await; - - // Force flush to generate manifest files in staging mode - engine - .handle_request( - region_id, - RegionRequest::Flush(RegionFlushRequest::default()), - ) - .await - .unwrap(); - // Check that manifest files are in staging directory let staging_manifest_dir = format!("{}/staging/manifest", region_dir); assert!( @@ -708,7 +694,7 @@ async fn test_staging_exit_success_with_manifests_with_format(flat_format: bool) .await .unwrap(); - // Add some data and flush in staging mode to generate staging manifests + // This data is flushed when entering staging mode. let rows_data = Rows { schema: column_schemas.clone(), rows: build_rows(0, 3), @@ -728,55 +714,55 @@ async fn test_staging_exit_success_with_manifests_with_format(flat_format: bool) .await .unwrap(); - // Add some data and flush in staging mode to generate staging manifests + // Rows written during staging stay in the memtable. let rows_data = Rows { schema: column_schemas.clone(), - rows: build_rows(3, 8), + rows: build_rows(3, 10), }; put_rows(&engine, region_id, rows_data).await; - // Force flush to generate staging manifests + // A source region provides the files that staging imports. + let source_region_id = RegionId::new(1024, 1); engine .handle_request( - region_id, - RegionRequest::Flush(RegionFlushRequest::default()), + source_region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), ) .await .unwrap(); - - // Add more data and flush again to generate multiple staging manifests - let rows_data2 = Rows { - schema: column_schemas.clone(), - rows: build_rows(8, 10), - }; - put_rows(&engine, region_id, rows_data2).await; - + for (start, end) in [(10, 13), (13, 15)] { + let rows_data = Rows { + schema: column_schemas.clone(), + rows: build_rows(start, end), + }; + put_rows(&engine, source_region_id, rows_data).await; + flush_region(&engine, source_region_id, None).await; + } + let source_files = engine + .get_region(source_region_id) + .unwrap() + .manifest_ctx + .manifest() + .await + .files + .values() + .cloned() + .collect::>(); engine .handle_request( - region_id, - RegionRequest::Flush(RegionFlushRequest::default()), + source_region_id, + RegionRequest::Close(RegionCloseRequest::default()), ) .await .unwrap(); let region = engine.get_region(region_id).unwrap(); - let file = { - let manager = region.manifest_ctx.manifest_manager.read().await; - manager - .staging_manifest() - .unwrap() - .files - .values() - .next() - .unwrap() - .clone() - }; - for expected in [11, 12] { + for (file, expected) in source_files.into_iter().zip([11, 12]) { engine .edit_region( region_id, RegionEdit { - files_to_add: vec![file.clone()], + files_to_add: vec![file], files_to_remove: vec![], timestamp_ms: None, compaction_time_window: None, @@ -803,37 +789,19 @@ async fn test_staging_exit_success_with_manifests_with_format(flat_format: bool) .unwrap(); assert_eq!( staging_files_before.len(), - // Two flushes, two edits and entering staging mode. - 5, - "Staging manifest directory should contain 5 files before exit, got: {:?}", + // Two edits and entering staging mode. + 3, + "Staging manifest directory should contain 3 files before exit, got: {:?}", staging_files_before ); { let manager = region.manifest_ctx.manifest_manager.read().await; let staging_manifest = manager.staging_manifest().unwrap(); assert_eq!(staging_manifest.files.len(), 3); - let parquet_bytes: u64 = staging_manifest - .files - .values() - .map(|file| file.file_size) - .sum(); - assert_eq!( - region.region_statistic().ingested_parquet_bytes, - parquet_bytes - ); assert_eq!( staging_manifest.metadata.partition_expr.as_ref().unwrap(), &partition_expr ); - let expr = parse_partition_expr(Some(partition_expr.as_str())) - .unwrap() - .unwrap(); - for file in staging_manifest.files.values() { - let Some(file_expr) = file.partition_expr.as_ref() else { - continue; - }; - assert_eq!(*file_expr, expr); - } } // Count normal manifest files before exit @@ -851,7 +819,7 @@ async fn test_staging_exit_success_with_manifests_with_format(flat_format: bool) "Normal manifest directory should initially contain 2 files" ); - // Try read data before exiting staging, SST files should be invisible + // Try read data before exiting staging, imported SST files should be invisible let request = ScanRequest::default(); let scanner = engine.scanner(region_id, request).await.unwrap(); assert_eq!( @@ -861,15 +829,15 @@ async fn test_staging_exit_success_with_manifests_with_format(flat_format: bool) ); assert_eq!( scanner.num_memtables(), - 0, - "Memtables should be removed in staging before exit" + 1, + "Memtables should be kept in staging before exit" ); let stream = scanner.scan().await.unwrap(); let batches = RecordBatches::try_collect(stream).await.unwrap(); let total_rows: usize = batches.iter().map(|rb| rb.num_rows()).sum(); assert_eq!( - total_rows, 3, - "3 rows should be readable before exit staging mode" + total_rows, 10, + "10 rows should be readable before exit staging mode" ); // Inspect SSTs from manifest @@ -1048,7 +1016,7 @@ async fn test_staging_exit_success_with_manifests_with_format(flat_format: bool) let stream = scanner.scan().await.unwrap(); let batches = RecordBatches::try_collect(stream).await.unwrap(); let total_rows: usize = batches.iter().map(|rb| rb.num_rows()).sum(); - assert_eq!(total_rows, 10, "Expected to read all staged rows"); + assert_eq!(total_rows, 15, "Expected to read all staged rows"); // Inspect SSTs from manifest let sst_entries = engine.all_ssts_from_manifest().await; @@ -1430,6 +1398,30 @@ impl Write for MockWriter { } } +/// Fails closing Parquet files while `fail_parquet` is set. +struct ParquetFailWriter { + path: String, + inner: Writer, + fail_parquet: Arc, +} + +impl Write for ParquetFailWriter { + async fn write(&mut self, bs: Buffer) -> MockResult<()> { + self.inner.write(bs).await + } + + async fn close(&mut self) -> MockResult { + if self.fail_parquet.load(Ordering::Relaxed) && self.path.ends_with(".parquet") { + return Err(MockError::new(ErrorKind::Unexpected, "mock error")); + } + self.inner.close().await + } + + async fn abort(&mut self) -> MockResult<()> { + self.inner.abort().await + } +} + struct MockDeleter { inner: Deleter, } @@ -1523,3 +1515,934 @@ async fn test_enter_staging_save_staging_manifest_error_with_format(flat_format: let mut env = TestEnv::new().await.with_mock_layer(mock_layer); test_enter_staging_error(&mut env, flat_format).await; } + +async fn enter_staging(engine: &MitoEngine, region_id: RegionId) { + engine + .handle_request( + region_id, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive: StagingPartitionDirective::UpdatePartitionExpr( + float_range_expr("field_0", 0., 100.).as_json_str().unwrap(), + ), + }), + ) + .await + .unwrap(); +} + +/// Creates a region and enters staging, returning its schema and table dir. +async fn create_staging_region( + engine: &MitoEngine, + region_id: RegionId, +) -> (Vec, String) { + let request = CreateRequestBuilder::new().build(); + let schema = rows_schema(&request); + let table_dir = request.table_dir.clone(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + enter_staging(engine, region_id).await; + (schema, table_dir) +} + +async fn create_writable_region(engine: &MitoEngine, region_id: RegionId) { + engine + .handle_request( + region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); +} + +async fn put_range( + engine: &MitoEngine, + region_id: RegionId, + schema: &[ColumnSchema], + range: (usize, usize), +) { + let rows = Rows { + schema: schema.to_vec(), + rows: build_rows(range.0, range.1), + }; + put_rows(engine, region_id, rows).await; +} + +/// Sends a put that the engine is expected to refuse with a retryable error. +async fn assert_put_rejected(engine: &MitoEngine, region_id: RegionId, schema: &[ColumnSchema]) { + let err = engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: schema.to_vec(), + rows: build_rows(50, 51), + }, + hint: None, + partition_expr_version: None, + skip_wal: false, + }), + ) + .await + .unwrap_err(); + assert_eq!(RetryHint::Retryable, err.retry_hint()); +} + +async fn put_range_skip_wal( + engine: &MitoEngine, + region_id: RegionId, + schema: &[ColumnSchema], + range: (usize, usize), +) { + engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: schema.to_vec(), + rows: build_rows(range.0, range.1), + }, + hint: None, + partition_expr_version: None, + skip_wal: true, + }), + ) + .await + .unwrap(); +} + +/// Returns one of the files in the region's manifest. +async fn first_file(region: &MitoRegionRef) -> FileMeta { + let manifest = region.manifest_ctx.manifest().await; + manifest.files.values().next().unwrap().clone() +} + +fn edit_adding(file: FileMeta) -> RegionEdit { + RegionEdit { + files_to_add: vec![file], + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + } +} + +/// Spawns an edit importing `file`. +fn spawn_edit( + engine: &MitoEngine, + region_id: RegionId, + file: FileMeta, +) -> tokio::task::JoinHandle> { + let edit_engine = engine.clone(); + tokio::spawn(async move { edit_engine.edit_region(region_id, edit_adding(file)).await }) +} + +/// Starts an edit importing `file` and waits until the region is `Editing`. +/// The caller holds the manifest write lock so that the edit cannot finish. +async fn start_staged_edit( + engine: &MitoEngine, + region: &MitoRegionRef, + file: FileMeta, +) -> tokio::task::JoinHandle> { + let edit = spawn_edit(engine, region.region_id, file); + tokio::time::timeout(Duration::from_secs(5), async { + while region.state() != RegionRoleState::Leader(RegionLeaderState::Editing) { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .unwrap(); + edit +} + +async fn count_rows(engine: &MitoEngine, region_id: RegionId) -> usize { + let scanner = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(scanner.scan().await.unwrap()) + .await + .unwrap(); + batches.iter().map(|batch| batch.num_rows()).sum() +} + +/// Returns the numbers of SST files and memtables that a scan reads. +async fn scan_inputs(engine: &MitoEngine, region_id: RegionId) -> (usize, usize) { + let scanner = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap(); + (scanner.num_files(), scanner.num_memtables()) +} + +/// Asserts that no flush has been scheduled for the rows in the region's memtable. +async fn assert_unflushed(engine: &MitoEngine, region_id: RegionId) { + let version = engine.get_region(region_id).unwrap().version(); + assert!(version.memtables.immutables().is_empty()); + assert_eq!((0, 1), scan_inputs(engine, region_id).await); +} + +#[tokio::test] +async fn test_memory_pressure_keeps_staging_rows_after_rollback() { + let mut write_buffer_manager = MockWriteBufferManager::default(); + write_buffer_manager.set_flush_limit(1 << 20); + write_buffer_manager.reserve_mem(1 << 25); + let write_buffer_manager = Arc::new(write_buffer_manager); + let listener = Arc::new(FlushListener::default()); + let mut env = TestEnv::new().await; + let engine = env + .create_engine_with( + MitoConfig { + num_workers: 1, + ..Default::default() + }, + Some(write_buffer_manager.clone()), + Some(listener.clone()), + None, + ) + .await; + let staging_region_id = RegionId::new(1024, 0); + let writable_region_id = RegionId::new(1024, 1); + let (schema, table_dir) = create_staging_region(&engine, staging_region_id).await; + create_writable_region(&engine, writable_region_id).await; + put_range(&engine, writable_region_id, &schema, (0, 3)).await; + put_range(&engine, staging_region_id, &schema, (0, 3)).await; + + write_buffer_manager.set_should_flush(true); + put_range(&engine, staging_region_id, &schema, (3, 5)).await; + listener.wait().await; + write_buffer_manager.set_should_flush(false); + + assert_eq!((1, 0), scan_inputs(&engine, writable_region_id).await); + assert_unflushed(&engine, staging_region_id).await; + + engine + .set_region_role(staging_region_id, RegionRole::Leader) + .unwrap(); + assert_eq!(5, count_rows(&engine, staging_region_id).await); + reopen_region(&engine, staging_region_id, table_dir, true, HashMap::new()).await; + assert_eq!(5, count_rows(&engine, staging_region_id).await); +} + +#[tokio::test] +async fn test_periodic_flush_skips_staging_region() { + let listener = Arc::new(FlushListener::default()); + let now = current_time_millis(); + let time_provider = Arc::new(MockTimeProvider::new(now)); + let mut env = TestEnv::new().await; + let engine = env + .create_engine_with_time( + MitoConfig { + auto_flush_interval: Duration::from_secs(300), + num_workers: 1, + ..Default::default() + }, + None, + Some(listener.clone()), + time_provider.clone(), + ) + .await; + let staging_region_id = RegionId::new(1024, 0); + let writable_region_id = RegionId::new(1024, 1); + let (schema, _) = create_staging_region(&engine, staging_region_id).await; + create_writable_region(&engine, writable_region_id).await; + put_range(&engine, staging_region_id, &schema, (0, 3)).await; + put_range(&engine, writable_region_id, &schema, (0, 3)).await; + + time_provider.set_now(now + 600_000); + time_provider.set_elapsed((MAX_INITIAL_CHECK_DELAY_SECS as i64 + 1) * 1000); + tokio::time::timeout(Duration::from_secs(3), listener.wait()) + .await + .unwrap(); + + assert_eq!((1, 0), scan_inputs(&engine, writable_region_id).await); + assert_unflushed(&engine, staging_region_id).await; +} + +#[tokio::test] +async fn test_region_write_buffer_limit_rejects_staging_writes_without_flush() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1024, 0); + let request = CreateRequestBuilder::new() + .insert_option(WRITE_BUFFER_SIZE_KEY, "1B") + .build(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + enter_staging(&engine, region_id).await; + + put_range(&engine, region_id, &schema, (0, 3)).await; + assert_put_rejected(&engine, region_id, &schema).await; + + assert_unflushed(&engine, region_id).await; +} + +#[tokio::test] +async fn test_region_write_buffer_limit_does_not_queue_staging_writes() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let probe_region_id = RegionId::new(1024, 0); + let (schema, _) = create_staging_region(&engine, probe_region_id).await; + put_range(&engine, probe_region_id, &schema, (0, 3)).await; + let usage = engine + .get_region(probe_region_id) + .unwrap() + .region_statistic() + .memtable_size; + + let region_id = RegionId::new(1024, 1); + let request = CreateRequestBuilder::new() + .insert_option(WRITE_BUFFER_SIZE_KEY, &format!("{usage}B")) + .build(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + enter_staging(&engine, region_id).await; + put_range(&engine, region_id, &schema, (0, 3)).await; + + tokio::time::timeout( + Duration::from_secs(5), + put_range(&engine, region_id, &schema, (3, 6)), + ) + .await + .unwrap(); + assert_unflushed(&engine, region_id).await; + assert_eq!(6, count_rows(&engine, region_id).await); +} + +#[tokio::test] +async fn test_staging_memtable_cap_rejects_only_staging_writes() { + let mut write_buffer_manager = MockWriteBufferManager::default(); + // Staging writes are rejected once the node's staging memtables hold one byte. + write_buffer_manager.set_flush_limit(2); + let mut env = TestEnv::new().await; + let engine = env + .create_engine_with( + MitoConfig { + num_workers: 2, + ..Default::default() + }, + Some(Arc::new(write_buffer_manager)), + None, + None, + ) + .await; + // The two staging regions belong to different workers. + let first_region_id = RegionId::new(1024, 0); + let second_region_id = RegionId::new(1024, 1); + let writable_region_id = RegionId::new(1024, 2); + let (schema, _) = create_staging_region(&engine, first_region_id).await; + create_staging_region(&engine, second_region_id).await; + create_writable_region(&engine, writable_region_id).await; + + put_range(&engine, first_region_id, &schema, (0, 3)).await; + assert_put_rejected(&engine, second_region_id, &schema).await; + put_range(&engine, writable_region_id, &schema, (0, 3)).await; + + engine + .set_region_role(first_region_id, RegionRole::Leader) + .unwrap(); + put_range(&engine, second_region_id, &schema, (0, 3)).await; +} + +#[tokio::test] +async fn test_staging_memtable_cap_counts_region_during_staged_edit() { + let mut write_buffer_manager = MockWriteBufferManager::default(); + write_buffer_manager.set_flush_limit(2); + let mut env = TestEnv::new().await; + let engine = env + .create_engine_with( + MitoConfig { + num_workers: 2, + ..Default::default() + }, + Some(Arc::new(write_buffer_manager)), + None, + None, + ) + .await; + let edited_region_id = RegionId::new(1024, 0); + let other_region_id = RegionId::new(1024, 1); + let request = CreateRequestBuilder::new().build(); + let schema = rows_schema(&request); + engine + .handle_request(edited_region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_range(&engine, edited_region_id, &schema, (0, 3)).await; + flush_region(&engine, edited_region_id, None).await; + enter_staging(&engine, edited_region_id).await; + put_range(&engine, edited_region_id, &schema, (3, 5)).await; + create_staging_region(&engine, other_region_id).await; + + let region = engine.get_region(edited_region_id).unwrap(); + let file = first_file(®ion).await; + let manager = region.manifest_ctx.manifest_manager.write().await; + let edit = start_staged_edit(&engine, ®ion, file).await; + assert_put_rejected(&engine, other_region_id, &schema).await; + + drop(manager); + edit.await.unwrap().unwrap(); + assert!(region.is_staging()); + assert_put_rejected(&engine, other_region_id, &schema).await; +} + +#[tokio::test] +async fn test_flush_request_on_staging_region_with_data_is_retryable() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1024, 0); + let (schema, _) = create_staging_region(&engine, region_id).await; + flush_region(&engine, region_id, None).await; + put_range(&engine, region_id, &schema, (0, 3)).await; + + for reason in [None, Some(RegionFlushReason::Repartition)] { + let err = engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest { + reason, + ..Default::default() + }), + ) + .await + .unwrap_err(); + assert_eq!(RetryHint::Retryable, err.retry_hint()); + assert_matches!( + err.into_inner().as_any().downcast_ref::().unwrap(), + Error::RegionBusy { .. } + ); + } + + assert_unflushed(&engine, region_id).await; +} + +#[tokio::test] +async fn test_migration_and_wal_prune_flushes_leave_staging_region_unflushed() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1024, 0); + let (schema, _) = create_staging_region(&engine, region_id).await; + put_range(&engine, region_id, &schema, (0, 3)).await; + + for reason in [ + RegionFlushReason::RegionMigration, + RegionFlushReason::RemoteWalPrune, + ] { + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest { + reason: Some(reason), + ..Default::default() + }), + ) + .await + .unwrap(); + } + + assert_unflushed(&engine, region_id).await; +} + +#[tokio::test] +async fn test_successful_staging_end_keeps_rows_of_rejected_flush() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1024, 0); + let request = CreateRequestBuilder::new().build(); + let schema = rows_schema(&request); + let table_dir = request.table_dir.clone(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_range(&engine, region_id, &schema, (0, 3)).await; + enter_staging(&engine, region_id).await; + put_range(&engine, region_id, &schema, (3, 5)).await; + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap_err(); + + let imported_file = first_file(&engine.get_region(region_id).unwrap()).await; + engine + .edit_region( + region_id, + RegionEdit { + files_to_add: vec![imported_file], + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }, + ) + .await + .unwrap(); + engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Leader) + .await + .unwrap(); + + assert_eq!(5, count_rows(&engine, region_id).await); + reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + assert_eq!(5, count_rows(&engine, region_id).await); +} + +#[tokio::test] +async fn test_downgrade_flushes_staging_rows_to_normal_manifest() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1024, 0); + let (schema, table_dir) = create_staging_region(&engine, region_id).await; + put_range(&engine, region_id, &schema, (0, 3)).await; + + engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::DowngradingLeader) + .await + .unwrap(); + flush_region(&engine, region_id, None).await; + + assert_eq!((1, 0), scan_inputs(&engine, region_id).await); + reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + assert_eq!(3, count_rows(&engine, region_id).await); +} + +#[tokio::test] +async fn test_staging_rows_survive_restart() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1024, 0); + let (schema, table_dir) = create_staging_region(&engine, region_id).await; + put_range(&engine, region_id, &schema, (0, 3)).await; + + let engine = env.reopen_engine(engine, MitoConfig::default()).await; + reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + + assert_eq!(3, count_rows(&engine, region_id).await); +} + +#[tokio::test] +async fn test_close_flushes_skip_wal_rows_of_staging_region() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1024, 0); + let (schema, table_dir) = create_staging_region(&engine, region_id).await; + put_range_skip_wal(&engine, region_id, &schema, (0, 3)).await; + + reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + + assert_eq!((1, 0), scan_inputs(&engine, region_id).await); + assert_eq!(3, count_rows(&engine, region_id).await); +} + +#[derive(Default)] +struct RequestCountListener { + count: AtomicUsize, + /// A flush waits here before it starts while a test holds the write side. + flush_gate: tokio::sync::RwLock<()>, +} + +impl RequestCountListener { + /// Waits until a worker has received a request batch after the count reached `received`. + async fn wait_received_after(&self, received: usize) { + tokio::time::timeout(Duration::from_secs(5), async { + while self.count.load(Ordering::Relaxed) == received { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .unwrap(); + } +} + +#[async_trait] +impl EventListener for RequestCountListener { + fn on_recv_requests(&self, request_num: usize) { + self.count.fetch_add(request_num, Ordering::Relaxed); + } + + async fn on_flush_begin(&self, _region_id: RegionId) { + let _open = self.flush_gate.read().await; + } +} + +#[tokio::test] +async fn test_edit_returns_to_staging_after_staging_leader_role() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1024, 0); + let request = CreateRequestBuilder::new().build(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_range(&engine, region_id, &schema, (0, 3)).await; + flush_region(&engine, region_id, None).await; + engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::StagingLeader) + .await + .unwrap(); + + let region = engine.get_region(region_id).unwrap(); + let file = first_file(®ion).await; + spawn_edit(&engine, region_id, file).await.unwrap().unwrap(); + + assert!(region.is_staging()); +} + +#[tokio::test] +async fn test_close_time_flush_rejects_queued_and_late_edits() { + let mut env = TestEnv::new().await; + let listener = Arc::new(RequestCountListener::default()); + let engine = env + .create_engine_with( + MitoConfig { + num_workers: 1, + ..Default::default() + }, + None, + Some(listener.clone()), + None, + ) + .await; + let region_id = RegionId::new(1024, 0); + let request = CreateRequestBuilder::new().build(); + let schema = rows_schema(&request); + let table_dir = request.table_dir.clone(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_range(&engine, region_id, &schema, (0, 3)).await; + flush_region(&engine, region_id, None).await; + enter_staging(&engine, region_id).await; + put_range_skip_wal(&engine, region_id, &schema, (3, 5)).await; + + let region = engine.get_region(region_id).unwrap(); + let file = first_file(®ion).await; + let queued_file = FileMeta { + file_id: FileId::random(), + ..file.clone() + }; + let late_file = FileMeta { + file_id: FileId::random(), + ..file.clone() + }; + let flush_gate = listener.flush_gate.write().await; + let manager = region.manifest_ctx.manifest_manager.write().await; + let edit = start_staged_edit(&engine, ®ion, file).await; + + let received = listener.count.load(Ordering::Relaxed); + let queued = spawn_edit(&engine, region_id, queued_file.clone()); + listener.wait_received_after(received).await; + + let received = listener.count.load(Ordering::Relaxed); + let close_engine = engine.clone(); + let close = tokio::spawn(async move { + close_engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + }); + listener.wait_received_after(received).await; + drop(manager); + + // The close-time flush is held back, so both edits arrive while the region waits for it. + let _ = edit.await; + let err = queued.await.unwrap().unwrap_err(); + assert_eq!(StatusCode::RegionNotFound, err.status_code()); + let err = spawn_edit(&engine, region_id, late_file.clone()) + .await + .unwrap() + .unwrap_err(); + assert_eq!(StatusCode::RegionBusy, err.status_code()); + + drop(flush_gate); + close.await.unwrap().unwrap(); + + reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + let region = engine.get_region(region_id).unwrap(); + let manifest = region.manifest_ctx.manifest_manager.read().await.manifest(); + assert!(!manifest.files.contains_key(&queued_file.file_id)); + assert!(!manifest.files.contains_key(&late_file.file_id)); + assert_eq!(5, count_rows(&engine, region_id).await); +} + +#[tokio::test] +async fn test_stale_staging_edit_rejected_after_close_time_flush_fails() { + let fail_parquet = Arc::new(AtomicBool::new(false)); + let writer_flag = fail_parquet.clone(); + let mock_layer = MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _args, inner| { + Box::new(ParquetFailWriter { + path: path.to_string(), + inner, + fail_parquet: writer_flag.clone(), + }) + })) + .build() + .unwrap(); + let mut env = TestEnv::new().await.with_mock_layer(mock_layer); + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1024, 0); + let request = CreateRequestBuilder::new().build(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_range(&engine, region_id, &schema, (0, 3)).await; + flush_region(&engine, region_id, None).await; + enter_staging(&engine, region_id).await; + put_range_skip_wal(&engine, region_id, &schema, (3, 5)).await; + + fail_parquet.store(true, Ordering::Relaxed); + engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap_err(); + + // The close failed, so the region is open, writable and no longer in staging. + let region = engine.get_region(region_id).unwrap(); + assert_eq!( + RegionRoleState::Leader(RegionLeaderState::Writable), + region.state() + ); + let file = FileMeta { + file_id: FileId::random(), + ..first_file(®ion).await + }; + let file_id = file.file_id; + let stale_expr = float_range_expr("field_0", 0., 100.).as_json_str().unwrap(); + let result = submit_staging_edit(&engine, region_id, file.clone(), stale_expr.clone()).await; + assert_matches!(result.unwrap_err(), Error::RegionState { .. }); + + // A later staging for another partition expression does not take the stale edit either. + fail_parquet.store(false, Ordering::Relaxed); + engine + .handle_request( + region_id, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive: StagingPartitionDirective::UpdatePartitionExpr( + float_range_expr("field_0", 0., 50.).as_json_str().unwrap(), + ), + }), + ) + .await + .unwrap(); + let result = submit_staging_edit(&engine, region_id, file, stale_expr).await; + assert_matches!(result.unwrap_err(), Error::RegionState { .. }); + + let manager = region.manifest_ctx.manifest_manager.read().await; + assert!(!manager.manifest().files.contains_key(&file_id)); + assert!( + manager + .staging_manifest() + .is_none_or(|staging| !staging.files.contains_key(&file_id)) + ); +} + +async fn submit_staging_edit( + engine: &MitoEngine, + region_id: RegionId, + file: FileMeta, + partition_expr: String, +) -> crate::error::Result<()> { + let (tx, rx) = tokio::sync::oneshot::channel(); + engine + .inner + .workers + .submit_to_worker( + region_id, + WorkerRequest::EditRegion( + RegionEditRequest::new(region_id, edit_adding(file), false, tx) + .for_staging(partition_expr), + ), + ) + .await + .unwrap(); + rx.await.unwrap() +} + +#[tokio::test] +async fn test_close_during_staged_edit_flushes_skip_wal_rows() { + let mut env = TestEnv::new().await; + let listener = Arc::new(RequestCountListener::default()); + let engine = env + .create_engine_with( + MitoConfig { + num_workers: 1, + ..Default::default() + }, + None, + Some(listener.clone()), + None, + ) + .await; + let region_id = RegionId::new(1024, 0); + let request = CreateRequestBuilder::new().build(); + let schema = rows_schema(&request); + let table_dir = request.table_dir.clone(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_range(&engine, region_id, &schema, (0, 3)).await; + flush_region(&engine, region_id, None).await; + enter_staging(&engine, region_id).await; + put_range_skip_wal(&engine, region_id, &schema, (3, 5)).await; + + let region = engine.get_region(region_id).unwrap(); + let file = first_file(®ion).await; + let manager = region.manifest_ctx.manifest_manager.write().await; + let edit = start_staged_edit(&engine, ®ion, file).await; + let received = listener.count.load(Ordering::Relaxed); + let close_engine = engine.clone(); + let close = tokio::spawn(async move { + close_engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + }); + // The worker has the close request before the edit, still waiting for the lock, can finish. + listener.wait_received_after(received).await; + drop(manager); + close.await.unwrap().unwrap(); + let _ = edit.await; + + reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + assert_eq!(5, count_rows(&engine, region_id).await); +} + +/// Waits until a writer is queued on the region's manifest lock while the caller holds it shared. +async fn wait_for_queued_manifest_writer(region: &MitoRegionRef) { + tokio::time::timeout(Duration::from_secs(5), async { + while region.manifest_ctx.manifest_manager.try_read().is_ok() { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_close_waits_for_manifest_lock_before_leaving_staging() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1024, 0); + create_staging_region(&engine, region_id).await; + let region = engine.get_region(region_id).unwrap(); + + let manager = region.manifest_ctx.manifest_manager.read().await; + let close_engine = engine.clone(); + let close = tokio::spawn(async move { + close_engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + }); + wait_for_queued_manifest_writer(®ion).await; + assert!(region.is_staging()); + + drop(manager); + close.await.unwrap().unwrap(); + assert!(!region.is_staging()); +} + +#[tokio::test] +async fn test_close_after_end_of_staging_is_queued_keeps_committed_result() { + let mut env = TestEnv::new().await; + let listener = Arc::new(RequestCountListener::default()); + let engine = env + .create_engine_with( + MitoConfig { + num_workers: 1, + ..Default::default() + }, + None, + Some(listener.clone()), + None, + ) + .await; + let region_id = RegionId::new(1024, 0); + let request = CreateRequestBuilder::new().build(); + let schema = rows_schema(&request); + let table_dir = request.table_dir.clone(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_range(&engine, region_id, &schema, (0, 3)).await; + enter_staging(&engine, region_id).await; + put_range(&engine, region_id, &schema, (3, 5)).await; + let region = engine.get_region(region_id).unwrap(); + let file = first_file(®ion).await; + engine + .edit_region( + region_id, + RegionEdit { + files_to_add: vec![file], + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }, + ) + .await + .unwrap(); + + let manager = region.manifest_ctx.manifest_manager.read().await; + let end_engine = engine.clone(); + let end = tokio::spawn(async move { + end_engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Leader) + .await + }); + wait_for_queued_manifest_writer(®ion).await; + let received = listener.count.load(Ordering::Relaxed); + let close_engine = engine.clone(); + let close = tokio::spawn(async move { + close_engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + }); + listener.wait_received_after(received).await; + drop(manager); + + end.await.unwrap().unwrap(); + close.await.unwrap().unwrap(); + reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + let reopened = engine.get_region(region_id).unwrap(); + assert_eq!( + Some(float_range_expr("field_0", 0., 100.).as_json_str().unwrap()), + reopened.metadata().partition_expr + ); + assert_eq!(5, count_rows(&engine, region_id).await); +} diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index ea0465322d6..b3369989346 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -60,7 +60,7 @@ use crate::read::flat_dedup::{FlatDedupIterator, FlatLastNonNull, FlatLastRow}; use crate::read::flat_merge::FlatMergeIterator; use crate::region::options::{IndexOptions, MergeMode, RegionOptions}; use crate::region::version::{VersionControlData, VersionControlRef, VersionRef}; -use crate::region::{ManifestContextRef, RegionLeaderState, RegionRoleState, parse_partition_expr}; +use crate::region::{ManifestContextRef, RegionLeaderState, parse_partition_expr}; use crate::request::{ BackgroundNotify, DdlRequest, FlushFailed, FlushFinished, OnFailure, OptionOutputTx, OutputTx, SenderBulkRequest, SenderDdlRequest, SenderWriteRequest, WorkerRequest, WorkerRequestWithTime, @@ -281,8 +281,6 @@ pub(crate) struct RegionFlushTask { pub(crate) index_options: IndexOptions, /// Semaphore to control flush concurrency. pub(crate) flush_semaphore: Arc, - /// Whether the region is in staging mode. - pub(crate) is_staging: bool, /// Partition expression of the region. /// /// This is used to generate the file meta. @@ -412,7 +410,6 @@ impl RegionFlushTask { _timer: timer, edit, memtables_to_remove, - is_staging: self.is_staging, }; WorkerRequest::Background { region_id: self.region_id, @@ -537,10 +534,7 @@ impl RegionFlushTask { flushed_sequence: Some(version_data.committed_sequence), committed_sequence: None, }; - info!( - "Applying {edit:?} to region {}, is_staging: {}", - self.region_id, self.is_staging - ); + info!("Applying {edit:?} to region {}", self.region_id); let action_list = RegionMetaActionList::with_action(RegionMetaAction::Edit(edit.clone())); @@ -553,17 +547,11 @@ impl RegionFlushTask { let expected_state = if matches!(self.reason, FlushReason::Downgrading) { RegionLeaderState::Downgrading } else { - // Check if region is in staging mode - let current_state = self.manifest_ctx.current_state(); - if current_state == RegionRoleState::Leader(RegionLeaderState::Staging) { - RegionLeaderState::Staging - } else { - RegionLeaderState::Writable - } + RegionLeaderState::Writable }; let manifest_version = match self .manifest_ctx - .update_manifest(expected_state, action_list, self.is_staging) + .update_manifest(expected_state, action_list, false) .await { Ok(manifest_version) => { @@ -584,9 +572,8 @@ impl RegionFlushTask { } }; info!( - "Successfully update manifest version to {manifest_version}, region: {}, is_staging: {}, reason: {}", + "Successfully update manifest version to {manifest_version}, region: {}, reason: {}", self.region_id, - self.is_staging, self.reason.as_str() ); @@ -1507,6 +1494,13 @@ impl FlushScheduler { }) } + /// Returns true if the region has a close-time flush. + pub(crate) fn has_close_time_flush(&self, region_id: RegionId) -> bool { + self.region_status + .get(®ion_id) + .is_some_and(|status| status.closing) + } + /// Returns true if the region has pending DDLs or a close-time flush. pub(crate) fn has_pending_ddls(&self, region_id: RegionId) -> bool { self.region_status @@ -1713,7 +1707,6 @@ mod tests { manifest_ctx, index_options: IndexOptions::default(), flush_semaphore: Arc::new(Semaphore::new(2)), - is_staging: false, partition_expr: None, durability_barrier: DurabilityBarrier::noop(), } @@ -1864,7 +1857,6 @@ mod tests { .await, index_options: IndexOptions::default(), flush_semaphore: Arc::new(Semaphore::new(2)), - is_staging: false, partition_expr: None, durability_barrier: DurabilityBarrier::noop(), }; @@ -1950,7 +1942,6 @@ mod tests { committed_sequence: None, }, memtables_to_remove: smallvec![], - is_staging: false, }), }; @@ -2166,7 +2157,6 @@ mod tests { manifest_ctx: manifest_ctx.clone(), index_options: IndexOptions::default(), flush_semaphore: Arc::new(Semaphore::new(2)), - is_staging: false, partition_expr: None, durability_barrier: DurabilityBarrier::noop(), }) @@ -2473,7 +2463,6 @@ mod tests { manifest_ctx: manifest_ctx.clone(), index_options: IndexOptions::default(), flush_semaphore: Arc::new(Semaphore::new(2)), - is_staging: false, partition_expr: None, durability_barrier: DurabilityBarrier::noop(), }) diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index 003a2240108..9584382e0ef 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -96,7 +96,7 @@ impl RegionUsage { pub enum RegionLeaderState { /// The region is opened and is writable. Writable, - /// The region is in staging mode - writable but no checkpoint/compaction. + /// The region is in staging mode - writable but no checkpoint/compaction/flush. Staging, /// The region is entering staging mode. - write requests will be stalled. EnteringStaging, @@ -407,6 +407,17 @@ impl MitoRegion { self.manifest_ctx.state.load() == RegionRoleState::Leader(RegionLeaderState::Staging) } + /// Returns whether the region is in staging mode, counting a staged edit in progress. + pub(crate) fn is_staging_or_staged_edit(&self) -> bool { + match self.manifest_ctx.state.load() { + RegionRoleState::Leader(RegionLeaderState::Staging) => true, + RegionRoleState::Leader(RegionLeaderState::Editing) => { + self.manifest_ctx.staging_partition_info().is_some() + } + _ => false, + } + } + /// Returns whether the region is entering staging mode. pub(crate) fn is_enter_staging(&self) -> bool { self.manifest_ctx.state.load() @@ -527,6 +538,29 @@ impl MitoRegion { ) } + /// Leaves staging without committing it, abandoning a staged edit in progress. + /// + /// You should call this method in the worker loop, which is what restores the state of an + /// edit, and only when [Self::is_staging_or_staged_edit] holds. + pub(crate) fn abort_staging(&self) { + let aborted = self + .manifest_ctx + .state + .fetch_update(|state| { + matches!( + state, + RegionRoleState::Leader( + RegionLeaderState::Staging | RegionLeaderState::Editing + ) + ) + .then_some(RegionRoleState::Leader(RegionLeaderState::Writable)) + }) + .is_ok(); + if aborted { + self.manifest_ctx.clear_staging_partition_info(); + } + } + /// Sets the region role state gracefully. This acquires the manifest write lock. pub(crate) async fn set_role_state_gracefully( &self, @@ -1142,30 +1176,6 @@ impl MitoRegion { Ok(Some(pending)) } - /// Returns the partition expression string for this region. - /// - /// If the region is currently in staging state, this returns the partition expression held in - /// the staging partition field. Otherwise, it returns the partition expression from the primary - /// region metadata (current committed version). - pub fn maybe_staging_partition_expr_str(&self) -> Option { - let is_staging = self.is_staging(); - if is_staging { - let staging_partition_info = self.manifest_ctx.staging_partition_info(); - if staging_partition_info.is_none() { - warn!( - "Staging partition expr is none for region {} in staging state", - self.region_id - ); - } - staging_partition_info - .as_ref() - .and_then(|info| info.partition_expr().map(ToString::to_string)) - } else { - let version = self.version(); - version.metadata.partition_expr.clone() - } - } - pub fn expected_partition_expr_version(&self) -> u64 { if self.is_staging() { self.manifest_ctx diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index 9c0178710f1..5f91a899005 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -1152,8 +1152,6 @@ pub(crate) struct FlushFinished { pub(crate) edit: RegionEdit, /// Memtables to remove. pub(crate) memtables_to_remove: SmallVec<[MemtableId; 2]>, - /// Whether the region is in staging mode. - pub(crate) is_staging: bool, } impl FlushFinished { @@ -1372,6 +1370,9 @@ pub(crate) struct RegionEditRequest { pub(crate) preload_sst_cache: bool, /// The waiters that are waiting for this region edit's result. pub(crate) waiters: Waiters, + /// The partition expression of the staging this edit belongs to, if any. + /// The edit must not run unless the region is staging for that expression. + pub(crate) staging_partition_expr: Option, } impl RegionEditRequest { @@ -1386,8 +1387,15 @@ impl RegionEditRequest { edit, preload_sst_cache, waiters: Waiters::one(waiter), + staging_partition_expr: None, } } + + /// Binds the edit to the staging for `partition_expr`. + pub(crate) fn for_staging(mut self, partition_expr: String) -> Self { + self.staging_partition_expr = Some(partition_expr); + self + } } /// Notifies the regin the result of editing region. diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 28026f12808..9d11e16785a 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -93,6 +93,7 @@ use crate::sst::index::puffin_manager::PuffinManagerFactory; use crate::temp_file_sweeper; use crate::time_provider::{StdTimeProvider, TimeProviderRef}; use crate::wal::Wal; +use crate::worker::handle_flush::StagingRegionsRef; use crate::worker::handle_manifest::RegionEditQueues; /// Identifier for a worker. @@ -258,6 +259,7 @@ impl WorkerGroup { let sweep_limiter = (config.gc.max_concurrent_gc_job > 0) .then(|| Arc::new(Semaphore::new(config.gc.max_concurrent_gc_job))); let node_lease = Arc::new(NodeLease::default()); + let staging_regions = StagingRegionsRef::default(); let workers = (0..config.num_workers) .map(|id| { @@ -290,6 +292,7 @@ impl WorkerGroup { #[cfg(feature = "hdfs-object-store")] sweep_limiter: sweep_limiter.clone(), node_lease: node_lease.clone(), + staging_regions: staging_regions.clone(), } .start() }) @@ -529,6 +532,7 @@ impl WorkerGroup { let sweep_limiter = (config.gc.max_concurrent_gc_job > 0) .then(|| Arc::new(Semaphore::new(config.gc.max_concurrent_gc_job))); let node_lease = Arc::new(NodeLease::default()); + let staging_regions = StagingRegionsRef::default(); let workers = (0..config.num_workers) .map(|id| { WorkerStarter { @@ -560,6 +564,7 @@ impl WorkerGroup { #[cfg(feature = "hdfs-object-store")] sweep_limiter: sweep_limiter.clone(), node_lease: node_lease.clone(), + staging_regions: staging_regions.clone(), } .start() }) @@ -677,6 +682,7 @@ struct WorkerStarter { /// no concurrent GC job. sweep_limiter: Option>, node_lease: Arc, + staging_regions: StagingRegionsRef, } impl WorkerStarter { @@ -782,6 +788,7 @@ impl WorkerStarter { #[cfg(feature = "hdfs-object-store")] sweep_limiter: self.sweep_limiter, node_lease: self.node_lease, + staging_regions: self.staging_regions, }; let handle = common_runtime::spawn_global(async move { worker_thread.run().await; @@ -1122,6 +1129,8 @@ struct RegionWorkerLoop { /// no concurrent GC job. sweep_limiter: Option>, node_lease: Arc, + /// Staging regions of this node that have received writes. + staging_regions: StagingRegionsRef, } impl RegionWorkerLoop { diff --git a/src/mito2/src/worker/handle_apply_staging.rs b/src/mito2/src/worker/handle_apply_staging.rs index 4ac8a8990f0..435734dfdf7 100644 --- a/src/mito2/src/worker/handle_apply_staging.rs +++ b/src/mito2/src/worker/handle_apply_staging.rs @@ -140,7 +140,8 @@ impl RegionWorkerLoop { region_id, edit, // we don't need to preload sst cache during repartition, as it may cause extra network overhead. false, tx, - ), + ) + .for_staging(request.partition_expr), ))) .await; diff --git a/src/mito2/src/worker/handle_close.rs b/src/mito2/src/worker/handle_close.rs index ae8b31fb861..c10ac761f43 100644 --- a/src/mito2/src/worker/handle_close.rs +++ b/src/mito2/src/worker/handle_close.rs @@ -36,6 +36,16 @@ impl RegionWorkerLoop { info!("Try to close region {}, worker: {}", region_id, self.id); + if region.is_staging_or_staged_edit() { + // A staging region cannot flush, so it leaves staging before the close-time flush + // below. The manifest lock lets a successful end of staging in progress finish first, + // after which the region is already out of staging. + let _manager = region.manifest_ctx.manifest_manager.write().await; + region.abort_staging(); + // Edits queued behind a staged edit belong to the staging that has just ended. + self.reject_region_edit_queue_as_not_found(region_id); + } + // Memtables containing skip-WAL writes must be flushed before closing. // Followers cannot flush; explicit flush-on-close and region WAL options // retain their existing behavior. diff --git a/src/mito2/src/worker/handle_flush.rs b/src/mito2/src/worker/handle_flush.rs index 09c0d66344e..984b551d7bb 100644 --- a/src/mito2/src/worker/handle_flush.rs +++ b/src/mito2/src/worker/handle_flush.rs @@ -14,9 +14,9 @@ //! Handling flush related requests. -use std::collections::HashSet; -use std::sync::Arc; +use std::collections::{HashMap, HashSet}; use std::sync::atomic::Ordering; +use std::sync::{Arc, Mutex, Weak}; use common_base::readable_size::ReadableSize; use common_telemetry::{debug, error, info}; @@ -25,10 +25,10 @@ use store_api::region_request::{RegionFlushReason, RegionFlushRequest}; use store_api::storage::RegionId; use crate::config::{IndexBuildMode, MitoConfig}; -use crate::error::{RegionNotFoundSnafu, Result}; +use crate::error::{RegionBusySnafu, RegionNotFoundSnafu, Result}; use crate::flush::{FlushReason, RegionFlushTask}; -use crate::region::MitoRegionRef; use crate::region::version::VersionRef; +use crate::region::{MitoRegion, MitoRegionRef}; use crate::request::{BuildIndexRequest, FlushFailed, FlushFinished, OnFailure, OptionOutputTx}; use crate::sst::index::IndexBuildType; use crate::worker::RegionWorkerLoop; @@ -39,6 +39,33 @@ pub(crate) struct RegionWriteBufferPressure { pub(crate) rejected_region_ids: HashSet, } +/// Staging regions of a node that have received writes. +/// +/// Workers share it to bound the memtables that stay unflushed until staging ends. +#[derive(Debug, Default)] +pub(crate) struct StagingRegions { + regions: Mutex>>, +} + +pub(crate) type StagingRegionsRef = Arc; + +impl StagingRegions { + /// Records `region` and returns the memtable usage of the regions still in staging. + fn record_and_memtable_usage(&self, region: &MitoRegionRef) -> usize { + let mut regions = self.regions.lock().unwrap(); + regions.insert(region.region_id, Arc::downgrade(region)); + let mut usage = 0; + regions.retain(|_, region| match region.upgrade() { + Some(region) if region.is_staging_or_staged_edit() => { + usage += region_memtable_usage(®ion.version()); + true + } + _ => false, + }); + usage + } +} + #[derive(Debug, Default, PartialEq, Eq)] struct RegionWriteBufferStatus { should_flush: bool, @@ -100,8 +127,11 @@ impl RegionWorkerLoop { let mut pending_regions = vec![]; for region in ®ions { - if self.flush_scheduler.is_flush_requested(region.region_id) || !region.is_writable() { - // Already flushing or not writable. + if self.flush_scheduler.is_flush_requested(region.region_id) + || !region.is_writable() + || region.is_staging() + { + // Already flushing, not writable, or in staging, whose rows must stay in the WAL. continue; } @@ -193,6 +223,19 @@ impl RegionWorkerLoop { self.stalled_requests.estimated_size(®ion_id), ); + if region.is_staging() { + // Staging regions don't flush, so waiting would not release their memory. Writes + // are rejected at twice the region limit or at half of the engine's flush limit, + // which reserves memory headroom for other regions. + let staging_usage = self.staging_regions.record_and_memtable_usage(®ion); + if status.should_reject + || staging_usage >= self.write_buffer_manager.flush_limit() / 2 + { + pressure.rejected_region_ids.insert(region_id); + } + continue; + } + if status.should_flush && !self.flush_scheduler.is_flush_requested(region.region_id) { let task = self.new_flush_task( ®ion, @@ -242,8 +285,7 @@ impl RegionWorkerLoop { manifest_ctx: region.manifest_ctx.clone(), index_options: region.version().options.index_options.clone(), flush_semaphore: self.flush_semaphore.clone(), - is_staging: region.is_staging(), - partition_expr: region.maybe_staging_partition_expr_str(), + partition_expr: region.metadata().partition_expr.clone(), durability_barrier: self .wal .durability_barrier(region.region_id, ®ion.provider), @@ -327,6 +369,18 @@ impl RegionWorkerLoop { // when handling flush request instead of in `schedule_flush` or `flush_finished`. self.update_topic_latest_entry_id(®ion); + // A staging region must keep its WAL until staging ends, so it flushes nothing. + if region.is_staging() && !region.version().memtables.is_empty() { + sender.send(match request.reason { + // These only ask to advance the flushed position, which can wait. + Some(RegionFlushReason::RegionMigration | RegionFlushReason::RemoteWalPrune) => { + Ok(0) + } + _ => RegionBusySnafu { region_id }.fail(), + }); + return; + } + let reason = resolve_flush_reason(request.reason, region.is_downgrading()); let mut task = self.new_flush_task(®ion, reason, request.row_group_size, self.config.clone()); @@ -350,6 +404,12 @@ impl RegionWorkerLoop { continue; } self.update_topic_latest_entry_id(region); + // The rows of a staging region must stay in the WAL until staging ends, so a staging + // region with rows pins its topic's prune point: the update above only keeps the id + // fresh while the memtable is empty. + if region.is_staging() { + continue; + } let auto_flush_interval = region .version() @@ -399,27 +459,14 @@ impl RegionWorkerLoop { } }; - if request.is_staging { - // Skip the region metadata update. - info!( - "Skipping region metadata update for region {} in staging mode", - region_id - ); - region.version_control.apply_edit( - None, - &request.memtables_to_remove, - region.file_purger.clone(), - ); - } else { - region.version_control.apply_edit( - Some(request.edit.clone()), - &request.memtables_to_remove, - region.file_purger.clone(), - ); - } + region.version_control.apply_edit( + Some(request.edit.clone()), + &request.memtables_to_remove, + region.file_purger.clone(), + ); - // The flush has committed its manifest, including in staging mode. Count only - // its Parquet output, before WAL cleanup (which can fail after the commit). + // The flush has committed its manifest. Count only its Parquet output, before + // WAL cleanup (which can fail after the commit). let ingested_parquet_bytes = request .edit .files_to_add diff --git a/src/mito2/src/worker/handle_manifest.rs b/src/mito2/src/worker/handle_manifest.rs index d7959880989..6eceba19586 100644 --- a/src/mito2/src/worker/handle_manifest.rs +++ b/src/mito2/src/worker/handle_manifest.rs @@ -30,7 +30,9 @@ use store_api::storage::{RegionId, SequenceNumber}; use crate::cache::CacheManagerRef; use crate::cache::file_cache::{FileType, IndexKey}; use crate::config::IndexBuildMode; -use crate::error::{EditRegionSnafu, RegionBusySnafu, RegionNotFoundSnafu, Result}; +use crate::error::{ + EditRegionSnafu, RegionBusySnafu, RegionNotFoundSnafu, RegionStateSnafu, Result, +}; use crate::manifest::action::{ RegionChange, RegionEdit, RegionMetaAction, RegionMetaActionList, RegionTruncate, TruncateKind, }; @@ -310,6 +312,36 @@ impl RegionWorkerLoop { return RegionEditDispatch::Rejected; }; + // A close-time flush ends with the region removed. A close that left staging keeps the + // region writable until then, so an edit meant for staging would run as a normal edit. + if self.flush_scheduler.has_close_time_flush(region_id) { + request + .waiters + .reply_with(|| RegionBusySnafu { region_id }.fail()); + return RegionEditDispatch::Rejected; + } + + // A close may have abandoned the staging this edit belongs to, even if its flush failed. + if let Some(expected) = &request.staging_partition_expr + && region + .manifest_ctx + .staging_partition_info() + .as_ref() + .and_then(|info| info.partition_expr()) + != Some(expected.as_str()) + { + let state = region.state(); + request.waiters.reply_with(|| { + RegionStateSnafu { + region_id, + state, + expect: RegionRoleState::Leader(RegionLeaderState::Staging), + } + .fail() + }); + return RegionEditDispatch::Rejected; + } + if !region.is_writable() || self.flush_scheduler.has_pending_alter(region_id) { if matches!( region.state(), @@ -334,6 +366,7 @@ impl RegionWorkerLoop { mut edit, waiters, preload_sst_cache, + staging_partition_expr: _, } = request; let file_sequence = region.version_control.current().last_sequence + 1; edit.committed_sequence = Some(file_sequence); @@ -427,7 +460,10 @@ impl RegionWorkerLoop { }; let need_compaction = if edit_result.is_staging { - if edit_result.update_region_state { + // A close may have abandoned this edit and left staging already. + if edit_result.update_region_state + && region.state() == RegionRoleState::Leader(RegionLeaderState::Editing) + { // For staging regions, edits are not applied immediately, // as they remain invisible until the region exits the staging state. region.switch_state_to_staging(RegionLeaderState::Editing); diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs index 48a802d0c98..3a740c79867 100644 --- a/src/store-api/src/region_engine.rs +++ b/src/store-api/src/region_engine.rs @@ -50,7 +50,7 @@ pub enum SettableRegionRoleState { DowngradingLeader, /// Exit staging mode and return to normal leader state. Only allowed from staging state. Leader, - /// Enter staging mode. Region remains writable but disables checkpoint and compaction. Only allowed from normal leader state. + /// Enter staging mode. Region remains writable but disables checkpoint, compaction and flush. Only allowed from normal leader state. StagingLeader, } @@ -606,7 +606,7 @@ pub struct RegionStatistic { /// Total ingested Parquet bytes since the region opened. /// Currently counts only committed flush output. /// - /// Includes staging flushes, but excludes index files, compaction and imported SSTs. + /// Excludes index files, compaction and imported SSTs. /// Resets when the region instance is recreated; manifest replay does not increment it. /// Accepts the legacy `flushed_parquet_bytes` key; upgrade metasrv before datanodes. #[serde(default, alias = "flushed_parquet_bytes")]