diff --git a/src/mito2/src/engine/skip_wal_test.rs b/src/mito2/src/engine/skip_wal_test.rs index 3b6cf89f07c..9af1f34982b 100644 --- a/src/mito2/src/engine/skip_wal_test.rs +++ b/src/mito2/src/engine/skip_wal_test.rs @@ -12,6 +12,9 @@ // 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}; @@ -21,7 +24,10 @@ use store_api::region_request::{ 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,71 @@ 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 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_begin().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_begin().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_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 b90e5a94040..ab1ee935c65 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, InvalidRequestSnafu, JoinSnafu, 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, FlushFailed, FlushFinished, OnFailure, OptionOutputTx, OutputTx, + SenderBulkRequest, SenderDdlRequest, SenderWriteRequest, WorkerRequest, WorkerRequestWithTime, }; use crate::schedule::scheduler::{Job, SchedulerRef}; use crate::sst::file::FileMeta; @@ -280,6 +280,32 @@ 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 RegionFlushTask { /// Push the sender if it is not none. pub(crate) fn push_sender(&mut self, mut sender: OptionOutputTx) { @@ -307,16 +333,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. @@ -719,10 +755,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 +1108,34 @@ 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( + InvalidRequestSnafu { + region_id, + reason: format!("Failed to freeze mutable memtable before flush: {e}"), + } + .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( + InvalidRequestSnafu { + region_id, + reason: format!("Failed to schedule flush job: {e}"), + } + .build(), + )); return Err(e); } @@ -1345,14 +1406,54 @@ mod tests { 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, + } + } + #[test] fn test_get_mutable_limit() { assert_eq!(4, WriteBufferManagerImpl::get_mutable_limit(8)); @@ -1457,6 +1558,85 @@ 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::Closing, + tx, + manifest_ctx, + ); + task.push_sender(OptionOutputTx::from(output_tx)); + + scheduler + .schedule_flush(builder.region_id(), &version_control, task) + .unwrap_err(); + + let output = output_rx.await.expect("waiter must receive explicit error"); + assert!(output.is_err()); + } + + #[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::Closing, + 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(), + 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, + flush_reason: FlushReason::Closing, + }), + }; + + 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_schedule_pending_request() { let job_scheduler = Arc::new(VecScheduler::default()); diff --git a/src/mito2/src/worker/handle_close.rs b/src/mito2/src/worker/handle_close.rs index 4f0e388f7fc..a616a0e32dc 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 {