diff --git a/src/metric-engine/src/engine.rs b/src/metric-engine/src/engine.rs index 7d5d74d0fa..6d5dace9c4 100644 --- a/src/metric-engine/src/engine.rs +++ b/src/metric-engine/src/engine.rs @@ -51,7 +51,7 @@ use store_api::region_engine::{ }; use store_api::region_request::{ AffectedRows, BatchRegionDdlRequest, RegionCatchupRequest, RegionOpenRequest, RegionPutRequest, - RegionRequest, + RegionRequest, RegionTruncateRequest, }; use store_api::storage::{RegionId, ScanRequest, SequenceNumber}; @@ -271,7 +271,28 @@ impl RegionEngine for MetricEngine { UnsupportedRegionRequestSnafu { request }.fail() } } - RegionRequest::Truncate(_) => UnsupportedRegionRequestSnafu { request }.fail(), + RegionRequest::Truncate(RegionTruncateRequest::Unflushed) => { + if self.inner.is_physical_region(region_id) { + self.inner + .mito + .handle_request( + utils::to_data_region_id(region_id), + RegionRequest::Truncate(RegionTruncateRequest::Unflushed), + ) + .await + .context(error::MitoTruncateOperationSnafu) + .map(|response| response.affected_rows) + } else { + UnsupportedRegionRequestSnafu { + request: RegionRequest::Truncate(RegionTruncateRequest::Unflushed), + } + .fail() + } + } + RegionRequest::Truncate(request) => UnsupportedRegionRequestSnafu { + request: RegionRequest::Truncate(request), + } + .fail(), RegionRequest::Delete(delete) => self.inner.delete_region(region_id, delete).await, RegionRequest::Catchup(_) => { let mut response = self @@ -574,6 +595,8 @@ mod test { use std::assert_matches; use std::collections::HashMap; + use api::v1::Rows; + use common_recordbatch::RecordBatches; use common_telemetry::info; use common_wal::options::{KafkaWalOptions, WalOptions}; use mito2::sst::location::region_dir_from_table_dir; @@ -582,12 +605,87 @@ mod test { use store_api::mito_engine_options::WAL_OPTIONS_KEY; use store_api::region_request::{ PathType, RegionCleanUpRequest, RegionCloseRequest, RegionDropRequest, RegionFlushRequest, - RegionOpenRequest, RegionRequest, + RegionOpenRequest, RegionPutRequest, RegionRequest, RegionTruncateRequest, }; use super::*; use crate::maybe_skip_kafka_log_store_integration_test; - use crate::test_util::{TestEnv, create_logical_region_request}; + use crate::test_util::{ + TestEnv, build_rows, create_logical_region_request, row_schema_with_tags, + }; + + #[tokio::test] + async fn test_discard_unflushed_data_only() { + let env = TestEnv::new().await; + env.init_metric_region().await; + + let engine = env.metric(); + let mito = env.mito(); + let physical_region_id = env.default_physical_region_id(); + let logical_region_id = env.default_logical_region_id(); + let data_region_id = utils::to_data_region_id(physical_region_id); + let metadata_region_id = utils::to_metadata_region_id(physical_region_id); + + let metadata_memtable_size = mito + .region_statistic(metadata_region_id) + .unwrap() + .memtable_size; + assert!(metadata_memtable_size > 0); + + engine + .handle_request( + logical_region_id, + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: row_schema_with_tags(&["job"]), + rows: build_rows(1, 5), + }, + hint: None, + partition_expr_version: None, + }), + ) + .await + .unwrap(); + assert!(mito.region_statistic(data_region_id).unwrap().memtable_size > 0); + + engine + .handle_request( + physical_region_id, + RegionRequest::Truncate(RegionTruncateRequest::Unflushed), + ) + .await + .unwrap(); + + assert_eq!( + 0, + mito.region_statistic(data_region_id).unwrap().memtable_size + ); + assert_eq!( + metadata_memtable_size, + mito.region_statistic(metadata_region_id) + .unwrap() + .memtable_size + ); + + let stream = engine + .scan_to_stream(logical_region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + 0, + batches.iter().map(|batch| batch.num_rows()).sum::() + ); + + // This maintenance operation applies to a physical region group only. + engine + .handle_request( + logical_region_id, + RegionRequest::Truncate(RegionTruncateRequest::Unflushed), + ) + .await + .unwrap_err(); + } #[tokio::test] async fn close_open_regions() { diff --git a/src/metric-engine/src/error.rs b/src/metric-engine/src/error.rs index 316588e6b3..0803223ba3 100644 --- a/src/metric-engine/src/error.rs +++ b/src/metric-engine/src/error.rs @@ -178,6 +178,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Mito truncate operation fails"))] + MitoTruncateOperation { + source: BoxedError, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Mito sync operation fails"))] MitoSyncOperation { source: BoxedError, @@ -451,6 +458,7 @@ impl ErrorExt for Error { | MitoReadOperation { source, .. } | MitoWriteOperation { source, .. } | MitoFlushOperation { source, .. } + | MitoTruncateOperation { source, .. } | MitoSyncOperation { source, .. } | MitoEnterStagingOperation { source, .. } | BatchOpenMitoRegion { source, .. } @@ -487,6 +495,7 @@ impl ErrorExt for Error { | MitoReadOperation { source, .. } | MitoWriteOperation { source, .. } | MitoFlushOperation { source, .. } + | MitoTruncateOperation { source, .. } | MitoSyncOperation { source, .. } | MitoEnterStagingOperation { source, .. } | MitoCopyRegionFromOperation { source, .. } diff --git a/src/mito2/src/engine/truncate_test.rs b/src/mito2/src/engine/truncate_test.rs index c23b9f770a..c90bb960f8 100644 --- a/src/mito2/src/engine/truncate_test.rs +++ b/src/mito2/src/engine/truncate_test.rs @@ -14,6 +14,7 @@ use std::collections::HashMap; use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; use api::v1::Rows; @@ -21,19 +22,80 @@ use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_recordbatch::RecordBatches; use common_telemetry::{info, init_default_ut_logging}; -use store_api::region_engine::RegionEngine; +use common_wal::options::{WAL_OPTIONS_KEY, WalOptions}; +use store_api::region_engine::{RegionEngine, RegionRole}; use store_api::region_request::{ PathType, RegionFlushRequest, RegionOpenRequest, RegionRequest, RegionTruncateRequest, }; use store_api::storage::RegionId; +use tokio::sync::Notify; use super::ScanRequest; use crate::config::MitoConfig; -use crate::engine::listener::FlushCancellationListener; +use crate::engine::listener::{EventListener, FlushCancellationListener}; use crate::test_util::{ CreateRequestBuilder, MockWriteBufferManager, TestEnv, build_rows, put_rows, rows_schema, }; +#[derive(Default)] +struct FlushCommitListener { + commit_started: Notify, + cancel_requested: Notify, + resume_flush: Notify, +} + +impl FlushCommitListener { + async fn wait_commit_started(&self) { + self.commit_started.notified().await; + } + + async fn wait_cancel_requested(&self) { + self.cancel_requested.notified().await; + } + + fn resume_flush(&self) { + self.resume_flush.notify_one(); + } +} + +#[async_trait::async_trait] +impl EventListener for FlushCommitListener { + async fn on_flush_commit_begin(&self, _region_id: RegionId) { + self.commit_started.notify_one(); + self.resume_flush.notified().await; + } + + fn on_flush_cancel_requested(&self, _region_id: RegionId) { + self.cancel_requested.notify_one(); + } +} + +#[derive(Default)] +struct CountingStallListener { + count: AtomicUsize, + notify: Notify, +} + +impl CountingStallListener { + async fn wait_for_count(&self, expected: usize) { + loop { + let notified = self.notify.notified(); + if self.count.load(Ordering::Relaxed) >= expected { + return; + } + notified.await; + } + } +} + +#[async_trait::async_trait] +impl EventListener for CountingStallListener { + fn on_write_stall(&self) { + self.count.fetch_add(1, Ordering::Relaxed); + self.notify.notify_waiters(); + } +} + #[tokio::test] async fn test_engine_truncate_region_basic() { test_engine_truncate_region_basic_with_format(false).await; @@ -344,12 +406,492 @@ async fn test_engine_truncate_reopen_with_format(flat_format: bool) { } #[tokio::test] -async fn test_engine_truncate_during_flush() { - test_engine_truncate_during_flush_with_format(false).await; - test_engine_truncate_during_flush_with_format(true).await; +async fn test_engine_discard_unflushed() { + test_engine_discard_unflushed_with_format(false).await; + test_engine_discard_unflushed_with_format(true).await; } -async fn test_engine_truncate_during_flush_with_format(flat_format: bool) { +async fn test_engine_discard_unflushed_with_format(flat_format: bool) { + let mut env = TestEnv::with_prefix("discard-unflushed").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(); + let table_dir = request.table_dir.clone(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + // Persist the first batch and leave the second batch in memtables and WAL. + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas.clone(), + rows: build_rows(0, 3), + }, + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas.clone(), + rows: build_rows(3, 6), + }, + ) + .await; + + let region = engine.get_region(region_id).unwrap(); + let version_data = region.version_control.current(); + let discarded_entry_id = version_data.last_entry_id; + let discarded_sequence = version_data.committed_sequence; + assert!(!version_data.version.memtables.is_empty()); + assert_eq!( + 1, + engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap() + .num_files() + ); + + engine + .handle_request( + region_id, + RegionRequest::Truncate(RegionTruncateRequest::Unflushed), + ) + .await + .unwrap(); + + let version = region.version(); + assert!(version.memtables.is_empty()); + assert_eq!(discarded_entry_id, version.flushed_entry_id); + assert_eq!(discarded_sequence, version.flushed_sequence); + assert_eq!( + 1, + engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap() + .num_files() + ); + + // Repeating the request is a no-op and also retries WAL obsoletion if needed. + engine + .handle_request( + region_id, + RegionRequest::Truncate(RegionTruncateRequest::Unflushed), + ) + .await + .unwrap(); + + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + let expected = "\ ++-------+---------+---------------------+ +| tag_0 | field_0 | ts | ++-------+---------+---------------------+ +| 0 | 0.0 | 1970-01-01T00:00:00 | +| 1 | 1.0 | 1970-01-01T00:00:01 | +| 2 | 2.0 | 1970-01-01T00:00:02 | ++-------+---------+---------------------+"; + assert_eq!(expected, batches.pretty_print().unwrap()); + + // The replay boundary is durable: reopening must not restore discarded rows. + 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, + path_type: PathType::Bare, + options: HashMap::default(), + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); + + let region = engine.get_region(region_id).unwrap(); + assert!(region.version().memtables.is_empty()); + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!(expected, batches.pretty_print().unwrap()); + + // The region remains writable after discarding its unflushed data. + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas, + rows: build_rows(6, 8), + }, + ) + .await; + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + 5, + batches.iter().map(|batch| batch.num_rows()).sum::() + ); +} + +#[tokio::test] +async fn test_engine_discard_unflushed_skip_wal() { + let mut env = TestEnv::with_prefix("discard-unflushed-skip-wal").await; + let engine = env.create_engine(MitoConfig::default()).await; + + let region_id = RegionId::new(1, 1); + let mut request = CreateRequestBuilder::new().build(); + request.options.insert( + WAL_OPTIONS_KEY.to_string(), + serde_json::to_string(&WalOptions::Noop).unwrap(), + ); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas.clone(), + rows: build_rows(0, 3), + }, + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas.clone(), + rows: build_rows(3, 6), + }, + ) + .await; + + let region = engine.get_region(region_id).unwrap(); + let version_data = region.version_control.current(); + let discarded_sequence = version_data.committed_sequence; + // Entry ids stay at 0 without a WAL, so only the sequence advances. + assert_eq!(0, version_data.last_entry_id); + assert!(discarded_sequence > version_data.version.flushed_sequence); + + engine + .handle_request( + region_id, + RegionRequest::Truncate(RegionTruncateRequest::Unflushed), + ) + .await + .unwrap(); + + let version = region.version(); + assert!(version.memtables.is_empty()); + assert_eq!(0, version.flushed_entry_id); + assert_eq!(discarded_sequence, version.flushed_sequence); + assert_eq!( + 1, + engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap() + .num_files() + ); + + // The flushed rows survive and the region stays writable. + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas, + rows: build_rows(6, 8), + }, + ) + .await; + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + 5, + batches.iter().map(|batch| batch.num_rows()).sum::() + ); +} + +#[tokio::test] +async fn test_engine_discard_unflushed_wakes_stalled_writes() { + let mut env = TestEnv::with_prefix("discard-unflushed-wake-stalled").await; + let write_buffer_manager = Arc::new(MockWriteBufferManager::default()); + let listener = Arc::new(CountingStallListener::default()); + let engine = env + .create_engine_with( + MitoConfig { + num_workers: 2, + ..Default::default() + }, + Some(write_buffer_manager.clone()), + Some(listener.clone()), + None, + ) + .await; + + // These region IDs are routed to different workers when num_workers is 2. + let discarded_region_id = RegionId::new(1, 1); + let peer_region_id = RegionId::new(2, 1); + let discarded_request = CreateRequestBuilder::new().build(); + let discarded_schema = rows_schema(&discarded_request); + engine + .handle_request( + discarded_region_id, + RegionRequest::Create(discarded_request), + ) + .await + .unwrap(); + let peer_request = CreateRequestBuilder::new().table_dir("peer").build(); + let peer_schema = rows_schema(&peer_request); + engine + .handle_request(peer_region_id, RegionRequest::Create(peer_request)) + .await + .unwrap(); + + put_rows( + &engine, + discarded_region_id, + Rows { + schema: discarded_schema.clone(), + rows: build_rows(0, 3), + }, + ) + .await; + + write_buffer_manager.set_should_stall(true); + let discarded_engine = engine.clone(); + let discarded_write = tokio::spawn(async move { + put_rows( + &discarded_engine, + discarded_region_id, + Rows { + schema: discarded_schema, + rows: build_rows(3, 6), + }, + ) + .await; + }); + let peer_engine = engine.clone(); + let peer_write = tokio::spawn(async move { + put_rows( + &peer_engine, + peer_region_id, + Rows { + schema: peer_schema, + rows: build_rows(6, 8), + }, + ) + .await; + }); + + tokio::time::timeout(Duration::from_secs(10), listener.wait_for_count(2)) + .await + .expect("writes should stall on both workers"); + assert!(!discarded_write.is_finished()); + assert!(!peer_write.is_finished()); + + write_buffer_manager.set_should_stall(false); + engine + .handle_request( + discarded_region_id, + RegionRequest::Truncate(RegionTruncateRequest::Unflushed), + ) + .await + .unwrap(); + + tokio::time::timeout(Duration::from_secs(10), async { + discarded_write.await.unwrap(); + peer_write.await.unwrap(); + }) + .await + .expect("discard should wake stalled writes on both workers"); + + let stream = engine + .scan_to_stream(discarded_region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + 3, + batches.iter().map(|batch| batch.num_rows()).sum::() + ); + + let stream = engine + .scan_to_stream(peer_region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + 2, + batches.iter().map(|batch| batch.num_rows()).sum::() + ); +} + +#[tokio::test] +async fn test_engine_truncate_during_flush() { + test_engine_truncate_during_flush_with_format(false, false).await; + test_engine_truncate_during_flush_with_format(true, false).await; +} + +#[tokio::test] +async fn test_engine_discard_unflushed_during_flush() { + test_engine_truncate_during_flush_with_format(false, true).await; + test_engine_truncate_during_flush_with_format(true, true).await; +} + +#[tokio::test] +async fn test_engine_discard_unflushed_after_flush_commit_started() { + test_engine_discard_unflushed_after_flush_commit_started_with_format(false).await; + test_engine_discard_unflushed_after_flush_commit_started_with_format(true).await; +} + +async fn test_engine_discard_unflushed_after_flush_commit_started_with_format(flat_format: bool) { + let mut env = TestEnv::with_prefix("discard-unflushed-after-flush-commit").await; + let listener = Arc::new(FlushCommitListener::default()); + let engine = env + .create_engine_with( + MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }, + None, + Some(listener.clone()), + None, + ) + .await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas, + rows: build_rows(0, 3), + }, + ) + .await; + + let flush_engine = engine.clone(); + let flush_task = tokio::spawn(async move { + flush_engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(5), listener.wait_commit_started()) + .await + .expect("flush did not reach the commit gate"); + + let discard_engine = engine.clone(); + let discard_task = tokio::spawn(async move { + discard_engine + .handle_request( + region_id, + RegionRequest::Truncate(RegionTruncateRequest::Unflushed), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(5), listener.wait_cancel_requested()) + .await + .expect("discard did not queue behind the committing flush"); + assert!(!discard_task.is_finished()); + + listener.resume_flush(); + tokio::time::timeout(Duration::from_secs(5), flush_task) + .await + .expect("flush did not finish") + .expect("flush task panicked") + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), discard_task) + .await + .expect("discard did not finish") + .expect("discard task panicked") + .unwrap(); + + let region = engine.get_region(region_id).unwrap(); + assert!(region.version().memtables.is_empty()); + assert_eq!( + 1, + engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap() + .num_files() + ); + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + 3, + batches.iter().map(|batch| batch.num_rows()).sum::() + ); +} + +async fn test_engine_truncate_during_flush_with_format(flat_format: bool, unflushed_only: bool) { init_default_ut_logging(); let mut env = TestEnv::with_prefix("truncate-during-flush").await; let write_buffer_manager = Arc::new(MockWriteBufferManager::default()); @@ -412,7 +954,11 @@ async fn test_engine_truncate_during_flush_with_format(flat_format: bool) { Duration::from_secs(5), engine.handle_request( region_id, - RegionRequest::Truncate(RegionTruncateRequest::All), + RegionRequest::Truncate(if unflushed_only { + RegionTruncateRequest::Unflushed + } else { + RegionTruncateRequest::All + }), ), ) .await @@ -434,7 +980,7 @@ async fn test_engine_truncate_during_flush_with_format(flat_format: bool) { let request = ScanRequest::default(); let scanner = engine.scanner(region_id, request.clone()).await.unwrap(); assert_eq!(0, scanner.num_files()); - assert_eq!(Some(entry_id), truncated_entry_id); + assert_eq!((!unflushed_only).then_some(entry_id), truncated_entry_id); assert_eq!(sequence, truncated_sequence); // Reopen the engine. @@ -465,5 +1011,20 @@ async fn test_engine_truncate_during_flush_with_format(flat_format: bool) { let region = engine.get_region(region_id).unwrap(); let current_version = region.version_control.current().version; - assert_eq!(current_version.truncated_entry_id, Some(entry_id)); + assert_eq!( + current_version.truncated_entry_id, + (!unflushed_only).then_some(entry_id) + ); + + // The cancelled flush wrote no files, so both kinds leave the region empty and the + // replay must not bring the discarded rows back. + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + 0, + batches.iter().map(|batch| batch.num_rows()).sum::() + ); } diff --git a/src/mito2/src/region/version.rs b/src/mito2/src/region/version.rs index 006f684a23..67db17048c 100644 --- a/src/mito2/src/region/version.rs +++ b/src/mito2/src/region/version.rs @@ -282,6 +282,28 @@ impl VersionControl { }; } + /// Discards all memtables while preserving persisted SST files. + pub(crate) fn discard_unflushed( + &self, + discarded_entry_id: EntryId, + discarded_sequence: SequenceNumber, + ) { + let version = self.current().version; + let memtable_builder = version.memtables.mutable.memtable_builder().clone(); + let new_mutable = + Self::new_mutable_from_version(&version, version.metadata.clone(), memtable_builder); + let new_version = Arc::new( + VersionBuilder::from_version(version) + .memtables(MemtableVersion::new(new_mutable)) + .flushed_entry_id(discarded_entry_id) + .flushed_sequence(discarded_sequence) + .build(), + ); + + let mut version_data = self.data.write().unwrap(); + version_data.version = new_version; + } + /// Overwrites the current version with a new version. pub(crate) fn overwrite_current(&self, version: VersionRef) { let mut version_data = self.data.write().unwrap(); diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index e0ea655be4..94d442e261 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -43,7 +43,7 @@ use store_api::region_request::{ RegionFlushRequest, RegionOpenRequest, RegionRequest, RegionTruncateRequest, StagingPartitionDirective, }; -use store_api::storage::{FileId, RegionId}; +use store_api::storage::{FileId, RegionId, SequenceNumber}; use tokio::sync::oneshot::{self, Receiver, Sender}; use crate::compaction::{CompactionExecution, CompactionPickFinished}; @@ -918,6 +918,8 @@ pub(crate) enum BackgroundNotify { CompactionFailed(CompactionFailed), /// Truncate result. Truncate(TruncateResult), + /// Discard unflushed data result. + DiscardUnflushed(DiscardUnflushedResult), /// Region change result. RegionChange(RegionChangeResult), /// Region edit result. @@ -1082,6 +1084,25 @@ pub(crate) struct TruncateResult { pub(crate) kind: TruncateKind, } +/// Notifies the result of discarding unflushed data from a region. +#[derive(Debug)] +pub(crate) struct DiscardUnflushedResult { + /// Region id. + pub(crate) region_id: RegionId, + /// Result sender. + pub(crate) sender: OptionOutputTx, + /// Manifest update result. + pub(crate) result: Result<()>, + /// Last WAL entry covered by the discard operation. + pub(crate) discarded_entry_id: EntryId, + /// Last sequence covered by the discard operation. + pub(crate) discarded_sequence: SequenceNumber, + /// Estimated number of discarded rows. + pub(crate) discarded_rows: u64, + /// Estimated number of discarded bytes. + pub(crate) discarded_bytes: u64, +} + /// Notifies the region the result of writing region change action. #[derive(Debug)] pub(crate) struct RegionChangeResult { diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 801ed5eb1b..63706cab4e 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -1264,6 +1264,9 @@ impl RegionWorkerLoop { } BackgroundNotify::CompactionFailed(req) => self.handle_compaction_failure(req).await, BackgroundNotify::Truncate(req) => self.handle_truncate_result(req).await, + BackgroundNotify::DiscardUnflushed(req) => { + self.handle_discard_unflushed_result(req).await + } BackgroundNotify::RegionChange(req) => { self.handle_manifest_region_change_result(req).await } diff --git a/src/mito2/src/worker/handle_manifest.rs b/src/mito2/src/worker/handle_manifest.rs index d932f89735..700debadad 100644 --- a/src/mito2/src/worker/handle_manifest.rs +++ b/src/mito2/src/worker/handle_manifest.rs @@ -25,7 +25,7 @@ use parquet::file::metadata::PageIndexPolicy; use snafu::ResultExt; use store_api::logstore::LogStore; use store_api::metadata::RegionMetadataRef; -use store_api::storage::RegionId; +use store_api::storage::{RegionId, SequenceNumber}; use crate::cache::CacheManagerRef; use crate::cache::file_cache::{FileType, IndexKey}; @@ -41,11 +41,13 @@ use crate::region::options::RegionOptions; use crate::region::version::VersionControlRef; use crate::region::{MitoRegionRef, RegionLeaderState, RegionRoleState}; use crate::request::{ - BackgroundNotify, BuildIndexRequest, OptionOutputTx, RegionChangeResult, RegionEditRequest, - RegionEditResult, RegionSyncRequest, TruncateResult, WorkerRequest, WorkerRequestWithTime, + BackgroundNotify, BuildIndexRequest, DiscardUnflushedResult, OptionOutputTx, + RegionChangeResult, RegionEditRequest, RegionEditResult, RegionSyncRequest, TruncateResult, + WorkerRequest, WorkerRequestWithTime, }; use crate::sst::index::IndexBuildType; use crate::sst::location; +use crate::wal::EntryId; use crate::worker::{RegionWorkerLoop, WorkerListener}; pub(crate) type RegionEditQueues = HashMap; @@ -506,6 +508,62 @@ impl RegionWorkerLoop { }); } + /// Advances the durable replay frontier before discarding a region's memtables. + pub(crate) fn handle_manifest_discard_unflushed_action( + &self, + region: MitoRegionRef, + discarded_entry_id: EntryId, + discarded_sequence: SequenceNumber, + discarded_rows: u64, + discarded_bytes: u64, + sender: OptionOutputTx, + ) { + if let Err(e) = region.set_truncating() { + sender.send(Err(e)); + return; + } + + let region_id = region.region_id; + let request_sender = self.sender.clone(); + let manifest_ctx = region.manifest_ctx.clone(); + + common_runtime::spawn_global(async move { + // The frontier moves to the last written entry and sequence, so replaying the + // WAL after a restart skips everything the memtables held. + let edit = RegionEdit { + files_to_add: Vec::new(), + files_to_remove: Vec::new(), + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: Some(discarded_entry_id), + flushed_sequence: Some(discarded_sequence), + committed_sequence: None, + }; + let action_list = RegionMetaActionList::with_action(RegionMetaAction::Edit(edit)); + let result = manifest_ctx + .update_manifest(RegionLeaderState::Truncating, action_list, false) + .await + .map(|_| ()); + + let result = DiscardUnflushedResult { + region_id, + sender, + result, + discarded_entry_id, + discarded_sequence, + discarded_rows, + discarded_bytes, + }; + let _ = request_sender + .send(WorkerRequestWithTime::new(WorkerRequest::Background { + region_id, + notify: BackgroundNotify::DiscardUnflushed(result), + })) + .await + .inspect_err(|_| warn!("failed to send discard unflushed result")); + }); + } + /// Writes region change action to the manifest and then applies it to the region in background. pub(crate) fn handle_manifest_region_change( &self, diff --git a/src/mito2/src/worker/handle_truncate.rs b/src/mito2/src/worker/handle_truncate.rs index 020bf04099..d62bf8369a 100644 --- a/src/mito2/src/worker/handle_truncate.rs +++ b/src/mito2/src/worker/handle_truncate.rs @@ -14,7 +14,7 @@ //! Handling truncate related requests. -use common_telemetry::{debug, info}; +use common_telemetry::{debug, info, warn}; use store_api::logstore::LogStore; use store_api::region_request::RegionTruncateRequest; use store_api::storage::RegionId; @@ -22,7 +22,7 @@ use store_api::storage::RegionId; use crate::error::RegionNotFoundSnafu; use crate::manifest::action::{RegionTruncate, TruncateKind}; use crate::region::RegionLeaderState; -use crate::request::{DdlRequest, OptionOutputTx, TruncateResult}; +use crate::request::{DdlRequest, DiscardUnflushedResult, OptionOutputTx, TruncateResult}; use crate::worker::RegionWorkerLoop; impl RegionWorkerLoop { @@ -87,6 +87,47 @@ impl RegionWorkerLoop { self.handle_manifest_truncate_action(region, truncate, sender); } + RegionTruncateRequest::Unflushed => { + let memtables = &version_data.version.memtables; + if memtables.is_empty() + && version_data.version.flushed_entry_id == version_data.last_entry_id + && version_data.version.flushed_sequence == version_data.committed_sequence + { + self.update_topic_latest_entry_id(®ion); + let result = self + .wal + .obsolete( + region_id, + version_data.version.flushed_entry_id, + ®ion.provider, + ) + .await + .map(|_| 0); + sender.send(result); + return; + } + + let discarded_rows = memtables.num_rows(); + let discarded_bytes = + (memtables.mutable_usage() + memtables.immutables_usage()) as u64; + warn!( + "Discarding unflushed data from region {}, entry_id: {}, sequence: {}, estimated rows: {}, estimated bytes: {}", + region_id, + version_data.last_entry_id, + version_data.committed_sequence, + discarded_rows, + discarded_bytes, + ); + + self.handle_manifest_discard_unflushed_action( + region, + version_data.last_entry_id, + version_data.committed_sequence, + discarded_rows, + discarded_bytes, + sender, + ); + } RegionTruncateRequest::ByTimeRanges { time_ranges } => { info!( "Try to partially truncate region {} by time ranges: {:?}", @@ -196,4 +237,61 @@ impl RegionWorkerLoop { truncate_result.sender.send(Ok(0)); } + + /// Handles the result of advancing the replay frontier for unflushed data. + pub(crate) async fn handle_discard_unflushed_result(&mut self, result: DiscardUnflushedResult) { + let region_id = result.region_id; + let Some(region) = self.regions.get_region(region_id) else { + result.sender.send(RegionNotFoundSnafu { region_id }.fail()); + return; + }; + + region.switch_state_to_writable(RegionLeaderState::Truncating); + + if let Err(e) = result.result { + result.sender.send(Err(e)); + return; + } + + region + .version_control + .discard_unflushed(result.discarded_entry_id, result.discarded_sequence); + + // The flush and compaction schedulers are already quiesced because the request + // queued itself behind any running job, so these two are usually no-ops. Index + // builds aren't queued that way, so pending builds for surviving files are + // retired here. They would abort anyway while the region sits in `Truncating`, + // and those files may need a later index rebuild. + self.flush_scheduler.on_region_truncated(region_id); + self.compaction_scheduler.on_region_truncated(region_id); + self.index_build_scheduler + .on_region_truncated(region_id) + .await; + self.update_topic_latest_entry_id(®ion); + + // Discarding memtables releases write buffer memory. Notify other workers and + // retry requests stalled on this worker before WAL cleanup, which may fail even + // though the memory has already been released. + self.notify_group(); + self.handle_stalled_requests().await; + + if let Err(e) = self + .wal + .obsolete(region_id, result.discarded_entry_id, ®ion.provider) + .await + { + result.sender.send(Err(e)); + return; + } + + warn!( + "Discarded unflushed data from region {}, entry_id: {}, sequence: {}, estimated rows: {}, estimated bytes: {}", + region_id, + result.discarded_entry_id, + result.discarded_sequence, + result.discarded_rows, + result.discarded_bytes, + ); + result.sender.send(Ok(0)); + } } diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index 5cbd4060e8..490e6c4d6e 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -1629,6 +1629,16 @@ pub struct RegionBuildIndexRequest {} pub enum RegionTruncateRequest { /// Truncate all data in the region. All, + /// Discard all unflushed data while preserving persisted SST files. + /// + /// This destroys the region's in-memory data irreversibly. Persisted SST files and + /// the writable state are preserved, and the WAL is obsoleted up to the discard point + /// so a restart won't replay the discarded data. + /// + /// An error may be returned after the data has already been discarded, because the + /// WAL is obsoleted last. Retrying is safe: a request against a region with nothing + /// left to discard only re-attempts the WAL obsoletion. + Unflushed, ByTimeRanges { /// Time ranges to truncate. Both bound are inclusive. /// only files that are fully contained in the time range will be truncated.