From eb8ae171cc38bebb43eef0ac9c18bf43253d966e Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Fri, 31 Jul 2026 10:46:05 +0800 Subject: [PATCH] refactor(mito2): split compaction scheduler modules (#8698) * refactor(mito2): split compaction module into scheduler/reader submodules Extract the compaction scheduler lifecycle (scheduler, status, phases, execution, SST reservations, pending requests) and its tests out of compaction.rs into compaction/scheduler.rs and compaction/scheduler_test.rs. Split the remaining helpers by responsibility: - estimate_compaction_bytes/refresh_picker_output move to scheduler.rs, the only call site - get_expired_ssts moves to picker.rs, shared by the TWCS and window pickers - CompactionSstReaderBuilder/time_range_to_predicate/ts_to_lit move to the new compaction/reader.rs The root compaction.rs keeps the shared output types and find_dynamic_options, and re-exports the moved types so existing call paths stay unchanged. Pure code motion, no behavior change. Signed-off-by: Lei, HUANG * refactor(mito2): separate compaction scheduling from execution details Turn compaction/scheduler.rs into a directory module to make the scheduling flow easier to review: - scheduler.rs keeps the pure scheduling core: the CompactionScheduler state machine, scheduling entry points, termination chaining, DDL coordination and region lifecycle events - scheduler/planning.rs holds the execution-facing parts: background planning dispatch, picker invocation, plan acceptance, remote/local submission and memory estimation - scheduler/state.rs holds the per-region lifecycle types: CompactionStatus, ActiveCompaction, CompactionPhase, CompactingFiles, LocalCompactionState, CompactionExecution and PendingCompaction Child modules keep access to the scheduler's private methods, so the split is pure code motion with minimal visibility changes (pub(super) only where the parent module or tests reach into child items). No behavior change. Signed-off-by: Lei, HUANG * refactor(mito2): use absolute scheduler imports Signed-off-by: Lei, HUANG * docs(mito2): document compaction scheduler modules Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/mito2/AGENTS.md | 2 +- src/mito2/src/compaction.rs | 4225 +---------------- src/mito2/src/compaction/compactor.rs | 3 +- src/mito2/src/compaction/picker.rs | 18 + src/mito2/src/compaction/reader.rs | 244 + src/mito2/src/compaction/scheduler.rs | 603 +++ .../src/compaction/scheduler/planning.rs | 656 +++ src/mito2/src/compaction/scheduler/state.rs | 587 +++ src/mito2/src/compaction/scheduler_test.rs | 2221 +++++++++ src/mito2/src/compaction/twcs.rs | 4 +- src/mito2/src/compaction/window.rs | 4 +- 11 files changed, 4347 insertions(+), 4220 deletions(-) create mode 100644 src/mito2/src/compaction/reader.rs create mode 100644 src/mito2/src/compaction/scheduler.rs create mode 100644 src/mito2/src/compaction/scheduler/planning.rs create mode 100644 src/mito2/src/compaction/scheduler/state.rs create mode 100644 src/mito2/src/compaction/scheduler_test.rs diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index c25ede60ed..73fc517fa7 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -23,7 +23,7 @@ snapshot isolation). It implements the `RegionEngine` trait from `store-api`. | `wal` | `src/mito2/src/wal.rs` | Write-ahead log wrapper over `log-store` | | `memtable` | `src/mito2/src/memtable/` | In-memory write buffers (time-series / bulk / partition) | | `flush` | `src/mito2/src/flush.rs` | `FlushScheduler`, `WriteBufferManager`, memtable → SST | -| `compaction` | `src/mito2/src/compaction/` | TWCS picker, strict-window manual picker, compactor, memory control | +| `compaction` | `src/mito2/src/compaction/` | Compaction scheduler (`scheduler.rs` + `scheduler/`), TWCS picker, strict-window manual picker, compactor, memory control | | `access_layer` | `src/mito2/src/access_layer.rs` | SST read/write over the object store | | `sst` | `src/mito2/src/sst/` | Parquet format, file metadata, index layout | | `read` | `src/mito2/src/read/` | `ScanRegion`, merge, dedup, projection, streaming | diff --git a/src/mito2/src/compaction.rs b/src/mito2/src/compaction.rs index 41139a252e..967c9b9cb7 100644 --- a/src/mito2/src/compaction.rs +++ b/src/mito2/src/compaction.rs @@ -16,1419 +16,31 @@ mod buckets; pub mod compactor; pub mod memory_manager; pub mod picker; +mod reader; pub mod run; +mod scheduler; mod task; #[cfg(test)] mod test_util; mod twcs; mod window; -use std::collections::{HashMap, HashSet}; -use std::fmt; -use std::future::Future; -use std::sync::{Arc, Mutex}; -use std::time::Instant; +use std::collections::HashMap; -use api::v1::region::compact_request; -use common_base::Plugins; -use common_base::cancellation::CancellationHandle; -use common_memory_manager::OnExhaustedPolicy; use common_meta::key::SchemaMetadataManagerRef; -use common_telemetry::{debug, error, info, warn}; +use common_telemetry::{debug, error}; +use common_time::TimeToLive; use common_time::range::TimestampRange; -use common_time::timestamp::TimeUnit; -use common_time::{TimeToLive, Timestamp}; -use datafusion_common::ScalarValue; -use datafusion_expr::Expr; -use datatypes::extension::json::is_structured_json_field; -use datatypes::types::json_type::JsonNativeType; -use futures::FutureExt; -use parquet::arrow::parquet_to_arrow_schema; -use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData}; +pub use scheduler::CompactionRequest; +pub(crate) use scheduler::{ + CompactionExecution, CompactionPickFinished, CompactionScheduler, LocalCompactionState, +}; use serde::{Deserialize, Serialize}; -use snafu::{OptionExt, ResultExt}; -use store_api::metadata::RegionMetadataRef; +use snafu::ResultExt; use store_api::storage::RegionId; -use task::MAX_PARALLEL_COMPACTION; -use tokio::sync::mpsc::{self, Sender}; -use crate::access_layer::AccessLayerRef; -use crate::cache::{CacheManagerRef, CacheStrategy}; -use crate::compaction::compactor::{CompactionRegion, CompactionVersion, DefaultCompactor}; -use crate::compaction::memory_manager::CompactionMemoryManager; -use crate::compaction::picker::{CompactionTask, PickerOutput, new_picker}; -use crate::compaction::task::CompactionTaskImpl; -use crate::config::MitoConfig; -use crate::error::{ - CompactRegionSnafu, CompactionCancelledSnafu, DataTypeMismatchSnafu, Error, - GetSchemaMetadataSnafu, JoinSnafu, ManualCompactionOverrideSnafu, ParquetToArrowSchemaSnafu, - RegionClosedSnafu, RegionDroppedSnafu, RegionTruncatedSnafu, RemoteCompactionSnafu, Result, - TimeRangePredicateOverflowSnafu, TimeoutSnafu, UnexpectedSnafu, -}; -use crate::metrics::{ - COMPACTION_MEMORY_REJECTED, COMPACTION_STAGE_ELAPSED, INFLIGHT_COMPACTION_COUNT, -}; -use crate::read::FlatSource; -use crate::read::flat_projection::FlatProjectionMapper; -use crate::read::read_columns::ReadColumns; -use crate::read::scan_region::{PredicateGroup, ScanInput}; -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::{ - BackgroundNotify, DdlRequest, OptionOutputTx, OutputTx, SenderDdlRequest, WorkerRequest, - WorkerRequestWithTime, -}; -use crate::schedule::remote_job_scheduler::{ - CompactionJob, DefaultNotifier, RemoteJob, RemoteJobSchedulerRef, -}; -use crate::schedule::scheduler::SchedulerRef; +use crate::error::{GetSchemaMetadataSnafu, Result, TimeoutSnafu}; use crate::sst::file::{FileHandle, FileMeta, Level}; -use crate::sst::parquet::reader::MetadataCacheMetrics; -use crate::sst::version::{LevelMeta, SstVersion}; -use crate::worker::WorkerListener; - -/// Region compaction request. -pub struct CompactionRequest { - pub(crate) engine_config: Arc, - pub(crate) current_version: CompactionVersion, - pub(crate) access_layer: AccessLayerRef, - /// Sender to send notification to the region worker. - pub(crate) request_sender: mpsc::Sender, - /// Start time of compaction task. - pub(crate) start_time: Instant, - pub(crate) cache_manager: CacheManagerRef, - pub(crate) manifest_ctx: ManifestContextRef, - pub(crate) listener: WorkerListener, - pub(crate) schema_metadata_manager: SchemaMetadataManagerRef, - pub(crate) max_parallelism: usize, -} - -impl CompactionRequest { - pub(crate) fn region_id(&self) -> RegionId { - self.current_version.metadata.region_id - } -} - -/// 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) result: CompactionPlanningResult, -} - -pub(crate) struct PreparedCompaction { - compaction_region: CompactionRegion, - picker_output: PickerOutput, - start_time: Instant, - ttl: TimeToLive, -} - -/// Identifies an accepted compaction attempt and keeps its SST reservations alive. -/// The plan id fences terminal notifications from superseded attempts. -#[derive(Debug, Clone)] -pub(crate) struct CompactionExecution { - plan_id: u64, - _files: CompactingFiles, -} - -impl CompactionExecution { - fn new(plan_id: u64, files: CompactingFiles) -> Self { - Self { - plan_id, - _files: files, - } - } - - pub(crate) fn matches(&self, other: &Self) -> bool { - self.plan_id == other.plan_id - } - - #[cfg(test)] - pub(crate) fn for_test(plan_id: u64) -> Self { - Self::new(plan_id, CompactingFiles::empty()) - } -} - -/// Compaction scheduler tracks and manages compaction tasks. -pub(crate) struct CompactionScheduler { - scheduler: SchedulerRef, - /// Compacting regions. - region_status: HashMap, - /// Request sender of the worker that this scheduler belongs to. - request_sender: Sender, - cache_manager: CacheManagerRef, - engine_config: Arc, - memory_manager: Arc, - memory_policy: OnExhaustedPolicy, - listener: WorkerListener, - /// Plugins for the compaction scheduler. - plugins: Plugins, - /// Scheduler-wide generation counter for compaction plans and executions. - /// It outlives region statuses so close/reopen cannot reuse an old identity. - next_plan_id: u64, -} - -fn requires_pending_compaction_slot( - options: &compact_request::Options, - time_range: Option, -) -> bool { - matches!(options, compact_request::Options::StrictWindow(_)) || time_range.is_some() -} - -impl CompactionScheduler { - #[allow(clippy::too_many_arguments)] - pub(crate) fn new( - scheduler: SchedulerRef, - request_sender: Sender, - cache_manager: CacheManagerRef, - engine_config: Arc, - listener: WorkerListener, - plugins: Plugins, - memory_manager: Arc, - memory_policy: OnExhaustedPolicy, - ) -> Self { - Self { - scheduler, - region_status: HashMap::new(), - request_sender, - cache_manager, - engine_config, - memory_manager, - memory_policy, - listener, - plugins, - next_plan_id: 0, - } - } - - /// Returns the current plan id and advances the counter. - /// - /// Takes the counter instead of `&mut self` so callers can bump it while - /// holding a mutable borrow of a region status. - fn next_plan_id(counter: &mut u64) -> u64 { - let plan_id = *counter; - *counter = counter.wrapping_add(1); - plan_id - } - - /// Schedules a compaction for the region. - /// Returns whether a compaction is scheduled. - #[allow(clippy::too_many_arguments)] - pub(crate) fn schedule_compaction( - &mut self, - region_id: RegionId, - compact_options: compact_request::Options, - version_control: &VersionControlRef, - access_layer: &AccessLayerRef, - waiter: OptionOutputTx, - manifest_ctx: &ManifestContextRef, - schema_metadata_manager: SchemaMetadataManagerRef, - max_parallelism: usize, - ) -> Result { - self.schedule_compaction_with_time_range( - region_id, - compact_options, - version_control, - access_layer, - waiter, - manifest_ctx, - schema_metadata_manager, - max_parallelism, - None, - ) - } - - /// Schedules a compaction constrained by an optional time range. - #[allow(clippy::too_many_arguments)] - pub(crate) fn schedule_compaction_with_time_range( - &mut self, - region_id: RegionId, - compact_options: compact_request::Options, - version_control: &VersionControlRef, - access_layer: &AccessLayerRef, - waiter: OptionOutputTx, - manifest_ctx: &ManifestContextRef, - schema_metadata_manager: SchemaMetadataManagerRef, - max_parallelism: usize, - time_range: Option, - ) -> Result { - // skip compaction if region is in staging state - let current_state = manifest_ctx.current_state(); - if current_state == RegionRoleState::Leader(RegionLeaderState::Staging) { - info!( - "Skipping compaction for region {} in staging mode, options: {:?}", - region_id, compact_options - ); - waiter.send(Ok(0)); - return Ok(false); - } - - if let Some(status) = self.region_status.get_mut(®ion_id) { - // Pending Truncate/EnterStaging requests form a scheduling fence. Any later - // compaction with a waiter is an explicit request and receives CompactionCancelled; - // automatic triggers have no waiter, so sending the error is a no-op and the trigger - // is simply ignored. - if !status.pending_ddl_requests.is_empty() { - waiter.send(CompactionCancelledSnafu.fail()); - info!( - "Region {} has pending DDL requests, ignoring compaction: {:?}", - region_id, compact_options - ); - return Ok(false); - } - - if requires_pending_compaction_slot(&compact_options, time_range) { - // Incoming compaction request is manually triggered. - status.set_pending_request(PendingCompaction { - options: compact_options, - waiter, - max_parallelism, - time_range, - }); - info!( - "Region {} is compacting, manually compaction will be re-scheduled.", - region_id - ); - } else { - status.merge_regular_trigger(waiter); - } - return Ok(false); - } - - // Publish the picking phase before dispatching background planning. - let mut status = - CompactionStatus::new(region_id, version_control.clone(), access_layer.clone()); - let request = status.new_compaction_request( - self.request_sender.clone(), - self.engine_config.clone(), - self.cache_manager.clone(), - manifest_ctx, - self.listener.clone(), - schema_metadata_manager, - max_parallelism, - ); - let plan_id = Self::next_plan_id(&mut self.next_plan_id); - status.start_picking_with_time_range(plan_id, time_range); - status.merge_waiter(waiter); - self.region_status.insert(region_id, status); - self.dispatch_compaction_planning(plan_id, request, compact_options, time_range); - self.listener.on_compaction_scheduled(region_id); - Ok(true) - } - - // Handle pending manual compaction request for the region. - // - // Returns true if should early return, false otherwise. - pub(crate) fn handle_pending_compaction_request( - &mut self, - region_id: RegionId, - manifest_ctx: &ManifestContextRef, - schema_metadata_manager: SchemaMetadataManagerRef, - ) -> bool { - let Some(status) = self.region_status.get_mut(®ion_id) else { - return true; - }; - - // If there is a pending manual compaction request, schedule it. - // and defer returning the pending DDL requests to the caller. - let Some(pending_request) = std::mem::take(&mut status.pending_request) else { - return false; - }; - - let PendingCompaction { - options, - waiter, - max_parallelism, - time_range, - } = pending_request; - - let request = status.new_compaction_request( - self.request_sender.clone(), - self.engine_config.clone(), - self.cache_manager.clone(), - manifest_ctx, - self.listener.clone(), - schema_metadata_manager, - max_parallelism, - ); - status.merge_waiter(waiter); - // Bump the counter through a disjoint field borrow so the `status` - // borrow stays alive; nothing could have removed the status since it - // was fetched above. - let plan_id = Self::next_plan_id(&mut self.next_plan_id); - status.start_picking_with_time_range(plan_id, time_range); - self.dispatch_compaction_planning(plan_id, request, options, time_range); - debug!( - "Successfully scheduled manual compaction planning for region id: {}", - region_id - ); - true - } - - /// Notifies the scheduler that the compaction job is finished successfully. - async fn on_compaction_finished( - &mut self, - region_id: RegionId, - manifest_ctx: &ManifestContextRef, - schema_metadata_manager: SchemaMetadataManagerRef, - ) -> Vec { - if !self - .region_status - .get(®ion_id) - .is_some_and(|s| s.is_busy()) - { - return Vec::new(); - } - - if self.handle_pending_compaction_request( - region_id, - manifest_ctx, - schema_metadata_manager.clone(), - ) { - return Vec::new(); - } - - // The region status might be removed by the previous steps. - // So we return empty DDL requests. - let Some(status) = self.region_status.get_mut(®ion_id) else { - return Vec::new(); - }; - let Some(mut active) = status.take_active() else { - return Vec::new(); - }; - - for waiter in std::mem::take(&mut active.waiters) { - waiter.send(Ok(0)); - } - - // A queued DDL was waiting for the current task to terminate; chaining - // another compaction ahead of it would delay the DDL by a whole extra - // plan/execution cycle, so dispatch the DDLs first. - let pending_ddl_requests = std::mem::take(&mut status.pending_ddl_requests); - if !pending_ddl_requests.is_empty() { - // The just-finished compaction satisfies any retained regular triggers. - for waiter in active.regular_followup_waiters.take().unwrap_or_default() { - waiter.send(Ok(0)); - } - self.region_status.remove(®ion_id); - // If there are pending DDL requests, we should return them to the caller. - // And skip try to schedule next compaction task. - return pending_ddl_requests; - } - - if active.regular_followup_waiters.is_some() { - self.schedule_next_compaction_with_active( - region_id, - manifest_ctx, - schema_metadata_manager, - Some(active), - None, - ); - return Vec::new(); - } - Vec::new() - } - - /// Returns whether a terminal notification belongs to the installed execution. - /// Background work may finish after its region status has been replaced, so - /// matching the region id alone is insufficient. - 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) async fn on_execution_finished( - &mut self, - region_id: RegionId, - execution: &CompactionExecution, - manifest_ctx: &ManifestContextRef, - schema_metadata_manager: SchemaMetadataManagerRef, - ) -> Vec { - // A stale finish must not clear the replacement phase or notify its waiters and DDLs. - 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(CompactionStatus::is_busy) - .unwrap_or(false) - } - - /// Removes the region status if it has no running task. - /// - /// A finished compaction leaves an idle status (`active = None`) behind when - /// there is nothing more to schedule. If the caller decides not to chain - /// the next compaction, it must remove the idle status; otherwise the - /// status becomes a zombie that makes `schedule_compaction` swallow all - /// future compaction triggers of the region. - pub(crate) fn remove_idle_status(&mut self, region_id: RegionId) { - if self - .region_status - .get(®ion_id) - .is_some_and(|status| !status.is_busy()) - { - self.region_status.remove(®ion_id); - } - } - - /// Schedules next compaction upon a finished compaction. - /// Returns whether the compaction is scheduled. - pub(crate) fn schedule_next_compaction( - &mut self, - region_id: RegionId, - manifest_ctx: &ManifestContextRef, - schema_metadata_manager: SchemaMetadataManagerRef, - ) -> bool { - let Some(status) = self.region_status.get_mut(®ion_id) else { - return false; - }; - // A plan is already in flight; treat it as scheduled instead of - // overwriting the current phase and orphaning the in-flight planning. - if status.is_busy() { - return true; - } - - let time_range = status.time_range; - self.schedule_next_compaction_with_active( - region_id, - manifest_ctx, - schema_metadata_manager, - None, - time_range, - ) - } - - fn schedule_next_compaction_with_active( - &mut self, - region_id: RegionId, - manifest_ctx: &ManifestContextRef, - schema_metadata_manager: SchemaMetadataManagerRef, - active: Option, - time_range: Option, - ) -> bool { - let Some(status) = self.region_status.get_mut(®ion_id) else { - return false; - }; - // We should always try to compact the region until picker returns None. - let request = status.new_compaction_request( - self.request_sender.clone(), - self.engine_config.clone(), - self.cache_manager.clone(), - manifest_ctx, - self.listener.clone(), - schema_metadata_manager, - MAX_PARALLEL_COMPACTION, - ); - // Bump the counter through a disjoint field borrow so the `status` - // borrow stays alive; nothing could have removed the status since it - // was fetched above. - let plan_id = Self::next_plan_id(&mut self.next_plan_id); - status.start_regular_picking(plan_id, active, time_range); - self.dispatch_compaction_planning( - plan_id, - request, - compact_request::Options::Regular(Default::default()), - time_range, - ); - debug!( - "Successfully scheduled next compaction planning for region id: {}", - region_id - ); - true - } - - /// Notifies the scheduler that the compaction job is cancelled cooperatively. - 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 { - // A stale cancellation must not remove a replacement execution's status. - 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. - 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, - ) { - // A stale failure must not tear down a replacement execution. - 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( - region_id, - Arc::new(RegionDroppedSnafu { region_id }.build()), - ); - } - - /// Notifies the scheduler that the region is closed. - pub(crate) fn on_region_closed(&mut self, region_id: RegionId) { - self.remove_region_on_failure(region_id, Arc::new(RegionClosedSnafu { region_id }.build())); - } - - /// Notifies the scheduler that the region is truncated. - pub(crate) fn on_region_truncated(&mut self, region_id: RegionId) { - self.remove_region_on_failure( - region_id, - Arc::new(RegionTruncatedSnafu { region_id }.build()), - ); - } - - /// Cancels the running compaction and queues its dependent DDL atomically. - /// - /// Production callers currently use this only for [`DdlRequest::Truncate`] and - /// [`DdlRequest::EnterStaging`]. If cancellation is still possible, the current picking or - /// local execution is asked to stop; otherwise the DDL waits for its terminal notification. - /// The worker dispatches the queued DDL only after that notification is handled, preventing - /// truncate or enter-staging from racing with compaction planning, execution, or commit. - /// Returns the sender and typed request unchanged if compaction is not running. - pub(crate) fn try_cancel_and_add_ddl( - &mut self, - region_id: RegionId, - sender: OptionOutputTx, - request: T, - into_ddl_request: impl FnOnce(T) -> DdlRequest, - ) -> std::result::Result<(), (OptionOutputTx, T)> { - let Some(status) = self.region_status.get_mut(®ion_id) else { - return Err((sender, request)); - }; - if status.request_cancel() == RequestCancelResult::NotRunning { - return Err((sender, request)); - } - - let request = SenderDdlRequest { - region_id, - sender, - request: into_ddl_request(request), - }; - debug!( - "Added pending DDL request for region: {}, ddl: {:?}", - request.region_id, request.request - ); - // The first queued Truncate/EnterStaging also fences later regular triggers from - // creating more follow-ups ahead of the DDL. - status.pending_ddl_requests.push(request); - Ok(()) - } - - #[cfg(test)] - fn add_ddl_request_to_pending(&mut self, request: SenderDdlRequest) { - self.region_status - .get_mut(&request.region_id) - .unwrap() - .pending_ddl_requests - .push(request); - } - - #[cfg(test)] - pub(crate) fn has_pending_ddls(&self, region_id: RegionId) -> bool { - let has_pending = self - .region_status - .get(®ion_id) - .map(|status| !status.pending_ddl_requests.is_empty()) - .unwrap_or(false); - debug!( - "Checked pending DDL requests for region: {}, has_pending: {}", - region_id, has_pending - ); - has_pending - } - - #[cfg(test)] - pub(crate) fn request_cancel(&mut self, region_id: RegionId) -> RequestCancelResult { - let Some(status) = self.region_status.get_mut(®ion_id) else { - return RequestCancelResult::NotRunning; - }; - - status.request_cancel() - } - - fn dispatch_compaction_planning( - &self, - plan_id: u64, - request: CompactionRequest, - options: compact_request::Options, - time_range: Option, - ) { - 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 planning = Self::prepare_compaction( - request, - options, - plugins, - max_background_compactions, - time_range, - ); - Self::notify_planning_result(region_id, plan_id, request_sender, planning).await; - }); - } - - /// Runs the planning future and always sends the planning result back to - /// the worker, even if the planning panics. - /// - /// The worker only leaves the picking phase after it receives the - /// `CompactionPickFinished` notification. If a panicked planning task - /// swallowed the notification, the region would be stuck in the picking - /// phase forever, blocking all future compactions and pending DDLs (e.g. - /// entering staging) of the region. - async fn notify_planning_result( - region_id: RegionId, - plan_id: u64, - request_sender: Sender, - planning: impl Future + Send, - ) { - // The idiomatic way to handle a panic result. - let result = std::panic::AssertUnwindSafe(planning).catch_unwind().await.unwrap_or_else(|payload| { - let reason = if let Some(message) = payload.as_ref().downcast_ref::<&str>() { - message.to_string() - } else if let Some(message) = payload.as_ref().downcast_ref::() { - message.clone() - } else { - "unknown panic".to_string() - }; - CompactionPlanningResult::Error(Arc::new( - UnexpectedSnafu { - reason: format!( - "Compaction planning panicked for region {region_id}, plan_id {plan_id}: {reason}" - ), - } - .build(), - )) - }); - 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, - 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, - time_range: Option, - ) -> CompactionPlanningResult { - let region_id = request.region_id(); - let (dynamic_compaction_opts, ttl) = find_dynamic_options( - region_id, - &request.current_version.options, - &request.schema_metadata_manager, - ) - .await - .unwrap_or_else(|e| { - warn!(e; "Failed to find dynamic options for region: {}", region_id); - ( - request.current_version.options.compaction.clone(), - request.current_version.options.ttl.unwrap_or_default(), - ) - }); - - let picker = new_picker( - &options, - &dynamic_compaction_opts, - request.current_version.options.append_mode, - Some(max_background_compactions), - time_range, - ); - let region_id = request.region_id(); - let CompactionRequest { - engine_config, - current_version, - access_layer, - request_sender: _, - start_time, - cache_manager, - manifest_ctx, - listener, - schema_metadata_manager: _, - max_parallelism, - } = request; - - debug!( - "Pick compaction strategy {:?} for region: {}, ttl: {:?}", - picker, region_id, ttl - ); - - let compaction_region = CompactionRegion { - region_id, - current_version: current_version.clone(), - region_options: RegionOptions { - compaction: dynamic_compaction_opts.clone(), - ..current_version.options.clone() - }, - engine_config: engine_config.clone(), - region_metadata: current_version.metadata.clone(), - cache_manager: cache_manager.clone(), - access_layer: access_layer.clone(), - manifest_ctx: manifest_ctx.clone(), - file_purger: None, - ttl: Some(ttl), - max_parallelism, - plugins, - }; - - 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(&picker_region) - }) - .await - .context(JoinSnafu) - { - Ok(output) => output, - Err(err) => return CompactionPlanningResult::Error(Arc::new(err)), - }; - - let Some(picker_output) = picker_output else { - return CompactionPlanningResult::NoPlan; - }; - - CompactionPlanningResult::Prepared(PreparedCompaction { - compaction_region, - picker_output, - start_time, - ttl, - }) - } - - pub(crate) async fn handle_compaction_pick_finished( - &mut self, - finished: CompactionPickFinished, - manifest_ctx: &ManifestContextRef, - schema_metadata_manager: SchemaMetadataManagerRef, - ) -> Vec { - let region_id = finished.region_id; - let plan_id = finished.plan_id; - let Some(status) = self.region_status.get(®ion_id) else { - return Vec::new(); - }; - // Picking runs detached from the worker. Its result may arrive after - // close/reopen or replanning installed another Picking phase for this region. - if !status.is_picking(finished.plan_id) { - 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 = status.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 = status.take_waiters(); - match self - .submit_prepared_compaction(prepared, files, waiters, plan_id) - .await - { - Ok(Some(phase)) => { - if let Some(status) = self.region_status.get_mut(®ion_id) { - status.set_phase(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 - } - } - } - - 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(); - }; - let Some(mut active) = status.take_active() else { - return Vec::new(); - }; - for waiter in std::mem::take(&mut active.waiters) { - if let Some(err) = &err { - waiter.send(Err(err.clone()).context(CompactRegionSnafu { region_id })); - } else { - waiter.send(Ok(0)); - } - } - - status.active = Some(active); - if self.handle_pending_compaction_request( - region_id, - manifest_ctx, - schema_metadata_manager.clone(), - ) { - return Vec::new(); - } - - let Some(active) = self - .region_status - .get_mut(®ion_id) - .and_then(CompactionStatus::take_active) - else { - return Vec::new(); - }; - if active.regular_followup_waiters.is_some() { - self.schedule_next_compaction_with_active( - region_id, - manifest_ctx, - schema_metadata_manager, - Some(active), - None, - ); - 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, - mut plan_id: u64, - ) -> 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, files.clone()); - let remote_compaction_job = CompactionJob { - compaction_region: compaction_region.clone(), - picker_output: picker_output.clone(), - start_time, - waiters, - ttl, - }; - - let result = remote_job_scheduler - .schedule( - RemoteJob::CompactionJob(remote_compaction_job), - Box::new(DefaultNotifier::new( - self.request_sender.clone(), - execution.clone(), - )), - ) - .await; - - match result { - Ok(job_id) => { - info!( - "Scheduled remote compaction job {} for region {}", - job_id, region_id - ); - INFLIGHT_COMPACTION_COUNT.inc(); - 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.extend_waiters(e.waiters); - } - return RemoteCompactionSnafu { - region_id, - job_id: None, - reason: e.reason, - } - .fail(); - } - - error!(e; "Failed to schedule remote compaction job for region {}, fallback to local compaction", region_id); - // An error may be ambiguous after the remote scheduler consumed - // the notifier. Fence a delayed remote callback from the local fallback. - plan_id = Self::next_plan_id(&mut self.next_plan_id); - e.waiters - } - } - } else { - debug!( - "Remote compaction is not enabled, fallback to local compaction for region {}", - region_id - ); - waiters - } - } else { - waiters - }; - - // Check whether this local compaction can ever fit before submitting it. - let estimated_bytes = estimate_compaction_bytes(&picker_output); - if let Some(limit_bytes) = self.exceeds_compaction_memory_limit(estimated_bytes) { - COMPACTION_MEMORY_REJECTED - .with_label_values(&["oversized"]) - .inc(); - warn!( - "Skip compaction for region {} because estimated memory {} bytes exceeds compaction memory limit {} bytes", - region_id, estimated_bytes, limit_bytes, - ); - for waiter in waiters { - waiter.send(Ok(0)); - } - return Ok(None); - } - - let cancel_handle = Arc::new(CancellationHandle::default()); - let state = LocalCompactionState::new(cancel_handle.clone()); - let execution = CompactionExecution::new(plan_id, files); - let local_compaction_task = Box::new(CompactionTaskImpl { - state: state.clone(), - execution: execution.clone(), - request_sender: self.request_sender.clone(), - waiters, - start_time, - listener: self.listener.clone(), - picker_output, - compaction_region, - compactor: Arc::new(DefaultCompactor::with_cancel_handle(cancel_handle.clone())), - memory_manager: self.memory_manager.clone(), - memory_policy: self.memory_policy, - estimated_memory_bytes: estimated_bytes, - }); - - 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.append_waiters(&mut task.waiters); - } - Err(err) - } - } - } - - fn submit_compaction_task( - &mut self, - task: Box, - region_id: RegionId, - ) -> 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(); - } 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 { - let limit_bytes = self.memory_manager.limit_bytes(); - if limit_bytes > 0 && estimated_bytes > limit_bytes { - Some(limit_bytes) - } else { - None - } - } - - fn remove_region_on_failure(&mut self, region_id: RegionId, err: Arc) { - // Remove this region. - let Some(status) = self.region_status.remove(®ion_id) else { - return; - }; - - // Notifies all pending tasks. - status.on_failure(err); - } - - fn remove_region_on_cancel(&mut self, region_id: RegionId) -> Vec { - let Some(status) = self.region_status.remove(®ion_id) else { - return Vec::new(); - }; - - status.on_cancel() - } -} - -#[derive(Debug, Clone)] -pub(crate) struct LocalCompactionState { - cancel_handle: Arc, - commit_started: Arc>, -} - -#[derive(Debug)] -enum CompactionPhase { - Picking { - plan_id: u64, - cancelled: bool, - }, - Local { - state: LocalCompactionState, - execution: CompactionExecution, - }, - Remote { - execution: CompactionExecution, - }, -} - -#[derive(Debug)] -struct ActiveCompaction { - phase: CompactionPhase, - /// Waiters satisfied by the current planning or execution cycle. Picking waiters move into - /// the submitted task; regular triggers coalesced during execution accumulate here. - waiters: Vec, - /// Requests one fresh regular picking cycle after the current cycle finishes. It is kept - /// separate because the current picker snapshot may predate the trigger; `Some(empty)` records - /// an automatic trigger without an explicit waiter. - regular_followup_waiters: Option>, -} - -impl ActiveCompaction { - fn picking(plan_id: u64, waiters: Vec) -> Self { - Self { - phase: CompactionPhase::Picking { - plan_id, - cancelled: false, - }, - waiters, - regular_followup_waiters: None, - } - } - - fn start_picking(&mut self, plan_id: u64) { - self.phase = CompactionPhase::Picking { - plan_id, - cancelled: false, - }; - } - - fn start_regular_picking(&mut self, plan_id: u64) { - self.waiters - .extend(self.regular_followup_waiters.take().unwrap_or_default()); - self.start_picking(plan_id); - } - - fn is_picking(&self, expected_plan_id: u64) -> bool { - matches!( - self.phase, - CompactionPhase::Picking { plan_id, .. } if plan_id == expected_plan_id - ) - } - - fn accept_plan(&self, expected_plan_id: u64) -> bool { - matches!( - self.phase, - CompactionPhase::Picking { - plan_id, - cancelled: false, - } if plan_id == expected_plan_id - ) - } - - fn matches_execution(&self, execution: &CompactionExecution) -> bool { - match &self.phase { - CompactionPhase::Picking { .. } => None, - CompactionPhase::Local { execution, .. } | CompactionPhase::Remote { execution } => { - Some(execution) - } - } - .is_some_and(|current| current.matches(execution)) - } - - fn request_cancel(&mut self) -> RequestCancelResult { - match &mut self.phase { - CompactionPhase::Picking { cancelled, .. } => { - if *cancelled { - RequestCancelResult::AlreadyCancelling - } else { - *cancelled = true; - RequestCancelResult::CancelIssued - } - } - CompactionPhase::Local { state, .. } => state.request_cancel(), - CompactionPhase::Remote { .. } => RequestCancelResult::TooLateToCancel, - } - } - - fn merge_regular_trigger(&mut self, mut waiter: OptionOutputTx) { - if matches!(self.phase, CompactionPhase::Picking { .. }) { - let regular_followup_waiters = self.regular_followup_waiters.get_or_insert_default(); - if let Some(waiter) = waiter.take_inner() { - regular_followup_waiters.push(waiter); - } - } else { - self.merge_waiter(waiter); - } - } - - fn merge_waiter(&mut self, mut waiter: OptionOutputTx) { - if let Some(waiter) = waiter.take_inner() { - self.waiters.push(waiter); - } - } -} - -/// 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 { - fn new(cancel_handle: Arc) -> Self { - Self { - cancel_handle, - commit_started: Arc::new(Mutex::new(false)), - } - } - - /// Returns the cancellation handle for this compaction task. - pub(crate) fn cancel_handle(&self) -> Arc { - self.cancel_handle.clone() - } - - /// Marks the compaction task as started to commit, - /// which means the compaction task is in the final stage and is about to update region version and manifest. - /// It will reject cancellation request after this method is called. - /// - /// Returns true if this is the first time to mark commit started, false otherwise. - pub(crate) fn mark_commit_started(&self) -> bool { - let mut commit_started = self.commit_started.lock().unwrap(); - if self.cancel_handle.is_cancelled() { - return false; - } - *commit_started = true; - true - } - - /// Request cancellation for this compaction task. - pub(crate) fn request_cancel(&self) -> RequestCancelResult { - // The cancel handle must under the lock of `commit_started` to avoid racing between cancellation and commit. - let commit_started = self.commit_started.lock().unwrap(); - if *commit_started { - return RequestCancelResult::TooLateToCancel; - } - if self.cancel_handle.is_cancelled() { - return RequestCancelResult::AlreadyCancelling; - } - - self.cancel_handle.cancel(); - RequestCancelResult::CancelIssued - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum RequestCancelResult { - CancelIssued, - AlreadyCancelling, - TooLateToCancel, - NotRunning, -} - -impl Drop for CompactionScheduler { - fn drop(&mut self) { - for (region_id, status) in self.region_status.drain() { - // We are shutting down so notify all pending tasks. - status.on_failure(Arc::new(RegionClosedSnafu { region_id }.build())); - } - } -} /// Finds compaction options and TTL together with a single metadata fetch to reduce RTT. async fn find_dynamic_options( @@ -1512,285 +124,6 @@ async fn find_dynamic_options( Ok((compaction, ttl)) } -/// Status of running and pending region compaction tasks. -struct CompactionStatus { - /// Id of the region. - region_id: RegionId, - /// Version control of the region. - version_control: VersionControlRef, - /// Access layer of the region. - access_layer: AccessLayerRef, - /// Current compaction lifecycle. `None` is the existing transient idle state. - // TODO: Remove idle statuses and make ActiveCompaction non-optional once chained - // scheduling can recreate the status from region context. - active: Option, - /// Optional range retained by automatic continuations of the current compaction. - time_range: Option, - /// Pending compactions that are supposed to run as soon as current compaction task finished. - /// - /// This holds strict-window requests and ranged regular requests. An unrestricted regular - /// request is instead merged into `ActiveCompaction::regular_followup_waiters` or `waiters`. - pending_request: Option, - /// Pending DDL requests that should run when compaction is done. - /// - /// Although [`SenderDdlRequest`] can wrap any DDL variant, production code only queues - /// [`DdlRequest::Truncate`] and [`DdlRequest::EnterStaging`] here. Both must serialize with - /// compaction so they observe the version after compaction terminates. - pending_ddl_requests: Vec, -} - -impl CompactionStatus { - /// Creates a new [CompactionStatus] - fn new( - region_id: RegionId, - version_control: VersionControlRef, - access_layer: AccessLayerRef, - ) -> CompactionStatus { - CompactionStatus { - region_id, - version_control, - access_layer, - active: None, - time_range: None, - pending_request: None, - pending_ddl_requests: Vec::new(), - } - } - - #[cfg(test)] - fn start_picking(&mut self, plan_id: u64) { - self.start_picking_with_time_range(plan_id, None); - } - - fn start_picking_with_time_range(&mut self, plan_id: u64, time_range: Option) { - self.time_range = time_range; - if let Some(active) = &mut self.active { - active.start_picking(plan_id); - } else { - self.active = Some(ActiveCompaction::picking(plan_id, Vec::new())); - } - } - - fn start_regular_picking( - &mut self, - plan_id: u64, - active: Option, - time_range: Option, - ) { - self.time_range = time_range; - self.active = Some(if let Some(mut active) = active { - active.start_regular_picking(plan_id); - active - } else { - ActiveCompaction::picking(plan_id, Vec::new()) - }); - } - - fn is_picking(&self, expected_plan_id: u64) -> bool { - self.active - .as_ref() - .is_some_and(|active| active.is_picking(expected_plan_id)) - } - - fn accept_plan(&self, expected_plan_id: u64) -> bool { - self.active - .as_ref() - .is_some_and(|active| active.accept_plan(expected_plan_id)) - } - - fn is_busy(&self) -> bool { - self.active.is_some() - } - - fn matches_execution(&self, execution: &CompactionExecution) -> bool { - self.active - .as_ref() - .is_some_and(|active| active.matches_execution(execution)) - } - - #[cfg(test)] - fn start_local_task(&mut self) -> LocalCompactionState { - let state = LocalCompactionState::new(Arc::new(CancellationHandle::default())); - let execution = CompactionExecution::new(0, CompactingFiles::empty()); - let phase = CompactionPhase::Local { - state: state.clone(), - execution, - }; - if let Some(active) = &mut self.active { - active.phase = phase; - } else { - self.active = Some(ActiveCompaction { - phase, - waiters: Vec::new(), - regular_followup_waiters: None, - }); - } - state - } - - #[cfg(test)] - fn start_remote_task(&mut self) { - let execution = CompactionExecution::new(0, CompactingFiles::empty()); - let phase = CompactionPhase::Remote { execution }; - if let Some(active) = &mut self.active { - active.phase = phase; - } else { - self.active = Some(ActiveCompaction { - phase, - waiters: Vec::new(), - regular_followup_waiters: None, - }); - } - } - - fn request_cancel(&mut self) -> RequestCancelResult { - let Some(active) = &mut self.active else { - return RequestCancelResult::NotRunning; - }; - active.request_cancel() - } - - #[cfg(test)] - fn clear_running_task(&mut self) -> bool { - self.active.take().is_some() - } - - fn merge_regular_trigger(&mut self, waiter: OptionOutputTx) { - if let Some(active) = &mut self.active { - active.merge_regular_trigger(waiter); - } - } - - /// Merge the waiter to the pending compaction. - fn merge_waiter(&mut self, waiter: OptionOutputTx) { - if let Some(active) = &mut self.active { - active.merge_waiter(waiter); - } - } - - fn take_active(&mut self) -> Option { - self.active.take() - } - - fn take_waiters(&mut self) -> Vec { - self.active - .as_mut() - .map(|active| std::mem::take(&mut active.waiters)) - .unwrap_or_default() - } - - fn extend_waiters(&mut self, waiters: Vec) { - if let Some(active) = &mut self.active { - active.waiters.extend(waiters); - } - } - - fn append_waiters(&mut self, waiters: &mut Vec) { - if let Some(active) = &mut self.active { - active.waiters.append(waiters); - } - } - - fn set_phase(&mut self, phase: CompactionPhase) { - if let Some(active) = &mut self.active { - active.phase = phase; - } - } - - /// Set pending compaction request or replace current value if already exist. - fn set_pending_request(&mut self, pending: PendingCompaction) { - if let Some(prev) = self.pending_request.replace(pending) { - debug!( - "Replace pending compaction options with new request {:?} for region: {}", - prev.options, self.region_id - ); - prev.waiter.send(ManualCompactionOverrideSnafu.fail()); - } - } - - fn on_failure(mut self, err: Arc) { - if let Some(mut active) = self.active.take() { - for waiter in active - .waiters - .drain(..) - .chain(active.regular_followup_waiters.take().unwrap_or_default()) - { - waiter.send(Err(err.clone()).context(CompactRegionSnafu { - region_id: self.region_id, - })); - } - } - - if let Some(pending_compaction) = self.pending_request { - pending_compaction - .waiter - .send(Err(err.clone()).context(CompactRegionSnafu { - region_id: self.region_id, - })); - } - - for pending_ddl in self.pending_ddl_requests { - pending_ddl - .sender - .send(Err(err.clone()).context(CompactRegionSnafu { - region_id: self.region_id, - })); - } - } - - #[must_use] - fn on_cancel(mut self) -> Vec { - if let Some(mut active) = self.active.take() { - for waiter in active - .waiters - .drain(..) - .chain(active.regular_followup_waiters.take().unwrap_or_default()) - { - waiter.send(CompactionCancelledSnafu.fail()); - } - } - - if let Some(pending_compaction) = self.pending_request { - pending_compaction.waiter.send( - Err(Arc::new(CompactionCancelledSnafu.build())).context(CompactRegionSnafu { - region_id: self.region_id, - }), - ); - } - - std::mem::take(&mut self.pending_ddl_requests) - } - - /// Creates an immutable request for background compaction planning. - #[allow(clippy::too_many_arguments)] - fn new_compaction_request( - &self, - request_sender: Sender, - engine_config: Arc, - cache_manager: CacheManagerRef, - manifest_ctx: &ManifestContextRef, - listener: WorkerListener, - schema_metadata_manager: SchemaMetadataManagerRef, - max_parallelism: usize, - ) -> CompactionRequest { - let current_version = CompactionVersion::from(self.version_control.current().version); - let start_time = Instant::now(); - - CompactionRequest { - engine_config, - current_version, - access_layer: self.access_layer.clone(), - request_sender: request_sender.clone(), - start_time, - cache_manager, - manifest_ctx: manifest_ctx.clone(), - listener, - schema_metadata_manager, - max_parallelism, - } - } -} - #[derive(Debug, Clone)] pub struct CompactionOutput { /// Compaction output file level. @@ -1811,2539 +144,3 @@ pub struct SerializedCompactionOutput { filter_deleted: bool, output_time_range: Option, } - -/// Builders to create [BoxedRecordBatchStream] for compaction. -struct CompactionSstReaderBuilder<'a> { - metadata: RegionMetadataRef, - sst_layer: AccessLayerRef, - cache: CacheManagerRef, - inputs: &'a [FileHandle], - append_mode: bool, - filter_deleted: bool, - time_range: Option, - merge_mode: MergeMode, -} - -impl CompactionSstReaderBuilder<'_> { - /// Build a [FlatSource] that yields Arrow `RecordBatch`s from reading all the input SST files, - /// for compaction. The schema of the [FlatSource] is unified. - async fn build_flat_sst_reader(self) -> Result { - let scan_input = self.build_scan_input().await?; - - let schema = scan_input.mapper.output_schema(); - let schema = schema.arrow_schema(); - - let stream = SeqScan::new(scan_input) - .build_flat_reader_for_compaction() - .await?; - Ok(FlatSource::new_stream(schema.clone(), stream)) - } - - async fn build_scan_input(self) -> Result { - let schema = self.metadata.schema.arrow_schema(); - let parquet_metadata = self.collect_parquet_metadata().await?; - let batch_size = crate::batch_size::estimate_batch_size( - parquet_metadata - .iter() - .flat_map(|metadata| metadata.row_groups()) - .map(|row_group| { - let uncompressed_bytes = row_group - .columns() - .iter() - .map(|column| column.uncompressed_size() as u64) - .sum(); - (row_group.num_rows() as u64, uncompressed_bytes) - }), - ); - let json_type_hint = if schema.fields().iter().any(is_structured_json_field) { - let mut json_type_hint = schema - .fields() - .iter() - .filter(|&field| is_structured_json_field(field)) - .map(|field| (field.name().clone(), JsonNativeType::Null)) - .collect::>(); - - for metadata in &parquet_metadata { - let file_metadata = metadata.file_metadata(); - let schema = parquet_to_arrow_schema( - file_metadata.schema_descr(), - file_metadata.key_value_metadata(), - ) - .context(ParquetToArrowSchemaSnafu { - file: "compaction input", - })?; - for field in schema.fields() { - let Some(merged) = json_type_hint.get_mut(field.name()) else { - continue; - }; - - let json_type = JsonNativeType::try_from(field.data_type()) - .context(DataTypeMismatchSnafu)?; - merged.merge(&json_type); - } - } - - Some(json_type_hint) - } else { - None - }; - - let projection = (0..self.metadata.column_metadatas.len()).collect(); - let read_columns = ReadColumns::from_deduped_column_ids( - self.metadata.column_metadatas.iter().map(|x| x.column_id), - ); - let mapper = FlatProjectionMapper::new_with_read_columns( - &self.metadata, - projection, - read_columns, - json_type_hint.as_ref(), - )?; - - let mut scan_input = ScanInput::new(self.sst_layer, mapper) - .with_files(self.inputs.to_vec()) - .with_compaction(true) - .with_batch_size(batch_size) - .with_append_mode(self.append_mode) - // We use special cache strategy for compaction. - .with_cache(CacheStrategy::Compaction(self.cache)) - .with_filter_deleted(self.filter_deleted) - // We ignore file not found error during compaction. - .with_ignore_file_not_found(true) - .with_merge_mode(self.merge_mode); - - // This serves as a workaround of https://github.com/GreptimeTeam/greptimedb/issues/3944 - // by converting time ranges into predicate. - if let Some(time_range) = self.time_range { - scan_input = - scan_input.with_predicate(time_range_to_predicate(time_range, &self.metadata)?); - } - - Ok(scan_input) - } - - async fn collect_parquet_metadata(&self) -> Result>> { - let mut metadata = Vec::with_capacity(self.inputs.len()); - - for file_handle in self.inputs { - let file_path = - file_handle.file_path(self.sst_layer.table_dir(), self.sst_layer.path_type()); - let file_size = file_handle.meta_ref().file_size; - let parquet_metadata = match self - .sst_layer - .read_sst(file_handle.clone()) - .cache(CacheStrategy::Compaction(self.cache.clone())) - .read_parquet_metadata( - &file_path, - file_size, - &mut MetadataCacheMetrics::default(), - PageIndexPolicy::default(), - ) - .await - .map(|x| x.0.parquet_metadata()) - { - Ok(x) => x, - Err(e) if e.is_object_not_found() => continue, - Err(e) => return Err(e), - }; - metadata.push(parquet_metadata); - } - Ok(metadata) - } -} - -/// Converts time range to predicates so that rows outside the range will be filtered. -fn time_range_to_predicate( - range: TimestampRange, - metadata: &RegionMetadataRef, -) -> Result { - let ts_col = metadata.time_index_column(); - - // safety: time index column's type must be a valid timestamp type. - let ts_col_unit = ts_col - .column_schema - .data_type - .as_timestamp() - .unwrap() - .unit(); - - let exprs = match (range.start(), range.end()) { - (Some(start), Some(end)) => { - vec![ - datafusion_expr::col(ts_col.column_schema.name.clone()) - .gt_eq(ts_to_lit(*start, ts_col_unit)?), - datafusion_expr::col(ts_col.column_schema.name.clone()) - .lt(ts_to_lit(*end, ts_col_unit)?), - ] - } - (Some(start), None) => { - vec![ - datafusion_expr::col(ts_col.column_schema.name.clone()) - .gt_eq(ts_to_lit(*start, ts_col_unit)?), - ] - } - - (None, Some(end)) => { - vec![ - datafusion_expr::col(ts_col.column_schema.name.clone()) - .lt(ts_to_lit(*end, ts_col_unit)?), - ] - } - (None, None) => { - return Ok(PredicateGroup::default()); - } - }; - - let predicate = PredicateGroup::new(metadata, &exprs)?; - Ok(predicate) -} - -fn ts_to_lit(ts: Timestamp, ts_col_unit: TimeUnit) -> Result { - let ts = ts - .convert_to(ts_col_unit) - .context(TimeRangePredicateOverflowSnafu { - timestamp: ts, - unit: ts_col_unit, - })?; - let val = ts.value(); - let scalar_value = match ts_col_unit { - TimeUnit::Second => ScalarValue::TimestampSecond(Some(val), None), - TimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(val), None), - TimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(val), None), - TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(val), None), - }; - Ok(datafusion_expr::lit(scalar_value)) -} - -/// Finds all expired SSTs across levels. -fn get_expired_ssts( - levels: &[LevelMeta], - ttl: Option, - now: Timestamp, -) -> Vec { - let Some(ttl) = ttl else { - return vec![]; - }; - - levels - .iter() - .flat_map(|l| l.get_expired_files(&now, &ttl).into_iter()) - .collect() -} - -/// Estimates compaction memory as the sum of all input files' maximum row-group -/// uncompressed sizes. -fn estimate_compaction_bytes(picker_output: &PickerOutput) -> u64 { - picker_output - .outputs - .iter() - .flat_map(|output| output.inputs.iter()) - .map(|file: &FileHandle| { - let meta = file.meta_ref(); - meta.max_row_group_uncompressed_size - }) - .sum() -} - -/// Rebuilds picker output with current SST handles while preserving the picker's grouping. -/// -/// Picking runs in background on a version snapshot that may be stale by the -/// time the plan is accepted: a concurrent flush, compaction or index rebuild -/// can replace a selected file with a new handle carrying updated metadata -/// (e.g. `index_version`), or remove the file entirely. The handles in the -/// picker output therefore cannot be used as-is; re-resolving them against the -/// current version both detects gone files (aborting the plan) and ensures the -/// execution reads and reserves the up-to-date handle. -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 { - /// Compaction options. - pub(crate) options: compact_request::Options, - /// Waiters of pending requests. - pub(crate) waiter: OptionOutputTx, - /// Max parallelism for pending compaction. - pub(crate) max_parallelism: usize, - /// Optional time range that constrains candidate compaction windows. - pub(crate) time_range: Option, -} - -#[cfg(test)] -mod tests { - use std::assert_matches; - use std::time::Duration; - - use api::v1::region::StrictWindow; - use api::v1::region::compact_request::Options; - 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::error::InvalidSchedulerStateSnafu; - use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions}; - use crate::region::ManifestContext; - use crate::schedule::scheduler::{Job, Scheduler}; - use crate::sst::FormatType; - use crate::test_util::mock_schema_metadata_manager; - use crate::test_util::scheduler_util::{SchedulerEnv, VecScheduler}; - use crate::test_util::version_util::{VersionControlBuilder, apply_edit}; - - #[test] - fn test_requires_pending_compaction_slot() { - let time_range = TimestampRange::new( - Timestamp::new_millisecond(1_000), - Timestamp::new_millisecond(2_000), - ) - .unwrap(); - - assert!(!requires_pending_compaction_slot( - &compact_request::Options::Regular(Default::default()), - None, - )); - assert!(requires_pending_compaction_slot( - &compact_request::Options::Regular(Default::default()), - Some(time_range), - )); - assert!(requires_pending_compaction_slot( - &compact_request::Options::StrictWindow(StrictWindow { window_seconds: 60 }), - None, - )); - } - - struct FailingScheduler; - - struct FailingRemoteScheduler; - - #[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, - }, - ) - } - } - - 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, - ) - .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, fallback_to_local: bool) { - 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 = fallback_to_local; - } - - 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<()> { - InvalidSchedulerStateSnafu.fail() - } - - async fn stop(&self, _await_termination: bool) -> Result<()> { - Ok(()) - } - } - - 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:?}"), - } - } - - #[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()); - } - - #[tokio::test] - async fn test_find_compaction_options_db_level() { - let builder = VersionControlBuilder::new(); - let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); - let region_id = builder.region_id(); - let table_id = region_id.table_id(); - // Register table without ttl but with db-level compaction options - let mut schema_value = SchemaNameValue { - ttl: Some(DatabaseTimeToLive::default()), - ..Default::default() - }; - schema_value - .extra_options - .insert("compaction.type".to_string(), "twcs".to_string()); - schema_value - .extra_options - .insert("compaction.twcs.time_window".to_string(), "2h".to_string()); - schema_metadata_manager - .register_region_table_info( - table_id, - "t", - "c", - "s", - Some(schema_value), - kv_backend.clone(), - ) - .await; - - let version_control = Arc::new(builder.build()); - let region_opts = version_control.current().version.options.clone(); - let (opts, _) = find_dynamic_options(region_id, ®ion_opts, &schema_metadata_manager) - .await - .unwrap(); - match opts { - crate::region::options::CompactionOptions::Twcs(t) => { - assert_eq!(t.time_window_seconds(), Some(2 * 3600)); - } - } - } - - #[tokio::test] - async fn test_find_compaction_options_priority() { - fn schema_value_with_twcs(time_window: &str) -> SchemaNameValue { - let mut schema_value = SchemaNameValue { - ttl: Some(DatabaseTimeToLive::default()), - ..Default::default() - }; - schema_value - .extra_options - .insert("compaction.type".to_string(), "twcs".to_string()); - schema_value.extra_options.insert( - "compaction.twcs.time_window".to_string(), - time_window.to_string(), - ); - schema_value - } - - let cases = [ - ( - "db options set and table override set", - Some(schema_value_with_twcs("2h")), - true, - Some(Duration::from_secs(5 * 3600)), - Some(5 * 3600), - ), - ( - "db options set and table override not set", - Some(schema_value_with_twcs("2h")), - false, - None, - Some(2 * 3600), - ), - ( - "db options not set and table override set", - None, - true, - Some(Duration::from_secs(4 * 3600)), - Some(4 * 3600), - ), - ( - "db options not set and table override not set", - None, - false, - None, - None, - ), - ]; - - for (case_name, schema_value, override_set, table_window, expected_window) in cases { - let builder = VersionControlBuilder::new(); - let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); - let region_id = builder.region_id(); - let table_id = region_id.table_id(); - schema_metadata_manager - .register_region_table_info( - table_id, - "t", - "c", - "s", - schema_value, - kv_backend.clone(), - ) - .await; - - let version_control = Arc::new(builder.build()); - let mut region_opts = version_control.current().version.options.clone(); - region_opts.compaction_override = override_set; - if let Some(window) = table_window { - let crate::region::options::CompactionOptions::Twcs(twcs) = - &mut region_opts.compaction; - twcs.time_window = Some(window); - } - - let (opts, _) = find_dynamic_options(region_id, ®ion_opts, &schema_metadata_manager) - .await - .unwrap(); - match opts { - crate::region::options::CompactionOptions::Twcs(t) => { - assert_eq!(t.time_window_seconds(), expected_window, "{case_name}"); - } - } - } - } - - #[tokio::test] - async fn test_schedule_empty() { - let env = SchedulerEnv::new().await; - 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(); - schema_metadata_manager - .register_region_table_info( - builder.region_id().table_id(), - "test_table", - "test_catalog", - "test_schema", - None, - kv_backend, - ) - .await; - // Nothing to compact. - let version_control = Arc::new(builder.build()); - let (output_tx, output_rx) = oneshot::channel(); - let waiter = OptionOutputTx::from(output_tx); - let manifest_ctx = env - .mock_manifest_context(version_control.current().version.metadata.clone()) - .await; - let scheduled = scheduler - .schedule_compaction( - builder.region_id(), - compact_request::Options::Regular(Default::default()), - &version_control, - &env.access_layer, - waiter, - &manifest_ctx, - schema_metadata_manager.clone(), - 1, - ) - .unwrap(); - 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()); - - // Only one file, picker won't compact it. - let version_control = Arc::new(builder.push_l0_file(0, 1000).build()); - let (output_tx, output_rx) = oneshot::channel(); - let waiter = OptionOutputTx::from(output_tx); - let scheduled = scheduler - .schedule_compaction( - builder.region_id(), - compact_request::Options::Regular(Default::default()), - &version_control, - &env.access_layer, - waiter, - &manifest_ctx, - schema_metadata_manager.clone(), - 1, - ) - .unwrap(); - 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()); - } - - #[tokio::test] - 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, 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; - // Five overlapping L0 files are enough for the regular picker to create a task. - 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 scheduled = scheduler - .schedule_compaction( - region_id, - Options::Regular(Default::default()), - &version_control, - &env.access_layer, - OptionOutputTx::none(), - &manifest_ctx, - schema_metadata_manager.clone(), - 1, - ) - .unwrap(); - - // 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_planning_panic_notifies_and_clears_status() { - let env = SchedulerEnv::new().await; - let (tx, mut rx) = mpsc::channel(4); - let mut scheduler = env.mock_compaction_scheduler(tx.clone()); - 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.start_picking(7); - status.merge_waiter(OptionOutputTx::from(waiter_tx)); - scheduler.region_status.insert(region_id, status); - - CompactionScheduler::notify_planning_result(region_id, 7, tx, async { - panic!("planning boom") - }) - .await; - - let finished = recv_compaction_pick_finished(&mut rx).await; - let CompactionPlanningResult::Error(err) = &finished.result else { - panic!("expected planning error, got {:?}", &finished.result); - }; - assert!(err.to_string().contains("planning boom")); - let pending_ddls = scheduler - .handle_compaction_pick_finished(finished, &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_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 mut status = - CompactionStatus::new(region_id, version_control.clone(), env.access_layer.clone()); - status.start_picking(7); - scheduler.region_status.insert(region_id, status); - - let (pre_fence_tx, pre_fence_rx) = oneshot::channel(); - assert!( - !scheduler - .schedule_compaction( - region_id, - compact_request::Options::Regular(Default::default()), - &version_control, - &env.access_layer, - OptionOutputTx::from(pre_fence_tx), - &manifest_ctx, - schema_metadata_manager.clone(), - 1, - ) - .unwrap() - ); - let (ddl_tx, _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 - .handle_compaction_pick_finished( - CompactionPickFinished { - region_id, - plan_id: 7, - result: CompactionPlanningResult::NoPlan, - }, - &manifest_ctx, - schema_metadata_manager.clone(), - ) - .await; - assert!(pending_ddls.is_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 (post_fence_tx, post_fence_rx) = oneshot::channel(); - assert!( - !scheduler - .schedule_compaction( - region_id, - compact_request::Options::Regular(Default::default()), - &version_control, - &env.access_layer, - OptionOutputTx::from(post_fence_tx), - &manifest_ctx, - schema_metadata_manager.clone(), - 1, - ) - .unwrap() - ); - assert_matches!( - post_fence_rx.await.unwrap().unwrap_err(), - Error::CompactionCancelled { .. } - ); - - followup_finished.result = CompactionPlanningResult::NoPlan; - let pending_ddls = scheduler - .handle_compaction_pick_finished( - followup_finished, - &manifest_ctx, - schema_metadata_manager, - ) - .await; - assert_eq!(pending_ddls.len(), 1); - assert_eq!(pre_fence_rx.await.unwrap().unwrap(), 0); - assert!(!scheduler.region_status.contains_key(®ion_id)); - assert!(rx.try_recv().is_err()); - } - - #[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] - .active - .as_ref() - .unwrap() - .waiters - .len(), - 1 - ); - assert_matches!( - waiter_rx.try_recv(), - Err(oneshot::error::TryRecvError::Empty) - ); - } - - #[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_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_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_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, false); - 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_remote_fallback_uses_new_execution_plan_id() { - 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); - 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; - let remote_plan_id = finished.plan_id; - use_remote_compaction(&mut finished, true); - - scheduler - .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) - .await; - - assert_eq!(job_scheduler.num_jobs(), 1); - assert!(matches!( - scheduler.region_status[®ion_id] - .active - .as_ref() - .map(|active| &active.phase), - Some(CompactionPhase::Local { .. }) - )); - assert!( - !scheduler.region_status[®ion_id] - .matches_execution(&CompactionExecution::for_test(remote_plan_id)) - ); - } - - #[tokio::test] - async fn test_stale_plan_execution_does_not_affect_replacement_status() { - let env = SchedulerEnv::new().await; - let (tx, _rx) = mpsc::channel(4); - let mut scheduler = env.mock_compaction_scheduler(tx); - let stale_execution = CompactionExecution::for_test(1); - let replacement_version_control = compactable_version(); - let region_id = replacement_version_control - .current() - .version - .metadata - .region_id; - 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 (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); - - let pending_ddls = scheduler - .on_execution_finished( - region_id, - &stale_execution, - &manifest_ctx, - schema_metadata_manager, - ) - .await; - assert!(pending_ddls.is_empty()); - assert!(scheduler.region_status[®ion_id].is_busy()); - scheduler.on_execution_failed( - region_id, - &stale_execution, - Arc::new(InvalidSchedulerStateSnafu.build()), - ); - assert!(scheduler.region_status[®ion_id].is_busy()); - assert_matches!( - waiter_rx.try_recv(), - Err(oneshot::error::TryRecvError::Empty) - ); - } - - #[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, mut rx) = mpsc::channel(4); - let mut scheduler = env.mock_compaction_scheduler(tx); - scheduler.memory_manager = Arc::new(new_compaction_memory_manager(1024 * 1024)); - - let mut builder = VersionControlBuilder::new(); - let region_id = builder.region_id(); - let end = 1000 * 1000; - let version_control = Arc::new( - builder - .push_l0_file_with_max_row_group_size(0, end, 1024 * 1024) - .push_l0_file_with_max_row_group_size(10, end, 1024 * 1024) - .push_l0_file_with_max_row_group_size(50, end, 1024 * 1024) - .push_l0_file_with_max_row_group_size(80, end, 1024 * 1024) - .push_l0_file_with_max_row_group_size(90, end, 1024 * 1024) - .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 (output_tx, output_rx) = oneshot::channel(); - let rejected = COMPACTION_MEMORY_REJECTED.with_label_values(&["oversized"]); - let rejected_before = rejected.get(); - - let scheduled = scheduler - .schedule_compaction( - region_id, - Options::Regular(Default::default()), - &version_control, - &env.access_layer, - OptionOutputTx::from(output_tx), - &manifest_ctx, - schema_metadata_manager.clone(), - 1, - ) - .unwrap(); - - 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] - async fn test_schedule_on_finished() { - common_telemetry::init_default_ut_logging(); - 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 purger = builder.file_purger(); - let region_id = builder.region_id(); - - let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); - schema_metadata_manager - .register_region_table_info( - builder.region_id().table_id(), - "test_table", - "test_catalog", - "test_schema", - None, - kv_backend, - ) - .await; - - // 5 files to compact. - 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 scheduled = 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, - ) - .unwrap(); - // 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] - .files - .values() - .map(|file| file.meta_ref().clone()) - .collect(); - - // 5 files for next compaction and removes old files. - apply_edit( - &version_control, - &[(0, end), (20, end), (40, end), (60, end), (80, end)], - &file_metas, - purger.clone(), - ); - // The task is pending. - let (tx, _rx) = oneshot::channel(); - let scheduled = scheduler - .schedule_compaction( - region_id, - compact_request::Options::Regular(Default::default()), - &version_control, - &env.access_layer, - OptionOutputTx::new(Some(OutputTx::new(tx))), - &manifest_ctx, - schema_metadata_manager.clone(), - 1, - ) - .unwrap(); - assert!(!scheduled); - assert_eq!(1, scheduler.region_status.len()); - assert_eq!(1, job_scheduler.num_jobs()); - assert!( - !scheduler - .region_status - .get(&builder.region_id()) - .unwrap() - .active - .as_ref() - .unwrap() - .waiters - .is_empty() - ); - - // On compaction finished and schedule next compaction. - scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) - .await; - let scheduled = scheduler.schedule_next_compaction( - region_id, - &manifest_ctx, - schema_metadata_manager.clone(), - ); - 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. - apply_edit( - &version_control, - &[(0, end), (20, end), (40, end), (60, end), (80, end)], - &[], - purger.clone(), - ); - let (tx, _rx) = oneshot::channel(); - // The task is pending. - let scheduled = scheduler - .schedule_compaction( - region_id, - compact_request::Options::Regular(Default::default()), - &version_control, - &env.access_layer, - OptionOutputTx::new(Some(OutputTx::new(tx))), - &manifest_ctx, - schema_metadata_manager, - 1, - ) - .unwrap(); - assert!(!scheduled); - assert_eq!(2, job_scheduler.num_jobs()); - assert!( - !scheduler - .region_status - .get(&builder.region_id()) - .unwrap() - .active - .as_ref() - .unwrap() - .waiters - .is_empty() - ); - } - - #[tokio::test] - async fn test_remove_idle_status_allows_rescheduling() { - 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(); - scheduler.region_status.insert( - region_id, - CompactionStatus::new(region_id, version_control.clone(), env.access_layer.clone()), - ); - - 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, - ) - .unwrap() - ); - scheduler.remove_idle_status(region_id); - assert!( - scheduler - .schedule_compaction( - region_id, - compact_request::Options::Regular(Default::default()), - &version_control, - &env.access_layer, - OptionOutputTx::none(), - &manifest_ctx, - schema_metadata_manager, - 1, - ) - .unwrap() - ); - } - - #[tokio::test] - async fn test_time_range_compaction_when_compaction_in_progress() { - common_telemetry::init_default_ut_logging(); - 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 purger = builder.file_purger(); - let region_id = builder.region_id(); - - let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); - schema_metadata_manager - .register_region_table_info( - builder.region_id().table_id(), - "test_table", - "test_catalog", - "test_schema", - None, - kv_backend, - ) - .await; - - // 5 files to compact. - 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 file_metas: Vec<_> = version_control.current().version.ssts.levels()[0] - .files - .values() - .map(|file| file.meta_ref().clone()) - .collect(); - - // 5 files for next compaction and removes old files. - apply_edit( - &version_control, - &[(0, end), (20, end), (40, end), (60, end), (80, end)], - &file_metas, - purger.clone(), - ); - - 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, - ) - .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 - .region_status - .get(®ion_id) - .unwrap() - .pending_request - .is_none() - ); - - // Schedule another manual compaction with a time range. - let time_range = TimestampRange::new( - Timestamp::new_millisecond(0), - Timestamp::new_millisecond(end + 1), - ) - .unwrap(); - let (tx, _rx) = oneshot::channel(); - scheduler - .schedule_compaction_with_time_range( - region_id, - compact_request::Options::Regular(Default::default()), - &version_control, - &env.access_layer, - OptionOutputTx::new(Some(OutputTx::new(tx))), - &manifest_ctx, - schema_metadata_manager.clone(), - 1, - Some(time_range), - ) - .unwrap(); - assert_eq!(1, scheduler.region_status.len()); - // Current job num should be 1 since compaction is in progress. - assert_eq!(1, job_scheduler.num_jobs()); - let status = scheduler.region_status.get(&builder.region_id()).unwrap(); - assert_eq!( - Some(time_range), - status - .pending_request - .as_ref() - .and_then(|pending| pending.time_range) - ); - - // On compaction finished and schedule next compaction. - scheduler - .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(); - assert!(status.pending_request.is_none()); - } - - #[tokio::test] - async fn test_ranged_compaction_continuation_preserves_time_range() { - 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(); - for offset in [0, 10, 20, 30] { - builder.push_l0_file(offset, 1_000); - } - for offset in [0, 10, 20, 30] { - builder.push_l0_file(2 * 3_600_000 + offset, 2 * 3_600_000 + 1_000); - } - 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 schema_value = SchemaNameValue::default(); - schema_value - .extra_options - .insert("compaction.type".to_string(), "twcs".to_string()); - schema_value - .extra_options - .insert("compaction.twcs.time_window".to_string(), "1h".to_string()); - schema_metadata_manager - .register_region_table_info( - region_id.table_id(), - "t", - "c", - "s", - Some(schema_value), - kv_backend, - ) - .await; - let time_range = TimestampRange::new( - Timestamp::new_millisecond(0), - Timestamp::new_millisecond(3_600_000), - ) - .unwrap(); - - scheduler - .schedule_compaction_with_time_range( - region_id, - compact_request::Options::StrictWindow(StrictWindow { - window_seconds: 3_600, - }), - &version_control, - &env.access_layer, - OptionOutputTx::none(), - &manifest_ctx, - schema_metadata_manager.clone(), - 1, - Some(time_range), - ) - .unwrap(); - let first = recv_compaction_pick_finished(&mut rx).await; - assert!( - selected_files(&first) - .iter() - .all(|file| file.time_range().1 < Timestamp::new_millisecond(3_600_000)) - ); - scheduler - .handle_compaction_pick_finished(first, &manifest_ctx, schema_metadata_manager.clone()) - .await; - assert_eq!(1, job_scheduler.num_jobs()); - - scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) - .await; - assert!(scheduler.schedule_next_compaction( - region_id, - &manifest_ctx, - schema_metadata_manager, - )); - - let continuation = recv_compaction_pick_finished(&mut rx).await; - match continuation.result { - CompactionPlanningResult::NoPlan => {} - CompactionPlanningResult::Prepared(prepared) => assert!( - prepared - .picker_output - .outputs - .iter() - .flat_map(|output| &output.inputs) - .all(|file| file.time_range().1 < Timestamp::new_millisecond(3_600_000)) - ), - CompactionPlanningResult::Error(err) => { - panic!("unexpected compaction planning error: {err}") - } - } - } - - #[tokio::test] - async fn test_compaction_bypass_in_staging_mode() { - let env = SchedulerEnv::new().await; - let (tx, _rx) = mpsc::channel(4); - let mut scheduler = env.mock_compaction_scheduler(tx); - - // Create version control and manifest context for staging mode - let builder = VersionControlBuilder::new(); - let version_control = Arc::new(builder.build()); - let region_id = version_control.current().version.metadata.region_id; - - // Create staging manifest context using the same pattern as SchedulerEnv - let staging_manifest_ctx = { - let manager = RegionManifestManager::new( - version_control.current().version.metadata.clone(), - 0, - RegionManifestOptions { - manifest_dir: "".to_string(), - object_store: env.access_layer.object_store().clone(), - compress_type: CompressionType::Uncompressed, - checkpoint_distance: 10, - remove_file_options: Default::default(), - manifest_cache: None, - }, - FormatType::PrimaryKey, - &Default::default(), - ) - .await - .unwrap(); - Arc::new(ManifestContext::new( - manager, - RegionRoleState::Leader(RegionLeaderState::Staging), - None, - )) - }; - - let (schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager(); - - // Test regular compaction bypass in staging mode - let (tx, rx) = oneshot::channel(); - scheduler - .schedule_compaction( - region_id, - compact_request::Options::Regular(Default::default()), - &version_control, - &env.access_layer, - OptionOutputTx::new(Some(OutputTx::new(tx))), - &staging_manifest_ctx, - schema_metadata_manager, - 1, - ) - .unwrap(); - - let result = rx.await.unwrap(); - assert_eq!(result.unwrap(), 0); // is there a better way to check this? - assert_eq!(0, scheduler.region_status.len()); - } - - #[tokio::test] - async fn test_add_ddl_request_to_pending() { - let env = SchedulerEnv::new().await; - let (tx, _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()), - ); - scheduler - .region_status - .get_mut(®ion_id) - .unwrap() - .start_local_task(); - - let (output_tx, _output_rx) = oneshot::channel(); - scheduler.add_ddl_request_to_pending(SenderDdlRequest { - region_id, - sender: OptionOutputTx::from(output_tx), - request: crate::request::DdlRequest::EnterStaging( - store_api::region_request::EnterStagingRequest { - partition_directive: - store_api::region_request::StagingPartitionDirective::RejectAllWrites, - }, - ), - }); - - assert!(scheduler.has_pending_ddls(region_id)); - } - - #[tokio::test] - async fn test_pending_ddl_fences_later_compaction_triggers() { - let env = SchedulerEnv::new().await; - let (tx, _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(); - 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_manual_tx, mut first_manual_rx) = oneshot::channel(); - let mut status = - CompactionStatus::new(region_id, version_control.clone(), env.access_layer.clone()); - status.start_local_task(); - status.set_pending_request(PendingCompaction { - options: compact_request::Options::StrictWindow(StrictWindow { window_seconds: 60 }), - waiter: OptionOutputTx::from(first_manual_tx), - max_parallelism: 1, - time_range: None, - }); - scheduler.region_status.insert(region_id, status); - - let (ddl_tx, _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, - }, - ), - }); - - // Automatic regular triggers have no waiter and are ignored by the DDL fence. - 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, - ) - .unwrap() - ); - let active = scheduler.region_status[®ion_id].active.as_ref().unwrap(); - assert!(active.waiters.is_empty()); - assert!(active.regular_followup_waiters.is_none()); - - // Explicit regular and strict-window requests both have waiters and are rejected. - for options in [ - compact_request::Options::Regular(Default::default()), - compact_request::Options::StrictWindow(StrictWindow { - window_seconds: 120, - }), - ] { - let (later_tx, later_rx) = oneshot::channel(); - assert!( - !scheduler - .schedule_compaction( - region_id, - options, - &version_control, - &env.access_layer, - OptionOutputTx::from(later_tx), - &manifest_ctx, - schema_metadata_manager.clone(), - 1, - ) - .unwrap() - ); - assert_matches!( - later_rx.await.unwrap().unwrap_err(), - Error::CompactionCancelled { .. } - ); - } - - assert_matches!( - first_manual_rx.try_recv(), - Err(oneshot::error::TryRecvError::Empty) - ); - let pending_request = scheduler.region_status[®ion_id] - .pending_request - .as_ref() - .expect("manual compaction queued before DDL was removed"); - assert_matches!( - &pending_request.options, - compact_request::Options::StrictWindow(StrictWindow { window_seconds: 60 }) - ); - } - - #[tokio::test] - async fn test_request_cancel_state_transitions() { - let env = SchedulerEnv::new().await; - let builder = VersionControlBuilder::new(); - let region_id = builder.region_id(); - let version_control = Arc::new(builder.build()); - let mut status = - CompactionStatus::new(region_id, version_control, env.access_layer.clone()); - let state = status.start_local_task(); - - assert_eq!(status.request_cancel(), RequestCancelResult::CancelIssued); - assert!(state.cancel_handle().is_cancelled()); - assert_eq!( - status.request_cancel(), - RequestCancelResult::AlreadyCancelling - ); - - assert!(!state.mark_commit_started()); - assert_eq!( - status.request_cancel(), - RequestCancelResult::AlreadyCancelling - ); - - assert!(status.clear_running_task()); - assert_eq!(status.request_cancel(), RequestCancelResult::NotRunning); - } - - #[tokio::test] - async fn test_request_cancel_remote_compaction_is_too_late() { - let env = SchedulerEnv::new().await; - let builder = VersionControlBuilder::new(); - let region_id = builder.region_id(); - let version_control = Arc::new(builder.build()); - let mut status = - CompactionStatus::new(region_id, version_control, env.access_layer.clone()); - - status.start_remote_task(); - - assert_eq!( - status.request_cancel(), - RequestCancelResult::TooLateToCancel - ); - assert!(status.is_busy()); - } - - #[tokio::test] - async fn test_try_cancel_and_add_ddl_returns_request_when_not_running() { - let env = SchedulerEnv::new().await; - let (tx, _rx) = mpsc::channel(4); - let mut scheduler = env.mock_compaction_scheduler(tx); - let region_id = RegionId::new(1, 1); - let (ddl_tx, ddl_rx) = oneshot::channel(); - - let result = scheduler.try_cancel_and_add_ddl( - region_id, - OptionOutputTx::from(ddl_tx), - 42_u64, - |_| { - crate::request::DdlRequest::EnterStaging( - store_api::region_request::EnterStagingRequest { - partition_directive: - store_api::region_request::StagingPartitionDirective::RejectAllWrites, - }, - ) - }, - ); - - let Err((sender, payload)) = result else { - panic!("DDL was queued without a running compaction"); - }; - assert_eq!(payload, 42); - sender.send(Ok(0)); - assert_eq!(ddl_rx.await.unwrap().unwrap(), 0); - } - - #[tokio::test] - async fn test_try_cancel_and_add_ddl_cancels_and_queues_atomically() { - 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 mut status = - CompactionStatus::new(region_id, version_control, env.access_layer.clone()); - status.start_picking(7); - scheduler.region_status.insert(region_id, status); - let (ddl_tx, _ddl_rx) = oneshot::channel(); - - let result = - scheduler.try_cancel_and_add_ddl(region_id, OptionOutputTx::from(ddl_tx), (), |_| { - crate::request::DdlRequest::EnterStaging( - store_api::region_request::EnterStagingRequest { - partition_directive: - store_api::region_request::StagingPartitionDirective::RejectAllWrites, - }, - ) - }); - - assert!(result.is_ok()); - assert!(scheduler.has_pending_ddls(region_id)); - assert_eq!( - scheduler.request_cancel(region_id), - RequestCancelResult::AlreadyCancelling - ); - } - - #[tokio::test] - async fn test_on_compaction_cancelled_returns_pending_ddl_requests() { - let job_scheduler = Arc::new(VecScheduler::default()); - let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); - let (tx, _rx) = mpsc::channel(4); - let mut scheduler = env.mock_compaction_scheduler(tx); - let builder = VersionControlBuilder::new(); - let version_control = Arc::new(builder.build()); - let region_id = builder.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(); - - 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 { - region_id, - sender: OptionOutputTx::from(output_tx), - request: crate::request::DdlRequest::EnterStaging( - store_api::region_request::EnterStagingRequest { - partition_directive: - store_api::region_request::StagingPartitionDirective::RejectAllWrites, - }, - ), - }); - - let pending_ddls = scheduler.on_compaction_cancelled(region_id).await; - - assert_eq!(pending_ddls.len(), 1); - 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] - async fn test_on_compaction_cancelled_prioritizes_pending_ddls_over_pending_compaction() { - let job_scheduler = Arc::new(VecScheduler::default()); - let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); - let (tx, _rx) = mpsc::channel(4); - let mut scheduler = env.mock_compaction_scheduler(tx); - let builder = VersionControlBuilder::new(); - let version_control = Arc::new(builder.build()); - let region_id = builder.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(); - - scheduler.region_status.insert( - region_id, - CompactionStatus::new(region_id, version_control, env.access_layer.clone()), - ); - let status = scheduler.region_status.get_mut(®ion_id).unwrap(); - status.start_local_task(); - let (manual_tx, manual_rx) = oneshot::channel(); - status.set_pending_request(PendingCompaction { - options: compact_request::Options::StrictWindow(StrictWindow { window_seconds: 60 }), - waiter: OptionOutputTx::from(manual_tx), - max_parallelism: 1, - time_range: None, - }); - - let (output_tx, _output_rx) = oneshot::channel(); - scheduler.add_ddl_request_to_pending(SenderDdlRequest { - region_id, - sender: OptionOutputTx::from(output_tx), - request: crate::request::DdlRequest::EnterStaging( - store_api::region_request::EnterStagingRequest { - partition_directive: - store_api::region_request::StagingPartitionDirective::RejectAllWrites, - }, - ), - }); - - let pending_ddls = scheduler.on_compaction_cancelled(region_id).await; - - assert_eq!(pending_ddls.len(), 1); - assert!(!scheduler.region_status.contains_key(®ion_id)); - assert_eq!(job_scheduler.num_jobs(), 0); - assert_matches!(manual_rx.await.unwrap(), Err(_)); - } - - #[tokio::test] - async fn test_pending_ddl_request_failed_on_compaction_failed() { - 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 version_control = Arc::new(builder.build()); - let region_id = builder.region_id(); - - 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 { - region_id, - sender: OptionOutputTx::from(output_tx), - request: crate::request::DdlRequest::EnterStaging( - store_api::region_request::EnterStagingRequest { - partition_directive: - store_api::region_request::StagingPartitionDirective::RejectAllWrites, - }, - ), - }); - - assert!(scheduler.has_pending_ddls(region_id)); - scheduler - .on_compaction_failed(region_id, Arc::new(RegionClosedSnafu { region_id }.build())); - - 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] - async fn test_pending_ddl_request_failed_on_region_closed() { - let env = SchedulerEnv::new().await; - let (tx, _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 (output_tx, output_rx) = oneshot::channel(); - scheduler.add_ddl_request_to_pending(SenderDdlRequest { - region_id, - sender: OptionOutputTx::from(output_tx), - request: crate::request::DdlRequest::EnterStaging( - store_api::region_request::EnterStagingRequest { - partition_directive: - store_api::region_request::StagingPartitionDirective::RejectAllWrites, - }, - ), - }); - - assert!(scheduler.has_pending_ddls(region_id)); - scheduler.on_region_closed(region_id); - - assert!(!scheduler.has_pending_ddls(region_id)); - let result = output_rx.await.unwrap(); - assert_matches!(result, Err(_)); - } - - #[tokio::test] - async fn test_pending_ddl_request_failed_on_region_dropped() { - let env = SchedulerEnv::new().await; - let (tx, _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 (output_tx, output_rx) = oneshot::channel(); - scheduler.add_ddl_request_to_pending(SenderDdlRequest { - region_id, - sender: OptionOutputTx::from(output_tx), - request: crate::request::DdlRequest::EnterStaging( - store_api::region_request::EnterStagingRequest { - partition_directive: - store_api::region_request::StagingPartitionDirective::RejectAllWrites, - }, - ), - }); - - assert!(scheduler.has_pending_ddls(region_id)); - scheduler.on_region_dropped(region_id); - - assert!(!scheduler.has_pending_ddls(region_id)); - let result = output_rx.await.unwrap(); - assert_matches!(result, Err(_)); - } - - #[tokio::test] - async fn test_pending_ddl_request_failed_on_region_truncated() { - let env = SchedulerEnv::new().await; - let (tx, _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 (output_tx, output_rx) = oneshot::channel(); - scheduler.add_ddl_request_to_pending(SenderDdlRequest { - region_id, - sender: OptionOutputTx::from(output_tx), - request: crate::request::DdlRequest::EnterStaging( - store_api::region_request::EnterStagingRequest { - partition_directive: - store_api::region_request::StagingPartitionDirective::RejectAllWrites, - }, - ), - }); - - assert!(scheduler.has_pending_ddls(region_id)); - scheduler.on_region_truncated(region_id); - - assert!(!scheduler.has_pending_ddls(region_id)); - let result = output_rx.await.unwrap(); - assert_matches!(result, Err(_)); - } - - #[tokio::test] - async fn test_on_compaction_finished_returns_pending_ddl_requests() { - let job_scheduler = Arc::new(VecScheduler::default()); - let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); - let (tx, _rx) = mpsc::channel(4); - let mut scheduler = env.mock_compaction_scheduler(tx); - let builder = VersionControlBuilder::new(); - let version_control = Arc::new(builder.build()); - let region_id = builder.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(); - - 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 (output_tx, _output_rx) = oneshot::channel(); - scheduler.add_ddl_request_to_pending(SenderDdlRequest { - region_id, - sender: OptionOutputTx::from(output_tx), - request: crate::request::DdlRequest::EnterStaging( - store_api::region_request::EnterStagingRequest { - partition_directive: - store_api::region_request::StagingPartitionDirective::RejectAllWrites, - }, - ), - }); - - let pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) - .await; - - assert_eq!(pending_ddls.len(), 1); - assert!(!scheduler.has_pending_ddls(region_id)); - assert!(!scheduler.region_status.contains_key(®ion_id)); - assert_eq!(job_scheduler.num_jobs(), 0); - } - - #[tokio::test] - async fn test_on_compaction_finished_replays_pending_ddl_after_manual_noop() { - 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 version_control = Arc::new(builder.build()); - let region_id = builder.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(); - - let (manual_tx, manual_rx) = oneshot::channel(); - let mut status = - CompactionStatus::new(region_id, version_control.clone(), env.access_layer.clone()); - status.start_local_task(); - status.set_pending_request(PendingCompaction { - options: compact_request::Options::Regular(Default::default()), - waiter: OptionOutputTx::from(manual_tx), - max_parallelism: 1, - time_range: None, - }); - scheduler.region_status.insert(region_id, status); - - let (ddl_tx, _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 - .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); - } - - #[tokio::test] - async fn test_on_compaction_finished_dispatches_pending_ddl_before_chained_regular() { - 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 version_control = Arc::new(builder.build()); - let region_id = builder.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(); - - // A regular trigger was retained while picking and the region is now - // executing; a DDL queued behind the task must be dispatched as soon - // as the task finishes instead of waiting for a whole extra cycle. - let (regular_tx, regular_rx) = oneshot::channel(); - let mut status = - CompactionStatus::new(region_id, version_control.clone(), 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 (ddl_tx, _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 - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) - .await; - - assert_eq!(pending_ddls.len(), 1); - assert_eq!(regular_rx.await.unwrap().unwrap(), 0); - assert!(!scheduler.region_status.contains_key(®ion_id)); - assert!(rx.try_recv().is_err()); - } - - #[tokio::test] - async fn test_on_compaction_finished_returns_empty_when_region_absent() { - 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 pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) - .await; - - assert!(pending_ddls.is_empty()); - } - - #[tokio::test] - async fn test_on_compaction_finished_manual_schedule_error_cleans_status() { - 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 mut builder = VersionControlBuilder::new(); - 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 region_id = builder.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(); - - let (manual_tx, manual_rx) = oneshot::channel(); - let mut status = - CompactionStatus::new(region_id, version_control.clone(), env.access_layer.clone()); - status.start_local_task(); - status.set_pending_request(PendingCompaction { - options: compact_request::Options::Regular(Default::default()), - waiter: OptionOutputTx::from(manual_tx), - max_parallelism: 1, - time_range: None, - }); - scheduler.region_status.insert(region_id, status); - - let (ddl_tx, 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 - .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_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, 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(); - let manifest_ctx = env - .mock_manifest_context(version_control.current().version.metadata.clone()) - .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 pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) - .await; - - assert!(pending_ddls.is_empty()); - assert!(scheduler.region_status.contains_key(®ion_id)); - - let (schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager(); - // 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.clone(), - ); - assert!(scheduled); - 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)); - } - - #[tokio::test] - async fn test_on_compaction_finished_next_schedule_error_cleans_status() { - 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 mut builder = VersionControlBuilder::new(); - 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 region_id = builder.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(); - - 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 pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) - .await; - - assert!(pending_ddls.is_empty()); - assert!(scheduler.region_status.contains_key(®ion_id)); - - 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.clone(), - ); - assert!(scheduled); - 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)); - } - - #[tokio::test] - async fn test_concurrent_memory_competition() { - let manager = Arc::new(new_compaction_memory_manager(3 * 1024 * 1024)); // 3MB - let barrier = Arc::new(Barrier::new(3)); - let mut handles = vec![]; - - // Spawn 3 tasks competing for memory, each trying to acquire 2MB - for _i in 0..3 { - let mgr = manager.clone(); - let bar = barrier.clone(); - let handle = tokio::spawn(async move { - bar.wait().await; // Synchronize start - mgr.try_acquire(2 * 1024 * 1024) - }); - handles.push(handle); - } - - let results: Vec> = futures::future::join_all(handles) - .await - .into_iter() - .map(|r| r.unwrap()) - .collect(); - - // Only 1 should succeed (3MB limit, 2MB request, can only fit one) - let succeeded = results.iter().filter(|r| r.is_some()).count(); - let failed = results.iter().filter(|r| r.is_none()).count(); - - assert_eq!(succeeded, 1, "Expected exactly 1 task to acquire memory"); - assert_eq!(failed, 2, "Expected 2 tasks to fail"); - - // Clean up - drop(results); - assert_eq!(manager.used_bytes(), 0); - } -} diff --git a/src/mito2/src/compaction/compactor.rs b/src/mito2/src/compaction/compactor.rs index bd18adbb1a..49ea25ffcd 100644 --- a/src/mito2/src/compaction/compactor.rs +++ b/src/mito2/src/compaction/compactor.rs @@ -37,7 +37,8 @@ use crate::access_layer::{ }; use crate::cache::{CacheManager, CacheManagerRef}; use crate::compaction::picker::PickerOutput; -use crate::compaction::{CompactionOutput, CompactionSstReaderBuilder, find_dynamic_options}; +use crate::compaction::reader::CompactionSstReaderBuilder; +use crate::compaction::{CompactionOutput, find_dynamic_options}; use crate::config::MitoConfig; use crate::engine::region_hook::{RegionHookRef, SstFileInfo}; use crate::error; diff --git a/src/mito2/src/compaction/picker.rs b/src/mito2/src/compaction/picker.rs index 403538606f..be82730507 100644 --- a/src/mito2/src/compaction/picker.rs +++ b/src/mito2/src/compaction/picker.rs @@ -17,6 +17,7 @@ use std::sync::Arc; use api::v1::region::compact_request; use common_time::range::TimestampRange; +use common_time::{TimeToLive, Timestamp}; use serde::{Deserialize, Serialize}; use crate::compaction::compactor::CompactionRegion; @@ -26,6 +27,7 @@ use crate::compaction::{CompactionOutput, SerializedCompactionOutput}; use crate::region::options::CompactionOptions; use crate::sst::file::{FileHandle, FileMeta}; use crate::sst::file_purger::FilePurger; +use crate::sst::version::LevelMeta; #[async_trait::async_trait] pub(crate) trait CompactionTask: Debug + Send + Sync + 'static { @@ -150,6 +152,22 @@ pub fn new_picker( } } +/// Finds all expired SSTs across levels. +pub(super) fn get_expired_ssts( + levels: &[LevelMeta], + ttl: Option, + now: Timestamp, +) -> Vec { + let Some(ttl) = ttl else { + return vec![]; + }; + + levels + .iter() + .flat_map(|l| l.get_expired_files(&now, &ttl).into_iter()) + .collect() +} + #[cfg(test)] mod tests { use store_api::storage::FileId; diff --git a/src/mito2/src/compaction/reader.rs b/src/mito2/src/compaction/reader.rs new file mode 100644 index 0000000000..04055ed9c1 --- /dev/null +++ b/src/mito2/src/compaction/reader.rs @@ -0,0 +1,244 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; +use std::sync::Arc; + +use common_time::Timestamp; +use common_time::range::TimestampRange; +use common_time::timestamp::TimeUnit; +use datafusion_common::ScalarValue; +use datafusion_expr::Expr; +use datatypes::extension::json::is_structured_json_field; +use datatypes::types::json_type::JsonNativeType; +use parquet::arrow::parquet_to_arrow_schema; +use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData}; +use snafu::{OptionExt, ResultExt}; +use store_api::metadata::RegionMetadataRef; + +use crate::access_layer::AccessLayerRef; +use crate::cache::{CacheManagerRef, CacheStrategy}; +use crate::error::{ + DataTypeMismatchSnafu, ParquetToArrowSchemaSnafu, Result, TimeRangePredicateOverflowSnafu, +}; +use crate::read::FlatSource; +use crate::read::flat_projection::FlatProjectionMapper; +use crate::read::read_columns::ReadColumns; +use crate::read::scan_region::{PredicateGroup, ScanInput}; +use crate::read::seq_scan::SeqScan; +use crate::region::options::MergeMode; +use crate::sst::file::FileHandle; +use crate::sst::parquet::reader::MetadataCacheMetrics; + +/// Builders to create [BoxedRecordBatchStream] for compaction. +pub(crate) struct CompactionSstReaderBuilder<'a> { + pub(crate) metadata: RegionMetadataRef, + pub(crate) sst_layer: AccessLayerRef, + pub(crate) cache: CacheManagerRef, + pub(crate) inputs: &'a [FileHandle], + pub(crate) append_mode: bool, + pub(crate) filter_deleted: bool, + pub(crate) time_range: Option, + pub(crate) merge_mode: MergeMode, +} + +impl CompactionSstReaderBuilder<'_> { + /// Build a [FlatSource] that yields Arrow `RecordBatch`s from reading all the input SST files, + /// for compaction. The schema of the [FlatSource] is unified. + pub(crate) async fn build_flat_sst_reader(self) -> Result { + let scan_input = self.build_scan_input().await?; + + let schema = scan_input.mapper.output_schema(); + let schema = schema.arrow_schema(); + + let stream = SeqScan::new(scan_input) + .build_flat_reader_for_compaction() + .await?; + Ok(FlatSource::new_stream(schema.clone(), stream)) + } + + async fn build_scan_input(self) -> Result { + let schema = self.metadata.schema.arrow_schema(); + let parquet_metadata = self.collect_parquet_metadata().await?; + let batch_size = crate::batch_size::estimate_batch_size( + parquet_metadata + .iter() + .flat_map(|metadata| metadata.row_groups()) + .map(|row_group| { + let uncompressed_bytes = row_group + .columns() + .iter() + .map(|column| column.uncompressed_size() as u64) + .sum(); + (row_group.num_rows() as u64, uncompressed_bytes) + }), + ); + let json_type_hint = if schema.fields().iter().any(is_structured_json_field) { + let mut json_type_hint = schema + .fields() + .iter() + .filter(|&field| is_structured_json_field(field)) + .map(|field| (field.name().clone(), JsonNativeType::Null)) + .collect::>(); + + for metadata in &parquet_metadata { + let file_metadata = metadata.file_metadata(); + let schema = parquet_to_arrow_schema( + file_metadata.schema_descr(), + file_metadata.key_value_metadata(), + ) + .context(ParquetToArrowSchemaSnafu { + file: "compaction input", + })?; + for field in schema.fields() { + let Some(merged) = json_type_hint.get_mut(field.name()) else { + continue; + }; + + let json_type = JsonNativeType::try_from(field.data_type()) + .context(DataTypeMismatchSnafu)?; + merged.merge(&json_type); + } + } + + Some(json_type_hint) + } else { + None + }; + + let projection = (0..self.metadata.column_metadatas.len()).collect(); + let read_columns = ReadColumns::from_deduped_column_ids( + self.metadata.column_metadatas.iter().map(|x| x.column_id), + ); + let mapper = FlatProjectionMapper::new_with_read_columns( + &self.metadata, + projection, + read_columns, + json_type_hint.as_ref(), + )?; + + let mut scan_input = ScanInput::new(self.sst_layer, mapper) + .with_files(self.inputs.to_vec()) + .with_compaction(true) + .with_batch_size(batch_size) + .with_append_mode(self.append_mode) + // We use special cache strategy for compaction. + .with_cache(CacheStrategy::Compaction(self.cache)) + .with_filter_deleted(self.filter_deleted) + // We ignore file not found error during compaction. + .with_ignore_file_not_found(true) + .with_merge_mode(self.merge_mode); + + // This serves as a workaround of https://github.com/GreptimeTeam/greptimedb/issues/3944 + // by converting time ranges into predicate. + if let Some(time_range) = self.time_range { + scan_input = + scan_input.with_predicate(time_range_to_predicate(time_range, &self.metadata)?); + } + + Ok(scan_input) + } + + async fn collect_parquet_metadata(&self) -> Result>> { + let mut metadata = Vec::with_capacity(self.inputs.len()); + + for file_handle in self.inputs { + let file_path = + file_handle.file_path(self.sst_layer.table_dir(), self.sst_layer.path_type()); + let file_size = file_handle.meta_ref().file_size; + let parquet_metadata = match self + .sst_layer + .read_sst(file_handle.clone()) + .cache(CacheStrategy::Compaction(self.cache.clone())) + .read_parquet_metadata( + &file_path, + file_size, + &mut MetadataCacheMetrics::default(), + PageIndexPolicy::default(), + ) + .await + .map(|x| x.0.parquet_metadata()) + { + Ok(x) => x, + Err(e) if e.is_object_not_found() => continue, + Err(e) => return Err(e), + }; + metadata.push(parquet_metadata); + } + Ok(metadata) + } +} + +/// Converts time range to predicates so that rows outside the range will be filtered. +fn time_range_to_predicate( + range: TimestampRange, + metadata: &RegionMetadataRef, +) -> Result { + let ts_col = metadata.time_index_column(); + + // safety: time index column's type must be a valid timestamp type. + let ts_col_unit = ts_col + .column_schema + .data_type + .as_timestamp() + .unwrap() + .unit(); + + let exprs = match (range.start(), range.end()) { + (Some(start), Some(end)) => { + vec![ + datafusion_expr::col(ts_col.column_schema.name.clone()) + .gt_eq(ts_to_lit(*start, ts_col_unit)?), + datafusion_expr::col(ts_col.column_schema.name.clone()) + .lt(ts_to_lit(*end, ts_col_unit)?), + ] + } + (Some(start), None) => { + vec![ + datafusion_expr::col(ts_col.column_schema.name.clone()) + .gt_eq(ts_to_lit(*start, ts_col_unit)?), + ] + } + + (None, Some(end)) => { + vec![ + datafusion_expr::col(ts_col.column_schema.name.clone()) + .lt(ts_to_lit(*end, ts_col_unit)?), + ] + } + (None, None) => { + return Ok(PredicateGroup::default()); + } + }; + + let predicate = PredicateGroup::new(metadata, &exprs)?; + Ok(predicate) +} + +fn ts_to_lit(ts: Timestamp, ts_col_unit: TimeUnit) -> Result { + let ts = ts + .convert_to(ts_col_unit) + .context(TimeRangePredicateOverflowSnafu { + timestamp: ts, + unit: ts_col_unit, + })?; + let val = ts.value(); + let scalar_value = match ts_col_unit { + TimeUnit::Second => ScalarValue::TimestampSecond(Some(val), None), + TimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(val), None), + TimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(val), None), + TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(val), None), + }; + Ok(datafusion_expr::lit(scalar_value)) +} diff --git a/src/mito2/src/compaction/scheduler.rs b/src/mito2/src/compaction/scheduler.rs new file mode 100644 index 0000000000..90d54d9f61 --- /dev/null +++ b/src/mito2/src/compaction/scheduler.rs @@ -0,0 +1,603 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +mod planning; +mod state; + +use std::collections::HashMap; +use std::sync::Arc; + +use api::v1::region::compact_request; +use common_base::Plugins; +use common_memory_manager::OnExhaustedPolicy; +use common_meta::key::SchemaMetadataManagerRef; +use common_telemetry::{debug, error, info}; +use common_time::range::TimestampRange; +pub(crate) use planning::CompactionPickFinished; +pub use planning::CompactionRequest; +use state::{ActiveCompaction, CompactionStatus, PendingCompaction, RequestCancelResult}; +pub(crate) use state::{CompactionExecution, LocalCompactionState}; +use store_api::storage::RegionId; +use tokio::sync::mpsc::Sender; + +use crate::access_layer::AccessLayerRef; +use crate::cache::CacheManagerRef; +use crate::compaction::memory_manager::CompactionMemoryManager; +use crate::compaction::task::MAX_PARALLEL_COMPACTION; +use crate::config::MitoConfig; +use crate::error::{ + CompactionCancelledSnafu, Error, RegionClosedSnafu, RegionDroppedSnafu, RegionTruncatedSnafu, + Result, +}; +use crate::region::version::VersionControlRef; +use crate::region::{ManifestContextRef, RegionLeaderState, RegionRoleState}; +use crate::request::{DdlRequest, OptionOutputTx, SenderDdlRequest, WorkerRequestWithTime}; +use crate::schedule::scheduler::SchedulerRef; +use crate::worker::WorkerListener; + +/// Compaction scheduler tracks and manages compaction tasks. +pub(crate) struct CompactionScheduler { + scheduler: SchedulerRef, + /// Compacting regions. + region_status: HashMap, + /// Request sender of the worker that this scheduler belongs to. + request_sender: Sender, + cache_manager: CacheManagerRef, + engine_config: Arc, + memory_manager: Arc, + memory_policy: OnExhaustedPolicy, + listener: WorkerListener, + /// Plugins for the compaction scheduler. + plugins: Plugins, + /// Scheduler-wide generation counter for compaction plans and executions. + /// It outlives region statuses so close/reopen cannot reuse an old identity. + next_plan_id: u64, +} + +fn requires_pending_compaction_slot( + options: &compact_request::Options, + time_range: Option, +) -> bool { + matches!(options, compact_request::Options::StrictWindow(_)) || time_range.is_some() +} + +impl CompactionScheduler { + #[allow(clippy::too_many_arguments)] + pub(crate) fn new( + scheduler: SchedulerRef, + request_sender: Sender, + cache_manager: CacheManagerRef, + engine_config: Arc, + listener: WorkerListener, + plugins: Plugins, + memory_manager: Arc, + memory_policy: OnExhaustedPolicy, + ) -> Self { + Self { + scheduler, + region_status: HashMap::new(), + request_sender, + cache_manager, + engine_config, + memory_manager, + memory_policy, + listener, + plugins, + next_plan_id: 0, + } + } + + /// Returns the current plan id and advances the counter. + /// + /// Takes the counter instead of `&mut self` so callers can bump it while + /// holding a mutable borrow of a region status. + fn next_plan_id(counter: &mut u64) -> u64 { + let plan_id = *counter; + *counter = counter.wrapping_add(1); + plan_id + } + + /// Schedules a compaction for the region. + /// Returns whether a compaction is scheduled. + #[allow(clippy::too_many_arguments)] + pub(crate) fn schedule_compaction( + &mut self, + region_id: RegionId, + compact_options: compact_request::Options, + version_control: &VersionControlRef, + access_layer: &AccessLayerRef, + waiter: OptionOutputTx, + manifest_ctx: &ManifestContextRef, + schema_metadata_manager: SchemaMetadataManagerRef, + max_parallelism: usize, + ) -> Result { + self.schedule_compaction_with_time_range( + region_id, + compact_options, + version_control, + access_layer, + waiter, + manifest_ctx, + schema_metadata_manager, + max_parallelism, + None, + ) + } + + /// Schedules a compaction constrained by an optional time range. + #[allow(clippy::too_many_arguments)] + pub(crate) fn schedule_compaction_with_time_range( + &mut self, + region_id: RegionId, + compact_options: compact_request::Options, + version_control: &VersionControlRef, + access_layer: &AccessLayerRef, + waiter: OptionOutputTx, + manifest_ctx: &ManifestContextRef, + schema_metadata_manager: SchemaMetadataManagerRef, + max_parallelism: usize, + time_range: Option, + ) -> Result { + // skip compaction if region is in staging state + let current_state = manifest_ctx.current_state(); + if current_state == RegionRoleState::Leader(RegionLeaderState::Staging) { + info!( + "Skipping compaction for region {} in staging mode, options: {:?}", + region_id, compact_options + ); + waiter.send(Ok(0)); + return Ok(false); + } + + if let Some(status) = self.region_status.get_mut(®ion_id) { + // Pending Truncate/EnterStaging requests form a scheduling fence. Any later + // compaction with a waiter is an explicit request and receives CompactionCancelled; + // automatic triggers have no waiter, so sending the error is a no-op and the trigger + // is simply ignored. + if !status.pending_ddl_requests.is_empty() { + waiter.send(CompactionCancelledSnafu.fail()); + info!( + "Region {} has pending DDL requests, ignoring compaction: {:?}", + region_id, compact_options + ); + return Ok(false); + } + + if requires_pending_compaction_slot(&compact_options, time_range) { + // Incoming compaction request is manually triggered. + status.set_pending_request(PendingCompaction { + options: compact_options, + waiter, + max_parallelism, + time_range, + }); + info!( + "Region {} is compacting, manually compaction will be re-scheduled.", + region_id + ); + } else { + status.merge_regular_trigger(waiter); + } + return Ok(false); + } + + // Publish the picking phase before dispatching background planning. + let mut status = + CompactionStatus::new(region_id, version_control.clone(), access_layer.clone()); + let request = status.new_compaction_request( + self.request_sender.clone(), + self.engine_config.clone(), + self.cache_manager.clone(), + manifest_ctx, + self.listener.clone(), + schema_metadata_manager, + max_parallelism, + ); + let plan_id = Self::next_plan_id(&mut self.next_plan_id); + status.start_picking_with_time_range(plan_id, time_range); + status.merge_waiter(waiter); + self.region_status.insert(region_id, status); + self.dispatch_compaction_planning(plan_id, request, compact_options, time_range); + self.listener.on_compaction_scheduled(region_id); + Ok(true) + } + + // Handle pending manual compaction request for the region. + // + // Returns true if should early return, false otherwise. + pub(crate) fn handle_pending_compaction_request( + &mut self, + region_id: RegionId, + manifest_ctx: &ManifestContextRef, + schema_metadata_manager: SchemaMetadataManagerRef, + ) -> bool { + let Some(status) = self.region_status.get_mut(®ion_id) else { + return true; + }; + + // If there is a pending manual compaction request, schedule it. + // and defer returning the pending DDL requests to the caller. + let Some(pending_request) = std::mem::take(&mut status.pending_request) else { + return false; + }; + + let PendingCompaction { + options, + waiter, + max_parallelism, + time_range, + } = pending_request; + + let request = status.new_compaction_request( + self.request_sender.clone(), + self.engine_config.clone(), + self.cache_manager.clone(), + manifest_ctx, + self.listener.clone(), + schema_metadata_manager, + max_parallelism, + ); + status.merge_waiter(waiter); + // Bump the counter through a disjoint field borrow so the `status` + // borrow stays alive; nothing could have removed the status since it + // was fetched above. + let plan_id = Self::next_plan_id(&mut self.next_plan_id); + status.start_picking_with_time_range(plan_id, time_range); + self.dispatch_compaction_planning(plan_id, request, options, time_range); + debug!( + "Successfully scheduled manual compaction planning for region id: {}", + region_id + ); + true + } + + /// Notifies the scheduler that the compaction job is finished successfully. + async fn on_compaction_finished( + &mut self, + region_id: RegionId, + manifest_ctx: &ManifestContextRef, + schema_metadata_manager: SchemaMetadataManagerRef, + ) -> Vec { + if !self + .region_status + .get(®ion_id) + .is_some_and(|s| s.is_busy()) + { + return Vec::new(); + } + + if self.handle_pending_compaction_request( + region_id, + manifest_ctx, + schema_metadata_manager.clone(), + ) { + return Vec::new(); + } + + // The region status might be removed by the previous steps. + // So we return empty DDL requests. + let Some(status) = self.region_status.get_mut(®ion_id) else { + return Vec::new(); + }; + let Some(mut active) = status.take_active() else { + return Vec::new(); + }; + + for waiter in std::mem::take(&mut active.waiters) { + waiter.send(Ok(0)); + } + + // A queued DDL was waiting for the current task to terminate; chaining + // another compaction ahead of it would delay the DDL by a whole extra + // plan/execution cycle, so dispatch the DDLs first. + let pending_ddl_requests = std::mem::take(&mut status.pending_ddl_requests); + if !pending_ddl_requests.is_empty() { + // The just-finished compaction satisfies any retained regular triggers. + for waiter in active.regular_followup_waiters.take().unwrap_or_default() { + waiter.send(Ok(0)); + } + self.region_status.remove(®ion_id); + // If there are pending DDL requests, we should return them to the caller. + // And skip try to schedule next compaction task. + return pending_ddl_requests; + } + + if active.regular_followup_waiters.is_some() { + self.schedule_next_compaction_with_active( + region_id, + manifest_ctx, + schema_metadata_manager, + Some(active), + None, + ); + return Vec::new(); + } + Vec::new() + } + + /// Returns whether a terminal notification belongs to the installed execution. + /// Background work may finish after its region status has been replaced, so + /// matching the region id alone is insufficient. + 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) async fn on_execution_finished( + &mut self, + region_id: RegionId, + execution: &CompactionExecution, + manifest_ctx: &ManifestContextRef, + schema_metadata_manager: SchemaMetadataManagerRef, + ) -> Vec { + // A stale finish must not clear the replacement phase or notify its waiters and DDLs. + 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(CompactionStatus::is_busy) + .unwrap_or(false) + } + + /// Removes the region status if it has no running task. + /// + /// A finished compaction leaves an idle status (`active = None`) behind when + /// there is nothing more to schedule. If the caller decides not to chain + /// the next compaction, it must remove the idle status; otherwise the + /// status becomes a zombie that makes `schedule_compaction` swallow all + /// future compaction triggers of the region. + pub(crate) fn remove_idle_status(&mut self, region_id: RegionId) { + if self + .region_status + .get(®ion_id) + .is_some_and(|status| !status.is_busy()) + { + self.region_status.remove(®ion_id); + } + } + + /// Schedules next compaction upon a finished compaction. + /// Returns whether the compaction is scheduled. + pub(crate) fn schedule_next_compaction( + &mut self, + region_id: RegionId, + manifest_ctx: &ManifestContextRef, + schema_metadata_manager: SchemaMetadataManagerRef, + ) -> bool { + let Some(status) = self.region_status.get_mut(®ion_id) else { + return false; + }; + // A plan is already in flight; treat it as scheduled instead of + // overwriting the current phase and orphaning the in-flight planning. + if status.is_busy() { + return true; + } + + let time_range = status.time_range; + self.schedule_next_compaction_with_active( + region_id, + manifest_ctx, + schema_metadata_manager, + None, + time_range, + ) + } + + fn schedule_next_compaction_with_active( + &mut self, + region_id: RegionId, + manifest_ctx: &ManifestContextRef, + schema_metadata_manager: SchemaMetadataManagerRef, + active: Option, + time_range: Option, + ) -> bool { + let Some(status) = self.region_status.get_mut(®ion_id) else { + return false; + }; + // We should always try to compact the region until picker returns None. + let request = status.new_compaction_request( + self.request_sender.clone(), + self.engine_config.clone(), + self.cache_manager.clone(), + manifest_ctx, + self.listener.clone(), + schema_metadata_manager, + MAX_PARALLEL_COMPACTION, + ); + // Bump the counter through a disjoint field borrow so the `status` + // borrow stays alive; nothing could have removed the status since it + // was fetched above. + let plan_id = Self::next_plan_id(&mut self.next_plan_id); + status.start_regular_picking(plan_id, active, time_range); + self.dispatch_compaction_planning( + plan_id, + request, + compact_request::Options::Regular(Default::default()), + time_range, + ); + debug!( + "Successfully scheduled next compaction planning for region id: {}", + region_id + ); + true + } + + /// Notifies the scheduler that the compaction job is cancelled cooperatively. + 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 { + // A stale cancellation must not remove a replacement execution's status. + 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. + 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, + ) { + // A stale failure must not tear down a replacement execution. + 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( + region_id, + Arc::new(RegionDroppedSnafu { region_id }.build()), + ); + } + + /// Notifies the scheduler that the region is closed. + pub(crate) fn on_region_closed(&mut self, region_id: RegionId) { + self.remove_region_on_failure(region_id, Arc::new(RegionClosedSnafu { region_id }.build())); + } + + /// Notifies the scheduler that the region is truncated. + pub(crate) fn on_region_truncated(&mut self, region_id: RegionId) { + self.remove_region_on_failure( + region_id, + Arc::new(RegionTruncatedSnafu { region_id }.build()), + ); + } + + /// Cancels the running compaction and queues its dependent DDL atomically. + /// + /// Production callers currently use this only for [`DdlRequest::Truncate`] and + /// [`DdlRequest::EnterStaging`]. If cancellation is still possible, the current picking or + /// local execution is asked to stop; otherwise the DDL waits for its terminal notification. + /// The worker dispatches the queued DDL only after that notification is handled, preventing + /// truncate or enter-staging from racing with compaction planning, execution, or commit. + /// Returns the sender and typed request unchanged if compaction is not running. + pub(crate) fn try_cancel_and_add_ddl( + &mut self, + region_id: RegionId, + sender: OptionOutputTx, + request: T, + into_ddl_request: impl FnOnce(T) -> DdlRequest, + ) -> std::result::Result<(), (OptionOutputTx, T)> { + let Some(status) = self.region_status.get_mut(®ion_id) else { + return Err((sender, request)); + }; + if status.request_cancel() == RequestCancelResult::NotRunning { + return Err((sender, request)); + } + + let request = SenderDdlRequest { + region_id, + sender, + request: into_ddl_request(request), + }; + debug!( + "Added pending DDL request for region: {}, ddl: {:?}", + request.region_id, request.request + ); + // The first queued Truncate/EnterStaging also fences later regular triggers from + // creating more follow-ups ahead of the DDL. + status.pending_ddl_requests.push(request); + Ok(()) + } + + #[cfg(test)] + fn add_ddl_request_to_pending(&mut self, request: SenderDdlRequest) { + self.region_status + .get_mut(&request.region_id) + .unwrap() + .pending_ddl_requests + .push(request); + } + + #[cfg(test)] + pub(crate) fn has_pending_ddls(&self, region_id: RegionId) -> bool { + let has_pending = self + .region_status + .get(®ion_id) + .map(|status| !status.pending_ddl_requests.is_empty()) + .unwrap_or(false); + debug!( + "Checked pending DDL requests for region: {}, has_pending: {}", + region_id, has_pending + ); + has_pending + } + + #[cfg(test)] + pub(crate) fn request_cancel(&mut self, region_id: RegionId) -> RequestCancelResult { + let Some(status) = self.region_status.get_mut(®ion_id) else { + return RequestCancelResult::NotRunning; + }; + + status.request_cancel() + } + + fn remove_region_on_failure(&mut self, region_id: RegionId, err: Arc) { + // Remove this region. + let Some(status) = self.region_status.remove(®ion_id) else { + return; + }; + + // Notifies all pending tasks. + status.on_failure(err); + } + + fn remove_region_on_cancel(&mut self, region_id: RegionId) -> Vec { + let Some(status) = self.region_status.remove(®ion_id) else { + return Vec::new(); + }; + + status.on_cancel() + } +} + +impl Drop for CompactionScheduler { + fn drop(&mut self) { + for (region_id, status) in self.region_status.drain() { + // We are shutting down so notify all pending tasks. + status.on_failure(Arc::new(RegionClosedSnafu { region_id }.build())); + } + } +} + +#[cfg(test)] +#[path = "scheduler_test.rs"] +mod tests; diff --git a/src/mito2/src/compaction/scheduler/planning.rs b/src/mito2/src/compaction/scheduler/planning.rs new file mode 100644 index 0000000000..386e9dfaeb --- /dev/null +++ b/src/mito2/src/compaction/scheduler/planning.rs @@ -0,0 +1,656 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::fmt; +use std::future::Future; +use std::sync::{Arc, Mutex}; +use std::time::Instant; + +use api::v1::region::compact_request; +use common_base::Plugins; +use common_base::cancellation::CancellationHandle; +use common_meta::key::SchemaMetadataManagerRef; +use common_telemetry::{debug, error, info, warn}; +use common_time::TimeToLive; +use common_time::range::TimestampRange; +use futures::FutureExt; +use snafu::ResultExt; +use store_api::storage::RegionId; +use tokio::sync::mpsc::{self, Sender}; + +use crate::access_layer::AccessLayerRef; +use crate::cache::CacheManagerRef; +use crate::compaction::compactor::{CompactionRegion, CompactionVersion, DefaultCompactor}; +use crate::compaction::picker::{CompactionTask, PickerOutput, new_picker}; +use crate::compaction::scheduler::CompactionScheduler; +use crate::compaction::scheduler::state::{ + CompactingFiles, CompactionExecution, CompactionPhase, CompactionStatus, LocalCompactionState, +}; +use crate::compaction::task::CompactionTaskImpl; +use crate::compaction::{CompactionOutput, find_dynamic_options}; +use crate::config::MitoConfig; +use crate::error::{ + CompactRegionSnafu, Error, JoinSnafu, RemoteCompactionSnafu, Result, UnexpectedSnafu, +}; +use crate::metrics::{ + COMPACTION_MEMORY_REJECTED, COMPACTION_STAGE_ELAPSED, INFLIGHT_COMPACTION_COUNT, +}; +use crate::region::ManifestContextRef; +use crate::region::options::RegionOptions; +use crate::request::{ + BackgroundNotify, OutputTx, SenderDdlRequest, WorkerRequest, WorkerRequestWithTime, +}; +use crate::schedule::remote_job_scheduler::{ + CompactionJob, DefaultNotifier, RemoteJob, RemoteJobSchedulerRef, +}; +use crate::sst::file::FileHandle; +use crate::sst::version::SstVersion; +use crate::worker::WorkerListener; + +/// Region compaction request. +pub struct CompactionRequest { + pub(crate) engine_config: Arc, + pub(crate) current_version: CompactionVersion, + pub(crate) access_layer: AccessLayerRef, + /// Sender to send notification to the region worker. + pub(crate) request_sender: mpsc::Sender, + /// Start time of compaction task. + pub(crate) start_time: Instant, + pub(crate) cache_manager: CacheManagerRef, + pub(crate) manifest_ctx: ManifestContextRef, + pub(crate) listener: WorkerListener, + pub(crate) schema_metadata_manager: SchemaMetadataManagerRef, + pub(crate) max_parallelism: usize, +} + +impl CompactionRequest { + pub(crate) fn region_id(&self) -> RegionId { + self.current_version.metadata.region_id + } +} + +/// 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) result: CompactionPlanningResult, +} + +pub(crate) struct PreparedCompaction { + pub(super) compaction_region: CompactionRegion, + pub(super) picker_output: PickerOutput, + start_time: Instant, + ttl: TimeToLive, +} + +impl CompactionScheduler { + pub(super) fn dispatch_compaction_planning( + &self, + plan_id: u64, + request: CompactionRequest, + options: compact_request::Options, + time_range: Option, + ) { + 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 planning = Self::prepare_compaction( + request, + options, + plugins, + max_background_compactions, + time_range, + ); + Self::notify_planning_result(region_id, plan_id, request_sender, planning).await; + }); + } + + /// Runs the planning future and always sends the planning result back to + /// the worker, even if the planning panics. + /// + /// The worker only leaves the picking phase after it receives the + /// `CompactionPickFinished` notification. If a panicked planning task + /// swallowed the notification, the region would be stuck in the picking + /// phase forever, blocking all future compactions and pending DDLs (e.g. + /// entering staging) of the region. + pub(super) async fn notify_planning_result( + region_id: RegionId, + plan_id: u64, + request_sender: Sender, + planning: impl Future + Send, + ) { + // The idiomatic way to handle a panic result. + let result = std::panic::AssertUnwindSafe(planning).catch_unwind().await.unwrap_or_else(|payload| { + let reason = if let Some(message) = payload.as_ref().downcast_ref::<&str>() { + message.to_string() + } else if let Some(message) = payload.as_ref().downcast_ref::() { + message.clone() + } else { + "unknown panic".to_string() + }; + CompactionPlanningResult::Error(Arc::new( + UnexpectedSnafu { + reason: format!( + "Compaction planning panicked for region {region_id}, plan_id {plan_id}: {reason}" + ), + } + .build(), + )) + }); + 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, + 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, + time_range: Option, + ) -> CompactionPlanningResult { + let region_id = request.region_id(); + let (dynamic_compaction_opts, ttl) = find_dynamic_options( + region_id, + &request.current_version.options, + &request.schema_metadata_manager, + ) + .await + .unwrap_or_else(|e| { + warn!(e; "Failed to find dynamic options for region: {}", region_id); + ( + request.current_version.options.compaction.clone(), + request.current_version.options.ttl.unwrap_or_default(), + ) + }); + + let picker = new_picker( + &options, + &dynamic_compaction_opts, + request.current_version.options.append_mode, + Some(max_background_compactions), + time_range, + ); + let region_id = request.region_id(); + let CompactionRequest { + engine_config, + current_version, + access_layer, + request_sender: _, + start_time, + cache_manager, + manifest_ctx, + listener, + schema_metadata_manager: _, + max_parallelism, + } = request; + + debug!( + "Pick compaction strategy {:?} for region: {}, ttl: {:?}", + picker, region_id, ttl + ); + + let compaction_region = CompactionRegion { + region_id, + current_version: current_version.clone(), + region_options: RegionOptions { + compaction: dynamic_compaction_opts.clone(), + ..current_version.options.clone() + }, + engine_config: engine_config.clone(), + region_metadata: current_version.metadata.clone(), + cache_manager: cache_manager.clone(), + access_layer: access_layer.clone(), + manifest_ctx: manifest_ctx.clone(), + file_purger: None, + ttl: Some(ttl), + max_parallelism, + plugins, + }; + + 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(&picker_region) + }) + .await + .context(JoinSnafu) + { + Ok(output) => output, + Err(err) => return CompactionPlanningResult::Error(Arc::new(err)), + }; + + let Some(picker_output) = picker_output else { + return CompactionPlanningResult::NoPlan; + }; + + CompactionPlanningResult::Prepared(PreparedCompaction { + compaction_region, + picker_output, + start_time, + ttl, + }) + } + + pub(crate) async fn handle_compaction_pick_finished( + &mut self, + finished: CompactionPickFinished, + manifest_ctx: &ManifestContextRef, + schema_metadata_manager: SchemaMetadataManagerRef, + ) -> Vec { + let region_id = finished.region_id; + let plan_id = finished.plan_id; + let Some(status) = self.region_status.get(®ion_id) else { + return Vec::new(); + }; + // Picking runs detached from the worker. Its result may arrive after + // close/reopen or replanning installed another Picking phase for this region. + if !status.is_picking(finished.plan_id) { + 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 = status.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 = status.take_waiters(); + match self + .submit_prepared_compaction(prepared, files, waiters, plan_id) + .await + { + Ok(Some(phase)) => { + if let Some(status) = self.region_status.get_mut(®ion_id) { + status.set_phase(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 + } + } + } + + 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(); + }; + let Some(mut active) = status.take_active() else { + return Vec::new(); + }; + for waiter in std::mem::take(&mut active.waiters) { + if let Some(err) = &err { + waiter.send(Err(err.clone()).context(CompactRegionSnafu { region_id })); + } else { + waiter.send(Ok(0)); + } + } + + status.active = Some(active); + if self.handle_pending_compaction_request( + region_id, + manifest_ctx, + schema_metadata_manager.clone(), + ) { + return Vec::new(); + } + + let Some(active) = self + .region_status + .get_mut(®ion_id) + .and_then(CompactionStatus::take_active) + else { + return Vec::new(); + }; + if active.regular_followup_waiters.is_some() { + self.schedule_next_compaction_with_active( + region_id, + manifest_ctx, + schema_metadata_manager, + Some(active), + None, + ); + 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, + mut plan_id: u64, + ) -> 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, files.clone()); + let remote_compaction_job = CompactionJob { + compaction_region: compaction_region.clone(), + picker_output: picker_output.clone(), + start_time, + waiters, + ttl, + }; + + let result = remote_job_scheduler + .schedule( + RemoteJob::CompactionJob(remote_compaction_job), + Box::new(DefaultNotifier::new( + self.request_sender.clone(), + execution.clone(), + )), + ) + .await; + + match result { + Ok(job_id) => { + info!( + "Scheduled remote compaction job {} for region {}", + job_id, region_id + ); + INFLIGHT_COMPACTION_COUNT.inc(); + 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.extend_waiters(e.waiters); + } + return RemoteCompactionSnafu { + region_id, + job_id: None, + reason: e.reason, + } + .fail(); + } + + error!(e; "Failed to schedule remote compaction job for region {}, fallback to local compaction", region_id); + // An error may be ambiguous after the remote scheduler consumed + // the notifier. Fence a delayed remote callback from the local fallback. + plan_id = Self::next_plan_id(&mut self.next_plan_id); + e.waiters + } + } + } else { + debug!( + "Remote compaction is not enabled, fallback to local compaction for region {}", + region_id + ); + waiters + } + } else { + waiters + }; + + // Check whether this local compaction can ever fit before submitting it. + let estimated_bytes = estimate_compaction_bytes(&picker_output); + if let Some(limit_bytes) = self.exceeds_compaction_memory_limit(estimated_bytes) { + COMPACTION_MEMORY_REJECTED + .with_label_values(&["oversized"]) + .inc(); + warn!( + "Skip compaction for region {} because estimated memory {} bytes exceeds compaction memory limit {} bytes", + region_id, estimated_bytes, limit_bytes, + ); + for waiter in waiters { + waiter.send(Ok(0)); + } + return Ok(None); + } + + let cancel_handle = Arc::new(CancellationHandle::default()); + let state = LocalCompactionState::new(cancel_handle.clone()); + let execution = CompactionExecution::new(plan_id, files); + let local_compaction_task = Box::new(CompactionTaskImpl { + state: state.clone(), + execution: execution.clone(), + request_sender: self.request_sender.clone(), + waiters, + start_time, + listener: self.listener.clone(), + picker_output, + compaction_region, + compactor: Arc::new(DefaultCompactor::with_cancel_handle(cancel_handle.clone())), + memory_manager: self.memory_manager.clone(), + memory_policy: self.memory_policy, + estimated_memory_bytes: estimated_bytes, + }); + + 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.append_waiters(&mut task.waiters); + } + Err(err) + } + } + } + + fn submit_compaction_task( + &mut self, + task: Box, + region_id: RegionId, + ) -> 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(); + } 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 { + let limit_bytes = self.memory_manager.limit_bytes(); + if limit_bytes > 0 && estimated_bytes > limit_bytes { + Some(limit_bytes) + } else { + None + } + } +} + +/// Estimates compaction memory as the sum of all input files' maximum row-group +/// uncompressed sizes. +fn estimate_compaction_bytes(picker_output: &PickerOutput) -> u64 { + picker_output + .outputs + .iter() + .flat_map(|output| output.inputs.iter()) + .map(|file: &FileHandle| { + let meta = file.meta_ref(); + meta.max_row_group_uncompressed_size + }) + .sum() +} + +/// Rebuilds picker output with current SST handles while preserving the picker's grouping. +/// +/// Picking runs in background on a version snapshot that may be stale by the +/// time the plan is accepted: a concurrent flush, compaction or index rebuild +/// can replace a selected file with a new handle carrying updated metadata +/// (e.g. `index_version`), or remove the file entirely. The handles in the +/// picker output therefore cannot be used as-is; re-resolving them against the +/// current version both detects gone files (aborting the plan) and ensures the +/// execution reads and reserves the up-to-date handle. +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, + }) +} diff --git a/src/mito2/src/compaction/scheduler/state.rs b/src/mito2/src/compaction/scheduler/state.rs new file mode 100644 index 0000000000..adcc915b79 --- /dev/null +++ b/src/mito2/src/compaction/scheduler/state.rs @@ -0,0 +1,587 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashSet; +use std::sync::{Arc, Mutex}; +use std::time::Instant; + +use api::v1::region::compact_request; +use common_base::cancellation::CancellationHandle; +use common_meta::key::SchemaMetadataManagerRef; +use common_telemetry::debug; +use common_time::range::TimestampRange; +use snafu::ResultExt; +use store_api::storage::RegionId; +use tokio::sync::mpsc::Sender; + +use crate::access_layer::AccessLayerRef; +use crate::cache::CacheManagerRef; +use crate::compaction::compactor::CompactionVersion; +use crate::compaction::picker::PickerOutput; +use crate::compaction::scheduler::planning::CompactionRequest; +use crate::config::MitoConfig; +use crate::error::{ + CompactRegionSnafu, CompactionCancelledSnafu, Error, ManualCompactionOverrideSnafu, +}; +use crate::region::ManifestContextRef; +use crate::region::version::VersionControlRef; +use crate::request::{OptionOutputTx, OutputTx, SenderDdlRequest, WorkerRequestWithTime}; +use crate::sst::file::FileHandle; +use crate::worker::WorkerListener; + +/// Identifies an accepted compaction attempt and keeps its SST reservations alive. +/// The plan id fences terminal notifications from superseded attempts. +#[derive(Debug, Clone)] +pub(crate) struct CompactionExecution { + plan_id: u64, + _files: CompactingFiles, +} + +impl CompactionExecution { + pub(super) fn new(plan_id: u64, files: CompactingFiles) -> Self { + Self { + plan_id, + _files: files, + } + } + + pub(crate) fn matches(&self, other: &Self) -> bool { + self.plan_id == other.plan_id + } + + #[cfg(test)] + pub(crate) fn for_test(plan_id: u64) -> Self { + Self::new(plan_id, CompactingFiles::empty()) + } +} + +#[derive(Debug, Clone)] +pub(crate) struct LocalCompactionState { + cancel_handle: Arc, + commit_started: Arc>, +} + +#[derive(Debug)] +pub(super) enum CompactionPhase { + Picking { + plan_id: u64, + cancelled: bool, + }, + Local { + state: LocalCompactionState, + execution: CompactionExecution, + }, + Remote { + execution: CompactionExecution, + }, +} + +#[derive(Debug)] +pub(super) struct ActiveCompaction { + pub(super) phase: CompactionPhase, + /// Waiters satisfied by the current planning or execution cycle. Picking waiters move into + /// the submitted task; regular triggers coalesced during execution accumulate here. + pub(super) waiters: Vec, + /// Requests one fresh regular picking cycle after the current cycle finishes. It is kept + /// separate because the current picker snapshot may predate the trigger; `Some(empty)` records + /// an automatic trigger without an explicit waiter. + pub(super) regular_followup_waiters: Option>, +} + +impl ActiveCompaction { + pub(super) fn picking(plan_id: u64, waiters: Vec) -> Self { + Self { + phase: CompactionPhase::Picking { + plan_id, + cancelled: false, + }, + waiters, + regular_followup_waiters: None, + } + } + + pub(super) fn start_picking(&mut self, plan_id: u64) { + self.phase = CompactionPhase::Picking { + plan_id, + cancelled: false, + }; + } + + pub(super) fn start_regular_picking(&mut self, plan_id: u64) { + self.waiters + .extend(self.regular_followup_waiters.take().unwrap_or_default()); + self.start_picking(plan_id); + } + + pub(super) fn is_picking(&self, expected_plan_id: u64) -> bool { + matches!( + self.phase, + CompactionPhase::Picking { plan_id, .. } if plan_id == expected_plan_id + ) + } + + pub(super) fn accept_plan(&self, expected_plan_id: u64) -> bool { + matches!( + self.phase, + CompactionPhase::Picking { + plan_id, + cancelled: false, + } if plan_id == expected_plan_id + ) + } + + pub(super) fn matches_execution(&self, execution: &CompactionExecution) -> bool { + match &self.phase { + CompactionPhase::Picking { .. } => None, + CompactionPhase::Local { execution, .. } | CompactionPhase::Remote { execution } => { + Some(execution) + } + } + .is_some_and(|current| current.matches(execution)) + } + + pub(super) fn request_cancel(&mut self) -> RequestCancelResult { + match &mut self.phase { + CompactionPhase::Picking { cancelled, .. } => { + if *cancelled { + RequestCancelResult::AlreadyCancelling + } else { + *cancelled = true; + RequestCancelResult::CancelIssued + } + } + CompactionPhase::Local { state, .. } => state.request_cancel(), + CompactionPhase::Remote { .. } => RequestCancelResult::TooLateToCancel, + } + } + + pub(super) fn merge_regular_trigger(&mut self, mut waiter: OptionOutputTx) { + if matches!(self.phase, CompactionPhase::Picking { .. }) { + let regular_followup_waiters = self.regular_followup_waiters.get_or_insert_default(); + if let Some(waiter) = waiter.take_inner() { + regular_followup_waiters.push(waiter); + } + } else { + self.merge_waiter(waiter); + } + } + + pub(super) fn merge_waiter(&mut self, mut waiter: OptionOutputTx) { + if let Some(waiter) = waiter.take_inner() { + self.waiters.push(waiter); + } + } +} + +/// Owns atomic reservations for every SST selected by a compaction plan. +#[derive(Debug, Clone)] +pub(super) struct CompactingFiles { + _inner: Arc, +} + +#[derive(Debug)] +struct CompactingFilesInner { + files: Vec, +} + +impl CompactingFiles { + pub(super) 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)] + pub(super) 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 { + pub(super) fn new(cancel_handle: Arc) -> Self { + Self { + cancel_handle, + commit_started: Arc::new(Mutex::new(false)), + } + } + + /// Returns the cancellation handle for this compaction task. + pub(crate) fn cancel_handle(&self) -> Arc { + self.cancel_handle.clone() + } + + /// Marks the compaction task as started to commit, + /// which means the compaction task is in the final stage and is about to update region version and manifest. + /// It will reject cancellation request after this method is called. + /// + /// Returns true if this is the first time to mark commit started, false otherwise. + pub(crate) fn mark_commit_started(&self) -> bool { + let mut commit_started = self.commit_started.lock().unwrap(); + if self.cancel_handle.is_cancelled() { + return false; + } + *commit_started = true; + true + } + + /// Request cancellation for this compaction task. + pub(crate) fn request_cancel(&self) -> RequestCancelResult { + // The cancel handle must under the lock of `commit_started` to avoid racing between cancellation and commit. + let commit_started = self.commit_started.lock().unwrap(); + if *commit_started { + return RequestCancelResult::TooLateToCancel; + } + if self.cancel_handle.is_cancelled() { + return RequestCancelResult::AlreadyCancelling; + } + + self.cancel_handle.cancel(); + RequestCancelResult::CancelIssued + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum RequestCancelResult { + CancelIssued, + AlreadyCancelling, + TooLateToCancel, + NotRunning, +} + +/// Status of running and pending region compaction tasks. +pub(super) struct CompactionStatus { + /// Id of the region. + pub(super) region_id: RegionId, + /// Version control of the region. + pub(super) version_control: VersionControlRef, + /// Access layer of the region. + pub(super) access_layer: AccessLayerRef, + /// Current compaction lifecycle. `None` is the existing transient idle state. + // TODO: Remove idle statuses and make ActiveCompaction non-optional once chained + // scheduling can recreate the status from region context. + pub(super) active: Option, + /// Optional range retained by automatic continuations of the current compaction. + pub(super) time_range: Option, + /// Pending compactions that are supposed to run as soon as current compaction task finished. + /// + /// This holds strict-window requests and ranged regular requests. An unrestricted regular + /// request is instead merged into `ActiveCompaction::regular_followup_waiters` or `waiters`. + pub(super) pending_request: Option, + /// Pending DDL requests that should run when compaction is done. + /// + /// Although [`SenderDdlRequest`] can wrap any DDL variant, production code only queues + /// [`crate::request::DdlRequest::Truncate`] and [`crate::request::DdlRequest::EnterStaging`] here. Both must serialize with + /// compaction so they observe the version after compaction terminates. + pub(super) pending_ddl_requests: Vec, +} + +impl CompactionStatus { + /// Creates a new [CompactionStatus] + pub(super) fn new( + region_id: RegionId, + version_control: VersionControlRef, + access_layer: AccessLayerRef, + ) -> CompactionStatus { + CompactionStatus { + region_id, + version_control, + access_layer, + active: None, + time_range: None, + pending_request: None, + pending_ddl_requests: Vec::new(), + } + } + + #[cfg(test)] + pub(super) fn start_picking(&mut self, plan_id: u64) { + self.start_picking_with_time_range(plan_id, None); + } + + pub(super) fn start_picking_with_time_range( + &mut self, + plan_id: u64, + time_range: Option, + ) { + self.time_range = time_range; + if let Some(active) = &mut self.active { + active.start_picking(plan_id); + } else { + self.active = Some(ActiveCompaction::picking(plan_id, Vec::new())); + } + } + + pub(super) fn start_regular_picking( + &mut self, + plan_id: u64, + active: Option, + time_range: Option, + ) { + self.time_range = time_range; + self.active = Some(if let Some(mut active) = active { + active.start_regular_picking(plan_id); + active + } else { + ActiveCompaction::picking(plan_id, Vec::new()) + }); + } + + pub(super) fn is_picking(&self, expected_plan_id: u64) -> bool { + self.active + .as_ref() + .is_some_and(|active| active.is_picking(expected_plan_id)) + } + + pub(super) fn accept_plan(&self, expected_plan_id: u64) -> bool { + self.active + .as_ref() + .is_some_and(|active| active.accept_plan(expected_plan_id)) + } + + pub(super) fn is_busy(&self) -> bool { + self.active.is_some() + } + + pub(super) fn matches_execution(&self, execution: &CompactionExecution) -> bool { + self.active + .as_ref() + .is_some_and(|active| active.matches_execution(execution)) + } + + #[cfg(test)] + pub(super) fn start_local_task(&mut self) -> LocalCompactionState { + let state = LocalCompactionState::new(Arc::new(CancellationHandle::default())); + let execution = CompactionExecution::new(0, CompactingFiles::empty()); + let phase = CompactionPhase::Local { + state: state.clone(), + execution, + }; + if let Some(active) = &mut self.active { + active.phase = phase; + } else { + self.active = Some(ActiveCompaction { + phase, + waiters: Vec::new(), + regular_followup_waiters: None, + }); + } + state + } + + #[cfg(test)] + pub(super) fn start_remote_task(&mut self) { + let execution = CompactionExecution::new(0, CompactingFiles::empty()); + let phase = CompactionPhase::Remote { execution }; + if let Some(active) = &mut self.active { + active.phase = phase; + } else { + self.active = Some(ActiveCompaction { + phase, + waiters: Vec::new(), + regular_followup_waiters: None, + }); + } + } + + pub(super) fn request_cancel(&mut self) -> RequestCancelResult { + let Some(active) = &mut self.active else { + return RequestCancelResult::NotRunning; + }; + active.request_cancel() + } + + #[cfg(test)] + pub(super) fn clear_running_task(&mut self) -> bool { + self.active.take().is_some() + } + + pub(super) fn merge_regular_trigger(&mut self, waiter: OptionOutputTx) { + if let Some(active) = &mut self.active { + active.merge_regular_trigger(waiter); + } + } + + /// Merge the waiter to the pending compaction. + pub(super) fn merge_waiter(&mut self, waiter: OptionOutputTx) { + if let Some(active) = &mut self.active { + active.merge_waiter(waiter); + } + } + + pub(super) fn take_active(&mut self) -> Option { + self.active.take() + } + + pub(super) fn take_waiters(&mut self) -> Vec { + self.active + .as_mut() + .map(|active| std::mem::take(&mut active.waiters)) + .unwrap_or_default() + } + + pub(super) fn extend_waiters(&mut self, waiters: Vec) { + if let Some(active) = &mut self.active { + active.waiters.extend(waiters); + } + } + + pub(super) fn append_waiters(&mut self, waiters: &mut Vec) { + if let Some(active) = &mut self.active { + active.waiters.append(waiters); + } + } + + pub(super) fn set_phase(&mut self, phase: CompactionPhase) { + if let Some(active) = &mut self.active { + active.phase = phase; + } + } + + /// Set pending compaction request or replace current value if already exist. + pub(super) fn set_pending_request(&mut self, pending: PendingCompaction) { + if let Some(prev) = self.pending_request.replace(pending) { + debug!( + "Replace pending compaction options with new request {:?} for region: {}", + prev.options, self.region_id + ); + prev.waiter.send(ManualCompactionOverrideSnafu.fail()); + } + } + + pub(super) fn on_failure(mut self, err: Arc) { + if let Some(mut active) = self.active.take() { + for waiter in active + .waiters + .drain(..) + .chain(active.regular_followup_waiters.take().unwrap_or_default()) + { + waiter.send(Err(err.clone()).context(CompactRegionSnafu { + region_id: self.region_id, + })); + } + } + + if let Some(pending_compaction) = self.pending_request { + pending_compaction + .waiter + .send(Err(err.clone()).context(CompactRegionSnafu { + region_id: self.region_id, + })); + } + + for pending_ddl in self.pending_ddl_requests { + pending_ddl + .sender + .send(Err(err.clone()).context(CompactRegionSnafu { + region_id: self.region_id, + })); + } + } + + #[must_use] + pub(super) fn on_cancel(mut self) -> Vec { + if let Some(mut active) = self.active.take() { + for waiter in active + .waiters + .drain(..) + .chain(active.regular_followup_waiters.take().unwrap_or_default()) + { + waiter.send(CompactionCancelledSnafu.fail()); + } + } + + if let Some(pending_compaction) = self.pending_request { + pending_compaction.waiter.send( + Err(Arc::new(CompactionCancelledSnafu.build())).context(CompactRegionSnafu { + region_id: self.region_id, + }), + ); + } + + std::mem::take(&mut self.pending_ddl_requests) + } + + /// Creates an immutable request for background compaction planning. + #[allow(clippy::too_many_arguments)] + pub(super) fn new_compaction_request( + &self, + request_sender: Sender, + engine_config: Arc, + cache_manager: CacheManagerRef, + manifest_ctx: &ManifestContextRef, + listener: WorkerListener, + schema_metadata_manager: SchemaMetadataManagerRef, + max_parallelism: usize, + ) -> CompactionRequest { + let current_version = CompactionVersion::from(self.version_control.current().version); + let start_time = Instant::now(); + + CompactionRequest { + engine_config, + current_version, + access_layer: self.access_layer.clone(), + request_sender: request_sender.clone(), + start_time, + cache_manager, + manifest_ctx: manifest_ctx.clone(), + listener, + schema_metadata_manager, + max_parallelism, + } + } +} + +/// Pending compaction request that is supposed to run after current task is finished, +/// typically used for manual compactions. +pub(super) struct PendingCompaction { + /// Compaction options. + pub(crate) options: compact_request::Options, + /// Waiters of pending requests. + pub(crate) waiter: OptionOutputTx, + /// Max parallelism for pending compaction. + pub(crate) max_parallelism: usize, + /// Optional time range that constrains candidate compaction windows. + pub(crate) time_range: Option, +} diff --git a/src/mito2/src/compaction/scheduler_test.rs b/src/mito2/src/compaction/scheduler_test.rs new file mode 100644 index 0000000000..5e97f60f70 --- /dev/null +++ b/src/mito2/src/compaction/scheduler_test.rs @@ -0,0 +1,2221 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::assert_matches; +use std::time::Duration; + +use api::v1::region::StrictWindow; +use api::v1::region::compact_request::Options; +use common_datasource::compression::CompressionType; +use common_meta::key::schema_name::SchemaNameValue; +use common_time::{DatabaseTimeToLive, Timestamp}; +use store_api::storage::FileId; +use tokio::sync::{Barrier, mpsc, oneshot}; + +use crate::compaction::memory_manager::{CompactionMemoryGuard, new_compaction_memory_manager}; +use crate::compaction::picker::PickerOutput; +use crate::compaction::scheduler::planning::CompactionPlanningResult; +use crate::compaction::scheduler::state::{CompactingFiles, CompactionPhase}; +use crate::compaction::scheduler::*; +use crate::compaction::test_util::new_file_handle; +use crate::compaction::{CompactionOutput, find_dynamic_options}; +use crate::error::InvalidSchedulerStateSnafu; +use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions}; +use crate::metrics::COMPACTION_MEMORY_REJECTED; +use crate::region::ManifestContext; +use crate::request::{BackgroundNotify, OutputTx, WorkerRequest}; +use crate::schedule::remote_job_scheduler::{RemoteJob, RemoteJobSchedulerRef}; +use crate::schedule::scheduler::{Job, Scheduler}; +use crate::sst::FormatType; +use crate::sst::file::FileHandle; +use crate::test_util::mock_schema_metadata_manager; +use crate::test_util::scheduler_util::{SchedulerEnv, VecScheduler}; +use crate::test_util::version_util::{VersionControlBuilder, apply_edit}; + +#[test] +fn test_requires_pending_compaction_slot() { + let time_range = TimestampRange::new( + Timestamp::new_millisecond(1_000), + Timestamp::new_millisecond(2_000), + ) + .unwrap(); + + assert!(!requires_pending_compaction_slot( + &compact_request::Options::Regular(Default::default()), + None, + )); + assert!(requires_pending_compaction_slot( + &compact_request::Options::Regular(Default::default()), + Some(time_range), + )); + assert!(requires_pending_compaction_slot( + &compact_request::Options::StrictWindow(StrictWindow { window_seconds: 60 }), + None, + )); +} + +struct FailingScheduler; + +struct FailingRemoteScheduler; + +#[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, + }, + ) + } +} + +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, + ) + .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, fallback_to_local: bool) { + 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 = fallback_to_local; +} + +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<()> { + InvalidSchedulerStateSnafu.fail() + } + + async fn stop(&self, _await_termination: bool) -> Result<()> { + Ok(()) + } +} + +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:?}"), + } +} + +#[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()); +} + +#[tokio::test] +async fn test_find_compaction_options_db_level() { + let builder = VersionControlBuilder::new(); + let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); + let region_id = builder.region_id(); + let table_id = region_id.table_id(); + // Register table without ttl but with db-level compaction options + let mut schema_value = SchemaNameValue { + ttl: Some(DatabaseTimeToLive::default()), + ..Default::default() + }; + schema_value + .extra_options + .insert("compaction.type".to_string(), "twcs".to_string()); + schema_value + .extra_options + .insert("compaction.twcs.time_window".to_string(), "2h".to_string()); + schema_metadata_manager + .register_region_table_info( + table_id, + "t", + "c", + "s", + Some(schema_value), + kv_backend.clone(), + ) + .await; + + let version_control = Arc::new(builder.build()); + let region_opts = version_control.current().version.options.clone(); + let (opts, _) = find_dynamic_options(region_id, ®ion_opts, &schema_metadata_manager) + .await + .unwrap(); + match opts { + crate::region::options::CompactionOptions::Twcs(t) => { + assert_eq!(t.time_window_seconds(), Some(2 * 3600)); + } + } +} + +#[tokio::test] +async fn test_find_compaction_options_priority() { + fn schema_value_with_twcs(time_window: &str) -> SchemaNameValue { + let mut schema_value = SchemaNameValue { + ttl: Some(DatabaseTimeToLive::default()), + ..Default::default() + }; + schema_value + .extra_options + .insert("compaction.type".to_string(), "twcs".to_string()); + schema_value.extra_options.insert( + "compaction.twcs.time_window".to_string(), + time_window.to_string(), + ); + schema_value + } + + let cases = [ + ( + "db options set and table override set", + Some(schema_value_with_twcs("2h")), + true, + Some(Duration::from_secs(5 * 3600)), + Some(5 * 3600), + ), + ( + "db options set and table override not set", + Some(schema_value_with_twcs("2h")), + false, + None, + Some(2 * 3600), + ), + ( + "db options not set and table override set", + None, + true, + Some(Duration::from_secs(4 * 3600)), + Some(4 * 3600), + ), + ( + "db options not set and table override not set", + None, + false, + None, + None, + ), + ]; + + for (case_name, schema_value, override_set, table_window, expected_window) in cases { + let builder = VersionControlBuilder::new(); + let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); + let region_id = builder.region_id(); + let table_id = region_id.table_id(); + schema_metadata_manager + .register_region_table_info(table_id, "t", "c", "s", schema_value, kv_backend.clone()) + .await; + + let version_control = Arc::new(builder.build()); + let mut region_opts = version_control.current().version.options.clone(); + region_opts.compaction_override = override_set; + if let Some(window) = table_window { + let crate::region::options::CompactionOptions::Twcs(twcs) = &mut region_opts.compaction; + twcs.time_window = Some(window); + } + + let (opts, _) = find_dynamic_options(region_id, ®ion_opts, &schema_metadata_manager) + .await + .unwrap(); + match opts { + crate::region::options::CompactionOptions::Twcs(t) => { + assert_eq!(t.time_window_seconds(), expected_window, "{case_name}"); + } + } + } +} + +#[tokio::test] +async fn test_schedule_empty() { + let env = SchedulerEnv::new().await; + 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(); + schema_metadata_manager + .register_region_table_info( + builder.region_id().table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + kv_backend, + ) + .await; + // Nothing to compact. + let version_control = Arc::new(builder.build()); + let (output_tx, output_rx) = oneshot::channel(); + let waiter = OptionOutputTx::from(output_tx); + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .await; + let scheduled = scheduler + .schedule_compaction( + builder.region_id(), + compact_request::Options::Regular(Default::default()), + &version_control, + &env.access_layer, + waiter, + &manifest_ctx, + schema_metadata_manager.clone(), + 1, + ) + .unwrap(); + 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()); + + // Only one file, picker won't compact it. + let version_control = Arc::new(builder.push_l0_file(0, 1000).build()); + let (output_tx, output_rx) = oneshot::channel(); + let waiter = OptionOutputTx::from(output_tx); + let scheduled = scheduler + .schedule_compaction( + builder.region_id(), + compact_request::Options::Regular(Default::default()), + &version_control, + &env.access_layer, + waiter, + &manifest_ctx, + schema_metadata_manager.clone(), + 1, + ) + .unwrap(); + 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()); +} + +#[tokio::test] +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, 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; + // Five overlapping L0 files are enough for the regular picker to create a task. + 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 scheduled = scheduler + .schedule_compaction( + region_id, + Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::none(), + &manifest_ctx, + schema_metadata_manager.clone(), + 1, + ) + .unwrap(); + + // 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_planning_panic_notifies_and_clears_status() { + let env = SchedulerEnv::new().await; + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx.clone()); + 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.start_picking(7); + status.merge_waiter(OptionOutputTx::from(waiter_tx)); + scheduler.region_status.insert(region_id, status); + + CompactionScheduler::notify_planning_result(region_id, 7, tx, async { + panic!("planning boom") + }) + .await; + + let finished = recv_compaction_pick_finished(&mut rx).await; + let CompactionPlanningResult::Error(err) = &finished.result else { + panic!("expected planning error, got {:?}", &finished.result); + }; + assert!(err.to_string().contains("planning boom")); + let pending_ddls = scheduler + .handle_compaction_pick_finished(finished, &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_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 mut status = + CompactionStatus::new(region_id, version_control.clone(), env.access_layer.clone()); + status.start_picking(7); + scheduler.region_status.insert(region_id, status); + + let (pre_fence_tx, pre_fence_rx) = oneshot::channel(); + assert!( + !scheduler + .schedule_compaction( + region_id, + compact_request::Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::from(pre_fence_tx), + &manifest_ctx, + schema_metadata_manager.clone(), + 1, + ) + .unwrap() + ); + let (ddl_tx, _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 + .handle_compaction_pick_finished( + CompactionPickFinished { + region_id, + plan_id: 7, + result: CompactionPlanningResult::NoPlan, + }, + &manifest_ctx, + schema_metadata_manager.clone(), + ) + .await; + assert!(pending_ddls.is_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 (post_fence_tx, post_fence_rx) = oneshot::channel(); + assert!( + !scheduler + .schedule_compaction( + region_id, + compact_request::Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::from(post_fence_tx), + &manifest_ctx, + schema_metadata_manager.clone(), + 1, + ) + .unwrap() + ); + assert_matches!( + post_fence_rx.await.unwrap().unwrap_err(), + Error::CompactionCancelled { .. } + ); + + followup_finished.result = CompactionPlanningResult::NoPlan; + let pending_ddls = scheduler + .handle_compaction_pick_finished(followup_finished, &manifest_ctx, schema_metadata_manager) + .await; + assert_eq!(pending_ddls.len(), 1); + assert_eq!(pre_fence_rx.await.unwrap().unwrap(), 0); + assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(rx.try_recv().is_err()); +} + +#[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] + .active + .as_ref() + .unwrap() + .waiters + .len(), + 1 + ); + assert_matches!( + waiter_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + ); +} + +#[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_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_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_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, false); + 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_remote_fallback_uses_new_execution_plan_id() { + 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); + 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; + let remote_plan_id = finished.plan_id; + use_remote_compaction(&mut finished, true); + + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) + .await; + + assert_eq!(job_scheduler.num_jobs(), 1); + assert!(matches!( + scheduler.region_status[®ion_id] + .active + .as_ref() + .map(|active| &active.phase), + Some(CompactionPhase::Local { .. }) + )); + assert!( + !scheduler.region_status[®ion_id] + .matches_execution(&CompactionExecution::for_test(remote_plan_id)) + ); +} + +#[tokio::test] +async fn test_stale_plan_execution_does_not_affect_replacement_status() { + let env = SchedulerEnv::new().await; + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let stale_execution = CompactionExecution::for_test(1); + let replacement_version_control = compactable_version(); + let region_id = replacement_version_control + .current() + .version + .metadata + .region_id; + 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 (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); + + let pending_ddls = scheduler + .on_execution_finished( + region_id, + &stale_execution, + &manifest_ctx, + schema_metadata_manager, + ) + .await; + assert!(pending_ddls.is_empty()); + assert!(scheduler.region_status[®ion_id].is_busy()); + scheduler.on_execution_failed( + region_id, + &stale_execution, + Arc::new(InvalidSchedulerStateSnafu.build()), + ); + assert!(scheduler.region_status[®ion_id].is_busy()); + assert_matches!( + waiter_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + ); +} + +#[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, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + scheduler.memory_manager = Arc::new(new_compaction_memory_manager(1024 * 1024)); + + let mut builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let end = 1000 * 1000; + let version_control = Arc::new( + builder + .push_l0_file_with_max_row_group_size(0, end, 1024 * 1024) + .push_l0_file_with_max_row_group_size(10, end, 1024 * 1024) + .push_l0_file_with_max_row_group_size(50, end, 1024 * 1024) + .push_l0_file_with_max_row_group_size(80, end, 1024 * 1024) + .push_l0_file_with_max_row_group_size(90, end, 1024 * 1024) + .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 (output_tx, output_rx) = oneshot::channel(); + let rejected = COMPACTION_MEMORY_REJECTED.with_label_values(&["oversized"]); + let rejected_before = rejected.get(); + + let scheduled = scheduler + .schedule_compaction( + region_id, + Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::from(output_tx), + &manifest_ctx, + schema_metadata_manager.clone(), + 1, + ) + .unwrap(); + + 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] +async fn test_schedule_on_finished() { + common_telemetry::init_default_ut_logging(); + 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 purger = builder.file_purger(); + let region_id = builder.region_id(); + + let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); + schema_metadata_manager + .register_region_table_info( + builder.region_id().table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + kv_backend, + ) + .await; + + // 5 files to compact. + 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 scheduled = 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, + ) + .unwrap(); + // 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] + .files + .values() + .map(|file| file.meta_ref().clone()) + .collect(); + + // 5 files for next compaction and removes old files. + apply_edit( + &version_control, + &[(0, end), (20, end), (40, end), (60, end), (80, end)], + &file_metas, + purger.clone(), + ); + // The task is pending. + let (tx, _rx) = oneshot::channel(); + let scheduled = scheduler + .schedule_compaction( + region_id, + compact_request::Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::new(Some(OutputTx::new(tx))), + &manifest_ctx, + schema_metadata_manager.clone(), + 1, + ) + .unwrap(); + assert!(!scheduled); + assert_eq!(1, scheduler.region_status.len()); + assert_eq!(1, job_scheduler.num_jobs()); + assert!( + !scheduler + .region_status + .get(&builder.region_id()) + .unwrap() + .active + .as_ref() + .unwrap() + .waiters + .is_empty() + ); + + // On compaction finished and schedule next compaction. + scheduler + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .await; + let scheduled = scheduler.schedule_next_compaction( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + ); + 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. + apply_edit( + &version_control, + &[(0, end), (20, end), (40, end), (60, end), (80, end)], + &[], + purger.clone(), + ); + let (tx, _rx) = oneshot::channel(); + // The task is pending. + let scheduled = scheduler + .schedule_compaction( + region_id, + compact_request::Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::new(Some(OutputTx::new(tx))), + &manifest_ctx, + schema_metadata_manager, + 1, + ) + .unwrap(); + assert!(!scheduled); + assert_eq!(2, job_scheduler.num_jobs()); + assert!( + !scheduler + .region_status + .get(&builder.region_id()) + .unwrap() + .active + .as_ref() + .unwrap() + .waiters + .is_empty() + ); +} + +#[tokio::test] +async fn test_remove_idle_status_allows_rescheduling() { + 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(); + scheduler.region_status.insert( + region_id, + CompactionStatus::new(region_id, version_control.clone(), env.access_layer.clone()), + ); + + 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, + ) + .unwrap() + ); + scheduler.remove_idle_status(region_id); + assert!( + scheduler + .schedule_compaction( + region_id, + compact_request::Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::none(), + &manifest_ctx, + schema_metadata_manager, + 1, + ) + .unwrap() + ); +} + +#[tokio::test] +async fn test_time_range_compaction_when_compaction_in_progress() { + common_telemetry::init_default_ut_logging(); + 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 purger = builder.file_purger(); + let region_id = builder.region_id(); + + let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); + schema_metadata_manager + .register_region_table_info( + builder.region_id().table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + kv_backend, + ) + .await; + + // 5 files to compact. + 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 file_metas: Vec<_> = version_control.current().version.ssts.levels()[0] + .files + .values() + .map(|file| file.meta_ref().clone()) + .collect(); + + // 5 files for next compaction and removes old files. + apply_edit( + &version_control, + &[(0, end), (20, end), (40, end), (60, end), (80, end)], + &file_metas, + purger.clone(), + ); + + 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, + ) + .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 + .region_status + .get(®ion_id) + .unwrap() + .pending_request + .is_none() + ); + + // Schedule another manual compaction with a time range. + let time_range = TimestampRange::new( + Timestamp::new_millisecond(0), + Timestamp::new_millisecond(end + 1), + ) + .unwrap(); + let (tx, _rx) = oneshot::channel(); + scheduler + .schedule_compaction_with_time_range( + region_id, + compact_request::Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::new(Some(OutputTx::new(tx))), + &manifest_ctx, + schema_metadata_manager.clone(), + 1, + Some(time_range), + ) + .unwrap(); + assert_eq!(1, scheduler.region_status.len()); + // Current job num should be 1 since compaction is in progress. + assert_eq!(1, job_scheduler.num_jobs()); + let status = scheduler.region_status.get(&builder.region_id()).unwrap(); + assert_eq!( + Some(time_range), + status + .pending_request + .as_ref() + .and_then(|pending| pending.time_range) + ); + + // On compaction finished and schedule next compaction. + scheduler + .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(); + assert!(status.pending_request.is_none()); +} + +#[tokio::test] +async fn test_ranged_compaction_continuation_preserves_time_range() { + 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(); + for offset in [0, 10, 20, 30] { + builder.push_l0_file(offset, 1_000); + } + for offset in [0, 10, 20, 30] { + builder.push_l0_file(2 * 3_600_000 + offset, 2 * 3_600_000 + 1_000); + } + 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 schema_value = SchemaNameValue::default(); + schema_value + .extra_options + .insert("compaction.type".to_string(), "twcs".to_string()); + schema_value + .extra_options + .insert("compaction.twcs.time_window".to_string(), "1h".to_string()); + schema_metadata_manager + .register_region_table_info( + region_id.table_id(), + "t", + "c", + "s", + Some(schema_value), + kv_backend, + ) + .await; + let time_range = TimestampRange::new( + Timestamp::new_millisecond(0), + Timestamp::new_millisecond(3_600_000), + ) + .unwrap(); + + scheduler + .schedule_compaction_with_time_range( + region_id, + compact_request::Options::StrictWindow(StrictWindow { + window_seconds: 3_600, + }), + &version_control, + &env.access_layer, + OptionOutputTx::none(), + &manifest_ctx, + schema_metadata_manager.clone(), + 1, + Some(time_range), + ) + .unwrap(); + let first = recv_compaction_pick_finished(&mut rx).await; + assert!( + selected_files(&first) + .iter() + .all(|file| file.time_range().1 < Timestamp::new_millisecond(3_600_000)) + ); + scheduler + .handle_compaction_pick_finished(first, &manifest_ctx, schema_metadata_manager.clone()) + .await; + assert_eq!(1, job_scheduler.num_jobs()); + + scheduler + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .await; + assert!(scheduler.schedule_next_compaction(region_id, &manifest_ctx, schema_metadata_manager,)); + + let continuation = recv_compaction_pick_finished(&mut rx).await; + match continuation.result { + CompactionPlanningResult::NoPlan => {} + CompactionPlanningResult::Prepared(prepared) => assert!( + prepared + .picker_output + .outputs + .iter() + .flat_map(|output| &output.inputs) + .all(|file| file.time_range().1 < Timestamp::new_millisecond(3_600_000)) + ), + CompactionPlanningResult::Error(err) => { + panic!("unexpected compaction planning error: {err}") + } + } +} + +#[tokio::test] +async fn test_compaction_bypass_in_staging_mode() { + let env = SchedulerEnv::new().await; + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + + // Create version control and manifest context for staging mode + let builder = VersionControlBuilder::new(); + let version_control = Arc::new(builder.build()); + let region_id = version_control.current().version.metadata.region_id; + + // Create staging manifest context using the same pattern as SchedulerEnv + let staging_manifest_ctx = { + let manager = RegionManifestManager::new( + version_control.current().version.metadata.clone(), + 0, + RegionManifestOptions { + manifest_dir: "".to_string(), + object_store: env.access_layer.object_store().clone(), + compress_type: CompressionType::Uncompressed, + checkpoint_distance: 10, + remove_file_options: Default::default(), + manifest_cache: None, + }, + FormatType::PrimaryKey, + &Default::default(), + ) + .await + .unwrap(); + Arc::new(ManifestContext::new( + manager, + RegionRoleState::Leader(RegionLeaderState::Staging), + None, + )) + }; + + let (schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager(); + + // Test regular compaction bypass in staging mode + let (tx, rx) = oneshot::channel(); + scheduler + .schedule_compaction( + region_id, + compact_request::Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::new(Some(OutputTx::new(tx))), + &staging_manifest_ctx, + schema_metadata_manager, + 1, + ) + .unwrap(); + + let result = rx.await.unwrap(); + assert_eq!(result.unwrap(), 0); // is there a better way to check this? + assert_eq!(0, scheduler.region_status.len()); +} + +#[tokio::test] +async fn test_add_ddl_request_to_pending() { + let env = SchedulerEnv::new().await; + let (tx, _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()), + ); + scheduler + .region_status + .get_mut(®ion_id) + .unwrap() + .start_local_task(); + + let (output_tx, _output_rx) = oneshot::channel(); + scheduler.add_ddl_request_to_pending(SenderDdlRequest { + region_id, + sender: OptionOutputTx::from(output_tx), + request: crate::request::DdlRequest::EnterStaging( + store_api::region_request::EnterStagingRequest { + partition_directive: + store_api::region_request::StagingPartitionDirective::RejectAllWrites, + }, + ), + }); + + assert!(scheduler.has_pending_ddls(region_id)); +} + +#[tokio::test] +async fn test_pending_ddl_fences_later_compaction_triggers() { + let env = SchedulerEnv::new().await; + let (tx, _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(); + 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_manual_tx, mut first_manual_rx) = oneshot::channel(); + let mut status = + CompactionStatus::new(region_id, version_control.clone(), env.access_layer.clone()); + status.start_local_task(); + status.set_pending_request(PendingCompaction { + options: compact_request::Options::StrictWindow(StrictWindow { window_seconds: 60 }), + waiter: OptionOutputTx::from(first_manual_tx), + max_parallelism: 1, + time_range: None, + }); + scheduler.region_status.insert(region_id, status); + + let (ddl_tx, _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, + }, + ), + }); + + // Automatic regular triggers have no waiter and are ignored by the DDL fence. + 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, + ) + .unwrap() + ); + let active = scheduler.region_status[®ion_id].active.as_ref().unwrap(); + assert!(active.waiters.is_empty()); + assert!(active.regular_followup_waiters.is_none()); + + // Explicit regular and strict-window requests both have waiters and are rejected. + for options in [ + compact_request::Options::Regular(Default::default()), + compact_request::Options::StrictWindow(StrictWindow { + window_seconds: 120, + }), + ] { + let (later_tx, later_rx) = oneshot::channel(); + assert!( + !scheduler + .schedule_compaction( + region_id, + options, + &version_control, + &env.access_layer, + OptionOutputTx::from(later_tx), + &manifest_ctx, + schema_metadata_manager.clone(), + 1, + ) + .unwrap() + ); + assert_matches!( + later_rx.await.unwrap().unwrap_err(), + Error::CompactionCancelled { .. } + ); + } + + assert_matches!( + first_manual_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + ); + let pending_request = scheduler.region_status[®ion_id] + .pending_request + .as_ref() + .expect("manual compaction queued before DDL was removed"); + assert_matches!( + &pending_request.options, + compact_request::Options::StrictWindow(StrictWindow { window_seconds: 60 }) + ); +} + +#[tokio::test] +async fn test_request_cancel_state_transitions() { + let env = SchedulerEnv::new().await; + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let mut status = CompactionStatus::new(region_id, version_control, env.access_layer.clone()); + let state = status.start_local_task(); + + assert_eq!(status.request_cancel(), RequestCancelResult::CancelIssued); + assert!(state.cancel_handle().is_cancelled()); + assert_eq!( + status.request_cancel(), + RequestCancelResult::AlreadyCancelling + ); + + assert!(!state.mark_commit_started()); + assert_eq!( + status.request_cancel(), + RequestCancelResult::AlreadyCancelling + ); + + assert!(status.clear_running_task()); + assert_eq!(status.request_cancel(), RequestCancelResult::NotRunning); +} + +#[tokio::test] +async fn test_request_cancel_remote_compaction_is_too_late() { + let env = SchedulerEnv::new().await; + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let mut status = CompactionStatus::new(region_id, version_control, env.access_layer.clone()); + + status.start_remote_task(); + + assert_eq!( + status.request_cancel(), + RequestCancelResult::TooLateToCancel + ); + assert!(status.is_busy()); +} + +#[tokio::test] +async fn test_try_cancel_and_add_ddl_returns_request_when_not_running() { + let env = SchedulerEnv::new().await; + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let region_id = RegionId::new(1, 1); + let (ddl_tx, ddl_rx) = oneshot::channel(); + + let result = + scheduler.try_cancel_and_add_ddl(region_id, OptionOutputTx::from(ddl_tx), 42_u64, |_| { + crate::request::DdlRequest::EnterStaging( + store_api::region_request::EnterStagingRequest { + partition_directive: + store_api::region_request::StagingPartitionDirective::RejectAllWrites, + }, + ) + }); + + let Err((sender, payload)) = result else { + panic!("DDL was queued without a running compaction"); + }; + assert_eq!(payload, 42); + sender.send(Ok(0)); + assert_eq!(ddl_rx.await.unwrap().unwrap(), 0); +} + +#[tokio::test] +async fn test_try_cancel_and_add_ddl_cancels_and_queues_atomically() { + 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 mut status = CompactionStatus::new(region_id, version_control, env.access_layer.clone()); + status.start_picking(7); + scheduler.region_status.insert(region_id, status); + let (ddl_tx, _ddl_rx) = oneshot::channel(); + + let result = + scheduler.try_cancel_and_add_ddl(region_id, OptionOutputTx::from(ddl_tx), (), |_| { + crate::request::DdlRequest::EnterStaging( + store_api::region_request::EnterStagingRequest { + partition_directive: + store_api::region_request::StagingPartitionDirective::RejectAllWrites, + }, + ) + }); + + assert!(result.is_ok()); + assert!(scheduler.has_pending_ddls(region_id)); + assert_eq!( + scheduler.request_cancel(region_id), + RequestCancelResult::AlreadyCancelling + ); +} + +#[tokio::test] +async fn test_on_compaction_cancelled_returns_pending_ddl_requests() { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let builder = VersionControlBuilder::new(); + let version_control = Arc::new(builder.build()); + let region_id = builder.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(); + + 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 { + region_id, + sender: OptionOutputTx::from(output_tx), + request: crate::request::DdlRequest::EnterStaging( + store_api::region_request::EnterStagingRequest { + partition_directive: + store_api::region_request::StagingPartitionDirective::RejectAllWrites, + }, + ), + }); + + let pending_ddls = scheduler.on_compaction_cancelled(region_id).await; + + assert_eq!(pending_ddls.len(), 1); + 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] +async fn test_on_compaction_cancelled_prioritizes_pending_ddls_over_pending_compaction() { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let builder = VersionControlBuilder::new(); + let version_control = Arc::new(builder.build()); + let region_id = builder.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(); + + scheduler.region_status.insert( + region_id, + CompactionStatus::new(region_id, version_control, env.access_layer.clone()), + ); + let status = scheduler.region_status.get_mut(®ion_id).unwrap(); + status.start_local_task(); + let (manual_tx, manual_rx) = oneshot::channel(); + status.set_pending_request(PendingCompaction { + options: compact_request::Options::StrictWindow(StrictWindow { window_seconds: 60 }), + waiter: OptionOutputTx::from(manual_tx), + max_parallelism: 1, + time_range: None, + }); + + let (output_tx, _output_rx) = oneshot::channel(); + scheduler.add_ddl_request_to_pending(SenderDdlRequest { + region_id, + sender: OptionOutputTx::from(output_tx), + request: crate::request::DdlRequest::EnterStaging( + store_api::region_request::EnterStagingRequest { + partition_directive: + store_api::region_request::StagingPartitionDirective::RejectAllWrites, + }, + ), + }); + + let pending_ddls = scheduler.on_compaction_cancelled(region_id).await; + + assert_eq!(pending_ddls.len(), 1); + assert!(!scheduler.region_status.contains_key(®ion_id)); + assert_eq!(job_scheduler.num_jobs(), 0); + assert_matches!(manual_rx.await.unwrap(), Err(_)); +} + +#[tokio::test] +async fn test_pending_ddl_request_failed_on_compaction_failed() { + 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 version_control = Arc::new(builder.build()); + let region_id = builder.region_id(); + + 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 { + region_id, + sender: OptionOutputTx::from(output_tx), + request: crate::request::DdlRequest::EnterStaging( + store_api::region_request::EnterStagingRequest { + partition_directive: + store_api::region_request::StagingPartitionDirective::RejectAllWrites, + }, + ), + }); + + assert!(scheduler.has_pending_ddls(region_id)); + scheduler.on_compaction_failed(region_id, Arc::new(RegionClosedSnafu { region_id }.build())); + + 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] +async fn test_pending_ddl_request_failed_on_region_closed() { + let env = SchedulerEnv::new().await; + let (tx, _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 (output_tx, output_rx) = oneshot::channel(); + scheduler.add_ddl_request_to_pending(SenderDdlRequest { + region_id, + sender: OptionOutputTx::from(output_tx), + request: crate::request::DdlRequest::EnterStaging( + store_api::region_request::EnterStagingRequest { + partition_directive: + store_api::region_request::StagingPartitionDirective::RejectAllWrites, + }, + ), + }); + + assert!(scheduler.has_pending_ddls(region_id)); + scheduler.on_region_closed(region_id); + + assert!(!scheduler.has_pending_ddls(region_id)); + let result = output_rx.await.unwrap(); + assert_matches!(result, Err(_)); +} + +#[tokio::test] +async fn test_pending_ddl_request_failed_on_region_dropped() { + let env = SchedulerEnv::new().await; + let (tx, _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 (output_tx, output_rx) = oneshot::channel(); + scheduler.add_ddl_request_to_pending(SenderDdlRequest { + region_id, + sender: OptionOutputTx::from(output_tx), + request: crate::request::DdlRequest::EnterStaging( + store_api::region_request::EnterStagingRequest { + partition_directive: + store_api::region_request::StagingPartitionDirective::RejectAllWrites, + }, + ), + }); + + assert!(scheduler.has_pending_ddls(region_id)); + scheduler.on_region_dropped(region_id); + + assert!(!scheduler.has_pending_ddls(region_id)); + let result = output_rx.await.unwrap(); + assert_matches!(result, Err(_)); +} + +#[tokio::test] +async fn test_pending_ddl_request_failed_on_region_truncated() { + let env = SchedulerEnv::new().await; + let (tx, _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 (output_tx, output_rx) = oneshot::channel(); + scheduler.add_ddl_request_to_pending(SenderDdlRequest { + region_id, + sender: OptionOutputTx::from(output_tx), + request: crate::request::DdlRequest::EnterStaging( + store_api::region_request::EnterStagingRequest { + partition_directive: + store_api::region_request::StagingPartitionDirective::RejectAllWrites, + }, + ), + }); + + assert!(scheduler.has_pending_ddls(region_id)); + scheduler.on_region_truncated(region_id); + + assert!(!scheduler.has_pending_ddls(region_id)); + let result = output_rx.await.unwrap(); + assert_matches!(result, Err(_)); +} + +#[tokio::test] +async fn test_on_compaction_finished_returns_pending_ddl_requests() { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let builder = VersionControlBuilder::new(); + let version_control = Arc::new(builder.build()); + let region_id = builder.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(); + + 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 (output_tx, _output_rx) = oneshot::channel(); + scheduler.add_ddl_request_to_pending(SenderDdlRequest { + region_id, + sender: OptionOutputTx::from(output_tx), + request: crate::request::DdlRequest::EnterStaging( + store_api::region_request::EnterStagingRequest { + partition_directive: + store_api::region_request::StagingPartitionDirective::RejectAllWrites, + }, + ), + }); + + let pending_ddls = scheduler + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) + .await; + + assert_eq!(pending_ddls.len(), 1); + assert!(!scheduler.has_pending_ddls(region_id)); + assert!(!scheduler.region_status.contains_key(®ion_id)); + assert_eq!(job_scheduler.num_jobs(), 0); +} + +#[tokio::test] +async fn test_on_compaction_finished_replays_pending_ddl_after_manual_noop() { + 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 version_control = Arc::new(builder.build()); + let region_id = builder.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(); + + let (manual_tx, manual_rx) = oneshot::channel(); + let mut status = + CompactionStatus::new(region_id, version_control.clone(), env.access_layer.clone()); + status.start_local_task(); + status.set_pending_request(PendingCompaction { + options: compact_request::Options::Regular(Default::default()), + waiter: OptionOutputTx::from(manual_tx), + max_parallelism: 1, + time_range: None, + }); + scheduler.region_status.insert(region_id, status); + + let (ddl_tx, _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 + .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); +} + +#[tokio::test] +async fn test_on_compaction_finished_dispatches_pending_ddl_before_chained_regular() { + 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 version_control = Arc::new(builder.build()); + let region_id = builder.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(); + + // A regular trigger was retained while picking and the region is now + // executing; a DDL queued behind the task must be dispatched as soon + // as the task finishes instead of waiting for a whole extra cycle. + let (regular_tx, regular_rx) = oneshot::channel(); + let mut status = + CompactionStatus::new(region_id, version_control.clone(), 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 (ddl_tx, _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 + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) + .await; + + assert_eq!(pending_ddls.len(), 1); + assert_eq!(regular_rx.await.unwrap().unwrap(), 0); + assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(rx.try_recv().is_err()); +} + +#[tokio::test] +async fn test_on_compaction_finished_returns_empty_when_region_absent() { + 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 pending_ddls = scheduler + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) + .await; + + assert!(pending_ddls.is_empty()); +} + +#[tokio::test] +async fn test_on_compaction_finished_manual_schedule_error_cleans_status() { + 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 mut builder = VersionControlBuilder::new(); + 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 region_id = builder.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(); + + let (manual_tx, manual_rx) = oneshot::channel(); + let mut status = + CompactionStatus::new(region_id, version_control.clone(), env.access_layer.clone()); + status.start_local_task(); + status.set_pending_request(PendingCompaction { + options: compact_request::Options::Regular(Default::default()), + waiter: OptionOutputTx::from(manual_tx), + max_parallelism: 1, + time_range: None, + }); + scheduler.region_status.insert(region_id, status); + + let (ddl_tx, 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 + .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_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, 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(); + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .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 pending_ddls = scheduler + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) + .await; + + assert!(pending_ddls.is_empty()); + assert!(scheduler.region_status.contains_key(®ion_id)); + + let (schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager(); + // 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.clone(), + ); + assert!(scheduled); + 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)); +} + +#[tokio::test] +async fn test_on_compaction_finished_next_schedule_error_cleans_status() { + 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 mut builder = VersionControlBuilder::new(); + 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 region_id = builder.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(); + + 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 pending_ddls = scheduler + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) + .await; + + assert!(pending_ddls.is_empty()); + assert!(scheduler.region_status.contains_key(®ion_id)); + + 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.clone(), + ); + assert!(scheduled); + 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)); +} + +#[tokio::test] +async fn test_concurrent_memory_competition() { + let manager = Arc::new(new_compaction_memory_manager(3 * 1024 * 1024)); // 3MB + let barrier = Arc::new(Barrier::new(3)); + let mut handles = vec![]; + + // Spawn 3 tasks competing for memory, each trying to acquire 2MB + for _i in 0..3 { + let mgr = manager.clone(); + let bar = barrier.clone(); + let handle = tokio::spawn(async move { + bar.wait().await; // Synchronize start + mgr.try_acquire(2 * 1024 * 1024) + }); + handles.push(handle); + } + + let results: Vec> = futures::future::join_all(handles) + .await + .into_iter() + .map(|r| r.unwrap()) + .collect(); + + // Only 1 should succeed (3MB limit, 2MB request, can only fit one) + let succeeded = results.iter().filter(|r| r.is_some()).count(); + let failed = results.iter().filter(|r| r.is_none()).count(); + + assert_eq!(succeeded, 1, "Expected exactly 1 task to acquire memory"); + assert_eq!(failed, 2, "Expected 2 tasks to fail"); + + // Clean up + drop(results); + assert_eq!(manager.used_bytes(), 0); +} diff --git a/src/mito2/src/compaction/twcs.rs b/src/mito2/src/compaction/twcs.rs index 9ec9fe647f..d41b9091ad 100644 --- a/src/mito2/src/compaction/twcs.rs +++ b/src/mito2/src/compaction/twcs.rs @@ -26,14 +26,14 @@ use common_time::timestamp_millis::BucketAligned; use rayon::prelude::*; use store_api::storage::RegionId; +use crate::compaction::CompactionOutput; use crate::compaction::buckets::infer_time_bucket; use crate::compaction::compactor::CompactionRegion; -use crate::compaction::picker::{Picker, PickerOutput}; +use crate::compaction::picker::{Picker, PickerOutput, get_expired_ssts}; use crate::compaction::run::{ FileGroup, Item, Ranged, find_sorted_runs, find_sorted_runs_by_time_range, merge_primary_key_ranges, merge_seq_files, primary_key_ranges_overlap, reduce_runs, }; -use crate::compaction::{CompactionOutput, get_expired_ssts}; use crate::sst::file::{FileHandle, Level, overlaps}; use crate::sst::version::LevelMeta; diff --git a/src/mito2/src/compaction/window.rs b/src/mito2/src/compaction/window.rs index af7a35412c..b5bf8e8dd6 100644 --- a/src/mito2/src/compaction/window.rs +++ b/src/mito2/src/compaction/window.rs @@ -22,10 +22,10 @@ use common_time::timestamp::TimeUnit; use common_time::timestamp_millis::BucketAligned; use store_api::storage::RegionId; +use crate::compaction::CompactionOutput; use crate::compaction::buckets::infer_time_bucket; use crate::compaction::compactor::{CompactionRegion, CompactionVersion}; -use crate::compaction::picker::{Picker, PickerOutput}; -use crate::compaction::{CompactionOutput, get_expired_ssts}; +use crate::compaction::picker::{Picker, PickerOutput, get_expired_ssts}; use crate::sst::file::FileHandle; /// Compaction picker that splits the time range of all involved files to windows, and merges