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