From 3796c478e2929b1239ff14cb2941af8927d704ee Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 9 Oct 2026 12:21:38 +0000 Subject: [PATCH] fix(mito): flush memtables containing skip-WAL writes on region close (#9440) * fix(mito): flush unlogged writes when closing regions Signed-off-by: WenyXu * fix(mito): align skip-WAL naming and metric recovery tests Signed-off-by: WenyXu --------- Signed-off-by: WenyXu (cherry picked from commit a026b3b280b4013488220e3d12c96fe3d0d02da1) Signed-off-by: WenyXu --- src/metric-engine/src/engine/bulk_insert.rs | 48 ++-- src/metric-engine/src/engine/put.rs | 56 +++-- src/metric-engine/src/test_util.rs | 11 + src/mito2/src/engine/basic_test.rs | 56 ++++- src/mito2/src/engine/scan_test.rs | 58 ++++- src/mito2/src/engine/skip_wal_test.rs | 192 ++++++++++++++- src/mito2/src/flush.rs | 2 +- src/mito2/src/memtable/version.rs | 254 +++++++++++++++++++- src/mito2/src/read/scan_region.rs | 86 +++---- src/mito2/src/region/opener.rs | 27 +-- src/mito2/src/region/version.rs | 10 +- src/mito2/src/region_write_ctx.rs | 196 ++++++++++++++- src/mito2/src/sst/index.rs | 4 +- src/mito2/src/sst/primary_key.rs | 4 +- src/mito2/src/test_util/version_util.rs | 7 +- src/mito2/src/worker/handle_close.rs | 18 +- src/mito2/src/worker/handle_manifest.rs | 11 +- src/mito2/src/worker/handle_write.rs | 27 +-- 18 files changed, 888 insertions(+), 179 deletions(-) diff --git a/src/metric-engine/src/engine/bulk_insert.rs b/src/metric-engine/src/engine/bulk_insert.rs index fa47ce48008..dd20f34adf8 100644 --- a/src/metric-engine/src/engine/bulk_insert.rs +++ b/src/metric-engine/src/engine/bulk_insert.rs @@ -354,17 +354,23 @@ mod tests { #[tokio::test] async fn test_bulk_insert_physical_region_passthrough() { - check_bulk_insert_physical_region_passthrough(false).await; - check_bulk_insert_physical_region_passthrough(true).await; + check_bulk_insert_physical_region_passthrough(false, false).await; + check_bulk_insert_physical_region_passthrough(true, false).await; } - async fn check_bulk_insert_physical_region_passthrough(skip_wal: bool) { + #[tokio::test] + async fn test_bulk_insert_physical_region_close_recovery() { + check_bulk_insert_physical_region_passthrough(false, true).await; + check_bulk_insert_physical_region_passthrough(true, true).await; + } + + async fn check_bulk_insert_physical_region_passthrough(skip_wal: bool, close: bool) { // Use flat format so that BulkMemtable is used (supports write_bulk). let mito_config = MitoConfig { default_flat_format: true, ..Default::default() }; - let env = TestEnv::with_mito_config("", mito_config, Default::default()).await; + let mut env = TestEnv::with_mito_config("", mito_config, Default::default()).await; env.init_metric_region().await; env.metric().inner.flush_task.stop().await.unwrap(); let physical_region_id = env.default_physical_region_id(); @@ -414,20 +420,23 @@ mod tests { let batches = RecordBatches::try_collect(stream).await.unwrap(); assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 6); - // Closing without a flush must recover only WAL-backed data. Recreate - // the wrapper too, so its metadata cache cannot hide missing metadata. let stat = env.mito().region_statistic(physical_region_id).unwrap(); assert_eq!(stat.sst_num, 0); - env.metric() - .handle_request( - physical_region_id, - RegionRequest::Close(RegionCloseRequest { - flush_on_close: false, - }), - ) - .await - .unwrap(); - let reopened = MetricEngine::try_new(env.mito(), Default::default()).unwrap(); + let reopened = if close { + env.metric() + .handle_request( + physical_region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + // Recreating the wrapper prevents its metadata cache from hiding data loss. + MetricEngine::try_new(env.mito(), Default::default()).unwrap() + } else { + // Stopping workers without Close exercises recovery from WAL alone. + env.reopen_engine(Default::default()).await; + env.metric() + }; reopened.inner.flush_task.stop().await.unwrap(); reopened .handle_request( @@ -446,6 +455,11 @@ mod tests { ) .await .unwrap(); + let stat = env + .mito() + .region_statistic(crate::utils::to_data_region_id(physical_region_id)) + .unwrap(); + assert_eq!(stat.sst_num > 0, skip_wal && close); let stream = reopened .scan_to_stream(logical_region_id, ScanRequest::default()) .await @@ -453,7 +467,7 @@ mod tests { let batches = RecordBatches::try_collect(stream).await.unwrap(); assert_eq!( batches.iter().map(|b| b.num_rows()).sum::(), - if skip_wal { 0 } else { 6 }, + if skip_wal && !close { 0 } else { 6 }, ); } diff --git a/src/metric-engine/src/engine/put.rs b/src/metric-engine/src/engine/put.rs index fdf9b02a5ff..b709d99d847 100644 --- a/src/metric-engine/src/engine/put.rs +++ b/src/metric-engine/src/engine/put.rs @@ -934,14 +934,22 @@ mod tests { #[tokio::test] async fn test_put_skip_wal_batch_recovery() { - check_put_skip_wal_batch_recovery("sparse", false).await; - check_put_skip_wal_batch_recovery("sparse", true).await; - check_put_skip_wal_batch_recovery("dense", false).await; - check_put_skip_wal_batch_recovery("dense", true).await; + check_put_skip_wal_batch_recovery("sparse", false, false).await; + check_put_skip_wal_batch_recovery("sparse", true, false).await; + check_put_skip_wal_batch_recovery("dense", false, false).await; + check_put_skip_wal_batch_recovery("dense", true, false).await; } - async fn check_put_skip_wal_batch_recovery(encoding: &str, skip_wal: bool) { - let env = TestEnv::new().await; + #[tokio::test] + async fn test_put_skip_wal_batch_close_recovery() { + check_put_skip_wal_batch_recovery("sparse", false, true).await; + check_put_skip_wal_batch_recovery("sparse", true, true).await; + check_put_skip_wal_batch_recovery("dense", false, true).await; + check_put_skip_wal_batch_recovery("dense", true, true).await; + } + + async fn check_put_skip_wal_batch_recovery(encoding: &str, skip_wal: bool, close: bool) { + let mut env = TestEnv::new().await; let engine = env.metric(); engine.inner.flush_task.stop().await.unwrap(); let physical_region_id = env.default_physical_region_id(); @@ -1009,18 +1017,21 @@ mod tests { assert!(stat.memtable_size > 0); assert_eq!(stat.sst_num, 0); } - engine - .handle_request( - physical_region_id, - RegionRequest::Close(RegionCloseRequest { - flush_on_close: false, - }), - ) - .await - .unwrap(); - - // Recreate the wrapper as well, discarding its metadata cache. - let reopened = MetricEngine::try_new(env.mito(), Default::default()).unwrap(); + let reopened = if close { + engine + .handle_request( + physical_region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + // Recreate the wrapper as well, discarding its metadata cache. + MetricEngine::try_new(env.mito(), Default::default()).unwrap() + } else { + // Bypass Close to retain coverage of rows lost without WAL protection. + env.reopen_engine(Default::default()).await; + env.metric() + }; reopened.inner.flush_task.stop().await.unwrap(); reopened .handle_request( @@ -1042,12 +1053,17 @@ mod tests { ) .await .unwrap(); + let stat = env + .mito() + .region_statistic(crate::utils::to_data_region_id(physical_region_id)) + .unwrap(); + assert_eq!(stat.sst_num > 0, skip_wal && close); let recovered_metadata = reopened.get_metadata(logical_region_id).await.unwrap(); assert_eq!( metadata_before.column_metadatas, recovered_metadata.column_metadatas ); - let expected = if skip_wal { + let expected = if skip_wal && !close { vec![] } else { vec![(0, 30.0), (1, 10.0), (2, 20.0), (3, 30.0)] @@ -1055,7 +1071,7 @@ mod tests { assert_eq!( scan_timestamp_values(&reopened, logical_region_id).await, expected, - "encoding={encoding}, skip_wal={skip_wal}" + "encoding={encoding}, skip_wal={skip_wal}, close={close}" ); } diff --git a/src/metric-engine/src/test_util.rs b/src/metric-engine/src/test_util.rs index 0671351ddb3..cf64d968182 100644 --- a/src/metric-engine/src/test_util.rs +++ b/src/metric-engine/src/test_util.rs @@ -121,6 +121,17 @@ impl TestEnv { self.metric.clone() } + /// Restarts Mito without closing regions and recreates the metric metadata cache. + /// Callers stop the old metric engine's metadata flush task before restarting. + pub async fn reopen_engine(&mut self, config: EngineConfig) { + let mito_config = self.mito.mito_config().clone(); + self.mito = self + .mito_env + .reopen_engine(self.mito.clone(), mito_config) + .await; + self.metric = MetricEngine::try_new(self.mito.clone(), config).unwrap(); + } + /// Creates a new follower engine with the same config as the leader engine. pub async fn create_follower_engine(&mut self) -> (MitoEngine, MetricEngine) { let mito = self diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index 776e1f78abe..35f15f2393c 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/src/mito2/src/engine/basic_test.rs @@ -1500,7 +1500,10 @@ async fn check_request_skip_wal_recovery(flat_format: bool, skip_wal: bool) { let before = RecordBatches::try_collect(stream).await.unwrap(); assert_eq!(before.iter().map(|b| b.num_rows()).sum::(), 4); - reopen_region(&engine, region_id, table_dir, false, HashMap::new()).await; + // Stop the engine directly, without region close, to exercise WAL-only recovery. + let engine = + restart_without_region_close(&mut env, engine, flat_format, region_id, table_dir).await; + let stream = engine .scan_to_stream(region_id, ScanRequest::default()) .await @@ -1625,7 +1628,7 @@ async fn check_request_skip_wal_mixed_batch_recovery( skip_wal: bool, flush_before_reopen: bool, ) { - use crate::region_write_ctx::RegionWriteCtx; + use crate::region_write_ctx::{RegionWriteCtx, WriteSource}; use crate::request::OptionOutputTx; use crate::test_util::LogStoreImpl; use crate::wal::Wal; @@ -1657,6 +1660,7 @@ async fn check_request_skip_wal_mixed_batch_recovery( ®ion.version_control, region.provider.clone(), None, + WriteSource::Request, ); let mut receivers = Vec::with_capacity(4); for (index, (skip, timestamp)) in [(false, 0), (skip_wal, 0), (false, 0), (skip_wal, 1)] @@ -1692,8 +1696,7 @@ async fn check_request_skip_wal_mixed_batch_recovery( ctx.add_wal_entry(&mut writer).unwrap(); let response = writer.write_to_wal().await.unwrap(); assert_eq!(response.last_entry_ids.get(®ion_id), Some(&1)); - ctx.write_memtable().await; - ctx.publish_sequence_and_entry_id(); + ctx.write_memtables().await; drop(ctx); for rx in receivers { assert_eq!(rx.await.unwrap().unwrap(), 2); @@ -1751,9 +1754,13 @@ async fn check_request_skip_wal_mixed_batch_recovery( ); } - // A no-flush close discards memtables. Only WAL-backed rows recover - // unless an explicit flush has already persisted all requests. - reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + // Bypass region close so only WAL-backed rows recover unless the explicit + // flush above persisted all requests. + let engine = + restart_without_region_close(&mut env, engine, flat_format, region_id, table_dir).await; + engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); let loses_skipped_rows = skip_wal && !flush_before_reopen; let mut expected_values = if loses_skipped_rows { vec![(0, 20.0), (1000, 1.0), (3000, 21.0)] @@ -1830,3 +1837,38 @@ async fn request_skip_wal_values(engine: &MitoEngine, region_id: RegionId) -> Ve values.sort_unstable_by_key(|(timestamp, _)| *timestamp); values } + +/// Restarts workers without closing regions, preserving WAL-only recovery coverage. +async fn restart_without_region_close( + env: &mut TestEnv, + engine: MitoEngine, + flat_format: bool, + region_id: RegionId, + table_dir: String, +) -> MitoEngine { + let engine = env + .reopen_engine( + engine, + MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }, + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir, + options: HashMap::new(), + skip_wal_replay: false, + path_type: PathType::Bare, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + engine +} diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index 39a1e66f6d8..f86b327735d 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -56,7 +56,7 @@ use store_api::metric_engine_consts::PRIMARY_KEY_ENCODING; use store_api::region_engine::{PrepareRequest, RegionEngine, RegionScanner}; use store_api::region_request::{ AlterKind, RegionAlterRequest, RegionBulkInsertsRequest, RegionCompactRequest, - RegionPutRequest, RegionRequest, SetRegionOption, + RegionOpenRequest, RegionPutRequest, RegionRequest, SetRegionOption, }; use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; use store_api::storage::{ @@ -2991,13 +2991,21 @@ fn build_bulk_insert_request( #[tokio::test] async fn test_bulk_skip_wal_recovery() { - check_bulk_skip_wal_recovery(false, false).await; - check_bulk_skip_wal_recovery(false, true).await; - check_bulk_skip_wal_recovery(true, false).await; - check_bulk_skip_wal_recovery(true, true).await; + check_bulk_skip_wal_recovery(false, false, false).await; + check_bulk_skip_wal_recovery(false, true, false).await; + check_bulk_skip_wal_recovery(true, false, false).await; + check_bulk_skip_wal_recovery(true, true, false).await; } -async fn check_bulk_skip_wal_recovery(skip_wal: bool, flush: bool) { +#[tokio::test] +async fn test_bulk_skip_wal_close_recovery() { + check_bulk_skip_wal_recovery(false, false, true).await; + check_bulk_skip_wal_recovery(false, true, true).await; + check_bulk_skip_wal_recovery(true, false, true).await; + check_bulk_skip_wal_recovery(true, true, true).await; +} + +async fn check_bulk_skip_wal_recovery(skip_wal: bool, flush: bool, close: bool) { let mut env = TestEnv::new().await; let engine = env.create_engine(MitoConfig::default()).await; let region_id = RegionId::new(1, 1); @@ -3048,7 +3056,41 @@ async fn check_bulk_skip_wal_recovery(skip_wal: bool, flush: bool) { u64::from(!skip_wal) ); } - reopen_region(&engine, region_id, table_dir, false, HashMap::new()).await; + let engine = if close { + reopen_region(&engine, region_id, table_dir, false, HashMap::new()).await; + assert_eq!( + region.version().flushed_sequence, + if flush || skip_wal { 4 } else { 0 } + ); + engine + } else { + // Stopping the engine directly exercises recovery without close-time flush. + let engine = env.reopen_engine(engine, MitoConfig::default()).await; + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir, + path_type: store_api::region_request::PathType::Bare, + options: HashMap::new(), + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + engine + }; + assert!( + !engine + .get_region(region_id) + .unwrap() + .version() + .memtables + .has_skip_wal_writes() + ); let stream = engine .scan_to_stream(region_id, ScanRequest::default()) .await @@ -3060,7 +3102,7 @@ async fn check_bulk_skip_wal_recovery(skip_wal: bool, flush: bool) { .iter() .map(|b| b.num_rows()) .sum::(), - if skip_wal && !flush { 0 } else { 4 } + if skip_wal && !flush && !close { 0 } else { 4 } ); } diff --git a/src/mito2/src/engine/skip_wal_test.rs b/src/mito2/src/engine/skip_wal_test.rs index 363318d1cf5..915099e7228 100644 --- a/src/mito2/src/engine/skip_wal_test.rs +++ b/src/mito2/src/engine/skip_wal_test.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; @@ -34,7 +35,7 @@ use crate::engine::listener::AlterFlushListener; use crate::test_util::{ CreateRequestBuilder, LogStoreFactory, TestEnv, build_rows, flush_region, kafka_log_store_factory, multiple_log_store_factories, prepare_test_for_kafka_log_store, - put_rows, raft_engine_log_store_factory, rows_schema, + put_rows, raft_engine_log_store_factory, reopen_region, rows_schema, }; fn set_skip_wal_request(skip_wal: bool) -> RegionRequest { @@ -920,3 +921,192 @@ async fn test_close_region_after_truncate_skip_wal() { let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); assert_eq!(3, total_rows); } + +#[rstest] +#[case(false)] +#[case(true)] +#[tokio::test] +async fn test_close_request_skip_wal( + #[case] skip_wal: bool, + #[values(false, true)] flat_format: bool, +) { + let mut env = TestEnv::with_prefix("close-request-skip-wal").await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + skip_wal, + rows: Rows { + schema: rows_schema(&request), + rows: build_rows(0, 3), + }, + hint: None, + partition_expr_version: None, + }), + ) + .await + .unwrap(); + let region = engine.get_region(region_id).unwrap(); + assert!(!region.skip_wal()); + assert_eq!(region.version().memtables.has_skip_wal_writes(), skip_wal); + let expected = common_recordbatch::RecordBatches::try_collect( + engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(), + ) + .await + .unwrap() + .pretty_print() + .unwrap(); + + reopen_region(&engine, region_id, request.table_dir, false, HashMap::new()).await; + // Logged writes are recovered without a close-time flush. + assert_eq!( + region.version().flushed_sequence, + if skip_wal { 3 } else { 0 } + ); + let reopened = engine.get_region(region_id).unwrap(); + assert!(!reopened.version().memtables.has_skip_wal_writes()); + let actual = common_recordbatch::RecordBatches::try_collect( + engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(), + ) + .await + .unwrap() + .pretty_print() + .unwrap(); + assert_eq!(actual, expected); +} + +#[rstest] +#[case(false)] +#[case(true)] +#[tokio::test] +async fn test_close_request_skip_wal_during_flush(#[case] new_skip_wal: bool) { + let mut env = TestEnv::with_prefix("close-request-skip-wal-during-flush").await; + let listener = Arc::new(AlterFlushListener::default()); + let engine = env + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + skip_wal: true, + rows: Rows { + schema: rows_schema(&request), + rows: build_rows(0, 3), + }, + hint: None, + partition_expr_version: None, + }), + ) + .await + .unwrap(); + let region = engine.get_region(region_id).unwrap(); + let engine_cloned = engine.clone(); + let flush_job = tokio::spawn(async move { + flush_region(&engine_cloned, region_id, None).await; + }); + tokio::time::timeout(Duration::from_secs(5), listener.wait_flush_begin()) + .await + .unwrap(); + + // These rows belong to the next mutable, while the skipped rows are frozen. + engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + skip_wal: new_skip_wal, + rows: Rows { + schema: rows_schema(&request), + rows: build_rows(3, 6), + }, + hint: None, + partition_expr_version: None, + }), + ) + .await + .unwrap(); + let expected = common_recordbatch::RecordBatches::try_collect( + engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(), + ) + .await + .unwrap() + .pretty_print() + .unwrap(); + let request_count = listener.request_count(); + let engine_cloned = engine.clone(); + let close_job = tokio::spawn(async move { + engine_cloned + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + }); + tokio::time::timeout( + Duration::from_secs(5), + listener.wait_request_count(request_count + 1), + ) + .await + .unwrap(); + listener.wake_flush(); + tokio::time::timeout(Duration::from_secs(5), flush_job) + .await + .unwrap() + .unwrap(); + if new_skip_wal { + tokio::time::timeout(Duration::from_secs(5), listener.wait_flush_begin()) + .await + .unwrap(); + listener.wake_flush(); + } + tokio::time::timeout(Duration::from_secs(5), close_job) + .await + .unwrap() + .unwrap(); + assert_eq!( + region.version().flushed_sequence, + if new_skip_wal { 6 } else { 3 } + ); + assert!(!region.version().memtables.has_skip_wal_writes()); + + reopen_region(&engine, region_id, request.table_dir, false, HashMap::new()).await; + let actual = common_recordbatch::RecordBatches::try_collect( + engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(), + ) + .await + .unwrap() + .pretty_print() + .unwrap(); + assert_eq!(actual, expected); +} diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index a5d9a7d0a73..295789aada4 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -231,7 +231,7 @@ pub enum FlushReason { Repartition, /// Flush triggered by remote WAL pruning. RemoteWalPrune, - /// Flush before closing a Noop WAL region. + /// Flush before closing a region. Closing, } diff --git a/src/mito2/src/memtable/version.rs b/src/mito2/src/memtable/version.rs index d707128c93a..cd5dae10315 100644 --- a/src/mito2/src/memtable/version.rs +++ b/src/mito2/src/memtable/version.rs @@ -15,6 +15,7 @@ //! Memtable version. use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; use common_time::Timestamp; @@ -23,8 +24,9 @@ use store_api::metadata::RegionMetadataRef; use store_api::storage::SequenceNumber; use crate::error::Result; -use crate::memtable::time_partition::TimePartitionsRef; -use crate::memtable::{MemtableId, MemtableRef}; +use crate::memtable::bulk::part::BulkPart; +use crate::memtable::time_partition::TimePartitions; +use crate::memtable::{KeyValues, MemtableBuilderRef, MemtableId, MemtableRef}; pub(crate) type SmallMemtableVec = SmallVec<[MemtableRef; 2]>; @@ -32,26 +34,35 @@ pub(crate) type SmallMemtableVec = SmallVec<[MemtableRef; 2]>; #[derive(Debug, Clone)] pub(crate) struct MemtableVersion { /// Mutable memtable. - pub(crate) mutable: TimePartitionsRef, + pub(crate) mutable: MutableMemtablesRef, /// Immutable memtables. /// /// We only allow one flush job per region but if a flush job failed, then we /// might need to store more than one immutable memtable on the next time we /// flush the region. immutables: SmallMemtableVec, + /// Whether immutable memtables may contain skip-WAL writes. + /// Flushes are serialized per region and cover all immutable memtables. + immutable_has_skip_wal_writes: bool, } pub(crate) type MemtableVersionRef = Arc; impl MemtableVersion { /// Returns a new [MemtableVersion] with specific mutable memtable. - pub(crate) fn new(mutable: TimePartitionsRef) -> MemtableVersion { + pub(crate) fn new(mutable: TimePartitions) -> MemtableVersion { MemtableVersion { - mutable, + mutable: Arc::new(MutableMemtables::new(mutable)), immutables: SmallVec::new(), + immutable_has_skip_wal_writes: false, } } + /// Returns whether current memtables may contain skip-WAL writes. + pub(crate) fn has_skip_wal_writes(&self) -> bool { + self.mutable.has_skip_wal_writes() || self.immutable_has_skip_wal_writes + } + /// Immutable memtables. pub(crate) fn immutables(&self) -> &[MemtableRef] { &self.immutables @@ -94,21 +105,25 @@ impl MemtableVersion { } // Update the time window. - let mutable = self.mutable.new_with_part_duration(time_window, None); + let mutable = self + .mutable + .partitions + .new_with_part_duration(time_window, None); common_telemetry::debug!( "Freeze empty memtable, update partition duration from {:?} to {:?}", self.mutable.part_duration(), time_window ); return Ok(Some(MemtableVersion { - mutable: Arc::new(mutable), + mutable: Arc::new(MutableMemtables::new(mutable)), immutables: self.immutables.clone(), + immutable_has_skip_wal_writes: self.immutable_has_skip_wal_writes, })); } // Marks the mutable memtable as immutable so it can free the memory usage from our // soft limit. - self.mutable.freeze()?; + self.mutable.partitions.freeze()?; // Fork the memtable. if Some(self.mutable.part_duration()) != time_window { common_telemetry::debug!( @@ -117,17 +132,23 @@ impl MemtableVersion { time_window ); } - let mutable = Arc::new(self.mutable.fork(metadata, time_window)); + let mutable = Arc::new(MutableMemtables::new( + self.mutable.partitions.fork(metadata, time_window), + )); let mut immutables = SmallVec::with_capacity(self.immutables.len() + self.mutable.num_partitions()); immutables.extend(self.immutables.iter().cloned()); // Pushes the mutable memtable to immutable list. - self.mutable.list_memtables_to_small_vec(&mut immutables); + self.mutable + .partitions + .list_memtables_to_small_vec(&mut immutables); Ok(Some(MemtableVersion { mutable, immutables, + immutable_has_skip_wal_writes: self.immutable_has_skip_wal_writes + || self.mutable.has_skip_wal_writes(), })) } @@ -139,6 +160,9 @@ impl MemtableVersion { .filter(|mem| !ids.contains(&mem.id())) .cloned() .collect(); + if self.immutables.is_empty() { + self.immutable_has_skip_wal_writes = false; + } } /// Returns the memory usage of the mutable memtable. @@ -182,3 +206,213 @@ impl MemtableVersion { self.mutable.is_empty() && self.immutables.is_empty() } } + +/// Mutable time partitions and their shared write state. +/// +/// Version clones share this object. Freezing creates a new object so writes to +/// the new mutable memtables cannot change the old version's state. +#[derive(Debug)] +pub(crate) struct MutableMemtables { + partitions: TimePartitions, + /// Set before installation because failed writes may install some rows. + has_skip_wal_writes: AtomicBool, +} + +pub(crate) type MutableMemtablesRef = Arc; + +impl MutableMemtables { + fn new(partitions: TimePartitions) -> Self { + Self { + partitions, + has_skip_wal_writes: AtomicBool::new(false), + } + } + + /// Marks that these memtables may contain skip-WAL writes. + pub(crate) fn mark_skip_wal_writes(&self) { + self.has_skip_wal_writes.store(true, Ordering::Relaxed); + } + + fn has_skip_wal_writes(&self) -> bool { + self.has_skip_wal_writes.load(Ordering::Relaxed) + } + + /// Writes key values to the mutable time partitions. + pub(crate) fn write(&self, kvs: &KeyValues) -> Result<()> { + self.partitions.write(kvs) + } + + /// Writes a bulk part to the mutable time partitions. + pub(crate) fn write_bulk(&self, part: BulkPart) -> Result<()> { + self.partitions.write_bulk(part) + } + + /// Returns whether all mutable time partitions are empty. + pub(crate) fn is_empty(&self) -> bool { + self.partitions.is_empty() + } + + /// Returns the memtable builder. + pub(crate) fn memtable_builder(&self) -> &MemtableBuilderRef { + self.partitions.memtable_builder() + } + + /// Returns the next memtable ID. + pub(crate) fn next_memtable_id(&self) -> MemtableId { + self.partitions.next_memtable_id() + } + + /// Returns the time partition duration. + pub(crate) fn part_duration(&self) -> Duration { + self.partitions.part_duration() + } + + /// Creates empty time partitions with an optional new duration and builder. + /// The returned partitions do not share data or write state with this object. + pub(crate) fn new_with_part_duration( + &self, + part_duration: Option, + memtable_builder: Option, + ) -> TimePartitions { + self.partitions + .new_with_part_duration(part_duration, memtable_builder) + } + + /// Returns the mutable memory usage. + pub(crate) fn memory_usage(&self) -> usize { + self.partitions.memory_usage() + } + + /// Returns the number of time series in mutable time partitions. + pub(crate) fn series_count(&self) -> usize { + self.partitions.series_count() + } + + fn num_rows(&self) -> u64 { + self.partitions.num_rows() + } + + fn num_partitions(&self) -> usize { + self.partitions.num_partitions() + } + + fn list_memtables(&self, memtables: &mut Vec) { + self.partitions.list_memtables(memtables); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::memtable::time_series::TimeSeriesMemtableBuilder; + use crate::test_util::memtable_util; + + #[test] + fn test_skip_wal_writes_follow_memtable_lifecycle() { + for new_skip_wal_writes in [false, true] { + let metadata = memtable_util::metadata_for_test(); + let partitions = TimePartitions::new( + metadata.clone(), + Arc::new(TimeSeriesMemtableBuilder::default()), + 0, + Some(Duration::from_secs(5)), + ); + let version = MemtableVersion::new(partitions); + let snapshot = version.clone(); + assert!(!version.has_skip_wal_writes()); + + let kvs = memtable_util::build_key_values( + &metadata, + "hello".to_string(), + 0, + &[1000, 7000], + 1, + ); + version.mutable.mark_skip_wal_writes(); + version.mutable.write(&kvs).unwrap(); + assert!(snapshot.has_skip_wal_writes()); + + let mut frozen = version + .freeze_mutable(&metadata, Some(Duration::from_secs(5))) + .unwrap() + .unwrap(); + assert_eq!(2, frozen.immutables().len()); + assert!(frozen.has_skip_wal_writes()); + assert!(!frozen.mutable.has_skip_wal_writes()); + assert!(!Arc::ptr_eq(&version.mutable, &frozen.mutable)); + + let during_flush = frozen.clone(); + if new_skip_wal_writes { + frozen.mutable.mark_skip_wal_writes(); + } + frozen.mutable.write(&kvs).unwrap(); + let ids: Vec<_> = frozen.immutables().iter().map(|mem| mem.id()).collect(); + frozen.remove_memtables(&ids); + assert_eq!(new_skip_wal_writes, frozen.has_skip_wal_writes()); + assert_eq!( + new_skip_wal_writes, + during_flush.mutable.has_skip_wal_writes() + ); + // Removing flushed memtables does not mutate the old flush snapshot. + assert_eq!(2, during_flush.immutables().len()); + assert!(during_flush.has_skip_wal_writes()); + + let mut next_flush = frozen + .freeze_mutable(&metadata, Some(Duration::from_secs(5))) + .unwrap() + .unwrap(); + let ids: Vec<_> = next_flush.immutables().iter().map(|mem| mem.id()).collect(); + next_flush.remove_memtables(&ids); + assert!(!next_flush.has_skip_wal_writes()); + } + } + + #[test] + fn test_skip_wal_writes_survive_flush_retry() { + let metadata = memtable_util::metadata_for_test(); + let partitions = TimePartitions::new( + metadata.clone(), + Arc::new(TimeSeriesMemtableBuilder::default()), + 0, + Some(Duration::from_secs(5)), + ); + let version = MemtableVersion::new(partitions); + let kvs = memtable_util::build_key_values(&metadata, "hello".to_string(), 0, &[1000], 1); + version.mutable.mark_skip_wal_writes(); + version.mutable.write(&kvs).unwrap(); + let frozen = version + .freeze_mutable(&metadata, Some(Duration::from_secs(5))) + .unwrap() + .unwrap(); + + // A failed flush leaves immutable memtables installed. Changing the + // empty mutable's partition duration must retain their write state. + let frozen = frozen + .freeze_mutable(&metadata, Some(Duration::from_secs(10))) + .unwrap() + .unwrap(); + assert!(frozen.has_skip_wal_writes()); + assert!(!frozen.mutable.has_skip_wal_writes()); + + // A retry includes the old immutable memtables and new logged writes. + frozen.mutable.write(&kvs).unwrap(); + let mut retried = frozen + .freeze_mutable(&metadata, Some(Duration::from_secs(10))) + .unwrap() + .unwrap(); + assert_eq!(2, retried.immutables().len()); + assert!(retried.has_skip_wal_writes()); + let ids: Vec<_> = retried.immutables().iter().map(|mem| mem.id()).collect(); + retried.remove_memtables(&ids); + assert!(!retried.has_skip_wal_writes()); + + let replacement = MemtableVersion::new(TimePartitions::new( + metadata, + frozen.mutable.memtable_builder().clone(), + frozen.mutable.next_memtable_id(), + None, + )); + assert!(!replacement.has_skip_wal_writes()); + assert!(frozen.has_skip_wal_writes()); + } +} diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 5cfd6a29844..0f9cea941b3 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -2888,12 +2888,12 @@ mod tests { }) .primary_key(vec![0]); let metadata = Arc::new(builder.build().unwrap()); - let mutable = Arc::new(crate::memtable::time_partition::TimePartitions::new( + let mutable = crate::memtable::time_partition::TimePartitions::new( metadata.clone(), Arc::new(crate::test_util::memtable_util::EmptyMemtableBuilder::default()), 0, None, - )); + ); let version = Arc::new( crate::region::version::VersionBuilder::new(metadata.clone(), mutable).build(), ); @@ -3167,12 +3167,12 @@ mod tests { ..Default::default() }; let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); - let mutable = Arc::new(crate::memtable::time_partition::TimePartitions::new( + let mutable = crate::memtable::time_partition::TimePartitions::new( metadata.clone(), Arc::new(crate::test_util::memtable_util::EmptyMemtableBuilder::default()), 0, None, - )); + ); let version = Arc::new( crate::region::version::VersionBuilder::new(metadata, mutable) .options(crate::region::options::RegionOptions { @@ -3207,12 +3207,12 @@ mod tests { ..Default::default() }; let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); - let mutable = Arc::new(crate::memtable::time_partition::TimePartitions::new( + let mutable = crate::memtable::time_partition::TimePartitions::new( metadata.clone(), Arc::new(crate::test_util::memtable_util::EmptyMemtableBuilder::default()), 0, None, - )); + ); let target_region_id = metadata.region_id; // Files at or below the cursor are excluded. Among the remaining files, @@ -3224,23 +3224,25 @@ mod tests { (None, false, true, false), (None, true, true, true), ] { - let version = - crate::region::version::VersionBuilder::new(metadata.clone(), mutable.clone()) - .options(crate::region::options::RegionOptions { - preserve_row_sequence: true, - ..Default::default() - }) - .add_files( - new_noop_file_purger(), - [FileMeta { - region_id: target_region_id, - sequence, - preserve_row_sequence: marked, - ..Default::default() - }] - .into_iter(), - ) - .build(); + let version = crate::region::version::VersionBuilder::new( + metadata.clone(), + mutable.new_with_part_duration(None, None), + ) + .options(crate::region::options::RegionOptions { + preserve_row_sequence: true, + ..Default::default() + }) + .add_files( + new_noop_file_purger(), + [FileMeta { + region_id: target_region_id, + sequence, + preserve_row_sequence: marked, + ..Default::default() + }] + .into_iter(), + ) + .build(); let (files, range) = exact_sequence_range(&request, &version).unwrap(); assert_eq!(expected_files, !files.is_empty()); assert_eq!(expected_range, range.is_some()); @@ -3255,24 +3257,26 @@ mod tests { (None, true, 1, true), ] { let file_id = store_api::storage::FileId::random(); - let version = - crate::region::version::VersionBuilder::new(metadata.clone(), mutable.clone()) - .options(crate::region::options::RegionOptions { - preserve_row_sequence: true, - ..Default::default() - }) - .add_files( - new_noop_file_purger(), - [FileMeta { - region_id: RegionId::new(1, 2), - file_id, - sequence, - preserve_row_sequence: marked, - ..Default::default() - }] - .into_iter(), - ) - .build(); + let version = crate::region::version::VersionBuilder::new( + metadata.clone(), + mutable.new_with_part_duration(None, None), + ) + .options(crate::region::options::RegionOptions { + preserve_row_sequence: true, + ..Default::default() + }) + .add_files( + new_noop_file_purger(), + [FileMeta { + region_id: RegionId::new(1, 2), + file_id, + sequence, + preserve_row_sequence: marked, + ..Default::default() + }] + .into_iter(), + ) + .build(); let result = exact_sequence_range(&request, &version); if expected_error { assert!(result.is_err()); diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 84a0e13833f..71c1192cbf2 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -58,7 +58,7 @@ use crate::error::{ use crate::manifest::action::RegionManifest; use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions}; use crate::memtable::bulk::part::BulkPart; -use crate::memtable::time_partition::{TimePartitions, TimePartitionsRef}; +use crate::memtable::time_partition::TimePartitions; use crate::memtable::{MemtableBuilderProvider, ensure_json2_not_use_time_series_memtable}; use crate::metrics::{CACHE_FILL_DOWNLOADED_FILES, CACHE_FILL_PENDING_FILES}; use crate::read::series_candidate::is_sparse_metric_metadata; @@ -68,7 +68,7 @@ use crate::region::{ ManifestContext, ManifestStats, MitoRegion, MitoRegionRef, RegionLeaderState, RegionRoleState, RegionStats, }; -use crate::region_write_ctx::RegionWriteCtx; +use crate::region_write_ctx::{RegionWriteCtx, WriteSource}; use crate::request::OptionOutputTx; use crate::schedule::scheduler::SchedulerRef; use crate::series_index::{IndexFilePurger, load_version_control}; @@ -410,12 +410,8 @@ impl RegionOpener { let memtable_builder = self.memtable_builder_provider.builder_for_options(&options); let part_duration = options.compaction.time_window(); // Initial memtable id is 0. - let mutable = Arc::new(TimePartitions::new( - metadata.clone(), - memtable_builder.clone(), - 0, - part_duration, - )); + let mutable = + TimePartitions::new(metadata.clone(), memtable_builder.clone(), 0, part_duration); debug!( "Create region {} with options: {:?}, default_flat_format: {}", @@ -587,12 +583,8 @@ impl RegionOpener { .time_window() .or(manifest.compaction_time_window); // Initial memtable id is 0. - let mutable = Arc::new(TimePartitions::new( - metadata.clone(), - memtable_builder.clone(), - 0, - part_duration, - )); + let mutable = + TimePartitions::new(metadata.clone(), memtable_builder.clone(), 0, part_duration); // Updates region options by manifest before creating version. let version_builder = version_builder_from_manifest( @@ -785,7 +777,7 @@ pub(crate) fn version_builder_from_manifest( manifest: &RegionManifest, metadata: RegionMetadataRef, file_purger: FilePurgerRef, - mutable: TimePartitionsRef, + mutable: TimePartitions, region_options: RegionOptions, ) -> VersionBuilder { VersionBuilder::new(metadata, mutable) @@ -961,6 +953,7 @@ where provider.clone(), // For WAL replay, we don't need to track the write bytes rate. None, + WriteSource::WalReplay, ); for mutation in entry.mutations { rows_replayed += mutation @@ -1009,11 +1002,9 @@ where // set next_entry_id and write to memtable. region_write_ctx.set_next_entry_id(last_entry_id + 1); - region_write_ctx.write_memtable().await; - region_write_ctx.write_bulk().await; // Publish the replayed sequences only after all rows (including bulk // parts) are installed, matching the write path ordering. - region_write_ctx.publish_sequence_and_entry_id(); + region_write_ctx.write_memtables().await; } // TODO(weny): We need to update `flushed_entry_id` in the region manifest diff --git a/src/mito2/src/region/version.rs b/src/mito2/src/region/version.rs index a56cef1f544..1c595fb0b9f 100644 --- a/src/mito2/src/region/version.rs +++ b/src/mito2/src/region/version.rs @@ -32,7 +32,7 @@ use store_api::storage::{RegionId, SequenceNumber}; use crate::error::Result; use crate::manifest::action::{RegionEdit, TruncateKind}; -use crate::memtable::time_partition::{TimePartitions, TimePartitionsRef}; +use crate::memtable::time_partition::TimePartitions; use crate::memtable::version::{MemtableVersion, MemtableVersionRef}; use crate::memtable::{MemtableBuilderRef, MemtableId}; use crate::region::options::RegionOptions; @@ -319,13 +319,13 @@ impl VersionControl { version: &Version, metadata: RegionMetadataRef, memtable_builder: MemtableBuilderRef, - ) -> TimePartitionsRef { - Arc::new(TimePartitions::new( + ) -> TimePartitions { + TimePartitions::new( metadata, memtable_builder, version.memtables.mutable.next_memtable_id(), Some(version.memtables.mutable.part_duration()), - )) + ) } } @@ -402,7 +402,7 @@ pub(crate) struct VersionBuilder { impl VersionBuilder { /// Returns a new builder. - pub(crate) fn new(metadata: RegionMetadataRef, mutable: TimePartitionsRef) -> Self { + pub(crate) fn new(metadata: RegionMetadataRef, mutable: TimePartitions) -> Self { VersionBuilder { metadata: metadata.clone(), memtables: Arc::new(MemtableVersion::new(mutable)), diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index a64c37dac7f..f8d7eb9c840 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -70,8 +70,19 @@ impl Drop for WriteNotify { } } +/// Source of writes installed into memtables. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum WriteSource { + Request, + WalReplay, +} + /// Context to keep region metadata and buffer write requests. pub(crate) struct RegionWriteCtx { + /// Source of the buffered writes. + source: WriteSource, + /// Whether accepted non-empty writes bypass WAL. WAL replay never sets this flag. + has_skip_wal_writes: bool, /// Id of region to write. region_id: RegionId, /// Version of the region while creating the context. @@ -120,6 +131,7 @@ impl RegionWriteCtx { version_control: &VersionControlRef, provider: Provider, written_bytes: Option>, + source: WriteSource, ) -> RegionWriteCtx { let VersionControlData { version, @@ -129,6 +141,8 @@ impl RegionWriteCtx { } = version_control.current(); RegionWriteCtx { + source, + has_skip_wal_writes: false, region_id, version, version_control: version_control.clone(), @@ -162,6 +176,7 @@ impl RegionWriteCtx { self.next_sequence = sequence; } let num_rows = rows.as_ref().map(|rows| rows.rows.len()).unwrap_or(0); + self.has_skip_wal_writes |= num_rows > 0 && self.is_skip_wal_write(skip_wal); let mutation = Mutation { op_type, sequence: self.next_sequence, @@ -255,8 +270,29 @@ impl RegionWriteCtx { self.next_entry_id } - /// Consumes mutations and writes them into mutable memtable. - pub(crate) async fn write_memtable(&mut self) { + /// Installs ordinary and bulk writes, then publishes their sequence and entry ID. + pub(crate) async fn write_memtables(&mut self) { + if self.failed { + return; + } + // A failed installation may leave some rows in memory. + if self.has_skip_wal_writes { + self.version.memtables.mutable.mark_skip_wal_writes(); + } + self.write_mutations().await; + self.write_bulk().await; + self.publish_sequence_and_entry_id(); + } + + fn is_skip_wal_write(&self, request_skip_wal: bool) -> bool { + self.source == WriteSource::Request + && (request_skip_wal + || self.provider == Provider::Noop + || self.version.options.skip_wal) + } + + /// Consumes mutations and writes them into mutable memtables. + async fn write_mutations(&mut self) { debug_assert_eq!(self.wal_notifiers.len(), self.wal_entry.mutations.len()); if self.failed { @@ -336,6 +372,7 @@ impl RegionWriteCtx { self.wal_entry.bulk_entries.push(entry); } + self.has_skip_wal_writes |= bulk.num_rows() > 0 && self.is_skip_wal_write(skip_wal); self.bulk_notifiers .push(WriteNotify::new(sender, bulk.num_rows())); @@ -344,7 +381,7 @@ impl RegionWriteCtx { true } - pub(crate) async fn write_bulk(&mut self) { + async fn write_bulk(&mut self) { if self.failed || self.bulk_parts.is_empty() { return; } @@ -398,10 +435,10 @@ impl RegionWriteCtx { } /// Publishes the assigned sequences and entry id to the region's committed - /// watermark. Only call after both [`write_memtable`](Self::write_memtable) + /// watermark. Only call after both [`write_mutations`](Self::write_mutations) /// and [`write_bulk`](Self::write_bulk) have completed; a failed context /// must not publish. - pub(crate) fn publish_sequence_and_entry_id(&self) { + fn publish_sequence_and_entry_id(&self) { if self.failed { return; } @@ -576,6 +613,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); for (op_type, skip, num_rows) in [ (OpType::Put, skip_wal, 2), @@ -657,6 +695,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); ctx.push_mutation( OpType::Delete as i32, @@ -693,6 +732,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); let (tx, rx) = oneshot::channel(); ctx.push_mutation( @@ -736,6 +776,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); // Alternate policies in one context, retaining all parts for the memtable. for skip in [skip_wal, false, skip_wal] { @@ -772,8 +813,13 @@ mod tests { let builder = VersionControlBuilder::new(); let region_id = builder.region_id(); let version_control = Arc::new(builder.build()); - let mut ctx = - RegionWriteCtx::new(region_id, &version_control, Provider::noop_provider(), None); + let mut ctx = RegionWriteCtx::new( + region_id, + &version_control, + Provider::noop_provider(), + None, + WriteSource::Request, + ); let (tx, rx) = oneshot::channel(); assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(), None, skip_wal)); @@ -791,6 +837,142 @@ mod tests { assert!(result.is_err(), "bulk notifier should report WAL error"); } + #[tokio::test] + async fn test_skip_wal_writes_source_and_policy() { + use api::v1::helper::{tag_column_schema, time_index_column_schema}; + use api::v1::value::ValueData; + use api::v1::{ColumnDataType, Row, Value}; + + for source in [WriteSource::Request, WriteSource::WalReplay] { + for skip_wal in [false, true] { + for bulk in [false, true] { + for fail_before_install in [false, true] { + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let mut ctx = RegionWriteCtx::new( + region_id, + &version_control, + Provider::raft_engine_provider(region_id.as_u64()), + None, + source, + ); + let (tx, rx) = oneshot::channel(); + if bulk { + let mut part = new_bulk_part(); + part.fill_missing_columns(&ctx.version.metadata).unwrap(); + assert!(ctx.push_bulk(OptionOutputTx::from(tx), part, None, skip_wal)); + } else { + ctx.push_mutation( + OpType::Put as i32, + Some(Rows { + schema: vec![ + time_index_column_schema( + "ts", + ColumnDataType::TimestampMillisecond, + ), + tag_column_schema("tag_0", ColumnDataType::String), + ], + rows: vec![Row { + values: vec![ + crate::test_util::ts_ms_value(0), + Value { + value_data: Some(ValueData::StringValue( + "a".to_string(), + )), + }, + ], + }], + }), + None, + OptionOutputTx::from(tx), + None, + skip_wal, + ); + } + if fail_before_install { + ctx.set_error(Arc::new( + UnexpectedSnafu { + reason: "WAL append failed".to_string(), + } + .build(), + )); + } + ctx.write_memtables().await; + assert_eq!( + version_control + .current() + .version + .memtables + .has_skip_wal_writes(), + source == WriteSource::Request && skip_wal && !fail_before_install, + "source={source:?}, skip_wal={skip_wal}, bulk={bulk}, failed={fail_before_install}", + ); + assert_eq!( + version_control.committed_sequence(), + if fail_before_install { + 0 + } else if bulk { + 2 + } else { + 1 + } + ); + drop(ctx); + assert_eq!(rx.await.unwrap().is_err(), fail_before_install); + } + } + } + } + } + + #[test] + fn test_skip_wal_writes_region_policy_and_empty_requests() { + for source in [WriteSource::Request, WriteSource::WalReplay] { + for noop in [false, true] { + for region_skip_wal in [false, true] { + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let mut version = builder.build_version(); + version.options.skip_wal = region_skip_wal; + let version_control = + Arc::new(crate::region::version::VersionControl::new(version)); + let provider = if noop { + Provider::Noop + } else { + Provider::raft_engine_provider(region_id.as_u64()) + }; + let mut ctx = + RegionWriteCtx::new(region_id, &version_control, provider, None, source); + ctx.push_mutation( + OpType::Put as i32, + None, + None, + OptionOutputTx::none(), + None, + true, + ); + assert!(!ctx.has_skip_wal_writes); + ctx.push_mutation( + OpType::Put as i32, + Some(Rows { + schema: vec![], + rows: vec![api::v1::Row::default()], + }), + None, + OptionOutputTx::none(), + None, + false, + ); + assert_eq!( + ctx.has_skip_wal_writes, + source == WriteSource::Request && (noop || region_skip_wal) + ); + } + } + } + } + fn new_bulk_part() -> BulkPart { let schema = Arc::new(Schema::new(vec![Field::new( "ts", diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 502f6bcdc05..be52170f713 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -1611,12 +1611,12 @@ mod tests { file_purger: FilePurgerRef, files: HashMap, ) -> VersionControlRef { - let mutable = Arc::new(TimePartitions::new( + let mutable = TimePartitions::new( metadata.clone(), Arc::new(EmptyMemtableBuilder::default()), 0, None, - )); + ); let version_builder = VersionBuilder::new(metadata, mutable) .add_files(file_purger, files.values().cloned()) .build(); diff --git a/src/mito2/src/sst/primary_key.rs b/src/mito2/src/sst/primary_key.rs index 5fb258052c7..7752a652f80 100644 --- a/src/mito2/src/sst/primary_key.rs +++ b/src/mito2/src/sst/primary_key.rs @@ -270,12 +270,12 @@ mod tests { } fn version_with_files(metadata: RegionMetadataRef, files: Vec) -> VersionRef { - let mutable = Arc::new(TimePartitions::new( + let mutable = TimePartitions::new( metadata.clone(), Arc::new(TimeSeriesMemtableBuilder::default()), 0, None, - )); + ); Arc::new( VersionBuilder::new(metadata, mutable) .add_files(new_noop_file_purger(), files.into_iter()) diff --git a/src/mito2/src/test_util/version_util.rs b/src/mito2/src/test_util/version_util.rs index 1c0d7bfebc8..44827b35cc1 100644 --- a/src/mito2/src/test_util/version_util.rs +++ b/src/mito2/src/test_util/version_util.rs @@ -132,12 +132,7 @@ impl VersionControlBuilder { pub(crate) fn build_version(&self) -> Version { let metadata = Arc::new(self.metadata.clone()); - let mutable = Arc::new(TimePartitions::new( - metadata.clone(), - self.memtable_builder.clone(), - 0, - None, - )); + let mutable = TimePartitions::new(metadata.clone(), self.memtable_builder.clone(), 0, None); VersionBuilder::new(metadata, mutable) .add_files(self.file_purger.clone(), self.files.values().cloned()) .build() diff --git a/src/mito2/src/worker/handle_close.rs b/src/mito2/src/worker/handle_close.rs index 5a0406c655f..ede17511ad7 100644 --- a/src/mito2/src/worker/handle_close.rs +++ b/src/mito2/src/worker/handle_close.rs @@ -36,16 +36,12 @@ impl RegionWorkerLoop { info!("Try to close region {}, worker: {}", region_id, self.id); - // If the close request asks for a flush, or the region skips WAL, - // and has data in memtable and region is flushable (like, not in follower state), - // we should flush it before closing to ensure durability. - if (request.flush_on_close || region.skip_wal()) - && !region - .version_control - .current() - .version - .memtables - .is_empty() + // 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. + let version = region.version(); + if (request.flush_on_close || region.skip_wal() || version.memtables.has_skip_wal_writes()) + && !version.memtables.is_empty() && region.is_flushable() { info!("Region {} has pending data, waiting for flush", region_id); @@ -69,7 +65,7 @@ impl RegionWorkerLoop { return; } - // WAL configured or memtable is empty, flush is not necessary. + // No close-time flush is required, or the region cannot flush. self.remove_region(region_id).await; info!("Region {} closed, worker: {}", region_id, self.id); sender.send(Ok(0)) diff --git a/src/mito2/src/worker/handle_manifest.rs b/src/mito2/src/worker/handle_manifest.rs index aa8475a49f1..a5761434e91 100644 --- a/src/mito2/src/worker/handle_manifest.rs +++ b/src/mito2/src/worker/handle_manifest.rs @@ -257,13 +257,10 @@ impl RegionWorkerLoop { } else { None }; - let new_mutable = Arc::new( - region - .version() - .memtables - .mutable - .new_with_part_duration(version.compaction_time_window, memtable_builder), - ); + let new_mutable = version + .memtables + .mutable + .new_with_part_duration(version.compaction_time_window, memtable_builder); // Here it assumes the leader has backfilled the partition_expr of the metadata. let metadata = manifest.metadata.clone(); diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 323d73c60ec..41af569c649 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -33,7 +33,7 @@ use crate::metrics::{ WRITE_REJECT_TOTAL, WRITE_ROWS_TOTAL, WRITE_STAGE_ELAPSED, WRITE_STALL_TOTAL, }; use crate::region::{RegionLeaderState, RegionRoleState}; -use crate::region_write_ctx::RegionWriteCtx; +use crate::region_write_ctx::{RegionWriteCtx, WriteSource}; use crate::request::{SenderBulkRequest, SenderWriteRequest, WriteRequest}; use crate::wal::Wal; use crate::worker::RegionWorkerLoop; @@ -126,9 +126,7 @@ impl RegionWorkerLoop { if region_ctxs.len() == 1 { // fast path for single region. let mut region_ctx = region_ctxs.into_values().next().unwrap(); - region_ctx.write_memtable().await; - region_ctx.write_bulk().await; - region_ctx.publish_sequence_and_entry_id(); + region_ctx.write_memtables().await; put_rows += region_ctx.put_num; delete_rows += region_ctx.delete_num; } else { @@ -137,9 +135,7 @@ impl RegionWorkerLoop { .map(|mut region_ctx| { // use tokio runtime to schedule tasks. common_runtime::spawn_global(async move { - region_ctx.write_memtable().await; - region_ctx.write_bulk().await; - region_ctx.publish_sequence_and_entry_id(); + region_ctx.write_memtables().await; (region_ctx.put_num, region_ctx.delete_num) }) }) @@ -348,6 +344,7 @@ impl RegionWorkerLoop { ®ion.version_control, region.provider.clone(), Some(region.region_stats.written_bytes.clone()), + WriteSource::Request, ); e.insert(region_ctx); @@ -492,6 +489,7 @@ impl RegionWorkerLoop { ®ion.version_control, region.provider.clone(), Some(region.region_stats.written_bytes.clone()), + WriteSource::Request, ); e.insert(region_ctx); @@ -897,6 +895,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); let (tx, rx) = oneshot::channel(); ctx.push_mutation( @@ -946,6 +945,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); let (tx, rx) = oneshot::channel(); assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(1), None, skip_wal)); @@ -959,9 +959,7 @@ mod tests { if skip_wal { let ctx = contexts.get_mut(®ion_id).unwrap(); assert_eq!(ctx.next_entry_id(), 1); - ctx.write_memtable().await; - ctx.write_bulk().await; - ctx.publish_sequence_and_entry_id(); + ctx.write_memtables().await; assert_eq!(version_control.committed_sequence(), 1); assert_eq!(version_control.current().last_entry_id, 0); } @@ -995,9 +993,7 @@ mod tests { assert_eq!(entry_id + 1, region_ctxs[&ok_region].next_entry_id()); for region_ctx in region_ctxs.values_mut() { - region_ctx.write_memtable().await; - region_ctx.write_bulk().await; - region_ctx.publish_sequence_and_entry_id(); + region_ctx.write_memtables().await; } assert_eq!( @@ -1034,6 +1030,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); let (tx, rx) = oneshot::channel(); assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(3), None, skip_wal)); @@ -1050,9 +1047,7 @@ mod tests { let write_handle = tokio::spawn(async move { let mut region_ctx = region_ctxs.remove(®ion_id).unwrap(); - region_ctx.write_memtable().await; - region_ctx.write_bulk().await; - region_ctx.publish_sequence_and_entry_id(); + region_ctx.write_memtables().await; }); tokio::time::timeout(