diff --git a/src/mito2/src/compaction.rs b/src/mito2/src/compaction.rs index 3d1d5f35b7..f055fa15d7 100644 --- a/src/mito2/src/compaction.rs +++ b/src/mito2/src/compaction.rs @@ -23,7 +23,8 @@ mod test_util; mod twcs; mod window; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; +use std::fmt; use std::sync::{Arc, Mutex}; use std::time::Instant; @@ -59,7 +60,7 @@ use crate::compaction::task::CompactionTaskImpl; use crate::config::MitoConfig; use crate::error::{ CompactRegionSnafu, CompactionCancelledSnafu, DataTypeMismatchSnafu, Error, - GetSchemaMetadataSnafu, ManualCompactionOverrideSnafu, ParquetToArrowSchemaSnafu, + GetSchemaMetadataSnafu, JoinSnafu, ManualCompactionOverrideSnafu, ParquetToArrowSchemaSnafu, RegionClosedSnafu, RegionDroppedSnafu, RegionTruncatedSnafu, RemoteCompactionSnafu, Result, TimeRangePredicateOverflowSnafu, TimeoutSnafu, }; @@ -74,14 +75,17 @@ use crate::read::seq_scan::SeqScan; use crate::region::options::{MergeMode, RegionOptions}; use crate::region::version::VersionControlRef; use crate::region::{ManifestContextRef, RegionLeaderState, RegionRoleState}; -use crate::request::{OptionOutputTx, OutputTx, SenderDdlRequest, WorkerRequestWithTime}; +use crate::request::{ + BackgroundNotify, OptionOutputTx, OutputTx, SenderDdlRequest, WorkerRequest, + WorkerRequestWithTime, +}; use crate::schedule::remote_job_scheduler::{ CompactionJob, DefaultNotifier, RemoteJob, RemoteJobSchedulerRef, }; use crate::schedule::scheduler::SchedulerRef; use crate::sst::file::{FileHandle, FileMeta, Level}; use crate::sst::parquet::reader::MetadataCacheMetrics; -use crate::sst::version::LevelMeta; +use crate::sst::version::{LevelMeta, SstVersion}; use crate::worker::WorkerListener; /// Region compaction request. @@ -91,8 +95,6 @@ pub struct CompactionRequest { pub(crate) access_layer: AccessLayerRef, /// Sender to send notification to the region worker. pub(crate) request_sender: mpsc::Sender, - /// Waiters of the compaction request. - pub(crate) waiters: Vec, /// Start time of compaction task. pub(crate) start_time: Instant, pub(crate) cache_manager: CacheManagerRef, @@ -108,6 +110,90 @@ impl CompactionRequest { } } +/// Result returned to the worker after background compaction planning. +pub(crate) enum CompactionPlanningResult { + Prepared(PreparedCompaction), + NoPlan, + Error(Arc), +} + +impl fmt::Debug for CompactionPlanningResult { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Prepared(prepared) => f + .debug_tuple("Prepared") + .field(&prepared.compaction_region.region_id) + .finish(), + Self::NoPlan => f.write_str("NoPlan"), + Self::Error(err) => f.debug_tuple("Error").field(err).finish(), + } + } +} + +/// Pure planning completion sent back to the owning region worker. +#[derive(Debug)] +pub(crate) struct CompactionPickFinished { + pub(crate) region_id: RegionId, + pub(crate) plan_id: u64, + pub(crate) version_control: VersionControlRef, + pub(crate) result: CompactionPlanningResult, +} + +pub(crate) struct PreparedCompaction { + compaction_region: CompactionRegion, + picker_output: PickerOutput, + start_time: Instant, + ttl: TimeToLive, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum CompactionExecutionKind { + Local, + Remote, +} + +#[derive(Debug, Clone)] +pub(crate) struct CompactionExecution { + plan_id: u64, + version_control: VersionControlRef, + kind: CompactionExecutionKind, + _files: CompactingFiles, +} + +impl CompactionExecution { + fn new( + plan_id: u64, + version_control: VersionControlRef, + kind: CompactionExecutionKind, + files: CompactingFiles, + ) -> Self { + Self { + plan_id, + version_control, + kind, + _files: files, + } + } + + pub(crate) fn matches(&self, other: &Self) -> bool { + self.plan_id == other.plan_id + && self.kind == other.kind + && Arc::ptr_eq(&self.version_control, &other.version_control) + } + + pub(crate) fn version_control(&self) -> &VersionControlRef { + &self.version_control + } + + #[cfg(test)] + pub(crate) fn for_test( + version_control: VersionControlRef, + kind: CompactionExecutionKind, + ) -> Self { + Self::new(0, version_control, kind, CompactingFiles::empty()) + } +} + /// Compaction scheduler tracks and manages compaction tasks. pub(crate) struct CompactionScheduler { scheduler: SchedulerRef, @@ -122,6 +208,8 @@ pub(crate) struct CompactionScheduler { listener: WorkerListener, /// Plugins for the compaction scheduler. plugins: Plugins, + /// Monotonically increasing token for compaction plans. + next_plan_id: u64, } impl CompactionScheduler { @@ -146,9 +234,16 @@ impl CompactionScheduler { memory_policy, listener, plugins, + next_plan_id: 0, } } + fn next_plan_id(&mut self) -> u64 { + let plan_id = self.next_plan_id; + self.next_plan_id = self.next_plan_id.wrapping_add(1); + plan_id + } + /// Schedules a compaction for the region. /// Returns whether a compaction is scheduled. #[allow(clippy::too_many_arguments)] @@ -177,8 +272,7 @@ impl CompactionScheduler { if let Some(status) = self.region_status.get_mut(®ion_id) { match compact_options { Options::Regular(_) => { - // Region is compacting. Add the waiter to pending list. - status.merge_waiter(waiter); + status.merge_regular_trigger(waiter); } options @ Options::StrictWindow(_) => { // Incoming compaction request is manually triggered. @@ -196,12 +290,12 @@ impl CompactionScheduler { return Ok(false); } - // The region can compact directly. + // Publish the picking phase before dispatching background planning. let mut status = CompactionStatus::new(region_id, version_control.clone(), access_layer.clone()); + status.merge_waiter(waiter); let request = status.new_compaction_request( self.request_sender.clone(), - waiter, self.engine_config.clone(), self.cache_manager.clone(), manifest_ctx, @@ -209,24 +303,17 @@ impl CompactionScheduler { schema_metadata_manager, max_parallelism, ); - - match self - .schedule_compaction_request(request, compact_options) - .await - { - Ok(Some(active_compaction)) => { - // Publish CompactionStatus only after a task has been accepted by the scheduler. - // This avoids exposing a half-initialized region status that could collect pending - // DDL/compaction state even though no compaction is actually running. - status.active_compaction = Some(active_compaction); - self.region_status.insert(region_id, status); - - self.listener.on_compaction_scheduled(region_id); - Ok(true) - } - Ok(None) => Ok(false), - Err(e) => Err(e), - } + let plan_id = self.next_plan_id(); + status.start_picking(plan_id); + self.region_status.insert(region_id, status); + self.dispatch_compaction_planning( + plan_id, + version_control.clone(), + request, + compact_options, + ); + self.listener.on_compaction_scheduled(region_id); + Ok(true) } // Handle pending manual compaction request for the region. @@ -257,7 +344,6 @@ impl CompactionScheduler { let request = { status.new_compaction_request( self.request_sender.clone(), - waiter, self.engine_config.clone(), self.cache_manager.clone(), manifest_ctx, @@ -266,32 +352,26 @@ impl CompactionScheduler { max_parallelism, ) }; - - match self.schedule_compaction_request(request, options).await { - Ok(Some(active_compaction)) => { - let status = self.region_status.get_mut(®ion_id).unwrap(); - status.active_compaction = Some(active_compaction); - debug!( - "Successfully scheduled manual compaction for region id: {}", - region_id - ); - true - } - Ok(None) => { - // We still need to handle the pending DDL requests. - // So we can't return early here. - false - } - Err(e) => { - error!(e; "Failed to continue pending manual compaction for region id: {}", region_id); - self.remove_region_on_failure(region_id, Arc::new(e)); - true - } - } + self.region_status + .get_mut(®ion_id) + .unwrap() + .merge_waiter(waiter); + let version_control = self.region_status[®ion_id].version_control.clone(); + let plan_id = self.next_plan_id(); + self.region_status + .get_mut(®ion_id) + .unwrap() + .start_picking(plan_id); + self.dispatch_compaction_planning(plan_id, version_control, request, options); + debug!( + "Successfully scheduled manual compaction planning for region id: {}", + region_id + ); + true } /// Notifies the scheduler that the compaction job is finished successfully. - pub(crate) async fn on_compaction_finished( + async fn on_compaction_finished( &mut self, region_id: RegionId, manifest_ctx: &ManifestContextRef, @@ -325,6 +405,12 @@ impl CompactionScheduler { waiter.send(Ok(0)); } + if status.regular_replan_pending { + self.schedule_next_compaction(region_id, manifest_ctx, schema_metadata_manager) + .await; + return Vec::new(); + } + // If there are pending DDL requests, run them. let pending_ddl_requests = std::mem::take(&mut status.pending_ddl_requests); if !pending_ddl_requests.is_empty() { @@ -336,10 +422,44 @@ impl CompactionScheduler { Vec::new() } + pub(crate) fn is_current_execution( + &self, + region_id: RegionId, + execution: &CompactionExecution, + ) -> bool { + self.region_status + .get(®ion_id) + .is_some_and(|status| status.matches_execution(execution)) + } + + pub(crate) fn is_current_region_execution( + &self, + region_id: RegionId, + current_version_control: &VersionControlRef, + execution: &CompactionExecution, + ) -> bool { + Arc::ptr_eq(current_version_control, execution.version_control()) + && self.is_current_execution(region_id, execution) + } + + pub(crate) async fn on_execution_finished( + &mut self, + region_id: RegionId, + execution: &CompactionExecution, + manifest_ctx: &ManifestContextRef, + schema_metadata_manager: SchemaMetadataManagerRef, + ) -> Vec { + if !self.is_current_execution(region_id, execution) { + return Vec::new(); + } + self.on_compaction_finished(region_id, manifest_ctx, schema_metadata_manager) + .await + } + pub(crate) fn is_compacting(&self, region_id: RegionId) -> bool { self.region_status .get(®ion_id) - .map(|status| status.active_compaction.is_some()) + .map(CompactionStatus::is_busy) .unwrap_or(false) } @@ -358,7 +478,6 @@ impl CompactionScheduler { // We should always try to compact the region until picker returns None. let request = status.new_compaction_request( self.request_sender.clone(), - OptionOutputTx::none(), self.engine_config.clone(), self.cache_manager.clone(), manifest_ctx, @@ -366,55 +485,59 @@ impl CompactionScheduler { schema_metadata_manager, MAX_PARALLEL_COMPACTION, ); - - // Try to schedule next compaction task for this region. - match self - .schedule_compaction_request( - request, - compact_request::Options::Regular(Default::default()), - ) - .await - { - Ok(Some(active_compaction)) => { - self.region_status - .get_mut(®ion_id) - .unwrap() - .active_compaction = Some(active_compaction); - debug!( - "Successfully scheduled next compaction for region id: {}", - region_id - ); - true - } - Ok(None) => { - // No further compaction tasks can be scheduled; cleanup the `CompactionStatus` for this region. - // All DDL requests and pending compaction requests have already been processed. - // Safe to remove the region from status tracking. - self.region_status.remove(®ion_id); - false - } - Err(e) => { - error!(e; "Failed to schedule next compaction for region {}", region_id); - self.remove_region_on_failure(region_id, Arc::new(e)); - false - } - } + let version_control = self.region_status[®ion_id].version_control.clone(); + let plan_id = self.next_plan_id(); + self.region_status + .get_mut(®ion_id) + .unwrap() + .start_regular_picking(plan_id); + self.dispatch_compaction_planning( + plan_id, + version_control, + request, + compact_request::Options::Regular(Default::default()), + ); + debug!( + "Successfully scheduled next compaction planning for region id: {}", + region_id + ); + true } /// Notifies the scheduler that the compaction job is cancelled cooperatively. - pub(crate) async fn on_compaction_cancelled( - &mut self, - region_id: RegionId, - ) -> Vec { + async fn on_compaction_cancelled(&mut self, region_id: RegionId) -> Vec { self.remove_region_on_cancel(region_id) } + pub(crate) async fn on_execution_cancelled( + &mut self, + region_id: RegionId, + execution: &CompactionExecution, + ) -> Vec { + if !self.is_current_execution(region_id, execution) { + return Vec::new(); + } + self.on_compaction_cancelled(region_id).await + } + /// Notifies the scheduler that the compaction job is failed. - pub(crate) fn on_compaction_failed(&mut self, region_id: RegionId, err: Arc) { + fn on_compaction_failed(&mut self, region_id: RegionId, err: Arc) { error!(err; "Region {} failed to compact, cancel all pending tasks", region_id); self.remove_region_on_failure(region_id, err); } + pub(crate) fn on_execution_failed( + &mut self, + region_id: RegionId, + execution: &CompactionExecution, + err: Arc, + ) { + if !self.is_current_execution(region_id, execution) { + return; + } + self.on_compaction_failed(region_id, err); + } + /// Notifies the scheduler that the region is dropped. pub(crate) fn on_region_dropped(&mut self, region_id: RegionId) { self.remove_region_on_failure( @@ -446,7 +569,7 @@ impl CompactionScheduler { request.region_id, request.request ); let status = self.region_status.get_mut(&request.region_id).unwrap(); - status.pending_ddl_requests.push(request); + status.queue_ddl(request); } #[cfg(test)] @@ -471,15 +594,45 @@ impl CompactionScheduler { status.request_cancel() } - /// Schedules a compaction request. - /// - /// Returns the active compaction state if the request is scheduled successfully. - /// Returns `None` if no compaction task can be scheduled for this region. - async fn schedule_compaction_request( - &mut self, + fn dispatch_compaction_planning( + &self, + plan_id: u64, + version_control: VersionControlRef, request: CompactionRequest, options: compact_request::Options, - ) -> Result> { + ) { + let plugins = self.plugins.clone(); + let max_background_compactions = self.engine_config.max_background_compactions; + common_runtime::spawn_compact(async move { + let region_id = request.region_id(); + let request_sender = request.request_sender.clone(); + let result = + Self::prepare_compaction(request, options, plugins, max_background_compactions) + .await; + if let CompactionPlanningResult::Error(err) = &result { + error!(err; "Compaction planning failed for region {}, plan_id: {}", region_id, plan_id); + } + let request = WorkerRequestWithTime::new(WorkerRequest::Background { + region_id, + notify: BackgroundNotify::CompactionPickFinished(CompactionPickFinished { + region_id, + plan_id, + version_control, + result, + }), + }); + if request_sender.send(request).await.is_err() { + warn!("Failed to send compaction planning result for region {region_id}"); + } + }); + } + + async fn prepare_compaction( + request: CompactionRequest, + options: compact_request::Options, + plugins: Plugins, + max_background_compactions: usize, + ) -> CompactionPlanningResult { let region_id = request.region_id(); let (dynamic_compaction_opts, ttl) = find_dynamic_options( region_id, @@ -499,15 +652,14 @@ impl CompactionScheduler { &options, &dynamic_compaction_opts, request.current_version.options.append_mode, - Some(self.engine_config.max_background_compactions), + Some(max_background_compactions), ); let region_id = request.region_id(); let CompactionRequest { engine_config, current_version, access_layer, - request_sender, - waiters, + request_sender: _, start_time, cache_manager, manifest_ctx, @@ -536,30 +688,226 @@ impl CompactionScheduler { file_purger: None, ttl: Some(ttl), max_parallelism, - plugins: self.plugins.clone(), + plugins, }; - let picker_output = { + listener.on_compaction_pick_begin(region_id).await; + let picker_region = compaction_region.clone(); + let picker_output = match common_runtime::spawn_blocking_compact(move || { let _pick_timer = COMPACTION_STAGE_ELAPSED .with_label_values(&["pick"]) .start_timer(); - picker.pick(&compaction_region) + picker.pick(&picker_region) + }) + .await + .context(JoinSnafu) + { + Ok(output) => output, + Err(err) => return CompactionPlanningResult::Error(Arc::new(err)), }; - let picker_output = if let Some(picker_output) = picker_output { - picker_output - } else { - // Nothing to compact, we are done. Notifies all waiters as we consume the compaction request. - for waiter in waiters { + let Some(picker_output) = picker_output else { + return CompactionPlanningResult::NoPlan; + }; + + CompactionPlanningResult::Prepared(PreparedCompaction { + compaction_region, + picker_output, + start_time, + ttl, + }) + } + + pub(crate) async fn accept_compaction_pick_finished( + &mut self, + finished: CompactionPickFinished, + current_version_control: &VersionControlRef, + manifest_ctx: &ManifestContextRef, + schema_metadata_manager: SchemaMetadataManagerRef, + ) -> Vec { + let region_id = finished.region_id; + let plan_id = finished.plan_id; + let version_control = finished.version_control.clone(); + if !Arc::ptr_eq(current_version_control, &finished.version_control) { + return Vec::new(); + } + let Some(status) = self.region_status.get(®ion_id) else { + return Vec::new(); + }; + if !status.is_picking(finished.plan_id) + || !Arc::ptr_eq(&status.version_control, &finished.version_control) + { + return Vec::new(); + } + if !status.accept_plan(finished.plan_id) { + return self.remove_region_on_cancel(region_id); + } + + match finished.result { + CompactionPlanningResult::Prepared(mut prepared) => { + let current = self.region_status[®ion_id] + .version_control + .current() + .version; + let Some(picker_output) = + refresh_picker_output(prepared.picker_output, ¤t.ssts) + else { + return self + .finish_compaction_planning( + region_id, + None, + manifest_ctx, + schema_metadata_manager, + ) + .await; + }; + let Some(files) = CompactingFiles::try_new(&picker_output) else { + return self + .finish_compaction_planning( + region_id, + None, + manifest_ctx, + schema_metadata_manager, + ) + .await; + }; + prepared.picker_output = picker_output; + let Some(status) = self.region_status.get_mut(®ion_id) else { + return Vec::new(); + }; + let waiters = std::mem::take(&mut status.waiters); + match self + .submit_prepared_compaction(prepared, files, waiters, plan_id, version_control) + .await + { + Ok(Some(phase)) => { + if let Some(status) = self.region_status.get_mut(®ion_id) { + status.phase = Some(phase); + } + Vec::new() + } + Ok(None) => { + self.finish_compaction_planning( + region_id, + None, + manifest_ctx, + schema_metadata_manager, + ) + .await + } + Err(err) => { + self.remove_region_on_failure(region_id, Arc::new(err)); + Vec::new() + } + } + } + CompactionPlanningResult::NoPlan => { + self.finish_compaction_planning( + region_id, + None, + manifest_ctx, + schema_metadata_manager, + ) + .await + } + CompactionPlanningResult::Error(err) => { + self.finish_compaction_planning( + region_id, + Some(err), + manifest_ctx, + schema_metadata_manager, + ) + .await + } + } + } + + #[cfg(test)] + async fn handle_compaction_pick_finished( + &mut self, + finished: CompactionPickFinished, + manifest_ctx: &ManifestContextRef, + schema_metadata_manager: SchemaMetadataManagerRef, + ) -> Vec { + let current_version_control = finished.version_control.clone(); + self.accept_compaction_pick_finished( + finished, + ¤t_version_control, + manifest_ctx, + schema_metadata_manager, + ) + .await + } + + async fn finish_compaction_planning( + &mut self, + region_id: RegionId, + err: Option>, + manifest_ctx: &ManifestContextRef, + schema_metadata_manager: SchemaMetadataManagerRef, + ) -> Vec { + let Some(status) = self.region_status.get_mut(®ion_id) else { + return Vec::new(); + }; + status.clear_running_task(); + for waiter in std::mem::take(&mut status.waiters) { + if let Some(err) = &err { + waiter.send(Err(err.clone()).context(CompactRegionSnafu { region_id })); + } else { waiter.send(Ok(0)); } - return Ok(None); - }; + } + + if self + .handle_pending_compaction_request( + region_id, + manifest_ctx, + schema_metadata_manager.clone(), + ) + .await + { + return Vec::new(); + } + + if self.region_status[®ion_id].regular_replan_pending { + self.schedule_next_compaction(region_id, manifest_ctx, schema_metadata_manager) + .await; + return Vec::new(); + } + + self.region_status + .remove(®ion_id) + .map(|mut status| std::mem::take(&mut status.pending_ddl_requests)) + .unwrap_or_default() + } + + async fn submit_prepared_compaction( + &mut self, + prepared: PreparedCompaction, + files: CompactingFiles, + waiters: Vec, + plan_id: u64, + version_control: VersionControlRef, + ) -> Result> { + let PreparedCompaction { + compaction_region, + picker_output, + start_time, + ttl, + } = prepared; + let region_id = compaction_region.region_id; + let dynamic_compaction_opts = &compaction_region.region_options.compaction; // If specified to run compaction remotely, we schedule the compaction job remotely. // It will fall back to local compaction if there is no remote job scheduler. let waiters = if dynamic_compaction_opts.remote_compaction() { if let Some(remote_job_scheduler) = &self.plugins.get::() { + let execution = CompactionExecution::new( + plan_id, + version_control.clone(), + CompactionExecutionKind::Remote, + files.clone(), + ); let remote_compaction_job = CompactionJob { compaction_region: compaction_region.clone(), picker_output: picker_output.clone(), @@ -571,9 +919,10 @@ impl CompactionScheduler { let result = remote_job_scheduler .schedule( RemoteJob::CompactionJob(remote_compaction_job), - Box::new(DefaultNotifier { - request_sender: request_sender.clone(), - }), + Box::new(DefaultNotifier::new( + self.request_sender.clone(), + execution.clone(), + )), ) .await; @@ -584,11 +933,14 @@ impl CompactionScheduler { job_id, region_id ); INFLIGHT_COMPACTION_COUNT.inc(); - return Ok(Some(ActiveCompaction::Remote)); + return Ok(Some(CompactionPhase::Remote { execution })); } Err(e) => { if !dynamic_compaction_opts.fallback_to_local() { error!(e; "Failed to schedule remote compaction job for region {}", region_id); + if let Some(status) = self.region_status.get_mut(®ion_id) { + status.waiters.extend(e.waiters); + } return RemoteCompactionSnafu { region_id, job_id: None, @@ -599,7 +951,6 @@ impl CompactionScheduler { error!(e; "Failed to schedule remote compaction job for region {}, fallback to local compaction", region_id); - // Return the waiters back to the caller for local compaction. e.waiters } } @@ -632,12 +983,19 @@ impl CompactionScheduler { let cancel_handle = Arc::new(CancellationHandle::default()); let state = LocalCompactionState::new(cancel_handle.clone()); + let execution = CompactionExecution::new( + plan_id, + version_control, + CompactionExecutionKind::Local, + files, + ); let local_compaction_task = Box::new(CompactionTaskImpl { state: state.clone(), - request_sender, + execution: execution.clone(), + request_sender: self.request_sender.clone(), waiters, start_time, - listener, + listener: self.listener.clone(), picker_output, compaction_region, compactor: Arc::new(DefaultCompactor::with_cancel_handle(cancel_handle.clone())), @@ -646,24 +1004,49 @@ impl CompactionScheduler { estimated_memory_bytes: estimated_bytes, }); - self.submit_compaction_task(local_compaction_task, region_id) - .map(|_| Some(ActiveCompaction::Local { state })) + match self.submit_compaction_task(local_compaction_task, region_id) { + Ok(()) => Ok(Some(CompactionPhase::Local { state, execution })), + Err((err, task)) => { + if let (Some(status), Some(mut task)) = + (self.region_status.get_mut(®ion_id), task) + { + status.waiters.append(&mut task.waiters); + } + Err(err) + } + } } fn submit_compaction_task( &mut self, - mut task: Box, + task: Box, region_id: RegionId, - ) -> Result<()> { - self.scheduler - .schedule(Box::pin(async move { + ) -> std::result::Result<(), (Error, Option>)> { + let task = Arc::new(Mutex::new(Some(task))); + let task_to_run = task.clone(); + match self.scheduler.schedule(Box::pin(async move { + let task = task_to_run + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .take(); + if let Some(mut task) = task { INFLIGHT_COMPACTION_COUNT.inc(); task.run().await; INFLIGHT_COMPACTION_COUNT.dec(); - })) - .inspect_err( - |e| error!(e; "Failed to submit compaction request for region {}", region_id), - ) + } else { + error!("Compaction task was missing when the scheduled job started"); + } + })) { + Ok(()) => Ok(()), + Err(err) => { + error!(err; "Failed to submit compaction request for region {}", region_id); + let task = task + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .take(); + Err((err, task)) + } + } } fn exceeds_compaction_memory_limit(&self, estimated_bytes: u64) -> Option { @@ -701,9 +1084,82 @@ pub(crate) struct LocalCompactionState { } #[derive(Debug)] -enum ActiveCompaction { - Local { state: LocalCompactionState }, - Remote, +enum CompactionPhase { + Picking { + plan_id: u64, + cancelled: bool, + }, + Local { + state: LocalCompactionState, + execution: CompactionExecution, + }, + Remote { + execution: CompactionExecution, + }, +} + +impl CompactionPhase { + fn execution(&self) -> Option<&CompactionExecution> { + match self { + Self::Picking { .. } => None, + Self::Local { execution, .. } | Self::Remote { execution } => Some(execution), + } + } +} + +/// Owns atomic reservations for every SST selected by a compaction plan. +#[derive(Debug, Clone)] +struct CompactingFiles { + _inner: Arc, +} + +#[derive(Debug)] +struct CompactingFilesInner { + files: Vec, +} + +impl CompactingFiles { + fn try_new(output: &PickerOutput) -> Option { + let mut seen = HashSet::new(); + let mut files: Vec = Vec::new(); + let selected_files = output + .outputs + .iter() + .flat_map(|output| output.inputs.iter()) + .chain(output.expired_ssts.iter()); + + for file in selected_files { + if !seen.insert(file.file_id()) { + continue; + } + if !file.try_set_compacting() { + for reserved in &files { + reserved.set_compacting(false); + } + return None; + } + files.push(file.clone()); + } + + Some(Self { + _inner: Arc::new(CompactingFilesInner { files }), + }) + } + + #[cfg(test)] + fn empty() -> Self { + Self { + _inner: Arc::new(CompactingFilesInner { files: Vec::new() }), + } + } +} + +impl Drop for CompactingFilesInner { + fn drop(&mut self) { + for file in &self.files { + file.set_compacting(false); + } + } } impl LocalCompactionState { @@ -862,8 +1318,12 @@ struct CompactionStatus { pending_request: Option, /// Pending DDL requests that should run when compaction is done. pending_ddl_requests: Vec, - /// Active compaction state. - active_compaction: Option, + /// Current compaction phase. + phase: Option, + /// Whether a regular trigger arrived while the current plan was being picked. + regular_replan_pending: bool, + /// Waiters owned by the pending regular follow-up. + pending_regular_waiters: Vec, } impl CompactionStatus { @@ -880,37 +1340,124 @@ impl CompactionStatus { waiters: Vec::new(), pending_request: None, pending_ddl_requests: Vec::new(), - active_compaction: None, + phase: None, + regular_replan_pending: false, + pending_regular_waiters: Vec::new(), } } + fn start_picking(&mut self, plan_id: u64) { + self.phase = Some(CompactionPhase::Picking { + plan_id, + cancelled: false, + }); + } + + fn start_regular_picking(&mut self, plan_id: u64) { + self.regular_replan_pending = false; + self.waiters.append(&mut self.pending_regular_waiters); + self.start_picking(plan_id); + } + + fn is_picking(&self, expected_plan_id: u64) -> bool { + matches!( + self.phase, + Some(CompactionPhase::Picking { plan_id, .. }) if plan_id == expected_plan_id + ) + } + + fn accept_plan(&self, expected_plan_id: u64) -> bool { + matches!( + self.phase, + Some(CompactionPhase::Picking { + plan_id, + cancelled: false, + }) if plan_id == expected_plan_id + ) + } + + fn is_busy(&self) -> bool { + self.phase.is_some() + } + + fn matches_execution(&self, execution: &CompactionExecution) -> bool { + Arc::ptr_eq(&self.version_control, execution.version_control()) + && self + .phase + .as_ref() + .and_then(CompactionPhase::execution) + .is_some_and(|current| current.matches(execution)) + } + #[cfg(test)] fn start_local_task(&mut self) -> LocalCompactionState { let state = LocalCompactionState::new(Arc::new(CancellationHandle::default())); - self.active_compaction = Some(ActiveCompaction::Local { + let execution = CompactionExecution::new( + 0, + self.version_control.clone(), + CompactionExecutionKind::Local, + CompactingFiles::empty(), + ); + self.phase = Some(CompactionPhase::Local { state: state.clone(), + execution, }); state } #[cfg(test)] fn start_remote_task(&mut self) { - self.active_compaction = Some(ActiveCompaction::Remote); + let execution = CompactionExecution::new( + 0, + self.version_control.clone(), + CompactionExecutionKind::Remote, + CompactingFiles::empty(), + ); + self.phase = Some(CompactionPhase::Remote { execution }); } fn request_cancel(&mut self) -> RequestCancelResult { - let Some(active_compaction) = &self.active_compaction else { + let Some(phase) = &mut self.phase else { return RequestCancelResult::NotRunning; }; - match active_compaction { - ActiveCompaction::Local { state, .. } => state.request_cancel(), - ActiveCompaction::Remote => RequestCancelResult::TooLateToCancel, + match phase { + CompactionPhase::Picking { cancelled, .. } => { + if *cancelled { + RequestCancelResult::AlreadyCancelling + } else { + *cancelled = true; + RequestCancelResult::CancelIssued + } + } + CompactionPhase::Local { state, .. } => state.request_cancel(), + CompactionPhase::Remote { .. } => RequestCancelResult::TooLateToCancel, } } fn clear_running_task(&mut self) -> bool { - self.active_compaction.take().is_some() + self.phase.take().is_some() + } + + fn merge_regular_trigger(&mut self, mut waiter: OptionOutputTx) { + if self.can_retain_regular_followup() { + self.regular_replan_pending = true; + if let Some(waiter) = waiter.take_inner() { + self.pending_regular_waiters.push(waiter); + } + } else { + self.merge_waiter(waiter); + } + } + + fn can_retain_regular_followup(&self) -> bool { + matches!(self.phase, Some(CompactionPhase::Picking { .. })) + && self.pending_ddl_requests.is_empty() + } + + fn queue_ddl(&mut self, request: SenderDdlRequest) { + // The first queued DDL fences later regular triggers from creating more follow-ups. + self.pending_ddl_requests.push(request); } /// Merge the waiter to the pending compaction. @@ -932,7 +1479,11 @@ impl CompactionStatus { } fn on_failure(mut self, err: Arc) { - for waiter in self.waiters.drain(..) { + for waiter in self + .waiters + .drain(..) + .chain(self.pending_regular_waiters.drain(..)) + { waiter.send(Err(err.clone()).context(CompactRegionSnafu { region_id: self.region_id, })); @@ -957,7 +1508,11 @@ impl CompactionStatus { #[must_use] fn on_cancel(mut self) -> Vec { - for waiter in self.waiters.drain(..) { + for waiter in self + .waiters + .drain(..) + .chain(self.pending_regular_waiters.drain(..)) + { waiter.send(CompactionCancelledSnafu.fail()); } @@ -972,14 +1527,11 @@ impl CompactionStatus { std::mem::take(&mut self.pending_ddl_requests) } - /// Creates a new compaction request for compaction picker. - /// - /// It consumes all pending compaction waiters. + /// Creates an immutable request for background compaction planning. #[allow(clippy::too_many_arguments)] fn new_compaction_request( - &mut self, + &self, request_sender: Sender, - mut waiter: OptionOutputTx, engine_config: Arc, cache_manager: CacheManagerRef, manifest_ctx: &ManifestContextRef, @@ -989,19 +1541,12 @@ impl CompactionStatus { ) -> CompactionRequest { let current_version = CompactionVersion::from(self.version_control.current().version); let start_time = Instant::now(); - let mut waiters = Vec::with_capacity(self.waiters.len() + 1); - waiters.extend(std::mem::take(&mut self.waiters)); - - if let Some(waiter) = waiter.take_inner() { - waiters.push(waiter); - } CompactionRequest { engine_config, current_version, access_layer: self.access_layer.clone(), request_sender: request_sender.clone(), - waiters, start_time, cache_manager, manifest_ctx: manifest_ctx.clone(), @@ -1265,6 +1810,40 @@ fn estimate_compaction_bytes(picker_output: &PickerOutput) -> u64 { .sum() } +/// Rebuilds picker output with current SST handles while preserving the picker's grouping. +fn refresh_picker_output(output: PickerOutput, current: &SstVersion) -> Option { + let refresh = |file: FileHandle| { + current + .file_for_compaction(&file) + .filter(|current| !current.is_deleted() && !current.compacting()) + .cloned() + }; + let outputs = output + .outputs + .into_iter() + .map(|output| { + let inputs = output + .inputs + .into_iter() + .map(&refresh) + .collect::>>()?; + Some(CompactionOutput { inputs, ..output }) + }) + .collect::>>()?; + let expired_ssts = output + .expired_ssts + .into_iter() + .map(refresh) + .collect::>>()?; + + Some(PickerOutput { + outputs, + expired_ssts, + time_window_size: output.time_window_size, + max_file_size: output.max_file_size, + }) +} + /// Pending compaction request that is supposed to run after current task is finished, /// typically used for manual compactions. struct PendingCompaction { @@ -1285,10 +1864,13 @@ mod tests { use common_datasource::compression::CompressionType; use common_meta::key::schema_name::SchemaNameValue; use common_time::DatabaseTimeToLive; + use store_api::storage::FileId; use tokio::sync::{Barrier, oneshot}; use super::*; use crate::compaction::memory_manager::{CompactionMemoryGuard, new_compaction_memory_manager}; + use crate::compaction::test_util::new_file_handle; + use crate::engine::listener::CompactionPlanningGate; use crate::error::InvalidSchedulerStateSnafu; use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions}; use crate::region::ManifestContext; @@ -1300,6 +1882,168 @@ mod tests { struct FailingScheduler; + struct SuccessfulRemoteScheduler; + + struct FailingRemoteScheduler; + + #[derive(Default)] + struct HoldingRemoteScheduler { + notifier: Mutex>>, + } + + impl HoldingRemoteScheduler { + fn drop_notifier(&self) { + self.notifier.lock().unwrap().take(); + } + } + + #[async_trait::async_trait] + impl crate::schedule::remote_job_scheduler::RemoteJobScheduler for SuccessfulRemoteScheduler { + async fn schedule( + &self, + _job: RemoteJob, + _notifier: Box, + ) -> std::result::Result< + crate::schedule::remote_job_scheduler::JobId, + crate::schedule::remote_job_scheduler::RemoteJobSchedulerError, + > { + Ok(crate::schedule::remote_job_scheduler::JobId::parse_str( + "00000000-0000-0000-0000-000000000001", + ) + .unwrap()) + } + } + + #[async_trait::async_trait] + impl crate::schedule::remote_job_scheduler::RemoteJobScheduler for FailingRemoteScheduler { + async fn schedule( + &self, + job: RemoteJob, + _notifier: Box, + ) -> std::result::Result< + crate::schedule::remote_job_scheduler::JobId, + crate::schedule::remote_job_scheduler::RemoteJobSchedulerError, + > { + let RemoteJob::CompactionJob(job) = job; + Err( + crate::schedule::remote_job_scheduler::RemoteJobSchedulerError { + location: snafu::location!(), + reason: "remote scheduler rejected job".to_string(), + waiters: job.waiters, + }, + ) + } + } + + #[async_trait::async_trait] + impl crate::schedule::remote_job_scheduler::RemoteJobScheduler for HoldingRemoteScheduler { + async fn schedule( + &self, + _job: RemoteJob, + notifier: Box, + ) -> std::result::Result< + crate::schedule::remote_job_scheduler::JobId, + crate::schedule::remote_job_scheduler::RemoteJobSchedulerError, + > { + self.notifier.lock().unwrap().replace(notifier); + Ok(crate::schedule::remote_job_scheduler::JobId::parse_str( + "00000000-0000-0000-0000-000000000002", + ) + .unwrap()) + } + } + + fn compactable_version() -> VersionControlRef { + let mut builder = VersionControlBuilder::new(); + let end = 1000 * 1000; + Arc::new( + builder + .push_l0_file(0, end) + .push_l0_file(10, end) + .push_l0_file(50, end) + .push_l0_file(80, end) + .push_l0_file(90, end) + .build(), + ) + } + + async fn begin_pick_result( + env: &SchedulerEnv, + scheduler: &mut CompactionScheduler, + rx: &mut mpsc::Receiver, + version_control: &VersionControlRef, + ) -> ( + CompactionPickFinished, + ManifestContextRef, + SchemaMetadataManagerRef, + ) { + let region_id = version_control.current().version.metadata.region_id; + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .await; + let (schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager(); + assert!( + scheduler + .schedule_compaction( + region_id, + Options::Regular(Default::default()), + version_control, + &env.access_layer, + OptionOutputTx::none(), + &manifest_ctx, + schema_metadata_manager.clone(), + 1, + ) + .await + .unwrap() + ); + let finished = recv_compaction_pick_finished(rx).await; + assert!(matches!( + &finished.result, + CompactionPlanningResult::Prepared(_) + )); + (finished, manifest_ctx, schema_metadata_manager) + } + + fn selected_files(finished: &CompactionPickFinished) -> Vec { + let CompactionPlanningResult::Prepared(prepared) = &finished.result else { + panic!("expected prepared compaction"); + }; + prepared + .picker_output + .outputs + .iter() + .flat_map(|output| output.inputs.iter().cloned()) + .chain(prepared.picker_output.expired_ssts.iter().cloned()) + .collect() + } + + fn use_remote_compaction(finished: &mut CompactionPickFinished) { + let CompactionPlanningResult::Prepared(prepared) = &mut finished.result else { + panic!("expected prepared compaction"); + }; + let crate::region::options::CompactionOptions::Twcs(options) = + &mut prepared.compaction_region.region_options.compaction; + options.remote_compaction = true; + options.fallback_to_local = false; + } + + fn picker_output_with_files( + output_files: Vec, + expired_ssts: Vec, + ) -> PickerOutput { + PickerOutput { + outputs: vec![CompactionOutput { + output_level: 1, + inputs: output_files, + filter_deleted: false, + output_time_range: None, + }], + expired_ssts, + ..Default::default() + } + } + #[async_trait::async_trait] impl Scheduler for FailingScheduler { fn schedule(&self, _job: Job) -> Result<()> { @@ -1311,9 +2055,233 @@ mod tests { } } + async fn recv_compaction_pick_finished( + rx: &mut mpsc::Receiver, + ) -> CompactionPickFinished { + let request = rx.recv().await.expect("worker request channel closed"); + match request.request { + WorkerRequest::Background { + notify: BackgroundNotify::CompactionPickFinished(finished), + .. + } => finished, + other => panic!("unexpected worker request: {other:?}"), + } + } + + fn assert_compaction_lifecycle_error(err: Error, lifecycle: &str) { + let Error::CompactRegion { source, .. } = err else { + panic!("expected compact-region error, got {err:?}"); + }; + let matches_lifecycle = matches!( + (lifecycle, source.as_ref()), + ("close", Error::RegionClosed { .. }) + | ("drop", Error::RegionDropped { .. }) + | ("truncate", Error::RegionTruncated { .. }) + ); + assert!( + matches_lifecycle, + "unexpected {lifecycle} error source: {source:?}" + ); + } + + #[tokio::test] + async fn test_picking_phase_tracks_plan_and_cancellation() { + let env = SchedulerEnv::new().await; + let builder = VersionControlBuilder::new(); + let mut status = CompactionStatus::new( + builder.region_id(), + Arc::new(builder.build()), + env.access_layer.clone(), + ); + + status.start_picking(7); + + assert!(status.is_picking(7)); + assert!(!status.is_picking(8)); + assert!(status.is_busy()); + assert!(status.accept_plan(7)); + assert!(!status.accept_plan(8)); + assert_eq!(status.request_cancel(), RequestCancelResult::CancelIssued); + assert_eq!( + status.request_cancel(), + RequestCancelResult::AlreadyCancelling + ); + assert!(status.is_picking(7)); + assert!(status.is_busy()); + assert!(!status.accept_plan(7)); + } + + #[tokio::test] + async fn test_picking_coalesces_regular_waiter_and_queues_manual_request() { + let env = SchedulerEnv::new().await; + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .await; + let (schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager(); + let mut status = + CompactionStatus::new(region_id, version_control.clone(), env.access_layer.clone()); + status.start_picking(1); + scheduler.region_status.insert(region_id, status); + + let (regular_tx, _regular_rx) = oneshot::channel(); + assert!( + !scheduler + .schedule_compaction( + region_id, + compact_request::Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::from(regular_tx), + &manifest_ctx, + schema_metadata_manager.clone(), + 1, + ) + .await + .unwrap() + ); + let (manual_tx, _manual_rx) = oneshot::channel(); + assert!( + !scheduler + .schedule_compaction( + region_id, + compact_request::Options::StrictWindow(StrictWindow { window_seconds: 60 }), + &version_control, + &env.access_layer, + OptionOutputTx::from(manual_tx), + &manifest_ctx, + schema_metadata_manager, + 1, + ) + .await + .unwrap() + ); + + let status = scheduler.region_status.get(®ion_id).unwrap(); + assert!(status.is_busy()); + assert!(status.waiters.is_empty()); + assert!(status.regular_replan_pending); + assert_eq!(status.pending_regular_waiters.len(), 1); + assert!(status.pending_request.is_some()); + } + + #[tokio::test] + async fn test_picking_plan_ids_are_monotonic() { + let env = SchedulerEnv::new().await; + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + + let first = scheduler.next_plan_id(); + let second = scheduler.next_plan_id(); + + assert_eq!(second, first + 1); + } + + #[test] + fn test_picking_compacting_files_deduplicates_handles() { + let file = new_file_handle(FileId::random(), 0, 10, 0); + let output = picker_output_with_files(vec![file.clone(), file.clone()], vec![file.clone()]); + + let files = CompactingFiles::try_new(&output).unwrap(); + + assert!(file.compacting()); + drop(files); + assert!(!file.compacting()); + } + + #[test] + fn test_picking_compacting_files_rolls_back_on_conflict() { + let first = new_file_handle(FileId::random(), 0, 10, 0); + let conflicting = new_file_handle(FileId::random(), 0, 10, 0); + conflicting.set_compacting(true); + let output = picker_output_with_files(vec![first.clone(), conflicting.clone()], vec![]); + + assert!(CompactingFiles::try_new(&output).is_none()); + assert!(!first.compacting()); + assert!(conflicting.compacting()); + } + + #[test] + fn test_picking_compacting_files_drop_clears_all_reservations() { + let output_file = new_file_handle(FileId::random(), 0, 10, 0); + let expired_file = new_file_handle(FileId::random(), 0, 10, 0); + let output = + picker_output_with_files(vec![output_file.clone()], vec![expired_file.clone()]); + + let files = CompactingFiles::try_new(&output).unwrap(); + assert!(output_file.compacting()); + assert!(expired_file.compacting()); + + drop(files); + assert!(!output_file.compacting()); + assert!(!expired_file.compacting()); + } + + #[test] + fn test_pick_result_refreshes_handles_and_preserves_output_options() { + let purger = crate::test_util::new_noop_file_purger(); + let stale = new_file_handle(FileId::random(), 0, 10, 0); + let expired = new_file_handle(FileId::random(), 20, 30, 0); + let mut current_meta = stale.meta_ref().clone(); + current_meta.index_version = 1; + current_meta.index_file_size = 128; + let mut current = crate::sst::version::SstVersion::new(); + current.add_files( + purger.clone(), + [current_meta.clone(), expired.meta_ref().clone()].into_iter(), + ); + let output_time_range = TimestampRange::new( + common_time::Timestamp::new_millisecond(0), + common_time::Timestamp::new_millisecond(10), + ); + let output = PickerOutput { + outputs: vec![CompactionOutput { + output_level: 2, + inputs: vec![stale.clone()], + filter_deleted: true, + output_time_range, + }], + expired_ssts: vec![expired], + time_window_size: 3600, + max_file_size: Some(4096), + }; + + let refreshed = refresh_picker_output(output, ¤t).unwrap(); + + assert_eq!(refreshed.outputs.len(), 1); + assert_eq!(refreshed.outputs[0].output_level, 2); + assert!(refreshed.outputs[0].filter_deleted); + assert_eq!(refreshed.outputs[0].output_time_range, output_time_range); + assert_eq!(refreshed.time_window_size, 3600); + assert_eq!(refreshed.max_file_size, Some(4096)); + assert_eq!(refreshed.outputs[0].inputs[0].meta_ref(), ¤t_meta); + assert_eq!(refreshed.expired_ssts.len(), 1); + stale.set_compacting(true); + assert!(!refreshed.outputs[0].inputs[0].compacting()); + } + + #[test] + fn test_pick_result_rejects_ambiguous_file_id_across_levels() { + let purger = crate::test_util::new_noop_file_purger(); + let selected = new_file_handle(FileId::random(), 0, 10, 1); + let mut wrong_level = selected.meta_ref().clone(); + wrong_level.level = 0; + let current_level = selected.meta_ref().clone(); + let mut current = crate::sst::version::SstVersion::new(); + current.add_files(purger.clone(), std::iter::once(wrong_level)); + let output = picker_output_with_files(vec![selected], Vec::new()); + + assert!(refresh_picker_output(output.clone(), ¤t).is_none()); + current.add_files(purger, std::iter::once(current_level)); + assert!(refresh_picker_output(output, ¤t).is_none()); + } + #[tokio::test] async fn test_find_compaction_options_db_level() { - let env = SchedulerEnv::new().await; let builder = VersionControlBuilder::new(); let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); let region_id = builder.region_id(); @@ -1350,39 +2318,6 @@ mod tests { assert_eq!(t.time_window_seconds(), Some(2 * 3600)); } } - let manifest_ctx = env - .mock_manifest_context(version_control.current().version.metadata.clone()) - .await; - let (tx, _rx) = mpsc::channel(4); - let mut scheduler = env.mock_compaction_scheduler(tx); - let (otx, _orx) = oneshot::channel(); - let request = scheduler - .region_status - .entry(region_id) - .or_insert_with(|| { - crate::compaction::CompactionStatus::new( - region_id, - version_control.clone(), - env.access_layer.clone(), - ) - }) - .new_compaction_request( - scheduler.request_sender.clone(), - OptionOutputTx::new(Some(OutputTx::new(otx))), - scheduler.engine_config.clone(), - scheduler.cache_manager.clone(), - &manifest_ctx, - scheduler.listener.clone(), - schema_metadata_manager.clone(), - 1, - ); - scheduler - .schedule_compaction_request( - request, - compact_request::Options::Regular(Default::default()), - ) - .await - .unwrap(); } #[tokio::test] @@ -1472,7 +2407,7 @@ mod tests { #[tokio::test] async fn test_schedule_empty() { let env = SchedulerEnv::new().await; - let (tx, _rx) = mpsc::channel(4); + let (tx, mut rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx); let mut builder = VersionControlBuilder::new(); let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); @@ -1506,7 +2441,16 @@ mod tests { ) .await .unwrap(); - assert!(!scheduled); + assert!(scheduled); + let finished = recv_compaction_pick_finished(&mut rx).await; + assert!(matches!(&finished.result, CompactionPlanningResult::NoPlan)); + scheduler + .handle_compaction_pick_finished( + finished, + &manifest_ctx, + schema_metadata_manager.clone(), + ) + .await; let output = output_rx.await.unwrap().unwrap(); assert_eq!(output, 0); assert!(scheduler.region_status.is_empty()); @@ -1523,12 +2467,17 @@ mod tests { &env.access_layer, waiter, &manifest_ctx, - schema_metadata_manager, + schema_metadata_manager.clone(), 1, ) .await .unwrap(); - assert!(!scheduled); + assert!(scheduled); + let finished = recv_compaction_pick_finished(&mut rx).await; + assert!(matches!(&finished.result, CompactionPlanningResult::NoPlan)); + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) + .await; let output = output_rx.await.unwrap().unwrap(); assert_eq!(output, 0); assert!(scheduler.region_status.is_empty()); @@ -1538,7 +2487,7 @@ mod tests { async fn test_schedule_compaction_returns_true_when_task_scheduled() { let job_scheduler = Arc::new(VecScheduler::default()); let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); - let (tx, _rx) = mpsc::channel(4); + let (tx, mut rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx); let mut builder = VersionControlBuilder::new(); let region_id = builder.region_id(); @@ -1576,7 +2525,7 @@ mod tests { &env.access_layer, OptionOutputTx::none(), &manifest_ctx, - schema_metadata_manager, + schema_metadata_manager.clone(), 1, ) .await @@ -1585,15 +2534,1467 @@ mod tests { // The boolean result is what the worker uses to decide whether to update // last_schedule_compaction_millis. assert!(scheduled); + assert_eq!(0, job_scheduler.num_jobs()); + let finished = recv_compaction_pick_finished(&mut rx).await; + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) + .await; assert_eq!(1, job_scheduler.num_jobs()); assert!(scheduler.region_status.contains_key(®ion_id)); } + #[tokio::test] + async fn test_picking_coalesces_duplicate_same_region_triggers() { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let mut builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let end = 1000 * 1000; + let version_control = Arc::new( + builder + .push_l0_file(0, end) + .push_l0_file(10, end) + .push_l0_file(50, end) + .push_l0_file(80, end) + .push_l0_file(90, end) + .build(), + ); + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .await; + let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); + schema_metadata_manager + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + kv_backend, + ) + .await; + let gate = Arc::new(CompactionPlanningGate::new(region_id)); + let gate_guard = gate.arm(); + scheduler.listener = WorkerListener::new(Some(gate.clone())); + let (first_tx, _first_rx) = oneshot::channel(); + + let mut schedule = tokio::spawn({ + let version_control = version_control.clone(); + let access_layer = env.access_layer.clone(); + let manifest_ctx = manifest_ctx.clone(); + let schema_metadata_manager = schema_metadata_manager.clone(); + async move { + let scheduled = scheduler + .schedule_compaction( + region_id, + Options::Regular(Default::default()), + &version_control, + &access_layer, + OptionOutputTx::from(first_tx), + &manifest_ctx, + schema_metadata_manager, + 1, + ) + .await + .unwrap(); + (scheduled, scheduler) + } + }); + + tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered()) + .await + .expect("planning did not reach the picker gate"); + let (scheduled, mut scheduler) = + match tokio::time::timeout(Duration::from_secs(5), &mut schedule).await { + Ok(result) => result.unwrap(), + Err(_) => { + panic!("schedule_compaction awaited picker planning") + } + }; + + assert!(scheduled); + let status = scheduler.region_status.get(®ion_id).unwrap(); + assert!(status.is_busy()); + assert_eq!(status.waiters.len(), 1); + assert_eq!(0, job_scheduler.num_jobs()); + + let (second_tx, _second_rx) = oneshot::channel(); + assert!( + !scheduler + .schedule_compaction( + region_id, + Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::from(second_tx), + &manifest_ctx, + schema_metadata_manager.clone(), + 1, + ) + .await + .unwrap() + ); + let (manual_tx, _manual_rx) = oneshot::channel(); + assert!( + !scheduler + .schedule_compaction( + region_id, + Options::StrictWindow(StrictWindow { window_seconds: 60 }), + &version_control, + &env.access_layer, + OptionOutputTx::from(manual_tx), + &manifest_ctx, + schema_metadata_manager, + 1, + ) + .await + .unwrap() + ); + let status = scheduler.region_status.get(®ion_id).unwrap(); + assert_eq!(status.waiters.len(), 1); + assert!(status.regular_replan_pending); + assert_eq!(status.pending_regular_waiters.len(), 1); + assert!(status.pending_request.is_some()); + assert_eq!(1, gate.invocation_count()); + + gate_guard.release(); + let finished = tokio::time::timeout( + Duration::from_secs(5), + recv_compaction_pick_finished(&mut rx), + ) + .await + .expect("planning did not send a terminal notification"); + assert_eq!(finished.region_id, region_id); + assert!(matches!( + finished.result, + CompactionPlanningResult::Prepared(_) + )); + } + + #[tokio::test] + async fn test_planning_no_plan_completion_clears_picking_and_notifies_waiter() { + let env = SchedulerEnv::new().await; + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .await; + let (schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager(); + let (waiter_tx, waiter_rx) = oneshot::channel(); + let mut status = + CompactionStatus::new(region_id, version_control.clone(), env.access_layer.clone()); + status.merge_waiter(OptionOutputTx::from(waiter_tx)); + status.start_picking(7); + scheduler.region_status.insert(region_id, status); + + let pending_ddls = scheduler + .handle_compaction_pick_finished( + CompactionPickFinished { + region_id, + plan_id: 7, + version_control, + result: CompactionPlanningResult::NoPlan, + }, + &manifest_ctx, + schema_metadata_manager.clone(), + ) + .await; + + assert!(pending_ddls.is_empty()); + assert_eq!(0, waiter_rx.await.unwrap().unwrap()); + assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(rx.try_recv().is_err()); + } + + #[tokio::test] + async fn test_planning_error_completion_clears_picking_and_notifies_waiter_once() { + let env = SchedulerEnv::new().await; + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .await; + let (schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager(); + let (waiter_tx, waiter_rx) = oneshot::channel(); + let mut status = + CompactionStatus::new(region_id, version_control.clone(), env.access_layer.clone()); + status.merge_waiter(OptionOutputTx::from(waiter_tx)); + status.start_picking(9); + scheduler.region_status.insert(region_id, status); + + let pending_ddls = scheduler + .handle_compaction_pick_finished( + CompactionPickFinished { + region_id, + plan_id: 9, + version_control, + result: CompactionPlanningResult::Error(Arc::new( + InvalidSchedulerStateSnafu.build(), + )), + }, + &manifest_ctx, + schema_metadata_manager, + ) + .await; + + assert!(pending_ddls.is_empty()); + assert!(waiter_rx.await.unwrap().is_err()); + assert!(!scheduler.region_status.contains_key(®ion_id)); + } + + #[tokio::test] + async fn test_picking_lifecycle_close_drop_truncate_fail_waiters_and_ignore_completion() { + for lifecycle in ["close", "drop", "truncate"] { + let env = SchedulerEnv::new().await; + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .await; + let (schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager(); + let (first_tx, first_rx) = oneshot::channel(); + let (second_tx, second_rx) = oneshot::channel(); + let (manual_tx, manual_rx) = oneshot::channel(); + let mut status = + CompactionStatus::new(region_id, version_control.clone(), env.access_layer.clone()); + status.merge_waiter(OptionOutputTx::from(first_tx)); + status.set_pending_request(PendingCompaction { + options: compact_request::Options::StrictWindow(StrictWindow { + window_seconds: 60, + }), + waiter: OptionOutputTx::from(manual_tx), + max_parallelism: 1, + }); + status.start_picking(7); + status.merge_regular_trigger(OptionOutputTx::from(second_tx)); + scheduler.region_status.insert(region_id, status); + + match lifecycle { + "close" => scheduler.on_region_closed(region_id), + "drop" => scheduler.on_region_dropped(region_id), + "truncate" => scheduler.on_region_truncated(region_id), + _ => unreachable!(), + } + + assert!(!scheduler.region_status.contains_key(®ion_id)); + for result in [ + first_rx.await.unwrap(), + second_rx.await.unwrap(), + manual_rx.await.unwrap(), + ] { + assert_compaction_lifecycle_error(result.unwrap_err(), lifecycle); + } + let pending_ddls = scheduler + .accept_compaction_pick_finished( + CompactionPickFinished { + region_id, + plan_id: 7, + version_control: version_control.clone(), + result: CompactionPlanningResult::NoPlan, + }, + &version_control, + &manifest_ctx, + schema_metadata_manager, + ) + .await; + assert!(pending_ddls.is_empty()); + assert!(!scheduler.region_status.contains_key(®ion_id)); + } + } + + #[tokio::test] + async fn test_picking_lifecycle_replacement_ignores_old_completion() { + let env = SchedulerEnv::new().await; + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let stale_version_control = compactable_version(); + let region_id = stale_version_control.current().version.metadata.region_id; + let (finished, _stale_manifest_ctx, schema_metadata_manager) = + begin_pick_result(&env, &mut scheduler, &mut rx, &stale_version_control).await; + let selected = selected_files(&finished); + let replacement_version_control = compactable_version(); + let manifest_ctx = env + .mock_manifest_context( + replacement_version_control + .current() + .version + .metadata + .clone(), + ) + .await; + let (stale_tx, stale_rx) = oneshot::channel(); + scheduler + .region_status + .get_mut(®ion_id) + .unwrap() + .merge_waiter(OptionOutputTx::from(stale_tx)); + scheduler.on_region_closed(region_id); + assert!(stale_rx.await.unwrap().is_err()); + + let (replacement_tx, mut replacement_rx) = oneshot::channel(); + let replacement_plan_id = finished.plan_id + 1; + let mut replacement_status = CompactionStatus::new( + region_id, + replacement_version_control.clone(), + env.access_layer.clone(), + ); + replacement_status.merge_waiter(OptionOutputTx::from(replacement_tx)); + replacement_status.start_picking(replacement_plan_id); + scheduler + .region_status + .insert(region_id, replacement_status); + + let pending_ddls = scheduler + .accept_compaction_pick_finished( + finished, + &replacement_version_control, + &manifest_ctx, + schema_metadata_manager, + ) + .await; + + assert!(pending_ddls.is_empty()); + let status = &scheduler.region_status[®ion_id]; + assert!(status.is_picking(replacement_plan_id)); + assert_eq!(status.waiters.len(), 1); + assert!(selected.iter().all(|file| !file.compacting())); + assert_matches!( + replacement_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + ); + } + + #[tokio::test] + async fn test_picking_lifecycle_staging_waits_for_cancellation_ack() { + let env = SchedulerEnv::new().await; + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .await; + let (schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager(); + let (waiter_tx, waiter_rx) = oneshot::channel(); + let mut status = + CompactionStatus::new(region_id, version_control.clone(), env.access_layer.clone()); + status.merge_waiter(OptionOutputTx::from(waiter_tx)); + status.start_picking(7); + scheduler.region_status.insert(region_id, status); + let (ddl_tx, mut ddl_rx) = oneshot::channel(); + scheduler.add_ddl_request_to_pending(SenderDdlRequest { + region_id, + sender: OptionOutputTx::from(ddl_tx), + request: crate::request::DdlRequest::EnterStaging( + store_api::region_request::EnterStagingRequest { + partition_directive: + store_api::region_request::StagingPartitionDirective::RejectAllWrites, + }, + ), + }); + + assert_eq!( + scheduler.request_cancel(region_id), + RequestCancelResult::CancelIssued + ); + assert!(scheduler.region_status[®ion_id].is_picking(7)); + assert!(scheduler.has_pending_ddls(region_id)); + assert_matches!(ddl_rx.try_recv(), Err(oneshot::error::TryRecvError::Empty)); + + let pending_ddls = scheduler + .accept_compaction_pick_finished( + CompactionPickFinished { + region_id, + plan_id: 7, + version_control: version_control.clone(), + result: CompactionPlanningResult::NoPlan, + }, + &version_control, + &manifest_ctx, + schema_metadata_manager, + ) + .await; + + assert_eq!(pending_ddls.len(), 1); + assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(waiter_rx.await.unwrap().is_err()); + assert_matches!(ddl_rx.try_recv(), Err(oneshot::error::TryRecvError::Empty)); + } + + #[tokio::test] + async fn test_picking_lifecycle_manual_noop_precedes_pending_ddl() { + let env = SchedulerEnv::new().await; + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .await; + let (schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager(); + let (regular_tx, regular_rx) = oneshot::channel(); + let (followup_tx, mut followup_rx) = oneshot::channel(); + let (manual_tx, mut manual_rx) = oneshot::channel(); + let mut status = + CompactionStatus::new(region_id, version_control.clone(), env.access_layer.clone()); + status.merge_waiter(OptionOutputTx::from(regular_tx)); + status.start_picking(7); + scheduler.region_status.insert(region_id, status); + assert!( + !scheduler + .schedule_compaction( + region_id, + compact_request::Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::from(followup_tx), + &manifest_ctx, + schema_metadata_manager.clone(), + 1, + ) + .await + .unwrap() + ); + assert!( + !scheduler + .schedule_compaction( + region_id, + compact_request::Options::StrictWindow(StrictWindow { window_seconds: 60 }), + &version_control, + &env.access_layer, + OptionOutputTx::from(manual_tx), + &manifest_ctx, + schema_metadata_manager.clone(), + 1, + ) + .await + .unwrap() + ); + let (ddl_tx, mut ddl_rx) = oneshot::channel(); + scheduler.add_ddl_request_to_pending(SenderDdlRequest { + region_id, + sender: OptionOutputTx::from(ddl_tx), + request: crate::request::DdlRequest::EnterStaging( + store_api::region_request::EnterStagingRequest { + partition_directive: + store_api::region_request::StagingPartitionDirective::RejectAllWrites, + }, + ), + }); + + let pending_ddls = scheduler + .accept_compaction_pick_finished( + CompactionPickFinished { + region_id, + plan_id: 7, + version_control: version_control.clone(), + result: CompactionPlanningResult::NoPlan, + }, + &version_control, + &manifest_ctx, + schema_metadata_manager.clone(), + ) + .await; + + assert!(pending_ddls.is_empty()); + assert_eq!( + tokio::time::timeout(Duration::from_secs(5), regular_rx) + .await + .expect("current regular waiter was not notified before manual planning") + .unwrap() + .unwrap(), + 0 + ); + assert!(scheduler.has_pending_ddls(region_id)); + assert_matches!( + followup_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + ); + assert_matches!( + manual_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + ); + assert_matches!(ddl_rx.try_recv(), Err(oneshot::error::TryRecvError::Empty)); + + let mut manual_finished = recv_compaction_pick_finished(&mut rx).await; + manual_finished.result = CompactionPlanningResult::NoPlan; + let pending_ddls = scheduler + .accept_compaction_pick_finished( + manual_finished, + &version_control, + &manifest_ctx, + schema_metadata_manager.clone(), + ) + .await; + assert!(pending_ddls.is_empty()); + assert_eq!( + tokio::time::timeout(Duration::from_secs(5), manual_rx) + .await + .expect("manual waiter was not notified before regular follow-up planning") + .unwrap() + .unwrap(), + 0 + ); + assert_matches!( + followup_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + ); + assert_matches!(ddl_rx.try_recv(), Err(oneshot::error::TryRecvError::Empty)); + + let mut regular_finished = tokio::time::timeout( + Duration::from_secs(5), + recv_compaction_pick_finished(&mut rx), + ) + .await + .expect("regular follow-up was not planned after manual completion"); + regular_finished.result = CompactionPlanningResult::NoPlan; + let pending_ddls = scheduler + .accept_compaction_pick_finished( + regular_finished, + &version_control, + &manifest_ctx, + schema_metadata_manager, + ) + .await; + assert_eq!(pending_ddls.len(), 1); + assert_eq!( + tokio::time::timeout(Duration::from_secs(5), followup_rx) + .await + .expect("regular follow-up waiter was not notified") + .unwrap() + .unwrap(), + 0 + ); + assert_matches!(ddl_rx.try_recv(), Err(oneshot::error::TryRecvError::Empty)); + } + + #[tokio::test] + async fn test_ddl_fence_prevents_repeated_regular_followups() { + let env = SchedulerEnv::new().await; + let (tx, mut rx) = mpsc::channel(8); + let mut scheduler = env.mock_compaction_scheduler(tx); + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .await; + let (schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager(); + let (current_tx, current_rx) = oneshot::channel(); + let mut status = + CompactionStatus::new(region_id, version_control.clone(), env.access_layer.clone()); + status.merge_waiter(OptionOutputTx::from(current_tx)); + status.start_picking(7); + scheduler.region_status.insert(region_id, status); + + let (prefence_tx, prefence_rx) = oneshot::channel(); + assert!( + !scheduler + .schedule_compaction( + region_id, + compact_request::Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::from(prefence_tx), + &manifest_ctx, + schema_metadata_manager.clone(), + 1, + ) + .await + .unwrap() + ); + let (ddl_tx, mut ddl_rx) = oneshot::channel(); + scheduler.add_ddl_request_to_pending(SenderDdlRequest { + region_id, + sender: OptionOutputTx::from(ddl_tx), + request: crate::request::DdlRequest::EnterStaging( + store_api::region_request::EnterStagingRequest { + partition_directive: + store_api::region_request::StagingPartitionDirective::RejectAllWrites, + }, + ), + }); + + let pending_ddls = scheduler + .accept_compaction_pick_finished( + CompactionPickFinished { + region_id, + plan_id: 7, + version_control: version_control.clone(), + result: CompactionPlanningResult::NoPlan, + }, + &version_control, + &manifest_ctx, + schema_metadata_manager.clone(), + ) + .await; + assert!(pending_ddls.is_empty()); + assert_eq!(current_rx.await.unwrap().unwrap(), 0); + assert_matches!(ddl_rx.try_recv(), Err(oneshot::error::TryRecvError::Empty)); + + let mut followup_finished = tokio::time::timeout( + Duration::from_secs(5), + recv_compaction_pick_finished(&mut rx), + ) + .await + .expect("pre-fence regular follow-up was not planned"); + let mut postfence_waiters = Vec::new(); + for _ in 0..3 { + let (waiter_tx, waiter_rx) = oneshot::channel(); + assert!( + !scheduler + .schedule_compaction( + region_id, + compact_request::Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::from(waiter_tx), + &manifest_ctx, + schema_metadata_manager.clone(), + 1, + ) + .await + .unwrap() + ); + postfence_waiters.push(waiter_rx); + } + assert!( + !scheduler + .schedule_compaction( + region_id, + compact_request::Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::none(), + &manifest_ctx, + schema_metadata_manager.clone(), + 1, + ) + .await + .unwrap() + ); + + followup_finished.result = CompactionPlanningResult::NoPlan; + let pending_ddls = scheduler + .accept_compaction_pick_finished( + followup_finished, + &version_control, + &manifest_ctx, + schema_metadata_manager, + ) + .await; + assert_eq!(pending_ddls.len(), 1); + assert_eq!(prefence_rx.await.unwrap().unwrap(), 0); + for waiter in postfence_waiters { + assert_eq!(waiter.await.unwrap().unwrap(), 0); + } + assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(rx.try_recv().is_err()); + assert_matches!(ddl_rx.try_recv(), Err(oneshot::error::TryRecvError::Empty)); + } + + #[tokio::test] + async fn test_picking_lifecycle_scheduler_drop_notifies_pending_requests() { + let env = SchedulerEnv::new().await; + let (waiter_rx, manual_rx, ddl_rx) = { + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let (waiter_tx, waiter_rx) = oneshot::channel(); + let (manual_tx, manual_rx) = oneshot::channel(); + let (ddl_tx, ddl_rx) = oneshot::channel(); + let mut status = + CompactionStatus::new(region_id, version_control, env.access_layer.clone()); + status.merge_waiter(OptionOutputTx::from(waiter_tx)); + status.set_pending_request(PendingCompaction { + options: compact_request::Options::StrictWindow(StrictWindow { + window_seconds: 60, + }), + waiter: OptionOutputTx::from(manual_tx), + max_parallelism: 1, + }); + status.pending_ddl_requests.push(SenderDdlRequest { + region_id, + sender: OptionOutputTx::from(ddl_tx), + request: crate::request::DdlRequest::EnterStaging( + store_api::region_request::EnterStagingRequest { + partition_directive: + store_api::region_request::StagingPartitionDirective::RejectAllWrites, + }, + ), + }); + status.start_picking(7); + scheduler.region_status.insert(region_id, status); + (waiter_rx, manual_rx, ddl_rx) + }; + + for result in [ + waiter_rx.await.unwrap(), + manual_rx.await.unwrap(), + ddl_rx.await.unwrap(), + ] { + assert_compaction_lifecycle_error(result.unwrap_err(), "close"); + } + } + + #[tokio::test] + async fn test_pick_result_matching_plan_submits_once_and_owns_reservations() { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let version_control = compactable_version(); + let region_id = version_control.current().version.metadata.region_id; + let (finished, manifest_ctx, schema_metadata_manager) = + begin_pick_result(&env, &mut scheduler, &mut rx, &version_control).await; + let selected = selected_files(&finished); + + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) + .await; + + assert_eq!(job_scheduler.num_jobs(), 1); + assert!(selected.iter().all(FileHandle::compacting)); + scheduler.on_compaction_failed(region_id, Arc::new(InvalidSchedulerStateSnafu.build())); + assert!(selected.iter().all(FileHandle::compacting)); + drop(scheduler); + drop(env); + drop(job_scheduler); + assert!(selected.iter().all(|file| !file.compacting())); + } + + #[tokio::test] + async fn test_pick_result_mismatched_token_keeps_status_and_waiter_untouched() { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let version_control = compactable_version(); + let region_id = version_control.current().version.metadata.region_id; + let (mut finished, manifest_ctx, schema_metadata_manager) = + begin_pick_result(&env, &mut scheduler, &mut rx, &version_control).await; + let (waiter_tx, mut waiter_rx) = oneshot::channel(); + scheduler + .region_status + .get_mut(®ion_id) + .unwrap() + .merge_waiter(OptionOutputTx::from(waiter_tx)); + finished.plan_id += 1; + + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) + .await; + + assert_eq!(job_scheduler.num_jobs(), 0); + assert!(scheduler.region_status[®ion_id].is_busy()); + assert_eq!(scheduler.region_status[®ion_id].waiters.len(), 1); + assert_matches!( + waiter_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + ); + } + + #[tokio::test] + async fn test_pick_result_different_version_control_keeps_status_untouched() { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let version_control = compactable_version(); + let region_id = version_control.current().version.metadata.region_id; + let (mut finished, manifest_ctx, schema_metadata_manager) = + begin_pick_result(&env, &mut scheduler, &mut rx, &version_control).await; + let (waiter_tx, mut waiter_rx) = oneshot::channel(); + scheduler + .region_status + .get_mut(®ion_id) + .unwrap() + .merge_waiter(OptionOutputTx::from(waiter_tx)); + finished.version_control = compactable_version(); + + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) + .await; + + assert_eq!(job_scheduler.num_jobs(), 0); + assert!(scheduler.region_status[®ion_id].is_busy()); + assert_eq!(scheduler.region_status[®ion_id].waiters.len(), 1); + assert_matches!( + waiter_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + ); + } + + #[tokio::test] + async fn test_pick_result_rejects_replaced_current_region_instance() { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let stale_version_control = compactable_version(); + let region_id = stale_version_control.current().version.metadata.region_id; + let (finished, manifest_ctx, schema_metadata_manager) = + begin_pick_result(&env, &mut scheduler, &mut rx, &stale_version_control).await; + let replacement_version_control = compactable_version(); + assert_eq!( + replacement_version_control + .current() + .version + .metadata + .region_id, + region_id + ); + assert!(!Arc::ptr_eq( + &stale_version_control, + &replacement_version_control + )); + assert!(Arc::ptr_eq( + &scheduler.region_status[®ion_id].version_control, + &finished.version_control + )); + let replacement_files: Vec<_> = replacement_version_control + .current() + .version + .ssts + .levels() + .iter() + .flat_map(LevelMeta::files) + .cloned() + .collect(); + let (waiter_tx, mut waiter_rx) = oneshot::channel(); + scheduler + .region_status + .get_mut(®ion_id) + .unwrap() + .merge_waiter(OptionOutputTx::from(waiter_tx)); + + scheduler + .accept_compaction_pick_finished( + finished, + &replacement_version_control, + &manifest_ctx, + schema_metadata_manager, + ) + .await; + + assert_eq!(job_scheduler.num_jobs(), 0); + assert!(scheduler.region_status[®ion_id].is_busy()); + assert_eq!(scheduler.region_status[®ion_id].waiters.len(), 1); + assert_matches!( + waiter_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + ); + assert!(replacement_files.iter().all(|file| !file.compacting())); + } + + #[tokio::test] + async fn test_pick_result_accepts_unrelated_concurrent_flush() { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let version_control = compactable_version(); + let (finished, manifest_ctx, schema_metadata_manager) = + begin_pick_result(&env, &mut scheduler, &mut rx, &version_control).await; + let selected = selected_files(&finished); + apply_edit( + &version_control, + &[(2_000_000, 3_000_000)], + &[], + selected[0].file_purger(), + ); + + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) + .await; + + assert_eq!(job_scheduler.num_jobs(), 1); + assert!(selected.iter().all(FileHandle::compacting)); + } + + #[tokio::test] + async fn test_pick_result_rejects_removed_selected_file_cleanly() { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let version_control = compactable_version(); + let (finished, manifest_ctx, schema_metadata_manager) = + begin_pick_result(&env, &mut scheduler, &mut rx, &version_control).await; + let selected = selected_files(&finished); + let (waiter_tx, waiter_rx) = oneshot::channel(); + let region_id = finished.region_id; + scheduler + .region_status + .get_mut(®ion_id) + .unwrap() + .merge_waiter(OptionOutputTx::from(waiter_tx)); + apply_edit( + &version_control, + &[], + &[selected[0].meta_ref().clone()], + selected[0].file_purger(), + ); + + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) + .await; + + assert_eq!(job_scheduler.num_jobs(), 0); + assert_eq!(waiter_rx.await.unwrap().unwrap(), 0); + assert!(!scheduler.region_status.contains_key(®ion_id)); + } + + #[tokio::test] + async fn test_pick_result_refreshes_replaced_selected_file() { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let version_control = compactable_version(); + let (finished, manifest_ctx, schema_metadata_manager) = + begin_pick_result(&env, &mut scheduler, &mut rx, &version_control).await; + let selected = selected_files(&finished); + let stale = selected[0].clone(); + let mut replacement = stale.meta_ref().clone(); + replacement.index_version = 1; + replacement.index_file_size = 128; + version_control.apply_edit( + Some(crate::manifest::action::RegionEdit { + files_to_add: vec![replacement], + files_to_remove: Vec::new(), + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }), + &[], + stale.file_purger(), + ); + let current = version_control + .current() + .version + .ssts + .file_for_compaction(&stale) + .unwrap() + .clone(); + + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) + .await; + + assert_eq!(job_scheduler.num_jobs(), 1); + assert!(!stale.compacting()); + assert!(current.compacting()); + assert_eq!(current.meta_ref().index_version, 1); + } + + #[tokio::test] + async fn test_pick_result_rejects_deleted_or_compacting_current_file() { + for deleted in [true, false] { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let version_control = compactable_version(); + let (finished, manifest_ctx, schema_metadata_manager) = + begin_pick_result(&env, &mut scheduler, &mut rx, &version_control).await; + let selected = selected_files(&finished); + if deleted { + selected[0].mark_deleted(); + } else { + selected[0].set_compacting(true); + } + + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) + .await; + + assert_eq!(job_scheduler.num_jobs(), 0); + assert!(scheduler.region_status.is_empty()); + if !deleted { + assert!(selected[0].compacting()); + selected[0].set_compacting(false); + } + } + } + + #[tokio::test] + async fn test_pick_result_local_submission_failure_releases_and_notifies_once() { + let env = SchedulerEnv::new() + .await + .scheduler(Arc::new(FailingScheduler)); + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let version_control = compactable_version(); + let region_id = version_control.current().version.metadata.region_id; + let (finished, manifest_ctx, schema_metadata_manager) = + begin_pick_result(&env, &mut scheduler, &mut rx, &version_control).await; + let selected = selected_files(&finished); + let (waiter_tx, waiter_rx) = oneshot::channel(); + scheduler + .region_status + .get_mut(®ion_id) + .unwrap() + .merge_waiter(OptionOutputTx::from(waiter_tx)); + + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) + .await; + + assert!(waiter_rx.await.unwrap().is_err()); + assert!(selected.iter().all(|file| !file.compacting())); + assert!(!scheduler.region_status.contains_key(®ion_id)); + } + + #[tokio::test] + async fn test_pick_result_rejects_removed_expired_file_cleanly() { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let version_control = compactable_version(); + let (mut finished, manifest_ctx, schema_metadata_manager) = + begin_pick_result(&env, &mut scheduler, &mut rx, &version_control).await; + let CompactionPlanningResult::Prepared(prepared) = &mut finished.result else { + unreachable!(); + }; + let expired = prepared.picker_output.outputs[0].inputs.pop().unwrap(); + prepared.picker_output.expired_ssts.push(expired.clone()); + apply_edit( + &version_control, + &[], + &[expired.meta_ref().clone()], + expired.file_purger(), + ); + + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) + .await; + + assert_eq!(job_scheduler.num_jobs(), 0); + assert!(scheduler.region_status.is_empty()); + } + + #[tokio::test] + async fn test_pick_result_local_completion_and_cancel_release_reservations() { + for cancel in [false, true] { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler); + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let version_control = compactable_version(); + let region_id = version_control.current().version.metadata.region_id; + let (finished, manifest_ctx, schema_metadata_manager) = + begin_pick_result(&env, &mut scheduler, &mut rx, &version_control).await; + let selected = selected_files(&finished); + scheduler + .handle_compaction_pick_finished( + finished, + &manifest_ctx, + schema_metadata_manager.clone(), + ) + .await; + assert!(selected.iter().all(FileHandle::compacting)); + + if cancel { + assert_eq!( + scheduler.request_cancel(region_id), + RequestCancelResult::CancelIssued + ); + scheduler.on_compaction_cancelled(region_id).await; + } else { + scheduler + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) + .await; + } + + assert!(selected.iter().all(FileHandle::compacting)); + drop(scheduler); + drop(env); + assert!(selected.iter().all(|file| !file.compacting())); + } + } + + #[tokio::test] + async fn test_pick_result_remote_completion_and_failure_release_reservations() { + for failed in [false, true] { + let env = SchedulerEnv::new().await; + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + scheduler + .plugins + .insert::(Arc::new(SuccessfulRemoteScheduler)); + let version_control = compactable_version(); + let region_id = version_control.current().version.metadata.region_id; + let (mut finished, manifest_ctx, schema_metadata_manager) = + begin_pick_result(&env, &mut scheduler, &mut rx, &version_control).await; + use_remote_compaction(&mut finished); + let selected = selected_files(&finished); + scheduler + .handle_compaction_pick_finished( + finished, + &manifest_ctx, + schema_metadata_manager.clone(), + ) + .await; + assert!(selected.iter().all(FileHandle::compacting)); + + if failed { + scheduler + .on_compaction_failed(region_id, Arc::new(InvalidSchedulerStateSnafu.build())); + } else { + scheduler + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) + .await; + } + + assert!(selected.iter().all(|file| !file.compacting())); + } + } + + #[tokio::test] + async fn test_pick_result_remote_submission_failure_releases_and_notifies_once() { + let env = SchedulerEnv::new().await; + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + scheduler + .plugins + .insert::(Arc::new(FailingRemoteScheduler)); + let version_control = compactable_version(); + let region_id = version_control.current().version.metadata.region_id; + let (mut finished, manifest_ctx, schema_metadata_manager) = + begin_pick_result(&env, &mut scheduler, &mut rx, &version_control).await; + use_remote_compaction(&mut finished); + let selected = selected_files(&finished); + let (waiter_tx, waiter_rx) = oneshot::channel(); + scheduler + .region_status + .get_mut(®ion_id) + .unwrap() + .merge_waiter(OptionOutputTx::from(waiter_tx)); + + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) + .await; + + assert!(waiter_rx.await.unwrap().is_err()); + assert!(selected.iter().all(|file| !file.compacting())); + assert!(!scheduler.region_status.contains_key(®ion_id)); + } + + #[tokio::test] + async fn test_local_reservations_survive_status_removal_until_execution_drops() { + let selected; + { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let version_control = compactable_version(); + let region_id = version_control.current().version.metadata.region_id; + let (finished, manifest_ctx, schema_metadata_manager) = + begin_pick_result(&env, &mut scheduler, &mut rx, &version_control).await; + selected = selected_files(&finished); + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) + .await; + assert!(selected.iter().all(FileHandle::compacting)); + + scheduler.on_region_closed(region_id); + + assert!(selected.iter().all(FileHandle::compacting)); + } + assert!(selected.iter().all(|file| !file.compacting())); + } + + #[tokio::test] + async fn test_remote_reservations_survive_status_removal_until_notifier_drops() { + let env = SchedulerEnv::new().await; + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let remote_scheduler = Arc::new(HoldingRemoteScheduler::default()); + scheduler + .plugins + .insert::(remote_scheduler.clone()); + let version_control = compactable_version(); + let region_id = version_control.current().version.metadata.region_id; + let (mut finished, manifest_ctx, schema_metadata_manager) = + begin_pick_result(&env, &mut scheduler, &mut rx, &version_control).await; + use_remote_compaction(&mut finished); + let selected = selected_files(&finished); + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) + .await; + assert!(selected.iter().all(FileHandle::compacting)); + + scheduler.on_region_closed(region_id); + + assert!(selected.iter().all(FileHandle::compacting)); + remote_scheduler.drop_notifier(); + assert!(selected.iter().all(|file| !file.compacting())); + } + + #[tokio::test] + async fn test_stale_local_success_does_not_clear_replacement_phase() { + let env = SchedulerEnv::new().await; + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let stale_version_control = compactable_version(); + let replacement_version_control = compactable_version(); + let region_id = replacement_version_control + .current() + .version + .metadata + .region_id; + assert!(!Arc::ptr_eq( + &stale_version_control, + &replacement_version_control + )); + let manifest_ctx = env + .mock_manifest_context( + replacement_version_control + .current() + .version + .metadata + .clone(), + ) + .await; + let (schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager(); + let stale_execution = + CompactionExecution::for_test(stale_version_control, CompactionExecutionKind::Local); + let mut status = CompactionStatus::new( + region_id, + replacement_version_control.clone(), + env.access_layer.clone(), + ); + status.start_local_task(); + scheduler.region_status.insert(region_id, status); + let files_before = replacement_version_control + .current() + .version + .ssts + .owned_num_files(region_id); + if scheduler.is_current_region_execution( + region_id, + &replacement_version_control, + &stale_execution, + ) { + apply_edit( + &replacement_version_control, + &[(2_000_000, 3_000_000)], + &[], + crate::test_util::new_noop_file_purger(), + ); + } + + scheduler + .on_execution_finished( + region_id, + &stale_execution, + &manifest_ctx, + schema_metadata_manager, + ) + .await; + + assert!(scheduler.region_status[®ion_id].is_busy()); + assert_eq!( + replacement_version_control + .current() + .version + .ssts + .owned_num_files(region_id), + files_before + ); + } + + #[tokio::test] + async fn test_stale_local_cancel_does_not_remove_replacement_status() { + let env = SchedulerEnv::new().await; + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let stale_version_control = compactable_version(); + let replacement_version_control = compactable_version(); + let region_id = replacement_version_control + .current() + .version + .metadata + .region_id; + let stale_execution = + CompactionExecution::for_test(stale_version_control, CompactionExecutionKind::Local); + let mut status = CompactionStatus::new( + region_id, + replacement_version_control, + env.access_layer.clone(), + ); + status.start_local_task(); + scheduler.region_status.insert(region_id, status); + + scheduler + .on_execution_cancelled(region_id, &stale_execution) + .await; + + assert!(scheduler.region_status.contains_key(®ion_id)); + } + + #[tokio::test] + async fn test_stale_local_failure_does_not_remove_replacement_status_or_waiter() { + let env = SchedulerEnv::new().await; + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let stale_version_control = compactable_version(); + let replacement_version_control = compactable_version(); + let region_id = replacement_version_control + .current() + .version + .metadata + .region_id; + let stale_execution = + CompactionExecution::for_test(stale_version_control, CompactionExecutionKind::Local); + let (waiter_tx, mut waiter_rx) = oneshot::channel(); + let mut status = CompactionStatus::new( + region_id, + replacement_version_control, + env.access_layer.clone(), + ); + status.start_local_task(); + status.merge_waiter(OptionOutputTx::from(waiter_tx)); + scheduler.region_status.insert(region_id, status); + + scheduler.on_execution_failed( + region_id, + &stale_execution, + Arc::new(InvalidSchedulerStateSnafu.build()), + ); + + assert!(scheduler.region_status.contains_key(®ion_id)); + assert_matches!( + waiter_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + ); + } + + #[tokio::test] + async fn test_stale_remote_success_does_not_clear_replacement_phase() { + let env = SchedulerEnv::new().await; + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let stale_version_control = compactable_version(); + let replacement_version_control = compactable_version(); + let region_id = replacement_version_control + .current() + .version + .metadata + .region_id; + let stale_execution = + CompactionExecution::for_test(stale_version_control, CompactionExecutionKind::Remote); + let manifest_ctx = env + .mock_manifest_context( + replacement_version_control + .current() + .version + .metadata + .clone(), + ) + .await; + let (schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager(); + let mut status = CompactionStatus::new( + region_id, + replacement_version_control.clone(), + env.access_layer.clone(), + ); + status.start_remote_task(); + scheduler.region_status.insert(region_id, status); + let files_before = replacement_version_control + .current() + .version + .ssts + .owned_num_files(region_id); + if scheduler.is_current_region_execution( + region_id, + &replacement_version_control, + &stale_execution, + ) { + apply_edit( + &replacement_version_control, + &[(2_000_000, 3_000_000)], + &[], + crate::test_util::new_noop_file_purger(), + ); + } + + scheduler + .on_execution_finished( + region_id, + &stale_execution, + &manifest_ctx, + schema_metadata_manager, + ) + .await; + + assert!(scheduler.region_status[®ion_id].is_busy()); + assert_eq!( + replacement_version_control + .current() + .version + .ssts + .owned_num_files(region_id), + files_before + ); + } + + #[tokio::test] + async fn test_stale_remote_failure_does_not_remove_replacement_status() { + let env = SchedulerEnv::new().await; + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let stale_version_control = compactable_version(); + let replacement_version_control = compactable_version(); + let region_id = replacement_version_control + .current() + .version + .metadata + .region_id; + let stale_execution = + CompactionExecution::for_test(stale_version_control, CompactionExecutionKind::Remote); + let mut status = CompactionStatus::new( + region_id, + replacement_version_control, + env.access_layer.clone(), + ); + status.start_remote_task(); + scheduler.region_status.insert(region_id, status); + + scheduler.on_execution_failed( + region_id, + &stale_execution, + Arc::new(InvalidSchedulerStateSnafu.build()), + ); + + assert!(scheduler.region_status.contains_key(®ion_id)); + } + #[tokio::test] async fn test_schedule_compaction_skips_task_exceeding_memory_limit() { let job_scheduler = Arc::new(VecScheduler::default()); let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); - let (tx, _rx) = mpsc::channel(4); + let (tx, mut rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx); scheduler.memory_manager = Arc::new(new_compaction_memory_manager(1024 * 1024)); @@ -1635,17 +4036,23 @@ mod tests { &env.access_layer, OptionOutputTx::from(output_tx), &manifest_ctx, - schema_metadata_manager, + schema_metadata_manager.clone(), 1, ) .await .unwrap(); - assert!(!scheduled); + assert!(scheduled); + let finished = recv_compaction_pick_finished(&mut rx).await; + let selected = selected_files(&finished); + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) + .await; assert_eq!(output_rx.await.unwrap().unwrap(), 0); assert_eq!(rejected_before + 1, rejected.get()); assert_eq!(0, job_scheduler.num_jobs()); assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(selected.iter().all(|file| !file.compacting())); } #[tokio::test] @@ -1653,7 +4060,7 @@ mod tests { 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 (tx, mut rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx); let mut builder = VersionControlBuilder::new(); let purger = builder.file_purger(); @@ -1701,6 +4108,15 @@ mod tests { // Should schedule 1 compaction. assert!(scheduled); assert_eq!(1, scheduler.region_status.len()); + assert_eq!(0, job_scheduler.num_jobs()); + let finished = recv_compaction_pick_finished(&mut rx).await; + scheduler + .handle_compaction_pick_finished( + finished, + &manifest_ctx, + schema_metadata_manager.clone(), + ) + .await; assert_eq!(1, job_scheduler.num_jobs()); let data = version_control.current(); let file_metas: Vec<_> = data.version.ssts.levels()[0] @@ -1752,6 +4168,15 @@ mod tests { .await; assert!(scheduled); assert_eq!(1, scheduler.region_status.len()); + assert_eq!(1, job_scheduler.num_jobs()); + let finished = recv_compaction_pick_finished(&mut rx).await; + scheduler + .handle_compaction_pick_finished( + finished, + &manifest_ctx, + schema_metadata_manager.clone(), + ) + .await; assert_eq!(2, job_scheduler.num_jobs()); // 5 files for next compaction. @@ -1789,12 +4214,12 @@ mod tests { } #[tokio::test] - async fn test_schedule_compaction_does_not_publish_status_when_schedule_fails() { + async fn test_schedule_compaction_clears_status_when_submission_fails() { common_telemetry::init_default_ut_logging(); let env = SchedulerEnv::new() .await .scheduler(Arc::new(FailingScheduler)); - let (tx, _rx) = mpsc::channel(4); + let (tx, mut rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx); let mut builder = VersionControlBuilder::new(); let end = 1000 * 1000; @@ -1831,12 +4256,17 @@ mod tests { &env.access_layer, OptionOutputTx::none(), &manifest_ctx, - schema_metadata_manager, + schema_metadata_manager.clone(), 1, ) .await; - assert!(result.is_err()); + assert!(result.unwrap()); + assert!(scheduler.region_status.contains_key(®ion_id)); + let finished = recv_compaction_pick_finished(&mut rx).await; + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) + .await; assert!(!scheduler.region_status.contains_key(®ion_id)); } @@ -1845,7 +4275,7 @@ mod tests { 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 (tx, mut rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx); let mut builder = VersionControlBuilder::new(); let purger = builder.file_purger(); @@ -1907,6 +4337,15 @@ mod tests { .unwrap(); // Should schedule 1 compaction. assert_eq!(1, scheduler.region_status.len()); + assert_eq!(0, job_scheduler.num_jobs()); + let finished = recv_compaction_pick_finished(&mut rx).await; + scheduler + .handle_compaction_pick_finished( + finished, + &manifest_ctx, + schema_metadata_manager.clone(), + ) + .await; assert_eq!(1, job_scheduler.num_jobs()); assert!( scheduler @@ -1943,6 +4382,15 @@ mod tests { .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) .await; assert_eq!(1, scheduler.region_status.len()); + assert_eq!(1, job_scheduler.num_jobs()); + let finished = recv_compaction_pick_finished(&mut rx).await; + scheduler + .handle_compaction_pick_finished( + finished, + &manifest_ctx, + schema_metadata_manager.clone(), + ) + .await; assert_eq!(2, job_scheduler.num_jobs()); let status = scheduler.region_status.get(&builder.region_id()).unwrap(); @@ -2084,7 +4532,7 @@ mod tests { status.request_cancel(), RequestCancelResult::TooLateToCancel ); - assert!(status.active_compaction.is_some()); + assert!(status.is_busy()); } #[tokio::test] @@ -2101,15 +4549,13 @@ mod tests { .await; let (_schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager(); - scheduler.region_status.insert( - region_id, - CompactionStatus::new(region_id, version_control, env.access_layer.clone()), - ); - scheduler - .region_status - .get_mut(®ion_id) - .unwrap() - .start_local_task(); + let (regular_tx, regular_rx) = oneshot::channel(); + let mut status = + CompactionStatus::new(region_id, version_control, env.access_layer.clone()); + status.start_picking(7); + status.merge_regular_trigger(OptionOutputTx::from(regular_tx)); + status.start_local_task(); + scheduler.region_status.insert(region_id, status); let (output_tx, _output_rx) = oneshot::channel(); scheduler.add_ddl_request_to_pending(SenderDdlRequest { @@ -2129,6 +4575,7 @@ mod tests { assert!(!scheduler.has_pending_ddls(region_id)); assert!(!scheduler.region_status.contains_key(®ion_id)); assert_eq!(job_scheduler.num_jobs(), 0); + assert!(regular_rx.await.unwrap().is_err()); } #[tokio::test] @@ -2181,16 +4628,19 @@ mod tests { #[tokio::test] async fn test_pending_ddl_request_failed_on_compaction_failed() { let env = SchedulerEnv::new().await; - let (tx, _rx) = mpsc::channel(4); + let (tx, mut rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx); let builder = VersionControlBuilder::new(); let version_control = Arc::new(builder.build()); let region_id = builder.region_id(); - scheduler.region_status.insert( - region_id, - CompactionStatus::new(region_id, version_control, env.access_layer.clone()), - ); + let (regular_tx, regular_rx) = oneshot::channel(); + let mut status = + CompactionStatus::new(region_id, version_control, env.access_layer.clone()); + status.start_picking(7); + status.merge_regular_trigger(OptionOutputTx::from(regular_tx)); + status.start_local_task(); + scheduler.region_status.insert(region_id, status); let (output_tx, output_rx) = oneshot::channel(); scheduler.add_ddl_request_to_pending(SenderDdlRequest { @@ -2211,6 +4661,8 @@ mod tests { assert!(!scheduler.has_pending_ddls(region_id)); let result = output_rx.await.unwrap(); assert_matches!(result, Err(_)); + assert!(regular_rx.await.unwrap().is_err()); + assert!(rx.try_recv().is_err()); } #[tokio::test] @@ -2364,7 +4816,7 @@ mod tests { #[tokio::test] async fn test_on_compaction_finished_replays_pending_ddl_after_manual_noop() { let env = SchedulerEnv::new().await; - let (tx, _rx) = mpsc::channel(4); + let (tx, mut rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx); let builder = VersionControlBuilder::new(); let version_control = Arc::new(builder.build()); @@ -2398,9 +4850,14 @@ mod tests { }); let pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) .await; + assert!(pending_ddls.is_empty()); + let finished = recv_compaction_pick_finished(&mut rx).await; + let pending_ddls = scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) + .await; assert_eq!(pending_ddls.len(), 1); assert!(!scheduler.region_status.contains_key(®ion_id)); assert_eq!(manual_rx.await.unwrap().unwrap(), 0); @@ -2431,7 +4888,7 @@ mod tests { let env = SchedulerEnv::new() .await .scheduler(Arc::new(FailingScheduler)); - let (tx, _rx) = mpsc::channel(4); + let (tx, mut rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx); let mut builder = VersionControlBuilder::new(); let end = 1000 * 1000; @@ -2474,19 +4931,24 @@ mod tests { }); let pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) .await; assert!(pending_ddls.is_empty()); + let finished = recv_compaction_pick_finished(&mut rx).await; + let pending_ddls = scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) + .await; + assert!(pending_ddls.is_empty()); assert!(!scheduler.region_status.contains_key(®ion_id)); - assert!(manual_rx.await.is_err()); + assert_matches!(manual_rx.await.unwrap(), Err(_)); assert_matches!(ddl_rx.await.unwrap(), Err(_)); } #[tokio::test] async fn test_on_compaction_finished_next_schedule_noop_removes_status() { let env = SchedulerEnv::new().await; - let (tx, _rx) = mpsc::channel(4); + let (tx, mut rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx); let builder = VersionControlBuilder::new(); let version_control = Arc::new(builder.build()); @@ -2517,9 +4979,13 @@ mod tests { // With no compactable files, next scheduling returns false and removes // the status without creating a background task. let scheduled = scheduler - .schedule_next_compaction(region_id, &manifest_ctx, schema_metadata_manager) + .schedule_next_compaction(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .await; + assert!(scheduled); + let finished = recv_compaction_pick_finished(&mut rx).await; + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) .await; - assert!(!scheduled); assert!(!scheduler.region_status.contains_key(®ion_id)); } @@ -2528,7 +4994,7 @@ mod tests { let env = SchedulerEnv::new() .await .scheduler(Arc::new(FailingScheduler)); - let (tx, _rx) = mpsc::channel(4); + let (tx, mut rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx); let mut builder = VersionControlBuilder::new(); let end = 1000 * 1000; @@ -2567,9 +5033,13 @@ mod tests { let (schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager(); // The failing scheduler simulates a submit error; callers must see false. let scheduled = scheduler - .schedule_next_compaction(region_id, &manifest_ctx, schema_metadata_manager) + .schedule_next_compaction(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .await; + assert!(scheduled); + let finished = recv_compaction_pick_finished(&mut rx).await; + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) .await; - assert!(!scheduled); assert!(!scheduler.region_status.contains_key(®ion_id)); } diff --git a/src/mito2/src/compaction/task.rs b/src/mito2/src/compaction/task.rs index a8da921968..54a7d74a99 100644 --- a/src/mito2/src/compaction/task.rs +++ b/src/mito2/src/compaction/task.rs @@ -24,10 +24,10 @@ use snafu::ResultExt; use store_api::ManifestVersion; use tokio::sync::mpsc; -use crate::compaction::LocalCompactionState; use crate::compaction::compactor::{CompactionRegion, Compactor, MergeOutput}; use crate::compaction::memory_manager::{CompactionMemoryGuard, CompactionMemoryManager}; use crate::compaction::picker::{CompactionTask, PickerOutput}; +use crate::compaction::{CompactionExecution, LocalCompactionState}; use crate::error::{CompactRegionSnafu, CompactionMemoryExhaustedSnafu}; use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; use crate::metrics::{COMPACTION_FAILURE_COUNT, COMPACTION_MEMORY_WAIT, COMPACTION_STAGE_ELAPSED}; @@ -46,6 +46,8 @@ pub const MAX_PARALLEL_COMPACTION: usize = 1; pub(crate) struct CompactionTaskImpl { /// Shared local-compaction state for cooperative cancellation. pub(crate) state: LocalCompactionState, + /// Identity and reservation lease of this accepted execution. + pub(crate) execution: CompactionExecution, pub compaction_region: CompactionRegion, /// Request sender to notify the worker. pub(crate) request_sender: mpsc::Sender, @@ -80,20 +82,7 @@ impl Debug for CompactionTaskImpl { } } -impl Drop for CompactionTaskImpl { - fn drop(&mut self) { - self.mark_files_compacting(false) - } -} - impl CompactionTaskImpl { - fn mark_files_compacting(&self, compacting: bool) { - self.picker_output - .outputs - .iter() - .for_each(|o| o.inputs.iter().for_each(|f| f.set_compacting(compacting))); - } - /// Acquires memory budget based on the configured policy. /// /// Returns an error if memory cannot be acquired according to the policy. @@ -301,6 +290,7 @@ impl CompactionTask for CompactionTaskImpl { self.on_failure(err.clone()); let notify = BackgroundNotify::CompactionFailed(CompactionFailed { region_id: self.compaction_region.region_id, + execution: self.execution.clone(), err, }); self.send_to_worker(WorkerRequest::Background { @@ -312,8 +302,6 @@ impl CompactionTask for CompactionTaskImpl { } }; - // Marks files compacting before compaction and unmark after compaction (even if compaction is cancelled or failed), so that they won't be picked by other compaction tasks. - self.mark_files_compacting(true); self.handle_expiration().await; let cancel_handle = self.state.cancel_handle(); @@ -331,14 +319,19 @@ impl CompactionTask for CompactionTaskImpl { let senders = std::mem::take(&mut self.waiters); BackgroundNotify::CompactionCancelled(CompactionCancelled { region_id: self.compaction_region.region_id, + execution: self.execution.clone(), senders, }) } else { + self.listener + .on_compaction_commit_begin(self.compaction_region.region_id) + .await; match self.update_manifest(merge_output).await { Ok((edit, _manifest_version)) => { let senders = std::mem::take(&mut self.waiters); BackgroundNotify::CompactionFinished(CompactionFinished { region_id: self.compaction_region.region_id, + execution: self.execution.clone(), senders, start_time: self.start_time, edit, @@ -350,6 +343,7 @@ impl CompactionTask for CompactionTaskImpl { self.on_failure(err.clone()); BackgroundNotify::CompactionFailed(CompactionFailed { region_id: self.compaction_region.region_id, + execution: self.execution.clone(), err, }) } @@ -364,6 +358,7 @@ impl CompactionTask for CompactionTaskImpl { let senders = std::mem::take(&mut self.waiters); BackgroundNotify::CompactionCancelled(CompactionCancelled { region_id: self.compaction_region.region_id, + execution: self.execution.clone(), senders, }) } @@ -374,6 +369,7 @@ impl CompactionTask for CompactionTaskImpl { self.on_failure(err.clone()); BackgroundNotify::CompactionFailed(CompactionFailed { region_id: self.compaction_region.region_id, + execution: self.execution.clone(), err, }) } diff --git a/src/mito2/src/compaction/twcs.rs b/src/mito2/src/compaction/twcs.rs index cfa44d5045..dd467a557a 100644 --- a/src/mito2/src/compaction/twcs.rs +++ b/src/mito2/src/compaction/twcs.rs @@ -13,7 +13,7 @@ // limitations under the License. use std::collections::hash_map::Entry; -use std::collections::{BTreeMap, HashMap}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::fmt::Debug; use std::num::NonZeroU64; @@ -246,9 +246,11 @@ impl Picker for TwcsPicker { get_expired_ssts(levels, compaction_region.ttl, Timestamp::current_millis()); if !expired_ssts.is_empty() { info!("Expired SSTs in region {}: {:?}", region_id, expired_ssts); - // here we mark expired SSTs as compacting to avoid them being picked. - expired_ssts.iter().for_each(|f| f.set_compacting(true)); } + let expired_file_ids = expired_ssts + .iter() + .map(|file| file.file_id()) + .collect::>(); let compaction_time_window = compaction_region .current_version @@ -268,8 +270,13 @@ impl Picker for TwcsPicker { // Find active window from files in level 0. let active_window = find_latest_window_in_seconds(levels[0].files(), time_window_size); // Assign files to windows - let mut windows = - assign_to_windows(levels.iter().flat_map(LevelMeta::files), time_window_size); + let mut windows = assign_to_windows( + levels + .iter() + .flat_map(LevelMeta::files) + .filter(|file| !expired_file_ids.contains(&file.file_id())), + time_window_size, + ); let outputs = self.build_output(region_id, &mut windows, active_window); if outputs.is_empty() && expired_ssts.is_empty() { @@ -411,16 +418,85 @@ fn find_latest_window_in_seconds<'a>( #[cfg(test)] mod tests { use std::collections::HashSet; + use std::num::NonZeroU64; + use std::sync::Arc; + use std::time::Duration; use bytes::Bytes; + use common_base::Plugins; use store_api::storage::FileId; use super::*; + use crate::cache::CacheManager; + use crate::compaction::compactor::CompactionVersion; use crate::compaction::test_util::{ new_file_handle, new_file_handle_with_sequence, new_file_handle_with_size_and_sequence, new_file_handle_with_size_sequence_and_primary_key_range, }; - use crate::sst::file::Level; + use crate::config::MitoConfig; + use crate::region::options::RegionOptions; + use crate::sst::file::{FileMeta, Level}; + use crate::sst::version::SstVersion; + use crate::test_util::memtable_util::metadata_for_test; + use crate::test_util::scheduler_util::SchedulerEnv; + + async fn compaction_region_with_expired_sst() -> CompactionRegion { + let env = SchedulerEnv::new().await; + let metadata = metadata_for_test(); + let manifest_ctx = env.mock_manifest_context(metadata.clone()).await; + let mut ssts = SstVersion::new(); + ssts.add_files( + Arc::new(crate::sst::file_purger::NoopFilePurger), + (1..=4).map(|sequence| FileMeta { + file_id: FileId::random(), + time_range: ( + Timestamp::new_millisecond(0), + Timestamp::new_millisecond(10), + ), + level: 0, + sequence: NonZeroU64::new(sequence), + ..Default::default() + }), + ); + + CompactionRegion { + region_id: metadata.region_id, + region_options: RegionOptions::default(), + engine_config: Arc::new(MitoConfig::default()), + region_metadata: metadata.clone(), + cache_manager: Arc::new(CacheManager::default()), + access_layer: env.access_layer, + manifest_ctx, + current_version: CompactionVersion { + metadata, + options: RegionOptions::default(), + ssts: Arc::new(ssts), + compaction_time_window: None, + }, + file_purger: None, + ttl: Some(Duration::from_millis(1).into()), + max_parallelism: 1, + plugins: Plugins::new(), + } + } + + #[tokio::test] + async fn test_pick_expired_ssts_without_marking_compacting() { + let picker = TwcsPicker { + trigger_file_num: 4, + time_window_seconds: Some(3), + max_output_file_size: None, + append_mode: false, + max_background_tasks: None, + }; + let compaction_region = compaction_region_with_expired_sst().await; + + let output = picker.pick(&compaction_region).unwrap(); + + assert!(output.outputs.is_empty()); + assert!(!output.expired_ssts.is_empty()); + assert!(output.expired_ssts.iter().all(|file| !file.compacting())); + } #[test] fn test_get_latest_window_in_seconds() { diff --git a/src/mito2/src/compaction/window.rs b/src/mito2/src/compaction/window.rs index f802601df3..a7b720c1a4 100644 --- a/src/mito2/src/compaction/window.rs +++ b/src/mito2/src/compaction/window.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashSet}; use std::fmt::Debug; use common_telemetry::info; @@ -86,9 +86,11 @@ impl WindowedCompactionPicker { ); if !expired_ssts.is_empty() { info!("Expired SSTs in region {}: {:?}", region_id, expired_ssts); - // here we mark expired SSTs as compacting to avoid them being picked. - expired_ssts.iter().for_each(|f| f.set_compacting(true)); } + let expired_file_ids = expired_ssts + .iter() + .map(|file| file.file_id()) + .collect::>(); let windows = assign_files_to_time_windows( time_window, @@ -96,7 +98,8 @@ impl WindowedCompactionPicker { .ssts .levels() .iter() - .flat_map(|level| level.files.values()), + .flat_map(|level| level.files.values()) + .filter(|file| !expired_file_ids.contains(&file.file_id())), ); (build_output(windows), expired_ssts, time_window) @@ -204,17 +207,22 @@ mod tests { use std::sync::Arc; use std::time::Duration; + use common_base::Plugins; use common_time::Timestamp; use common_time::range::TimestampRange; use store_api::storage::{FileId, RegionId}; - use crate::compaction::compactor::CompactionVersion; + use crate::cache::CacheManager; + use crate::compaction::compactor::{CompactionRegion, CompactionVersion}; + use crate::compaction::picker::Picker; use crate::compaction::window::{WindowedCompactionPicker, file_time_bucket_span}; + use crate::config::MitoConfig; use crate::region::options::RegionOptions; use crate::sst::file::{FileMeta, Level}; use crate::sst::file_purger::NoopFilePurger; use crate::sst::version::SstVersion; use crate::test_util::memtable_util::metadata_for_test; + use crate::test_util::scheduler_util::SchedulerEnv; fn build_version( files: &[(FileId, i64, i64, Level)], @@ -260,19 +268,33 @@ mod tests { } } - #[test] - fn test_pick_expired() { + #[tokio::test] + async fn test_pick_expired_ssts_without_marking_compacting() { let picker = WindowedCompactionPicker::new(None); let files = vec![(FileId::random(), 0, 10, 0)]; - let version = build_version(&files, Some(Duration::from_millis(1))); - let (outputs, expired_ssts, _window) = picker.pick_inner( - RegionId::new(0, 0), - &version, - Timestamp::new_millisecond(12), - ); - assert!(outputs.is_empty()); - assert_eq!(1, expired_ssts.len()); + let env = SchedulerEnv::new().await; + let manifest_ctx = env.mock_manifest_context(version.metadata.clone()).await; + let compaction_region = CompactionRegion { + region_id: version.metadata.region_id, + region_options: RegionOptions::default(), + engine_config: Arc::new(MitoConfig::default()), + region_metadata: version.metadata.clone(), + cache_manager: Arc::new(CacheManager::default()), + access_layer: env.access_layer, + manifest_ctx, + current_version: version, + file_purger: None, + ttl: None, + max_parallelism: 1, + plugins: Plugins::new(), + }; + + let output = picker.pick(&compaction_region).unwrap(); + + assert!(output.outputs.is_empty()); + assert!(!output.expired_ssts.is_empty()); + assert!(output.expired_ssts.iter().all(|file| !file.compacting())); } const HOUR: i64 = 60 * 60 * 1000; diff --git a/src/mito2/src/engine/compaction_test.rs b/src/mito2/src/engine/compaction_test.rs index fd0982b7e5..d5cbc94aec 100644 --- a/src/mito2/src/engine/compaction_test.rs +++ b/src/mito2/src/engine/compaction_test.rs @@ -17,6 +17,7 @@ use std::ops::Range; use std::sync::Arc; use std::time::Duration; +use api::v1::region::{StrictWindow, compact_request}; use api::v1::{ColumnSchema, Rows}; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; @@ -26,8 +27,8 @@ use datatypes::arrow::datatypes::TimestampMillisecondType; use store_api::region_engine::{RegionEngine, RegionRole}; use store_api::region_request::AlterKind::SetRegionOptions; use store_api::region_request::{ - EnterStagingRequest, PathType, RegionAlterRequest, RegionCompactRequest, RegionDeleteRequest, - RegionFlushRequest, RegionOpenRequest, RegionRequest, SetRegionOption, + EnterStagingRequest, PathType, RegionAlterRequest, RegionCloseRequest, RegionCompactRequest, + RegionDeleteRequest, RegionFlushRequest, RegionOpenRequest, RegionRequest, SetRegionOption, StagingPartitionDirective, }; use store_api::storage::{RegionId, ScanRequest}; @@ -35,7 +36,7 @@ use tokio::sync::Notify; use crate::config::MitoConfig; use crate::engine::MitoEngine; -use crate::engine::listener::CompactionListener; +use crate::engine::listener::{CompactionListener, CompactionPlanningGate}; use crate::test_util::{ CreateRequestBuilder, TestEnv, build_rows_for_key, column_metadata_to_column_schema, put_rows, }; @@ -132,6 +133,656 @@ async fn collect_stream_ts(stream: SendableRecordBatchStream) -> Vec { res } +struct CompactionListenerGuard(Option>); + +impl CompactionListenerGuard { + fn new(listener: Arc) -> Self { + Self(Some(listener)) + } + + fn release(mut self) { + self.0.take().unwrap().wake(); + } +} + +impl Drop for CompactionListenerGuard { + fn drop(&mut self) { + if let Some(listener) = self.0.take() { + listener.wake(); + } + } +} + +#[tokio::test] +async fn test_region_b_progresses_while_same_worker_region_a_is_picking() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::new().await; + let region_a = RegionId::new(1, 1); + let region_b = RegionId::new(2, 1); + let gate = Arc::new(CompactionPlanningGate::new(region_a)); + let engine = env + .create_engine_with( + MitoConfig { + num_workers: 1, + min_compaction_interval: Duration::ZERO, + ..Default::default() + }, + None, + Some(gate.clone()), + None, + ) + .await; + + for (region_id, table_name) in [(region_a, "region_a"), (region_b, "region_b")] { + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + table_name, + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Create( + CreateRequestBuilder::new() + .insert_option("compaction.type", "twcs") + .build(), + ), + ) + .await + .unwrap(); + } + + let request = CreateRequestBuilder::new().build(); + let column_schemas = request + .column_metadatas + .iter() + .map(column_metadata_to_column_schema) + .collect::>(); + let gate_guard = gate.arm(); + let engine_for_compaction = engine.clone(); + let region_a_compaction = tokio::spawn(async move { + engine_for_compaction + .handle_request( + region_a, + RegionRequest::Compact(RegionCompactRequest::default()), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered()) + .await + .expect("region A planning did not reach the gate"); + + tokio::time::timeout( + Duration::from_secs(5), + put_and_flush(&engine, region_a, &column_schemas, 10..20), + ) + .await + .expect("region A automatic compaction trigger did not finish"); + tokio::time::timeout(Duration::from_secs(5), gate.wait_until_schedule_attempts(2)) + .await + .expect("region A automatic trigger did not reach the scheduler"); + assert_eq!(2, gate.schedule_attempt_count()); + assert_eq!(1, gate.invocation_count()); + + let engine_for_region_b = engine.clone(); + let mut region_b_work = tokio::spawn(async move { + put_and_flush(&engine_for_region_b, region_b, &column_schemas, 0..10).await; + }); + tokio::time::timeout(Duration::from_secs(5), &mut region_b_work) + .await + .expect("region B was blocked by region A compaction planning") + .expect("region B work task panicked"); + let followup_guard = gate.arm(); + gate_guard.release(); + tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered()) + .await + .expect("coalesced region A trigger did not start its follow-up plan"); + tokio::time::timeout(Duration::from_secs(5), region_a_compaction) + .await + .expect("region A compaction task did not finish after gate release") + .expect("region A compaction task panicked") + .expect("region A compaction failed"); + assert_eq!(2, gate.invocation_count()); + followup_guard.release(); +} + +#[tokio::test] +async fn test_regular_trigger_while_picking_replans_after_no_plan() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::new().await; + let region_id = RegionId::new(7, 1); + let gate = Arc::new(CompactionPlanningGate::new(region_id)); + let engine = env + .create_engine_with( + MitoConfig { + min_compaction_interval: Duration::ZERO, + ..Default::default() + }, + None, + Some(gate.clone()), + None, + ) + .await; + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "replan_after_no_plan", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + let create = CreateRequestBuilder::new() + .insert_option("compaction.type", "twcs") + .build(); + let column_schemas = create + .column_metadatas + .iter() + .map(column_metadata_to_column_schema) + .collect::>(); + engine + .handle_request(region_id, RegionRequest::Create(create)) + .await + .unwrap(); + + let first_plan_guard = gate.arm(); + put_and_flush(&engine, region_id, &column_schemas, 0..10).await; + tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered()) + .await + .expect("first automatic compaction did not reach the planning gate"); + + put_and_flush(&engine, region_id, &column_schemas, 5..20).await; + tokio::time::timeout(Duration::from_secs(5), gate.wait_until_schedule_attempts(2)) + .await + .expect("second flush did not trigger automatic compaction"); + assert_eq!(1, gate.invocation_count()); + + let second_plan_guard = gate.arm(); + first_plan_guard.release(); + tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered()) + .await + .expect("coalesced regular trigger was lost after the first plan returned no plan"); + assert_eq!(2, gate.invocation_count()); + second_plan_guard.release(); +} + +#[tokio::test] +async fn test_regular_trigger_while_picking_replans_after_prepared_execution() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::new().await; + let region_id = RegionId::new(8, 1); + let gate = Arc::new(CompactionPlanningGate::new(region_id)); + let engine = env + .create_engine_with( + MitoConfig { + min_compaction_interval: Duration::from_secs(60 * 60), + ..Default::default() + }, + None, + Some(gate.clone()), + None, + ) + .await; + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "replan_after_prepared", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + let create = CreateRequestBuilder::new() + .insert_option("compaction.type", "twcs") + .build(); + let column_schemas = create + .column_metadatas + .iter() + .map(column_metadata_to_column_schema) + .collect::>(); + engine + .handle_request(region_id, RegionRequest::Create(create)) + .await + .unwrap(); + put_and_flush(&engine, region_id, &column_schemas, 0..10).await; + put_and_flush(&engine, region_id, &column_schemas, 5..20).await; + + let first_plan_guard = gate.arm(); + let first_engine = engine.clone(); + let first = tokio::spawn(async move { + first_engine + .handle_request( + region_id, + RegionRequest::Compact(RegionCompactRequest::default()), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered()) + .await + .expect("first regular compaction did not reach the planning gate"); + + let expected_attempts = gate.schedule_attempt_count() + 1; + let second_engine = engine.clone(); + let second = tokio::spawn(async move { + second_engine + .handle_request( + region_id, + RegionRequest::Compact(RegionCompactRequest::default()), + ) + .await + }); + tokio::time::timeout( + Duration::from_secs(5), + gate.wait_until_schedule_attempts(expected_attempts), + ) + .await + .expect("second regular compaction did not reach the scheduler"); + + let commit_guard = gate.arm_commit(); + first_plan_guard.release(); + tokio::time::timeout(Duration::from_secs(5), gate.wait_until_commit_entered()) + .await + .expect("first regular compaction did not produce a prepared execution"); + + let followup_plan_guard = gate.arm(); + commit_guard.release(); + tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered()) + .await + .expect("prepared execution did not immediately admit the retained regular follow-up"); + tokio::time::timeout(Duration::from_secs(5), first) + .await + .expect("first regular compaction waiter was not notified") + .expect("first regular compaction task panicked") + .expect("first regular compaction failed"); + assert!(!second.is_finished()); + + followup_plan_guard.release(); + tokio::time::timeout(Duration::from_secs(5), second) + .await + .expect("retained regular compaction waiter was not notified") + .expect("retained regular compaction task panicked") + .expect("retained regular compaction failed"); +} + +#[tokio::test] +async fn test_pending_manual_compaction_finishes_before_queued_ddl() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::new().await; + let region_id = RegionId::new(6, 1); + let gate = Arc::new(CompactionPlanningGate::new(region_id)); + let engine = env + .create_engine_with( + MitoConfig { + num_workers: 1, + min_compaction_interval: Duration::from_secs(60 * 60), + ..Default::default() + }, + None, + Some(gate.clone()), + None, + ) + .await; + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "pending_manual_before_ddl", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + let create = CreateRequestBuilder::new() + .insert_option("compaction.type", "twcs") + .build(); + let column_schemas = create + .column_metadatas + .iter() + .map(column_metadata_to_column_schema) + .collect::>(); + engine + .handle_request(region_id, RegionRequest::Create(create)) + .await + .unwrap(); + put_and_flush(&engine, region_id, &column_schemas, 0..10).await; + put_and_flush(&engine, region_id, &column_schemas, 5..20).await; + + let commit_guard = gate.arm_commit(); + let regular_engine = engine.clone(); + let regular_task = tokio::spawn(async move { + regular_engine + .handle_request( + region_id, + RegionRequest::Compact(RegionCompactRequest::default()), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(5), gate.wait_until_commit_entered()) + .await + .expect("regular compaction did not reach its non-cancellable commit gate"); + + let manual_plan_guard = gate.arm(); + let manual_engine = engine.clone(); + let manual_task = tokio::spawn(async move { + manual_engine + .handle_request( + region_id, + RegionRequest::Compact(RegionCompactRequest { + options: compact_request::Options::StrictWindow(StrictWindow { + window_seconds: 60, + }), + ..Default::default() + }), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(5), gate.wait_until_schedule_attempts(2)) + .await + .expect("manual compaction did not reach the scheduler"); + assert!(!manual_task.is_finished()); + + let ddl_engine = engine.clone(); + let ddl_task = tokio::spawn(async move { + ddl_engine + .handle_request( + region_id, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive: StagingPartitionDirective::RejectAllWrites, + }), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(5), gate.wait_until_cancel_requested()) + .await + .expect("enter-staging DDL was not queued behind regular compaction"); + assert!(!ddl_task.is_finished()); + + commit_guard.release(); + tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered()) + .await + .expect("pending manual compaction was not planned after regular completion"); + assert_eq!(2, gate.invocation_count()); + assert!(!manual_task.is_finished()); + assert!(!ddl_task.is_finished()); + + let pending_ddl_guard = gate.arm_pending_ddl_dispatch(); + manual_plan_guard.release(); + tokio::time::timeout( + Duration::from_secs(5), + gate.wait_until_pending_ddl_dispatch(), + ) + .await + .expect("manual result was not notified before pending DDL dispatch"); + tokio::time::timeout(Duration::from_secs(5), manual_task) + .await + .expect("manual compaction did not notify its waiter") + .expect("manual compaction task panicked") + .expect("manual compaction failed"); + assert!(!ddl_task.is_finished()); + + pending_ddl_guard.release(); + tokio::time::timeout(Duration::from_secs(5), regular_task) + .await + .expect("regular compaction did not finish after commit release") + .expect("regular compaction task panicked") + .expect("regular compaction failed"); + tokio::time::timeout(Duration::from_secs(5), ddl_task) + .await + .expect("queued DDL did not finish after manual compaction") + .expect("queued DDL task panicked") + .expect("queued DDL failed"); + assert!(engine.get_region(region_id).unwrap().is_staging()); +} + +#[tokio::test] +async fn test_picking_close_reopen_ignores_old_plan() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::new().await; + let region_id = RegionId::new(3, 1); + let gate = Arc::new(CompactionPlanningGate::new(region_id)); + let engine = env + .create_engine_with( + MitoConfig { + num_workers: 1, + ..Default::default() + }, + None, + Some(gate.clone()), + None, + ) + .await; + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "close_reopen", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + let create = CreateRequestBuilder::new() + .insert_option("compaction.type", "twcs") + .build(); + let table_dir = create.table_dir.clone(); + let options = create.options.clone(); + engine + .handle_request(region_id, RegionRequest::Create(create)) + .await + .unwrap(); + + let gate_guard = gate.arm(); + let compact_engine = engine.clone(); + let compact_task = tokio::spawn(async move { + compact_engine + .handle_request( + region_id, + RegionRequest::Compact(RegionCompactRequest::default()), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered()) + .await + .expect("planning did not reach the gate"); + + engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + let compact_err = tokio::time::timeout(Duration::from_secs(5), compact_task) + .await + .expect("closed region compaction waiter was not released") + .expect("closed region compaction task panicked") + .unwrap_err(); + assert_eq!(compact_err.status_code(), StatusCode::Cancelled); + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir, + path_type: PathType::Bare, + options, + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); + + gate_guard.release(); + tokio::time::timeout(Duration::from_secs(5), compact(&engine, region_id)) + .await + .expect("replacement compaction was blocked by the stale plan"); + assert!(engine.is_region_exists(region_id)); + assert_eq!(2, gate.invocation_count()); +} + +#[tokio::test] +async fn test_enter_staging_waits_for_picking_logical_cancellation_ack() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::new().await; + let region_id = RegionId::new(4, 1); + let gate = Arc::new(CompactionPlanningGate::new(region_id)); + let engine = env + .create_engine_with( + MitoConfig { + num_workers: 1, + ..Default::default() + }, + None, + Some(gate.clone()), + None, + ) + .await; + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "enter_staging", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Create( + CreateRequestBuilder::new() + .insert_option("compaction.type", "twcs") + .build(), + ), + ) + .await + .unwrap(); + + let gate_guard = gate.arm(); + let compact_engine = engine.clone(); + let compact_task = tokio::spawn(async move { + compact_engine + .handle_request( + region_id, + RegionRequest::Compact(RegionCompactRequest::default()), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered()) + .await + .expect("planning did not reach the gate"); + let staging_engine = engine.clone(); + let staging_task = tokio::spawn(async move { + staging_engine + .handle_request( + region_id, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive: StagingPartitionDirective::RejectAllWrites, + }), + ) + .await + }); + + tokio::time::timeout(Duration::from_secs(5), gate.wait_until_cancel_requested()) + .await + .expect("enter-staging did not request picking cancellation"); + assert!(!compact_task.is_finished()); + assert!(!staging_task.is_finished()); + + gate_guard.release(); + let compact_err = tokio::time::timeout(Duration::from_secs(5), compact_task) + .await + .expect("cancelled compaction waiter was not released") + .expect("cancelled compaction task panicked") + .unwrap_err(); + assert_eq!(compact_err.status_code(), StatusCode::Cancelled); + tokio::time::timeout(Duration::from_secs(5), staging_task) + .await + .expect("enter-staging did not finish after cancellation acknowledgment") + .expect("enter-staging task panicked") + .expect("enter-staging request failed"); + assert!(engine.get_region(region_id).unwrap().is_staging()); +} + +#[tokio::test] +async fn test_worker_shutdown_fails_picking_waiter() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::new().await; + let region_id = RegionId::new(5, 1); + let gate = Arc::new(CompactionPlanningGate::new(region_id)); + let engine = env + .create_engine_with( + MitoConfig { + num_workers: 1, + ..Default::default() + }, + None, + Some(gate.clone()), + None, + ) + .await; + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "worker_shutdown", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Create( + CreateRequestBuilder::new() + .insert_option("compaction.type", "twcs") + .build(), + ), + ) + .await + .unwrap(); + + let gate_guard = gate.arm(); + let compact_engine = engine.clone(); + let compact_task = tokio::spawn(async move { + compact_engine + .handle_request( + region_id, + RegionRequest::Compact(RegionCompactRequest::default()), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered()) + .await + .expect("planning did not reach the gate"); + + tokio::time::timeout(Duration::from_secs(5), engine.stop()) + .await + .expect("worker shutdown blocked on picking") + .unwrap(); + let compact_err = tokio::time::timeout(Duration::from_secs(5), compact_task) + .await + .expect("worker shutdown did not release the compaction waiter") + .expect("compaction task panicked during worker shutdown") + .unwrap_err(); + assert_eq!(compact_err.status_code(), StatusCode::Cancelled); + gate_guard.release(); +} + #[tokio::test] async fn test_compaction_region() { test_compaction_region_with_format(false).await; @@ -603,6 +1254,7 @@ async fn test_readonly_during_compaction_with_format(flat_format: bool) { .handle_request(region_id, RegionRequest::Create(request)) .await .unwrap(); + let listener_guard = CompactionListenerGuard::new(listener.clone()); // Flush 2 SSTs for compaction. put_and_flush(&engine, region_id, &column_schemas, 0..10).await; put_and_flush(&engine, region_id, &column_schemas, 5..20).await; @@ -615,7 +1267,7 @@ async fn test_readonly_during_compaction_with_format(flat_format: bool) { .set_region_role(region_id, RegionRole::Follower) .unwrap(); // Wakes up the listener. - listener.wake(); + listener_guard.release(); let notify = Arc::new(Notify::new()); // We already sets max background purges to 1, so we can submit a task to the @@ -686,28 +1338,28 @@ async fn test_enter_staging_cancels_inflight_local_compaction_before_commit() { .handle_request(region_id, RegionRequest::Create(request)) .await .unwrap(); + let _listener_guard = CompactionListenerGuard::new(listener.clone()); put_and_flush(&engine, region_id, &column_schemas, 0..10).await; put_and_flush(&engine, region_id, &column_schemas, 5..20).await; - listener.wait_handle_finished().await; + tokio::time::timeout(Duration::from_secs(5), listener.wait_handle_finished()) + .await + .expect("local compaction did not reach its pre-commit gate"); - let engine_cloned = engine.clone(); - let enter_staging = tokio::spawn(async move { - engine_cloned - .handle_request( - region_id, - RegionRequest::EnterStaging(EnterStagingRequest { - partition_directive: StagingPartitionDirective::RejectAllWrites, - }), - ) - .await - }); - - tokio::time::sleep(Duration::from_millis(100)).await; - // The enter staging should finished, and the compaction should be cancelled. - assert!(enter_staging.is_finished()); - let _ = enter_staging.await.unwrap().unwrap(); + tokio::time::timeout( + Duration::from_secs(5), + engine.handle_request( + region_id, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive: StagingPartitionDirective::RejectAllWrites, + }), + ), + ) + .await + .expect("enter-staging waited for the blocked local compaction") + .expect("enter-staging request failed"); + assert!(engine.get_region(region_id).unwrap().is_staging()); } #[tokio::test] @@ -751,6 +1403,7 @@ async fn test_manual_compaction_returns_cancelled_when_enter_staging_cancels_it( .handle_request(region_id, RegionRequest::Create(request)) .await .unwrap(); + let _listener_guard = CompactionListenerGuard::new(listener.clone()); put_and_flush(&engine, region_id, &column_schemas, 0..10).await; put_and_flush(&engine, region_id, &column_schemas, 5..20).await; @@ -765,28 +1418,29 @@ async fn test_manual_compaction_returns_cancelled_when_enter_staging_cancels_it( .await }); - listener.wait_handle_finished().await; + tokio::time::timeout(Duration::from_secs(5), listener.wait_handle_finished()) + .await + .expect("manual compaction did not reach its pre-commit gate"); - let engine_cloned = engine.clone(); - let enter_staging = tokio::spawn(async move { - engine_cloned - .handle_request( - region_id, - RegionRequest::EnterStaging(EnterStagingRequest { - partition_directive: StagingPartitionDirective::RejectAllWrites, - }), - ) - .await - }); + tokio::time::timeout( + Duration::from_secs(5), + engine.handle_request( + region_id, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive: StagingPartitionDirective::RejectAllWrites, + }), + ), + ) + .await + .expect("enter-staging waited for the blocked manual compaction") + .expect("enter-staging request failed"); - tokio::time::sleep(Duration::from_millis(100)).await; - assert!(compact.is_finished()); - assert!(enter_staging.is_finished()); - - let err = compact.await.unwrap().unwrap_err(); + let err = tokio::time::timeout(Duration::from_secs(5), compact) + .await + .expect("cancelled manual compaction waiter was not released") + .expect("manual compaction task panicked") + .unwrap_err(); assert_eq!(err.status_code(), StatusCode::Cancelled); - - let _ = enter_staging.await.unwrap(); } #[tokio::test] diff --git a/src/mito2/src/engine/listener.rs b/src/mito2/src/engine/listener.rs index 224268d5a6..818d60b7cb 100644 --- a/src/mito2/src/engine/listener.rs +++ b/src/mito2/src/engine/listener.rs @@ -15,7 +15,7 @@ //! Engine event listener for tests. use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::time::Duration; use async_trait::async_trait; @@ -71,6 +71,21 @@ pub trait EventListener: Send + Sync { /// Notifies the listener that the compaction is scheduled. fn on_compaction_scheduled(&self, _region_id: RegionId) {} + /// Notifies the listener immediately before a worker asks the scheduler for compaction. + fn on_compaction_schedule_attempt(&self, _region_id: RegionId) {} + + /// Notifies the listener immediately before compaction planning invokes the picker. + async fn on_compaction_pick_begin(&self, _region_id: RegionId) {} + + /// Notifies the listener after local compaction becomes non-cancellable and before commit. + async fn on_compaction_commit_begin(&self, _region_id: RegionId) {} + + /// Notifies the listener after compaction results are sent and before pending DDL dispatch. + async fn on_compaction_result_notified(&self, _region_id: RegionId) {} + + /// Notifies the listener after compaction cancellation is requested and its DDL is queued. + fn on_compaction_cancel_requested(&self, _region_id: RegionId) {} + /// Notifies the listener that region starts to send a region change result to worker. async fn on_notify_region_change_result_begin(&self, _region_id: RegionId) {} @@ -89,6 +104,213 @@ pub trait EventListener: Send + Sync { pub type EventListenerRef = Arc; +/// Test gate that blocks compaction planning for one region before picker execution. +pub struct CompactionPlanningGate { + region_id: RegionId, + armed: AtomicBool, + entered: Notify, + cancel_requested: Notify, + permits: Semaphore, + invocation_count: AtomicUsize, + schedule_attempted: Notify, + schedule_attempt_count: AtomicUsize, + commit_armed: AtomicBool, + commit_entered: Notify, + commit_permits: Semaphore, + pending_ddl_armed: AtomicBool, + pending_ddl_entered: Notify, + pending_ddl_permits: Semaphore, +} + +/// Releases an armed [`CompactionPlanningGate`] when a test exits unexpectedly. +pub struct CompactionPlanningGateGuard { + gate: Option>, +} + +impl CompactionPlanningGateGuard { + pub fn release(mut self) { + self.gate.take().unwrap().release(); + } +} + +impl Drop for CompactionPlanningGateGuard { + fn drop(&mut self) { + if let Some(gate) = self.gate.take() { + gate.release(); + } + } +} + +/// Releases an armed commit gate when a test exits unexpectedly. +pub struct CompactionCommitGateGuard { + gate: Option>, +} + +impl CompactionCommitGateGuard { + pub fn release(mut self) { + self.gate.take().unwrap().release_commit(); + } +} + +impl Drop for CompactionCommitGateGuard { + fn drop(&mut self) { + if let Some(gate) = self.gate.take() { + gate.release_commit(); + } + } +} + +/// Releases an armed pending-DDL dispatch gate when a test exits unexpectedly. +pub struct CompactionPendingDdlGateGuard { + gate: Option>, +} + +impl CompactionPendingDdlGateGuard { + pub fn release(mut self) { + self.gate.take().unwrap().release_pending_ddl_dispatch(); + } +} + +impl Drop for CompactionPendingDdlGateGuard { + fn drop(&mut self) { + if let Some(gate) = self.gate.take() { + gate.release_pending_ddl_dispatch(); + } + } +} + +impl CompactionPlanningGate { + pub fn new(region_id: RegionId) -> Self { + Self { + region_id, + armed: AtomicBool::new(false), + entered: Notify::new(), + cancel_requested: Notify::new(), + permits: Semaphore::new(0), + invocation_count: AtomicUsize::new(0), + schedule_attempted: Notify::new(), + schedule_attempt_count: AtomicUsize::new(0), + commit_armed: AtomicBool::new(false), + commit_entered: Notify::new(), + commit_permits: Semaphore::new(0), + pending_ddl_armed: AtomicBool::new(false), + pending_ddl_entered: Notify::new(), + pending_ddl_permits: Semaphore::new(0), + } + } + + pub fn arm(self: &Arc) -> CompactionPlanningGateGuard { + self.armed.store(true, Ordering::Relaxed); + CompactionPlanningGateGuard { + gate: Some(self.clone()), + } + } + + pub async fn wait_until_entered(&self) { + self.entered.notified().await; + } + + pub async fn wait_until_cancel_requested(&self) { + self.cancel_requested.notified().await; + } + + pub async fn wait_until_schedule_attempts(&self, expected: usize) { + while self.schedule_attempt_count() < expected { + self.schedule_attempted.notified().await; + } + } + + pub fn schedule_attempt_count(&self) -> usize { + self.schedule_attempt_count.load(Ordering::Relaxed) + } + + pub fn arm_commit(self: &Arc) -> CompactionCommitGateGuard { + self.commit_armed.store(true, Ordering::Relaxed); + CompactionCommitGateGuard { + gate: Some(self.clone()), + } + } + + pub async fn wait_until_commit_entered(&self) { + self.commit_entered.notified().await; + } + + pub fn arm_pending_ddl_dispatch(self: &Arc) -> CompactionPendingDdlGateGuard { + self.pending_ddl_armed.store(true, Ordering::Relaxed); + CompactionPendingDdlGateGuard { + gate: Some(self.clone()), + } + } + + pub async fn wait_until_pending_ddl_dispatch(&self) { + self.pending_ddl_entered.notified().await; + } + + pub fn release(&self) { + self.permits.add_permits(1); + } + + pub fn invocation_count(&self) -> usize { + self.invocation_count.load(Ordering::Relaxed) + } + + fn release_commit(&self) { + self.commit_permits.add_permits(1); + } + + fn release_pending_ddl_dispatch(&self) { + self.pending_ddl_permits.add_permits(1); + } +} + +#[async_trait] +impl EventListener for CompactionPlanningGate { + fn on_compaction_schedule_attempt(&self, region_id: RegionId) { + if region_id == self.region_id { + self.schedule_attempt_count.fetch_add(1, Ordering::Relaxed); + self.schedule_attempted.notify_one(); + } + } + + async fn on_compaction_pick_begin(&self, region_id: RegionId) { + if region_id != self.region_id { + return; + } + + self.invocation_count.fetch_add(1, Ordering::Relaxed); + if !self.armed.swap(false, Ordering::Relaxed) { + return; + } + + self.entered.notify_one(); + self.permits.acquire().await.unwrap().forget(); + } + + async fn on_compaction_commit_begin(&self, region_id: RegionId) { + if region_id != self.region_id || !self.commit_armed.swap(false, Ordering::Relaxed) { + return; + } + + self.commit_entered.notify_one(); + self.commit_permits.acquire().await.unwrap().forget(); + } + + async fn on_compaction_result_notified(&self, region_id: RegionId) { + if region_id != self.region_id || !self.pending_ddl_armed.swap(false, Ordering::Relaxed) { + return; + } + + self.pending_ddl_entered.notify_one(); + self.pending_ddl_permits.acquire().await.unwrap().forget(); + } + + fn on_compaction_cancel_requested(&self, region_id: RegionId) { + if region_id == self.region_id { + self.cancel_requested.notify_one(); + } + } +} + /// Listener to watch flush events. #[derive(Default)] pub struct FlushListener { @@ -566,3 +788,20 @@ impl EventListener for GateIndexBuildListener { self.stop_notify.notify_one(); } } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_compaction_planning_gate_counts_every_matching_callback() { + let region_id = RegionId::new(1, 1); + let gate = CompactionPlanningGate::new(region_id); + + gate.on_compaction_pick_begin(region_id).await; + gate.on_compaction_pick_begin(RegionId::new(2, 1)).await; + gate.on_compaction_pick_begin(region_id).await; + + assert_eq!(2, gate.invocation_count()); + } +} diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index 3fdd15db4f..2cfe007652 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -46,6 +46,7 @@ use store_api::region_request::{ use store_api::storage::{FileId, RegionId}; use tokio::sync::oneshot::{self, Receiver, Sender}; +use crate::compaction::{CompactionExecution, CompactionPickFinished}; use crate::error::{ CompactRegionSnafu, CompactionCancelledSnafu, ConvertColumnDataTypeSnafu, CreateDefaultSnafu, Error, FillDefaultSnafu, FlushRegionSnafu, InvalidPartitionExprSnafu, InvalidRequestSnafu, @@ -895,6 +896,8 @@ pub(crate) struct SenderDdlRequest { /// Notification from a background job. #[derive(Debug)] pub(crate) enum BackgroundNotify { + /// Compaction planning has finished. + CompactionPickFinished(CompactionPickFinished), /// Flush has finished. FlushFinished(FlushFinished), /// Flush has failed. @@ -997,6 +1000,8 @@ pub(crate) struct IndexBuildFailed { pub(crate) struct CompactionFinished { /// Region id. pub(crate) region_id: RegionId, + /// Identity and reservation lease of the accepted execution. + pub(crate) execution: CompactionExecution, /// Compaction result senders. pub(crate) senders: Vec, /// Start time of compaction task. @@ -1010,6 +1015,8 @@ pub(crate) struct CompactionFinished { pub(crate) struct CompactionCancelled { /// Region id. pub(crate) region_id: RegionId, + /// Identity and reservation lease of the accepted execution. + pub(crate) execution: CompactionExecution, /// Waiters to wake once the cancellation has been observed by the worker. pub(crate) senders: Vec, } @@ -1051,6 +1058,8 @@ impl OnFailure for CompactionFinished { #[derive(Debug)] pub(crate) struct CompactionFailed { pub(crate) region_id: RegionId, + /// Identity and reservation lease of the accepted execution. + pub(crate) execution: CompactionExecution, /// The error source of the failure. pub(crate) err: Arc, } @@ -1341,9 +1350,15 @@ mod tests { #[test] fn test_compaction_cancelled_sends_cancelled_error() { + let version_control = + Arc::new(crate::test_util::version_util::VersionControlBuilder::new().build()); let (tx, rx) = oneshot::channel(); let request = CompactionCancelled { region_id: RegionId::new(1, 1), + execution: crate::compaction::CompactionExecution::for_test( + version_control, + crate::compaction::CompactionExecutionKind::Local, + ), senders: vec![OutputTx::new(tx)], }; diff --git a/src/mito2/src/schedule/remote_job_scheduler.rs b/src/mito2/src/schedule/remote_job_scheduler.rs index 2eca2a248b..487b2c9365 100644 --- a/src/mito2/src/schedule/remote_job_scheduler.rs +++ b/src/mito2/src/schedule/remote_job_scheduler.rs @@ -13,7 +13,7 @@ // limitations under the License. use std::fmt; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::time::Instant; use common_telemetry::error; @@ -24,6 +24,7 @@ use store_api::storage::RegionId; use tokio::sync::mpsc::Sender; use uuid::Uuid; +use crate::compaction::CompactionExecution; use crate::compaction::compactor::CompactionRegion; use crate::compaction::picker::PickerOutput; use crate::error::{CompactRegionSnafu, Error, ParseJobIdSnafu, Result}; @@ -134,9 +135,20 @@ pub struct CompactionJobResult { pub(crate) struct DefaultNotifier { /// The sender to send WorkerRequest to the mito engine. This is used to notify the mito engine when a remote job is completed. pub(crate) request_sender: Sender, + execution: Mutex>, } impl DefaultNotifier { + pub(crate) fn new( + request_sender: Sender, + execution: CompactionExecution, + ) -> Self { + Self { + request_sender, + execution: Mutex::new(Some(execution)), + } + } + fn on_failure(&self, err: Arc, region_id: RegionId, mut waiters: Vec) { COMPACTION_FAILURE_COUNT.inc(); for waiter in waiters.drain(..) { @@ -149,12 +161,22 @@ impl DefaultNotifier { impl Notifier for DefaultNotifier { async fn notify(&self, result: RemoteJobResult, waiters: Vec) { INFLIGHT_COMPACTION_COUNT.dec(); + let Some(execution) = self + .execution + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .take() + else { + error!("Remote compaction notifier invoked more than once"); + return; + }; match result { RemoteJobResult::CompactionJobResult(result) => { let notify = { match result.region_edit { Ok(edit) => BackgroundNotify::CompactionFinished(CompactionFinished { region_id: result.region_id, + execution, senders: waiters, start_time: result.start_time, edit, @@ -168,6 +190,7 @@ impl Notifier for DefaultNotifier { self.on_failure(err.clone(), result.region_id, waiters); BackgroundNotify::CompactionFailed(CompactionFailed { region_id: result.region_id, + execution, err, }) } @@ -195,6 +218,9 @@ impl Notifier for DefaultNotifier { #[cfg(test)] mod tests { use super::*; + use crate::compaction::{CompactionExecution, CompactionExecutionKind}; + use crate::error::InvalidSchedulerStateSnafu; + use crate::test_util::version_util::VersionControlBuilder; #[test] fn test_job_id() { @@ -202,4 +228,80 @@ mod tests { let job_id = JobId::parse_str(&id).unwrap(); assert_eq!(job_id.to_string(), id); } + + #[tokio::test] + async fn test_default_notifier_carries_remote_execution_on_success() { + let (tx, mut rx) = tokio::sync::mpsc::channel(1); + let execution = CompactionExecution::for_test( + Arc::new(VersionControlBuilder::new().build()), + CompactionExecutionKind::Remote, + ); + let expected = execution.clone(); + let notifier = DefaultNotifier::new(tx, execution); + let region_id = RegionId::new(1, 1); + + notifier + .notify( + RemoteJobResult::CompactionJobResult(CompactionJobResult { + job_id: JobId::parse_str("00000000-0000-0000-0000-000000000003").unwrap(), + region_id, + start_time: Instant::now(), + region_edit: Ok(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, + }), + }), + Vec::new(), + ) + .await; + + let request = rx.recv().await.unwrap(); + let WorkerRequest::Background { + notify: BackgroundNotify::CompactionFinished(finished), + .. + } = request.request + else { + panic!("expected remote compaction success notification"); + }; + assert!(finished.execution.matches(&expected)); + } + + #[tokio::test] + async fn test_default_notifier_carries_remote_execution_on_failure() { + let (tx, mut rx) = tokio::sync::mpsc::channel(1); + let execution = CompactionExecution::for_test( + Arc::new(VersionControlBuilder::new().build()), + CompactionExecutionKind::Remote, + ); + let expected = execution.clone(); + let notifier = DefaultNotifier::new(tx, execution); + let region_id = RegionId::new(1, 1); + + notifier + .notify( + RemoteJobResult::CompactionJobResult(CompactionJobResult { + job_id: JobId::parse_str("00000000-0000-0000-0000-000000000004").unwrap(), + region_id, + start_time: Instant::now(), + region_edit: Err(InvalidSchedulerStateSnafu.build()), + }), + Vec::new(), + ) + .await; + + let request = rx.recv().await.unwrap(); + let WorkerRequest::Background { + notify: BackgroundNotify::CompactionFailed(failed), + .. + } = request.request + else { + panic!("expected remote compaction failure notification"); + }; + assert!(failed.execution.matches(&expected)); + } } diff --git a/src/mito2/src/sst/file.rs b/src/mito2/src/sst/file.rs index 213cd4f714..e226bd967a 100644 --- a/src/mito2/src/sst/file.rs +++ b/src/mito2/src/sst/file.rs @@ -538,6 +538,14 @@ impl FileHandle { self.inner.compacting.store(compacting, Ordering::Relaxed); } + /// Atomically marks this file as compacting if it is currently available. + pub fn try_set_compacting(&self) -> bool { + self.inner + .compacting + .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed) + .is_ok() + } + pub fn index_outdated(&self) -> bool { self.inner.index_outdated.load(Ordering::Relaxed) } @@ -830,6 +838,19 @@ mod tests { } } + #[test] + fn test_try_set_compacting() { + let file = FileHandle::new( + create_file_meta(FileId::random(), 0), + crate::test_util::new_noop_file_purger(), + ); + + assert!(file.try_set_compacting()); + assert!(!file.try_set_compacting()); + file.set_compacting(false); + assert!(file.try_set_compacting()); + } + #[test] fn test_deserialize_file_meta() { let file_meta = create_file_meta(FileId::random(), 0); diff --git a/src/mito2/src/sst/version.rs b/src/mito2/src/sst/version.rs index 67d41a3b82..e07d0b5ba9 100644 --- a/src/mito2/src/sst/version.rs +++ b/src/mito2/src/sst/version.rs @@ -45,6 +45,22 @@ impl SstVersion { &self.levels } + /// Returns the unique current handle matching the selected file's identity and level. + pub(crate) fn file_for_compaction(&self, selected: &FileHandle) -> Option<&FileHandle> { + let mut files = self + .levels + .iter() + .filter_map(|level| level.files.get(&selected.file_id().file_id())); + let current = files.next()?; + if files.next().is_some() + || current.file_id() != selected.file_id() + || current.level() != selected.level() + { + return None; + } + Some(current) + } + /// Add files to the version. If a file with the same `file_id` already exists, /// it will be overwritten with the new file. /// @@ -270,6 +286,38 @@ mod tests { }); } + #[test] + fn test_file_for_compaction_returns_unambiguous_matching_level() { + let purger = new_noop_file_purger(); + let file_id = FileId::random(); + let file = FileMeta { + file_id, + level: 1, + ..Default::default() + }; + let selected = FileHandle::new(file.clone(), purger.clone()); + let mut version = SstVersion::new(); + version.add_files(purger, std::iter::once(file)); + + assert_eq!( + version + .file_for_compaction(&selected) + .unwrap() + .file_id() + .file_id(), + file_id + ); + let missing = FileHandle::new( + FileMeta { + file_id: FileId::random(), + level: 1, + ..Default::default() + }, + new_noop_file_purger(), + ); + assert!(version.file_for_compaction(&missing).is_none()); + } + #[test] fn test_usage_only_counts_owned_files() { let purger = new_noop_file_purger(); diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index a3602d29ff..4493bb1bbb 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -1235,6 +1235,9 @@ impl RegionWorkerLoop { /// Handles region background request async fn handle_background_notify(&mut self, region_id: RegionId, notify: BackgroundNotify) { match notify { + BackgroundNotify::CompactionPickFinished(req) => { + self.handle_compaction_pick_finished(region_id, req).await + } BackgroundNotify::FlushFinished(req) => { self.handle_flush_finished(region_id, req).await } @@ -1415,6 +1418,41 @@ impl WorkerListener { } } + pub(crate) fn on_compaction_schedule_attempt(&self, _region_id: RegionId) { + #[cfg(any(test, feature = "test"))] + if let Some(listener) = &self.listener { + listener.on_compaction_schedule_attempt(_region_id); + } + } + + pub(crate) async fn on_compaction_pick_begin(&self, _region_id: RegionId) { + #[cfg(any(test, feature = "test"))] + if let Some(listener) = &self.listener { + listener.on_compaction_pick_begin(_region_id).await; + } + } + + pub(crate) async fn on_compaction_commit_begin(&self, _region_id: RegionId) { + #[cfg(any(test, feature = "test"))] + if let Some(listener) = &self.listener { + listener.on_compaction_commit_begin(_region_id).await; + } + } + + pub(crate) async fn on_compaction_result_notified(&self, _region_id: RegionId) { + #[cfg(any(test, feature = "test"))] + if let Some(listener) = &self.listener { + listener.on_compaction_result_notified(_region_id).await; + } + } + + pub(crate) fn on_compaction_cancel_requested(&self, _region_id: RegionId) { + #[cfg(any(test, feature = "test"))] + if let Some(listener) = &self.listener { + listener.on_compaction_cancel_requested(_region_id); + } + } + pub(crate) async fn on_notify_region_change_result_begin(&self, _region_id: RegionId) { #[cfg(any(test, feature = "test"))] if let Some(listener) = &self.listener { diff --git a/src/mito2/src/worker/handle_compaction.rs b/src/mito2/src/worker/handle_compaction.rs index 3abdf0e349..ef789300d8 100644 --- a/src/mito2/src/worker/handle_compaction.rs +++ b/src/mito2/src/worker/handle_compaction.rs @@ -18,8 +18,9 @@ use store_api::logstore::LogStore; use store_api::region_request::RegionCompactRequest; use store_api::storage::RegionId; +use crate::compaction::{CompactionExecution, CompactionPickFinished}; use crate::config::IndexBuildMode; -use crate::error::RegionNotFoundSnafu; +use crate::error::{RegionClosedSnafu, RegionNotFoundSnafu}; use crate::metrics::COMPACTION_REQUEST_COUNT; use crate::region::MitoRegionRef; use crate::request::{ @@ -30,6 +31,43 @@ use crate::sst::index::IndexBuildType; use crate::worker::RegionWorkerLoop; impl RegionWorkerLoop { + fn is_current_compaction_execution( + &self, + region: &MitoRegionRef, + execution: &CompactionExecution, + ) -> bool { + self.compaction_scheduler.is_current_region_execution( + region.region_id, + ®ion.version_control, + execution, + ) + } + + pub(crate) async fn handle_compaction_pick_finished( + &mut self, + region_id: RegionId, + request: CompactionPickFinished, + ) where + S: LogStore, + { + let Some(region) = self.regions.get_region(region_id) else { + return; + }; + let mut pending_ddls = self + .compaction_scheduler + .accept_compaction_pick_finished( + request, + ®ion.version_control, + ®ion.manifest_ctx, + self.schema_metadata_manager.clone(), + ) + .await; + if !pending_ddls.is_empty() { + self.listener.on_compaction_result_notified(region_id).await; + } + self.handle_ddl_requests(&mut pending_ddls).await; + } + /// Handles compaction request submitted to region worker. pub(crate) async fn handle_compaction_request( &mut self, @@ -42,6 +80,7 @@ impl RegionWorkerLoop { }; COMPACTION_REQUEST_COUNT.inc(); let parallelism = req.parallelism.unwrap_or(1) as usize; + self.listener.on_compaction_schedule_attempt(region_id); if let Err(e) = self .compaction_scheduler .schedule_compaction( @@ -80,6 +119,11 @@ impl RegionWorkerLoop { return; } }; + if !self.is_current_compaction_execution(®ion, &request.execution) { + request.on_failure(RegionClosedSnafu { region_id }.build()); + return; + } + let execution = request.execution.clone(); region.version_control.apply_edit( Some(request.edit.clone()), @@ -91,6 +135,7 @@ impl RegionWorkerLoop { // compaction finished. request.on_success(); + self.listener.on_compaction_result_notified(region_id).await; // In async mode, create indexes after compact if new files are created. if self.config.index.build_mode == IndexBuildMode::Async @@ -110,8 +155,9 @@ impl RegionWorkerLoop { // Schedule next compaction if necessary. let mut pending_ddls = self .compaction_scheduler - .on_compaction_finished( + .on_execution_finished( region_id, + &execution, ®ion.manifest_ctx, self.schema_metadata_manager.clone(), ) @@ -151,17 +197,22 @@ impl RegionWorkerLoop { ) where S: LogStore, { + let execution = request.execution.clone(); + let is_current = self + .regions + .get_region(region_id) + .is_some_and(|region| self.is_current_compaction_execution(®ion, &execution)); request.on_success(); + if !is_current { + return; + } + // Reuse the scheduler's finish path to wake pending DDLs after a cooperative stop. - let mut pending_ddls = match self.regions.get_region(region_id) { - Some(_) => { - self.compaction_scheduler - .on_compaction_cancelled(region_id) - .await - } - None => Vec::new(), - }; + let mut pending_ddls = self + .compaction_scheduler + .on_execution_cancelled(region_id, &execution) + .await; self.handle_ddl_requests(&mut pending_ddls).await; } @@ -170,8 +221,14 @@ impl RegionWorkerLoop { pub(crate) async fn handle_compaction_failure(&mut self, req: CompactionFailed) { error!(req.err; "Failed to compact region: {}", req.region_id); + let Some(region) = self.regions.get_region(req.region_id) else { + return; + }; + if !self.is_current_compaction_execution(®ion, &req.execution) { + return; + } self.compaction_scheduler - .on_compaction_failed(req.region_id, req.err); + .on_execution_failed(req.region_id, &req.execution, req.err); } /// Schedule compaction for the region if necessary. @@ -191,6 +248,8 @@ impl RegionWorkerLoop { "minimal compaction interval time {:?} has passed, scheduling next compaction", self.config.min_compaction_interval ); + self.listener + .on_compaction_schedule_attempt(region.region_id); match self .compaction_scheduler .schedule_compaction( diff --git a/src/mito2/src/worker/handle_enter_staging.rs b/src/mito2/src/worker/handle_enter_staging.rs index 54bca8e4d8..8bc91ac96d 100644 --- a/src/mito2/src/worker/handle_enter_staging.rs +++ b/src/mito2/src/worker/handle_enter_staging.rs @@ -113,6 +113,7 @@ impl RegionWorkerLoop { partition_directive, }), }); + self.listener.on_compaction_cancel_requested(region_id); return; }