diff --git a/src/mito2/src/engine/listener.rs b/src/mito2/src/engine/listener.rs index a8470dec6f..224268d5a6 100644 --- a/src/mito2/src/engine/listener.rs +++ b/src/mito2/src/engine/listener.rs @@ -253,6 +253,7 @@ pub struct AlterFlushListener { flush_begin_notify: Notify, block_flush_notify: Notify, request_begin_notify: Notify, + request_count: AtomicUsize, } impl AlterFlushListener { @@ -266,6 +267,18 @@ impl AlterFlushListener { self.request_begin_notify.notified().await; } + /// Waits until the worker has received at least `times` request batches. + pub async fn wait_request_count(&self, times: usize) { + while self.request_count.load(Ordering::Relaxed) < times { + self.request_begin_notify.notified().await; + } + } + + /// Returns received request batch count. + pub fn request_count(&self) -> usize { + self.request_count.load(Ordering::Relaxed) + } + /// Continue the flush job. pub fn wake_flush(&self) { self.block_flush_notify.notify_one(); @@ -286,6 +299,7 @@ impl EventListener for AlterFlushListener { fn on_recv_requests(&self, request_num: usize) { info!("receive {} request", request_num); + self.request_count.fetch_add(1, Ordering::Relaxed); self.request_begin_notify.notify_one(); } } diff --git a/src/mito2/src/engine/skip_wal_test.rs b/src/mito2/src/engine/skip_wal_test.rs index 3b6cf89f07..eca19f22d7 100644 --- a/src/mito2/src/engine/skip_wal_test.rs +++ b/src/mito2/src/engine/skip_wal_test.rs @@ -12,16 +12,22 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; +use std::time::Duration; + use api::v1::Rows; use common_wal::options::{WAL_OPTIONS_KEY, WalOptions}; use store_api::region_engine::{RegionEngine, RegionRole}; use store_api::region_request::{ - RegionCloseRequest, RegionOpenRequest, RegionRequest, RegionTruncateRequest, + RegionCloseRequest, RegionOpenRequest, RegionPutRequest, RegionRequest, RegionTruncateRequest, }; use store_api::storage::{RegionId, ScanRequest}; use crate::config::MitoConfig; -use crate::test_util::{CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema}; +use crate::engine::listener::AlterFlushListener; +use crate::test_util::{ + CreateRequestBuilder, TestEnv, build_rows, flush_region, put_rows, rows_schema, +}; #[tokio::test] async fn test_close_region_skip_wal_with_pending_data() { @@ -215,6 +221,259 @@ async fn test_close_follower_region_skip_wal_with_pending_data() { assert!(!engine.is_region_exists(region_id)); } +#[tokio::test] +async fn test_close_region_skip_wal_while_flush_in_flight_closes_region() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("close-skip-wal-while-flush-in-flight").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 mut request = CreateRequestBuilder::new().build(); + let wal_options = WalOptions::Noop; + request.options.insert( + WAL_OPTIONS_KEY.to_string(), + serde_json::to_string(&wal_options).unwrap(), + ); + + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + + let rows = Rows { + schema: rows_schema(&request), + rows: build_rows(0, 3), + }; + put_rows(&engine, region_id, rows).await; + assert!( + !engine + .get_region(region_id) + .unwrap() + .version() + .memtables + .is_empty() + ); + + let request_count = listener.request_count(); + let engine_cloned = engine.clone(); + let flush_job = tokio::spawn(async move { + flush_region(&engine_cloned, region_id, None).await; + }); + listener.wait_flush_begin().await; + listener.wait_request_count(request_count + 1).await; + + let engine_cloned = engine.clone(); + let close_job = tokio::spawn(async move { + engine_cloned + .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .await + .unwrap(); + }); + listener.wait_request_count(request_count + 2).await; + + listener.wake_flush(); + tokio::time::timeout(Duration::from_secs(5), flush_job) + .await + .unwrap() + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), close_job) + .await + .unwrap() + .unwrap(); + + assert!(!engine.is_region_exists(region_id)); +} + +#[tokio::test] +async fn test_close_region_skip_wal_rejects_writes_queued_after_close() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("close-skip-wal-rejects-writes-after-close").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 mut request = CreateRequestBuilder::new().build(); + request.options.insert( + WAL_OPTIONS_KEY.to_string(), + serde_json::to_string(&WalOptions::Noop).unwrap(), + ); + + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: rows_schema(&request), + rows: build_rows(0, 3), + }, + ) + .await; + + let request_count = listener.request_count(); + let engine_cloned = engine.clone(); + let flush_job = tokio::spawn(async move { + flush_region(&engine_cloned, region_id, None).await; + }); + listener.wait_flush_begin().await; + listener.wait_request_count(request_count + 1).await; + + let pre_close_write = engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: rows_schema(&request), + rows: build_rows(3, 4), + }, + hint: None, + partition_expr_version: None, + }), + ) + .await + .unwrap(); + assert_eq!(1, pre_close_write.affected_rows); + + let engine_cloned = engine.clone(); + let close_job = tokio::spawn(async move { + engine_cloned + .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .await + .unwrap(); + }); + listener.wait_request_count(request_count + 3).await; + + let engine_cloned = engine.clone(); + let request_cloned = request.clone(); + let write_job = tokio::spawn(async move { + engine_cloned + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: rows_schema(&request_cloned), + rows: build_rows(4, 5), + }, + hint: None, + partition_expr_version: None, + }), + ) + .await + }); + listener.wait_request_count(request_count + 4).await; + + listener.wake_flush(); + tokio::time::timeout(Duration::from_secs(5), flush_job) + .await + .unwrap() + .unwrap(); + listener.wait_flush_begin().await; + + let engine_cloned = engine.clone(); + let request_cloned = request.clone(); + let late_write_job = tokio::spawn(async move { + engine_cloned + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: rows_schema(&request_cloned), + rows: build_rows(5, 6), + }, + hint: None, + partition_expr_version: None, + }), + ) + .await + }); + listener.wait_request_count(request_count + 5).await; + + listener.wake_flush(); + tokio::time::timeout(Duration::from_secs(5), close_job) + .await + .unwrap() + .unwrap(); + let write_result = tokio::time::timeout(Duration::from_secs(5), write_job) + .await + .unwrap() + .unwrap(); + let late_write_result = tokio::time::timeout(Duration::from_secs(5), late_write_job) + .await + .unwrap() + .unwrap(); + + assert!(write_result.is_err()); + assert!(late_write_result.is_err()); + assert!(!engine.is_region_exists(region_id)); +} + +#[tokio::test] +async fn test_concurrent_close_region_skip_wal_while_flush_in_flight_succeeds() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("concurrent-close-skip-wal-while-flush-in-flight").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 mut request = CreateRequestBuilder::new().build(); + request.options.insert( + WAL_OPTIONS_KEY.to_string(), + serde_json::to_string(&WalOptions::Noop).unwrap(), + ); + + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + + let rows = Rows { + schema: rows_schema(&request), + rows: build_rows(0, 3), + }; + put_rows(&engine, region_id, rows).await; + + let request_count = listener.request_count(); + let engine_cloned = engine.clone(); + let first_close = tokio::spawn(async move { + engine_cloned + .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .await + .unwrap(); + }); + listener.wait_flush_begin().await; + listener.wait_request_count(request_count + 1).await; + + let engine_cloned = engine.clone(); + let second_close = tokio::spawn(async move { + engine_cloned + .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .await + .unwrap(); + }); + listener.wait_request_count(request_count + 2).await; + + listener.wake_flush(); + tokio::time::timeout(Duration::from_secs(5), first_close) + .await + .unwrap() + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), second_close) + .await + .unwrap() + .unwrap(); + + assert!(!engine.is_region_exists(region_id)); +} + #[tokio::test] async fn test_close_region_after_truncate_skip_wal() { common_telemetry::init_default_ut_logging(); diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index b90e5a9404..f087d8a554 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -16,8 +16,8 @@ use std::collections::HashMap; use std::num::NonZeroU64; -use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; use std::time::Instant; use bytes::Bytes; @@ -39,7 +39,7 @@ use crate::cache::CacheManagerRef; use crate::config::MitoConfig; use crate::engine::region_hook::SstFileInfo; use crate::error::{ - Error, FlushRegionSnafu, JoinSnafu, RegionClosedSnafu, RegionDroppedSnafu, + Error, FlushRegionSnafu, JoinSnafu, RegionBusySnafu, RegionClosedSnafu, RegionDroppedSnafu, RegionTruncatedSnafu, Result, }; use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; @@ -57,8 +57,8 @@ 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::request::{ - BackgroundNotify, FlushFailed, FlushFinished, OptionOutputTx, OutputTx, SenderBulkRequest, - SenderDdlRequest, SenderWriteRequest, WorkerRequest, WorkerRequestWithTime, + BackgroundNotify, DdlRequest, FlushFailed, FlushFinished, OnFailure, OptionOutputTx, OutputTx, + SenderBulkRequest, SenderDdlRequest, SenderWriteRequest, WorkerRequest, WorkerRequestWithTime, }; use crate::schedule::scheduler::{Job, SchedulerRef}; use crate::sst::file::FileMeta; @@ -221,14 +221,14 @@ pub enum FlushReason { Downgrading, /// Enter staging mode. EnterStaging, - /// Flush when region is closing. - Closing, /// Flush triggered before region migration. RegionMigration, /// Flush triggered by repartition procedure. Repartition, /// Flush triggered by remote WAL pruning. RemoteWalPrune, + /// Flush before closing a Noop WAL region. + Closing, } impl FlushReason { @@ -280,6 +280,43 @@ pub(crate) struct RegionFlushTask { pub(crate) partition_expr: Option, } +struct FlushTaskWaiters { + region_id: RegionId, + senders: Mutex>, +} + +impl FlushTaskWaiters { + fn new(region_id: RegionId, senders: Vec) -> Self { + Self { + region_id, + senders: Mutex::new(senders), + } + } + + fn take(&self) -> Vec { + std::mem::take(&mut *self.senders.lock().unwrap()) + } + + fn on_failure(&self, err: Arc) { + for sender in self.take() { + sender.send(Err(err.clone()).context(FlushRegionSnafu { + region_id: self.region_id, + })); + } + } +} + +impl Drop for FlushTaskWaiters { + fn drop(&mut self) { + self.on_failure(Arc::new( + RegionBusySnafu { + region_id: self.region_id, + } + .build(), + )); + } +} + impl RegionFlushTask { /// Push the sender if it is not none. pub(crate) fn push_sender(&mut self, mut sender: OptionOutputTx) { @@ -307,16 +344,26 @@ impl RegionFlushTask { /// Converts the flush task into a background job. /// /// We must call this in the region worker. - fn into_flush_job(mut self, version_control: &VersionControlRef) -> Job { + fn into_flush_job( + mut self, + version_control: &VersionControlRef, + ) -> (Job, Arc) { // Get a version of this region before creating a job to get current // wal entry id, sequence and immutable memtables. let version_data = version_control.current(); + let waiters = Arc::new(FlushTaskWaiters::new( + self.region_id, + std::mem::take(&mut self.senders), + )); + let job_waiters = waiters.clone(); - Box::pin(async move { + let job = Box::pin(async move { + self.senders = job_waiters.take(); INFLIGHT_FLUSH_COUNT.inc(); self.do_flush(version_data).await; INFLIGHT_FLUSH_COUNT.dec(); - }) + }); + (job, waiters) } /// Runs the flush task. @@ -335,6 +382,7 @@ impl RegionFlushTask { .collect(); let flush_finished = FlushFinished { region_id: self.region_id, + flush_reason: self.reason, // The last entry has been flushed. flushed_entry_id: version_data.last_entry_id, senders: std::mem::take(&mut self.senders), @@ -342,7 +390,6 @@ impl RegionFlushTask { edit, memtables_to_remove, is_staging: self.is_staging, - flush_reason: self.reason, }; WorkerRequest::Background { region_id: self.region_id, @@ -719,10 +766,23 @@ impl RegionFlushTask { .send(WorkerRequestWithTime::new(request)) .await { + let request = e.0.request; error!( "Failed to notify flush job status for region {}, request: {:?}", - self.region_id, e.0 + self.region_id, request ); + if let WorkerRequest::Background { + notify: BackgroundNotify::FlushFinished(mut finished), + .. + } = request + { + finished.on_failure( + RegionClosedSnafu { + region_id: self.region_id, + } + .build(), + ); + } } } @@ -1059,22 +1119,22 @@ impl FlushScheduler { fn schedule_flush_task( &mut self, version_control: &VersionControlRef, - task: RegionFlushTask, + mut task: RegionFlushTask, ) -> Result<()> { let region_id = task.region_id; // If current region doesn't have flush status, we can flush the region directly. if let Err(e) = version_control.freeze_mutable() { error!(e; "Failed to freeze the mutable memtable for region {}", region_id); + task.on_failure(Arc::new(RegionBusySnafu { region_id }.build())); return Err(e); } // Submit a flush job. - let job = task.into_flush_job(version_control); + let (job, waiters) = task.into_flush_job(version_control); if let Err(e) = self.scheduler.schedule(job) { - // If scheduler returns error, senders in the job will be dropped and waiters - // can get recv errors. error!(e; "Failed to schedule flush job for region {}", region_id); + waiters.on_failure(Arc::new(RegionBusySnafu { region_id }.build())); return Err(e); } @@ -1111,12 +1171,13 @@ impl FlushScheduler { return Ok(()); } + let closing = task.reason == FlushReason::Closing; self.schedule_flush_task(version_control, task)?; // Add this region to status map. let _ = self.region_status.insert( region_id, - FlushStatus::new(region_id, version_control.clone()), + FlushStatus::new(region_id, version_control.clone(), closing), ); Ok(()) @@ -1181,6 +1242,9 @@ impl FlushScheduler { err; "Flush succeeded for region {region_id}, but failed to schedule next flush for it." ); + let flush_status = self.region_status.remove(®ion_id).unwrap(); + flush_status.on_failure(Arc::new(RegionBusySnafu { region_id }.build())); + return None; } // We can flush the region again, keep it in the region status. None @@ -1211,7 +1275,11 @@ impl FlushScheduler { /// Notifies the scheduler that the region is closed. pub(crate) fn on_region_closed(&mut self, region_id: RegionId) { - self.remove_region_on_failure(region_id, Arc::new(RegionClosedSnafu { region_id }.build())); + let Some(flush_status) = self.region_status.remove(®ion_id) else { + return; + }; + + flush_status.on_region_closed(Arc::new(RegionClosedSnafu { region_id }.build())); } /// Notifies the scheduler that the region is truncated. @@ -1262,11 +1330,11 @@ impl FlushScheduler { status.pending_bulk_writes.push(request); } - /// Returns true if the region has pending DDLs. + /// 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 .get(®ion_id) - .map(|status| !status.pending_ddls.is_empty()) + .map(|status| !status.pending_ddls.is_empty() || status.closing) .unwrap_or(false) } } @@ -1290,6 +1358,8 @@ struct FlushStatus { version_control: VersionControlRef, /// Task waiting for next flush. pending_task: Option, + /// Whether a close-time flush is in progress or pending. + closing: bool, /// Pending ddl requests. pending_ddls: Vec, /// Requests waiting to write after altering the region. @@ -1299,11 +1369,12 @@ struct FlushStatus { } impl FlushStatus { - fn new(region_id: RegionId, version_control: VersionControlRef) -> FlushStatus { + fn new(region_id: RegionId, version_control: VersionControlRef, closing: bool) -> FlushStatus { FlushStatus { region_id, version_control, pending_task: None, + closing, pending_ddls: Vec::new(), pending_writes: Vec::new(), pending_bulk_writes: Vec::new(), @@ -1312,6 +1383,7 @@ impl FlushStatus { /// Merges the task to pending task. fn merge_task(&mut self, task: RegionFlushTask) { + self.closing |= task.reason == FlushReason::Closing; if let Some(pending) = &mut self.pending_task { pending.merge(task); } else { @@ -1335,24 +1407,140 @@ impl FlushStatus { region_id: self.region_id, })); } + for bulk_req in self.pending_bulk_writes { + bulk_req + .sender + .send(Err(err.clone()).context(FlushRegionSnafu { + region_id: self.region_id, + })); + } + } + + fn on_region_closed(self, err: Arc) { + if let Some(mut task) = self.pending_task { + task.on_failure(err.clone()); + } + + for ddl in self.pending_ddls { + if matches!(ddl.request, DdlRequest::Close(_)) { + ddl.sender.send(Ok(0)); + } else { + ddl.sender.send(Err(err.clone()).context(FlushRegionSnafu { + region_id: self.region_id, + })); + } + } + + for write_req in self.pending_writes { + write_req + .sender + .send(Err(err.clone()).context(FlushRegionSnafu { + region_id: self.region_id, + })); + } + for bulk_req in self.pending_bulk_writes { + bulk_req + .sender + .send(Err(err.clone()).context(FlushRegionSnafu { + region_id: self.region_id, + })); + } } } #[cfg(test)] mod tests { + use common_error::ext::ErrorExt; + use common_error::status_code::StatusCode; use mito_codec::row_converter::build_primary_key_codec; use tokio::sync::oneshot; use super::*; use crate::cache::CacheManager; + use crate::error::InvalidSchedulerStateSnafu; use crate::memtable::bulk::part::BulkPartConverter; use crate::memtable::time_series::TimeSeriesMemtableBuilder; use crate::memtable::{Memtable, RangesOptions}; + use crate::schedule::scheduler::Scheduler; use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; use crate::test_util::memtable_util::{build_key_values_with_ts_seq_values, metadata_for_test}; use crate::test_util::scheduler_util::{SchedulerEnv, VecScheduler}; use crate::test_util::version_util::{VersionControlBuilder, write_rows_to_version}; + struct FailingScheduler; + + #[async_trait::async_trait] + impl Scheduler for FailingScheduler { + fn schedule(&self, _job: Job) -> Result<()> { + InvalidSchedulerStateSnafu.fail() + } + + async fn stop(&self, _await_termination: bool) -> Result<()> { + Ok(()) + } + } + + fn new_test_flush_task( + env: &SchedulerEnv, + region_id: RegionId, + reason: FlushReason, + request_sender: mpsc::Sender, + manifest_ctx: ManifestContextRef, + ) -> RegionFlushTask { + RegionFlushTask { + region_id, + reason, + senders: Vec::new(), + request_sender, + access_layer: env.access_layer.clone(), + listener: WorkerListener::default(), + engine_config: Arc::new(MitoConfig::default()), + row_group_size: None, + cache_manager: Arc::new(CacheManager::default()), + manifest_ctx, + index_options: IndexOptions::default(), + flush_semaphore: Arc::new(Semaphore::new(2)), + is_staging: false, + partition_expr: None, + } + } + + fn new_test_bulk_request( + region_id: RegionId, + ) -> ( + SenderBulkRequest, + oneshot::Receiver>, + ) { + let metadata = metadata_for_test(); + let schema = to_flat_sst_arrow_schema( + &metadata, + &FlatSchemaOptions::from_encoding(metadata.primary_key_encoding), + ); + let pk_codec = build_primary_key_codec(&metadata); + let mut converter = BulkPartConverter::new(&metadata, schema, 16, pk_codec, true); + let kvs = build_key_values_with_ts_seq_values( + &metadata, + "bulk_key".to_string(), + 1, + std::iter::once(1000i64), + std::iter::once(Some(1.0f64)), + 1, + ); + converter.append_key_values(&kvs).unwrap(); + let (sender, receiver) = oneshot::channel(); + + ( + SenderBulkRequest { + sender: OptionOutputTx::from(sender), + region_id, + request: converter.convert().unwrap(), + region_metadata: Some(metadata), + partition_expr_version: None, + }, + receiver, + ) + } + #[test] fn test_get_mutable_limit() { assert_eq!(4, WriteBufferManagerImpl::get_mutable_limit(8)); @@ -1457,6 +1645,151 @@ mod tests { assert!(scheduler.region_status.is_empty()); } + #[tokio::test] + async fn test_schedule_flush_failure_notifies_waiter() { + let env = SchedulerEnv::new() + .await + .scheduler(Arc::new(FailingScheduler)); + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_flush_scheduler(); + let mut builder = VersionControlBuilder::new(); + builder.set_memtable_builder(Arc::new(TimeSeriesMemtableBuilder::default())); + let version_control = Arc::new(builder.build()); + let version_data = version_control.current(); + write_rows_to_version(&version_data.version, "host0", 0, 10); + let manifest_ctx = env + .mock_manifest_context(version_data.version.metadata.clone()) + .await; + let (output_tx, output_rx) = oneshot::channel(); + let mut task = new_test_flush_task( + &env, + builder.region_id(), + FlushReason::Manual, + tx, + manifest_ctx, + ); + task.push_sender(OptionOutputTx::from(output_tx)); + + scheduler + .schedule_flush(builder.region_id(), &version_control, task) + .unwrap_err(); + + let err = output_rx + .await + .expect("waiter must receive explicit error") + .unwrap_err(); + assert_eq!(err.status_code(), StatusCode::RegionBusy); + } + + #[tokio::test] + async fn test_send_worker_request_failure_notifies_flush_finished_waiter() { + let env = SchedulerEnv::new().await; + let (tx, rx) = mpsc::channel(1); + drop(rx); + let builder = VersionControlBuilder::new(); + let version_control = Arc::new(builder.build()); + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .await; + let task = new_test_flush_task( + &env, + builder.region_id(), + FlushReason::Manual, + tx, + manifest_ctx, + ); + let (output_tx, output_rx) = oneshot::channel(); + let request = WorkerRequest::Background { + region_id: builder.region_id(), + notify: BackgroundNotify::FlushFinished(FlushFinished { + region_id: builder.region_id(), + flush_reason: FlushReason::Manual, + flushed_entry_id: 0, + senders: vec![OutputTx::new(output_tx)], + _timer: FLUSH_ELAPSED.with_label_values(&["total"]).start_timer(), + edit: RegionEdit { + files_to_add: Vec::new(), + files_to_remove: Vec::new(), + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }, + memtables_to_remove: smallvec![], + is_staging: false, + }), + }; + + task.send_worker_request(request).await; + + let output = output_rx.await.expect("waiter must receive explicit error"); + assert!(output.is_err()); + } + + #[tokio::test] + async fn test_flush_waiters_drop_notifies_waiter() { + let region_id = RegionId::new(1, 1); + let (output_tx, output_rx) = oneshot::channel(); + let waiters = FlushTaskWaiters::new(region_id, vec![OutputTx::new(output_tx)]); + + drop(waiters); + + let err = output_rx + .await + .expect("waiter must receive explicit error") + .unwrap_err(); + assert_eq!(err.status_code(), StatusCode::RegionBusy); + } + + #[tokio::test] + async fn test_flush_failure_notifies_pending_bulk_writes() { + let region_id = RegionId::new(1, 1); + let version_control = Arc::new(VersionControlBuilder::new().build()); + let (bulk_req, output_rx) = new_test_bulk_request(region_id); + let status = FlushStatus { + region_id, + version_control, + pending_task: None, + closing: false, + pending_ddls: Vec::new(), + pending_writes: Vec::new(), + pending_bulk_writes: vec![bulk_req], + }; + + status.on_failure(Arc::new(RegionClosedSnafu { region_id }.build())); + + let err = output_rx + .await + .expect("pending bulk write must receive explicit error") + .unwrap_err(); + assert_eq!(err.status_code(), StatusCode::Cancelled); + } + + #[tokio::test] + async fn test_region_closed_notifies_pending_bulk_writes() { + let region_id = RegionId::new(1, 1); + let version_control = Arc::new(VersionControlBuilder::new().build()); + let (bulk_req, output_rx) = new_test_bulk_request(region_id); + let status = FlushStatus { + region_id, + version_control, + pending_task: None, + closing: false, + pending_ddls: Vec::new(), + pending_writes: Vec::new(), + pending_bulk_writes: vec![bulk_req], + }; + + status.on_region_closed(Arc::new(RegionClosedSnafu { region_id }.build())); + + let err = output_rx + .await + .expect("pending bulk write must receive explicit error") + .unwrap_err(); + assert_eq!(err.status_code(), StatusCode::Cancelled); + } + #[tokio::test] async fn test_schedule_pending_request() { let job_scheduler = Arc::new(VecScheduler::default()); @@ -1730,4 +2063,77 @@ mod tests { .is_none() ); } + + #[tokio::test] + async fn test_schedule_pending_request_failure_drains_pending_ddls() { + common_telemetry::init_default_ut_logging(); + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_flush_scheduler(); + let mut builder = VersionControlBuilder::new(); + builder.set_memtable_builder(Arc::new(TimeSeriesMemtableBuilder::default())); + let version_control = Arc::new(builder.build()); + + let version_data = version_control.current(); + write_rows_to_version(&version_data.version, "host0", 0, 10); + let manifest_ctx = env + .mock_manifest_context(version_data.version.metadata.clone()) + .await; + + let task = new_test_flush_task( + &env, + builder.region_id(), + FlushReason::Manual, + tx.clone(), + manifest_ctx.clone(), + ); + scheduler + .schedule_flush(builder.region_id(), &version_control, task) + .unwrap(); + + let task = new_test_flush_task( + &env, + builder.region_id(), + FlushReason::Closing, + tx, + manifest_ctx, + ); + scheduler + .schedule_flush(builder.region_id(), &version_control, task) + .unwrap(); + + let (sender, receiver) = oneshot::channel(); + scheduler.add_ddl_request_to_pending(SenderDdlRequest { + sender: OptionOutputTx::from(sender), + region_id: builder.region_id(), + request: DdlRequest::Close(store_api::region_request::RegionCloseRequest {}), + }); + + let version_data = version_control.current(); + version_control.apply_edit( + Some(RegionEdit { + files_to_add: Vec::new(), + files_to_remove: Vec::new(), + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }), + &[0], + builder.file_purger(), + ); + write_rows_to_version(&version_data.version, "host1", 0, 10); + + scheduler.scheduler = Arc::new(FailingScheduler); + scheduler.on_flush_success(builder.region_id()); + + assert!(scheduler.region_status.is_empty()); + let err = receiver + .await + .expect("pending DDL must be notified") + .unwrap_err(); + assert_eq!(err.status_code(), StatusCode::RegionBusy); + } } diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index 7e071965f2..e9df654f83 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -921,6 +921,8 @@ pub(crate) enum BackgroundNotify { pub(crate) struct FlushFinished { /// Region id. pub(crate) region_id: RegionId, + /// Reason to flush. + pub(crate) flush_reason: FlushReason, /// Entry id of flushed data. pub(crate) flushed_entry_id: EntryId, /// Flush result senders. @@ -933,8 +935,6 @@ pub(crate) struct FlushFinished { pub(crate) memtables_to_remove: SmallVec<[MemtableId; 2]>, /// Whether the region is in staging mode. pub(crate) is_staging: bool, - /// Reason for flush. - pub(crate) flush_reason: FlushReason, } impl FlushFinished { diff --git a/src/mito2/src/worker/handle_close.rs b/src/mito2/src/worker/handle_close.rs index 4f0e388f7f..a616a0e32d 100644 --- a/src/mito2/src/worker/handle_close.rs +++ b/src/mito2/src/worker/handle_close.rs @@ -17,10 +17,10 @@ use common_telemetry::info; use store_api::logstore::LogStore; use store_api::logstore::provider::Provider; -use store_api::region_request::{RegionFlushReason, RegionFlushRequest}; +use store_api::region_request::{RegionCloseRequest, RegionFlushReason, RegionFlushRequest}; use store_api::storage::RegionId; -use crate::request::OptionOutputTx; +use crate::request::{DdlRequest, OptionOutputTx, SenderDdlRequest}; use crate::worker::RegionWorkerLoop; impl RegionWorkerLoop { @@ -48,6 +48,15 @@ impl RegionWorkerLoop { && region.is_flushable() { info!("Region {} has pending data, waiting for flush", region_id); + if self.flush_scheduler.is_flush_requested(region_id) { + self.flush_scheduler + .add_ddl_request_to_pending(SenderDdlRequest { + region_id, + sender, + request: DdlRequest::Close(RegionCloseRequest {}), + }); + return; + } self.handle_flush_request( region_id, RegionFlushRequest { diff --git a/src/mito2/src/worker/handle_flush.rs b/src/mito2/src/worker/handle_flush.rs index ef7d9da9e2..c8e86699ea 100644 --- a/src/mito2/src/worker/handle_flush.rs +++ b/src/mito2/src/worker/handle_flush.rs @@ -312,6 +312,7 @@ impl RegionWorkerLoop { } let flush_on_close = request.flush_reason == FlushReason::Closing; + let index_build_file_metas = std::mem::take(&mut request.edit.files_to_add); // In async mode, create indexes after flush. @@ -328,31 +329,54 @@ impl RegionWorkerLoop { } if flush_on_close { - // Remove region from server for flush on closing, - // no need to handle requests and schedule compactions. self.remove_region(region_id).await; info!("Region {} closed after flush", region_id); request.on_success(); - } else { - // Notifies waiters and observes the flush timer. - request.on_success(); - // Handle pending requests for the region. - if let Some((mut ddl_requests, mut write_requests, mut bulk_writes)) = - self.flush_scheduler.on_flush_success(region_id) - { - // Perform DDLs first because they require empty memtables. - self.handle_ddl_requests(&mut ddl_requests).await; - // Handle pending write requests, we don't stall these requests. + self.listener.on_flush_success(region_id); + return; + } + + // Notifies waiters and observes the flush timer. + request.on_success(); + // Handle pending requests for the region. + if let Some((mut ddl_requests, mut write_requests, mut bulk_writes)) = + self.flush_scheduler.on_flush_success(region_id) + { + // Perform DDLs first because they require empty memtables. + self.handle_ddl_requests(&mut ddl_requests).await; + if self.flush_scheduler.is_flush_requested(region_id) { + // The DDL may schedule another flush, e.g. a close-time flush after writes + // arrived in the mutable memtable during the previous flush. Keep pending + // writes fenced until that flush reaches its terminal state instead of + // accepting them while the DDL is still in progress. + for write_request in write_requests { + self.flush_scheduler + .add_write_request_to_pending(write_request); + } + for bulk_write in bulk_writes { + self.flush_scheduler.add_bulk_request_to_pending(bulk_write); + } + self.listener.on_flush_success(region_id); + return; + } + // A pending close DDL may have removed the region. Reject queued writes as + // not found, then stop instead of scheduling compaction for a closed region. + if !self.regions.is_region_exists(region_id) { self.handle_write_requests(&mut write_requests, &mut bulk_writes, false) .await; + self.listener.on_flush_success(region_id); + return; } - // Maybe flush worker again. - self.maybe_flush_worker(); - // Handle stalled requests. - self.handle_stalled_requests().await; - // Schedules compaction. - self.schedule_compaction(®ion).await; + // Handle pending write requests, we don't stall these requests. + self.handle_write_requests(&mut write_requests, &mut bulk_writes, false) + .await; } + // Maybe flush worker again. + self.maybe_flush_worker(); + // Handle stalled requests. + self.handle_stalled_requests().await; + // Schedules compaction. + self.schedule_compaction(®ion).await; self.listener.on_flush_success(region_id); }