mito2/compaction/

scheduler.rs

// 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 coordinator;
pub(crate) mod history;
mod local;
mod planning;
mod state;

use std::collections::HashMap;
use std::sync::{Arc, LazyLock};

use api::v1::region::compact_request;
use common_base::Plugins;
use common_memory_manager::OnExhaustedPolicy;
use common_meta::key::SchemaMetadataManagerRef;
#[cfg(test)]
use common_recordbatch::filter::MetadataFilter;
use common_telemetry::{debug, error, info, warn};
use common_time::range::TimestampRange;
pub(crate) use coordinator::{CompactionCoordinator, ReadyUnit};
#[cfg(test)]
use datafusion_common::ScalarValue;
pub(crate) use planning::CompactionPickFinished;
pub use planning::CompactionRequest;
#[cfg(test)]
pub(crate) use planning::pick_automatic_for_test;
pub(crate) use state::{CompactionExecution, CompactionState, CompactionStateRef};
use state::{CompactionStatus, CompactionTrigger, DdlExecution, PendingCompaction};
#[cfg(test)]
use store_api::compaction_task::{self, CompactionTaskEntry};
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::config::MitoConfig;
use crate::error::{
    CompactionCancelledSnafu, Error, ManualCompactionAlreadyRunningSnafu, RegionClosedSnafu,
    RegionDroppedSnafu, RegionTruncatedSnafu, Result, UnsupportedOperationSnafu,
};
use crate::region::version::VersionControlRef;
use crate::region::{ManifestContextRef, RegionLeaderState, RegionRoleState};
use crate::request::{DdlRequest, OptionOutputTx, SenderDdlRequest, WorkerRequestWithTime};
#[cfg(test)]
use crate::schedule::RequestCancelResult;
use crate::worker::WorkerListener;

/// Internal testing switch to prevent SST compaction from changing file layouts
/// during query regression tests. This is not a public configuration option.
const DISABLE_COMPACTION_ENV: &str = "GREPTIME_DISABLE_COMPACTION";

/// Internal query-regression switch, read once when the first scheduler starts.
/// Applies to SST compaction only; ingestion and flushing remain enabled.
static DISABLE_COMPACTION: LazyLock<bool> = LazyLock::new(|| {
    let value = std::env::var(DISABLE_COMPACTION_ENV).ok();
    let disabled = parse_disable_compaction(value.as_deref());
    if disabled {
        warn!(
            "SST compaction is disabled by {DISABLE_COMPACTION_ENV}; restart without it to resume"
        );
    }
    disabled
});

fn parse_disable_compaction(value: Option<&str>) -> bool {
    match value {
        None | Some("false") => false,
        Some("true") => true,
        Some(_) => {
            warn!(
                "Invalid {DISABLE_COMPACTION_ENV}: expected true or false; compaction remains enabled"
            );
            false
        }
    }
}

/// Compaction scheduler tracks and manages compaction tasks.
pub(crate) struct CompactionScheduler {
    /// Disables SST compaction for internal query regression tests.
    /// Initialized from [`DISABLE_COMPACTION`] and fixed for the scheduler's lifetime
    /// in production; unit tests can override it without changing the process environment.
    disable_compaction: bool,
    coordinator: Arc<CompactionCoordinator>,
    /// All loaded regions, including idle regions and regions executing DDLs.
    region_status: HashMap<RegionId, CompactionStatus>,
    /// Request sender of the worker that this scheduler belongs to.
    request_sender: Sender<WorkerRequestWithTime>,
    cache_manager: CacheManagerRef,
    engine_config: Arc<MitoConfig>,
    memory_manager: Arc<CompactionMemoryManager>,
    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,
}

/// Describes the immediate action produced by a compaction terminal transition.
#[derive(Debug)]
pub(crate) enum CompactionTransition {
    /// No follow-up work was dispatched and no DDL became ready.
    NoAction,
    /// An automatic follow-up planning cycle was dispatched.
    AutomaticFollowupScheduled,
    /// DDL requests became ready after the compaction lifecycle terminated.
    DdlReady(Vec<SenderDdlRequest>),
}

impl CompactionTransition {
    fn from_pending_ddls(pending_ddls: Vec<SenderDdlRequest>) -> Self {
        if pending_ddls.is_empty() {
            Self::NoAction
        } else {
            Self::DdlReady(pending_ddls)
        }
    }

    #[cfg(test)]
    fn is_empty(&self) -> bool {
        !matches!(self, Self::DdlReady(pending_ddls) if !pending_ddls.is_empty())
    }

    #[cfg(test)]
    fn len(&self) -> usize {
        match self {
            Self::DdlReady(pending_ddls) => pending_ddls.len(),
            Self::NoAction | Self::AutomaticFollowupScheduled => 0,
        }
    }
}

// API used by callers outside the compaction scheduler module.
impl CompactionScheduler {
    /// Creates an empty scheduler bound to one region worker.
    ///
    /// # Effects
    ///
    /// Stores the worker request sender and shared scheduling resources. No
    /// compaction is dispatched until a scheduling method is called.
    ///
    /// # Constraints
    ///
    /// All lifecycle methods on the returned scheduler must be driven by the
    /// same worker so region state transitions remain serialized.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        coordinator: Arc<CompactionCoordinator>,
        request_sender: Sender<WorkerRequestWithTime>,
        cache_manager: CacheManagerRef,
        engine_config: Arc<MitoConfig>,
        listener: WorkerListener,
        plugins: Plugins,
        memory_manager: Arc<CompactionMemoryManager>,
        memory_policy: OnExhaustedPolicy,
    ) -> Self {
        Self {
            disable_compaction: *DISABLE_COMPACTION,
            coordinator,
            region_status: HashMap::new(),
            request_sender,
            cache_manager,
            engine_config,
            memory_manager,
            memory_policy,
            listener,
            plugins,
            next_plan_id: 0,
        }
    }

    /// Registers a loaded region; replacing its VersionControl invalidates the old incarnation.
    pub(crate) fn register_region(
        &mut self,
        version_control: &VersionControlRef,
        access_layer: &AccessLayerRef,
        shared_state: &CompactionStateRef,
    ) {
        let region_id = version_control.region_id();
        if let Some(status) = self.region_status.get_mut(&region_id) {
            if Arc::ptr_eq(&status.version_control, version_control) {
                return;
            }
            status.on_failure(Arc::new(RegionClosedSnafu { region_id }.build()));
        }
        let mut status =
            CompactionStatus::new(region_id, version_control.clone(), access_layer.clone());
        status.shared_state = shared_state.clone();
        self.region_status.insert(region_id, status);
    }

    /// Reports active work rather than the existence of a resident region record.
    #[cfg(test)]
    fn is_compacting(&self, region_id: RegionId) -> bool {
        self.region_status
            .get(&region_id)
            .is_some_and(CompactionStatus::is_compacting)
    }

    /// Accepts an automatic compaction trigger for a region.
    ///
    /// # Effects
    ///
    /// Starts planning when the registered region is idle.
    /// Otherwise, an active status coalesces the trigger into one unrestricted
    /// follow-up cycle. Returns `true` only when this call dispatches planning.
    ///
    /// # Constraints
    ///
    /// The owning region worker must call this method serially. Automatic
    /// triggers have no waiter and may be ignored while a DDL fence is active.
    /// The region identity comes from `version_control`; `access_layer` and
    /// `manifest_ctx` must belong to the same region.
    pub(crate) fn schedule_automatic_compaction(
        &mut self,
        compact_options: compact_request::Options,
        version_control: &VersionControlRef,
        access_layer: &AccessLayerRef,
        manifest_ctx: &ManifestContextRef,
        schema_metadata_manager: SchemaMetadataManagerRef,
    ) -> Result<bool> {
        self.schedule_compaction(
            CompactionTrigger::Automatic,
            compact_options,
            version_control,
            access_layer,
            OptionOutputTx::none(),
            manifest_ctx,
            schema_metadata_manager,
            self.engine_config.max_background_compactions,
            None,
        )
    }

    /// Accepts a manual compaction request for a region.
    ///
    /// # Effects
    ///
    /// Starts planning when the registered region is idle. During
    /// an automatic compaction, queues the request and replaces any older
    /// pending manual request; during a manual compaction, rejects the new
    /// request through `waiter`. Returns `true` only when this call dispatches
    /// planning.
    ///
    /// # Constraints
    ///
    /// The owning region worker must call this method serially. The supplied
    /// waiter is completed exactly once by cycle completion, replacement,
    /// cancellation, failure, or region teardown. The region identity comes
    /// from `version_control`; `access_layer` and `manifest_ctx` must belong to
    /// the same region.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn schedule_manual_compaction(
        &mut self,
        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<TimestampRange>,
    ) -> Result<bool> {
        self.schedule_compaction(
            CompactionTrigger::Manual,
            compact_options,
            version_control,
            access_layer,
            waiter,
            manifest_ctx,
            schema_metadata_manager,
            max_parallelism,
            time_range,
        )
    }

    /// Applies a background planning result to the current region state.
    ///
    /// # Effects
    ///
    /// Rejects stale plans, submits a prepared local or remote execution, or
    /// completes the cycle when planning produced no executable plan. The
    /// returned transition reports a dispatched automatic follow-up or DDL
    /// requests released by the resulting terminal transition.
    ///
    /// # Constraints
    ///
    /// The owning worker must call this method for planning notifications and
    /// must execute every returned DDL request. The notification may be stale;
    /// stale notifications intentionally have no effect.
    pub(crate) async fn handle_compaction_pick_finished(
        &mut self,
        finished: CompactionPickFinished,
        manifest_ctx: &ManifestContextRef,
        schema_metadata_manager: SchemaMetadataManagerRef,
    ) -> CompactionTransition {
        self.handle_compaction_pick_finished_inner(finished, manifest_ctx, schema_metadata_manager)
            .await
    }

    /// Returns whether a terminal notification belongs to the installed execution.
    ///
    /// # Effects
    ///
    /// Performs a read-only identity check against the region's current plan.
    ///
    /// # Constraints
    ///
    /// Callers must check this before applying a compaction edit. Matching only
    /// the region id is insufficient because close/reopen can replace a status.
    pub(crate) fn is_current_execution(
        &self,
        region_id: RegionId,
        execution: &CompactionExecution,
    ) -> bool {
        self.region_status
            .get(&region_id)
            .is_some_and(|status| status.matches_execution(execution))
    }

    /// Completes the installed execution after its edit has been applied.
    ///
    /// # Effects
    ///
    /// Notifies waiters, schedules a pending manual or automatic follow-up, or
    /// returns the region to idle. A follow-up is scheduled whenever the cycle
    /// latched an automatic trigger, or when `made_progress` is true (the
    /// execution produced output or reduced the file count): successful compaction
    /// keeps draining the region until the picker returns no plan. The returned
    /// transition reports a dispatched automatic follow-up or DDLs that are now
    /// safe to execute.
    ///
    /// # Constraints
    ///
    /// The owning worker must call this only after accepting and applying the
    /// execution's result, and must execute every returned DDL request. Stale
    /// executions are ignored.
    pub(crate) async fn on_execution_finished(
        &mut self,
        region_id: RegionId,
        execution: &CompactionExecution,
        manifest_ctx: &ManifestContextRef,
        schema_metadata_manager: SchemaMetadataManagerRef,
        made_progress: bool,
    ) -> CompactionTransition {
        if !self.is_current_execution(region_id, execution) {
            return CompactionTransition::NoAction;
        }
        self.on_compaction_finished(
            region_id,
            manifest_ctx,
            schema_metadata_manager,
            made_progress,
        )
        .await
    }

    /// Records failure of the installed execution.
    ///
    /// # Effects
    ///
    /// Resets the matching region to idle and fails all compaction waiters and
    /// dependent DDL requests with the supplied error.
    ///
    /// # Constraints
    ///
    /// A stale execution failure is ignored so it cannot tear down replacement
    /// state.
    pub(crate) fn on_execution_failed(
        &mut self,
        region_id: RegionId,
        execution: &CompactionExecution,
        err: Arc<Error>,
    ) {
        if !self.is_current_execution(region_id, execution) {
            return;
        }
        self.on_compaction_failed(region_id, err);
    }

    /// Removes compaction state because the region was dropped.
    ///
    /// # Effects
    ///
    /// Fails all compaction and DDL waiters owned by the status.
    ///
    /// # Constraints
    ///
    /// The owning worker must invoke this as part of serialized region teardown.
    pub(crate) fn on_region_dropped(&mut self, region_id: RegionId) {
        self.remove_region(
            region_id,
            Arc::new(RegionDroppedSnafu { region_id }.build()),
        );
    }

    /// Removes compaction state because the region was closed.
    ///
    /// # Effects
    ///
    /// Fails all compaction and DDL waiters owned by the status.
    ///
    /// # Constraints
    ///
    /// The owning worker must invoke this as part of serialized region teardown.
    pub(crate) fn on_region_closed(&mut self, region_id: RegionId) {
        self.remove_region(region_id, Arc::new(RegionClosedSnafu { region_id }.build()));
    }

    /// Resets the cycle after truncate, retaining registration and outstanding DDL replies.
    ///
    /// # Effects
    ///
    /// Fails all compaction and DDL waiters owned by the status.
    ///
    /// # Constraints
    ///
    /// The owning worker must invoke this after truncate completes.
    pub(crate) fn on_region_truncated(&mut self, region_id: RegionId) {
        self.fail_compaction(
            region_id,
            Arc::new(RegionTruncatedSnafu { region_id }.build()),
        );
    }

    /// Cancels a running compaction and queues its dependent DDL atomically.
    ///
    /// # Effects
    ///
    /// Requests cancellation when the active phase can still stop, then queues
    /// the DDL. The queued DDL forms a scheduling fence for subsequent
    /// compaction triggers. If no compaction is running, returns the sender and
    /// typed request unchanged.
    ///
    /// # Constraints
    ///
    /// Production callers use this only for [`DdlRequest::Truncate`] and
    /// [`DdlRequest::EnterStaging`]. `Ok(())` means the caller must wait for a
    /// terminal callback to receive and execute the queued DDL; `Err` means the
    /// caller must execute the returned request directly.
    pub(crate) fn try_cancel_and_add_ddl<T>(
        &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(&region_id) else {
            return Err((sender, request));
        };
        if !status.is_compacting() {
            return Err((sender, request));
        }
        status.request_cancel();

        let request = SenderDdlRequest {
            region_id,
            sender,
            request: into_ddl_request(request),
        };
        debug!(
            "Added pending DDL request for region: {}, ddl: {:?}",
            request.region_id, request.request
        );
        status.pending_ddl_requests.push(request);
        Ok(())
    }
}

// Internal state-machine helpers.
impl CompactionScheduler {
    /// 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
    }

    #[allow(clippy::too_many_arguments)]
    fn schedule_compaction(
        &mut self,
        trigger: CompactionTrigger,
        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<TimestampRange>,
    ) -> Result<bool> {
        if self.disable_compaction {
            if matches!(trigger, CompactionTrigger::Manual) {
                waiter.send(
                    UnsupportedOperationSnafu {
                        err_msg: format!("SST compaction is disabled by {DISABLE_COMPACTION_ENV}"),
                    }
                    .fail(),
                );
            }
            return Ok(false);
        }
        let region_id = version_control.region_id();
        let status = self.region_status.entry(region_id).or_insert_with(|| {
            CompactionStatus::new(region_id, version_control.clone(), access_layer.clone())
        });
        if status.has_ddl() {
            info!(
                "Region {} has pending DDL requests, ignoring compaction: queued={}, executing={}",
                region_id,
                status.pending_ddl_requests.len(),
                status.executing_ddl.is_some(),
            );
            waiter.send(CompactionCancelledSnafu.fail());
            return Ok(false);
        }
        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 status.is_compacting() {
            match trigger {
                CompactionTrigger::Automatic => status.mark_automatic_trigger(),
                CompactionTrigger::Manual if status.is_manual_compaction() => {
                    waiter.send(ManualCompactionAlreadyRunningSnafu { region_id }.fail());
                    info!(
                        "Region {} already has a manually triggered compaction running",
                        region_id
                    );
                }
                CompactionTrigger::Manual => {
                    status.set_pending_request(PendingCompaction {
                        options: compact_options,
                        waiter,
                        max_parallelism,
                        time_range,
                    });
                    info!(
                        "Region {} is running an automatic compaction; manual compaction will be re-scheduled",
                        region_id
                    );
                }
            }
            return Ok(false);
        }

        // Publish the picking phase before dispatching background planning.
        let plan_id = Self::next_plan_id(&mut self.next_plan_id);
        status.start_picking_with_trigger(plan_id, trigger);
        let pick_history = (trigger == CompactionTrigger::Automatic
            && matches!(compact_options, compact_request::Options::Regular(_)))
        .then(|| status.pick_history.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,
            pick_history,
        );
        status.merge_waiter(waiter);
        self.dispatch_compaction_planning(plan_id, request, compact_options, time_range, trigger);
        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.
    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(&region_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,
            None,
        );
        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_trigger(plan_id, CompactionTrigger::Manual);
        status.active.task_state = status.pending_task_state.take();
        self.dispatch_compaction_planning(
            plan_id,
            request,
            options,
            time_range,
            CompactionTrigger::Manual,
        );
        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,
        made_progress: bool,
    ) -> CompactionTransition {
        if let Some(ddls) = self.take_ready_ddls(region_id) {
            return CompactionTransition::DdlReady(ddls);
        }

        if self.handle_pending_compaction_request(
            region_id,
            manifest_ctx,
            schema_metadata_manager.clone(),
        ) {
            return CompactionTransition::NoAction;
        }

        let Some(status) = self.region_status.get_mut(&region_id) else {
            return CompactionTransition::NoAction;
        };
        for waiter in status.take_waiters() {
            waiter.send(Ok(0));
        }

        // Keep draining when the cycle latched an automatic trigger, or when the
        // execution produced output or reduced the file count. An empty edit stops
        // here; the next flush trigger resumes.
        let should_continue = status.active.reset_automatic_followup() || made_progress;
        if should_continue
            && self.schedule_automatic_followup(region_id, manifest_ctx, schema_metadata_manager)
        {
            return CompactionTransition::AutomaticFollowupScheduled;
        }
        if let Some(status) = self.region_status.get_mut(&region_id) {
            status.become_idle();
        }
        CompactionTransition::NoAction
    }

    fn schedule_automatic_followup(
        &mut self,
        region_id: RegionId,
        manifest_ctx: &ManifestContextRef,
        schema_metadata_manager: SchemaMetadataManagerRef,
    ) -> bool {
        if self.disable_compaction {
            return false;
        }
        let Some(status) = self.region_status.get_mut(&region_id) else {
            return false;
        };
        // An external automatic trigger requires one unrestricted follow-up pick.
        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,
            self.engine_config.max_background_compactions,
            Some(status.pick_history.clone()),
        );
        // 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);
        self.dispatch_compaction_planning(
            plan_id,
            request,
            compact_request::Options::Regular(Default::default()),
            None,
            CompactionTrigger::Automatic,
        );
        debug!(
            "Successfully scheduled next compaction planning for region id: {}",
            region_id
        );
        true
    }

    /// Notifies the scheduler that the compaction job is cancelled cooperatively.
    #[cfg(test)]
    async fn on_compaction_cancelled(&mut self, region_id: RegionId) -> Vec<SenderDdlRequest> {
        self.finish_compaction_on_cancel(region_id)
    }

    /// Notifies the scheduler that the compaction job is failed.
    fn on_compaction_failed(&mut self, region_id: RegionId, err: Arc<Error>) {
        error!(err; "Region {} failed to compact, cancel all pending tasks", region_id);
        self.fail_compaction(region_id, err);
    }

    #[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)]
    fn has_pending_ddls(&self, region_id: RegionId) -> bool {
        let has_pending = self
            .region_status
            .get(&region_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)]
    fn request_cancel(&mut self, region_id: RegionId) -> RequestCancelResult {
        self.region_status
            .get_mut(&region_id)
            .unwrap()
            .request_cancel()
    }

    /// Unregisters a closed or dropped region and fails its queued requests.
    fn remove_region(&mut self, region_id: RegionId, err: Arc<Error>) {
        if let Some(mut status) = self.region_status.remove(&region_id) {
            status.on_failure(err);
        }
        self.coordinator.forget_region(region_id);
    }

    /// Fails a cycle without discarding the loaded region's registration.
    fn fail_compaction(&mut self, region_id: RegionId, err: Arc<Error>) {
        if let Some(status) = self.region_status.get_mut(&region_id) {
            status.on_failure(err);
        }
    }

    /// Returns the cycle to idle and records any DDL handoff in the same resident entry.
    fn finish_compaction_on_cancel(&mut self, region_id: RegionId) -> Vec<SenderDdlRequest> {
        let Some(status) = self.region_status.get_mut(&region_id) else {
            return Vec::new();
        };

        let ddls = status.on_cancel();
        self.fence_ddl_replies(region_id, ddls)
    }

    /// Hands off pending DDLs after compaction ends and cancels remaining compaction waiters.
    fn take_ready_ddls(&mut self, region_id: RegionId) -> Option<Vec<SenderDdlRequest>> {
        if self
            .region_status
            .get(&region_id)?
            .pending_ddl_requests
            .is_empty()
        {
            return None;
        }
        Some(self.finish_compaction_on_cancel(region_id))
    }

    /// Hands off a drained batch of DDLs and fences new compaction until every reply lands.
    ///
    /// Installs the fence, intercepts each DDL's reply, and returns the same batch so the caller
    /// can execute it. The real result of each DDL is still forwarded to its original sender.
    ///
    /// # Constraints
    /// Must only be called from a compaction terminal transition, and only by the worker that owns
    /// `region_id`.
    fn fence_ddl_replies(
        &mut self,
        region_id: RegionId,
        mut ddls: Vec<SenderDdlRequest>,
    ) -> Vec<SenderDdlRequest> {
        if ddls.is_empty() {
            return ddls;
        }
        let generation = Self::next_plan_id(&mut self.next_plan_id);
        if let Some(status) = self.region_status.get_mut(&region_id) {
            // Set region status to "executing DDLs" so future compaction won't be scheduled.
            status.executing_ddl = Some(DdlExecution::new(generation, ddls.len()));
        }

        for (request_index, ddl) in ddls.iter_mut().enumerate() {
            // Swap the DDL's real reply sender for a temporary one so we can observe it. The
            // spawned task forwards the observed result to the original caller and, in between,
            // notifies the scheduler that this slot of the batch has completed.
            let (tx, rx) = tokio::sync::oneshot::channel();
            let reply = std::mem::replace(&mut ddl.sender, tx.into());
            let sender = self.request_sender.clone();
            common_runtime::spawn_global(async move {
                use snafu::ResultExt;
                let result = rx
                    .await
                    .context(crate::error::RecvSnafu)
                    .and_then(|result| result);
                // Route the completion through the worker channel so `on_ddl_complete` sees it
                // serialized with other region state transitions.
                //
                // At most one batch can be outstanding per region: [`CompactionStatus::has_ddl`] blocks
                // compaction admission while `executing_ddl` is set, so `fence_ddl_replies` cannot run again
                // for the same region before the current batch drains. A single `(generation, request_index)`
                // pair therefore identifies every reply. `generation` comes from the scheduler-wide counter,
                // which outlives region statuses, so a late reply from a previous incarnation cannot clear the
                // current fence. Inside the batch `request_index` makes completion idempotent: only the last
                // distinct reply releases the fence.
                let _ = sender
                    .send(WorkerRequestWithTime::new(
                        crate::request::WorkerRequest::Background {
                            region_id,
                            notify: crate::request::BackgroundNotify::CompactionDdlComplete {
                                generation,
                                request_index,
                            },
                        },
                    ))
                    .await;
                reply.send(result);
            });
        }
        ddls
    }

    /// Releases the DDL fence after its last reply, ignoring replies from an older generation.
    pub(crate) fn on_ddl_complete(
        &mut self,
        region_id: RegionId,
        generation: u64,
        request_index: usize,
    ) {
        let Some(status) = self.region_status.get_mut(&region_id) else {
            return;
        };
        // `complete` returns false for a stale generation or a duplicate `request_index`, so a late or
        // repeated notify cannot end the batch early. It only clears the fence on the last distinct
        // reply, which is what lets `has_ddl` unblock compaction admission again.
        if let Some(ddl) = &mut status.executing_ddl
            && ddl.complete(generation, request_index)
        {
            status.executing_ddl = None;
        }
    }
}

impl Drop for CompactionScheduler {
    fn drop(&mut self) {
        for (region_id, mut status) in self.region_status.drain() {
            // We are shutting down so notify all pending tasks.
            status.on_failure(Arc::new(RegionClosedSnafu { region_id }.build()));
            self.coordinator.forget_region(region_id);
        }
    }
}

impl CompactionScheduler {
    /// Reads the same published task state as engine queries.
    #[cfg(test)]
    pub(crate) fn inspect_tasks(
        &self,
        filters: &MetadataFilter,
        node_id: Option<u64>,
    ) -> Vec<CompactionTaskEntry> {
        let mut entries = Vec::new();
        for (region_id, status) in &self.region_status {
            if !filters.matches(|name| match name {
                compaction_task::REGION_ID => Some(ScalarValue::UInt64(Some(region_id.as_u64()))),
                compaction_task::TABLE_ID => Some(ScalarValue::UInt32(Some(region_id.table_id()))),
                compaction_task::NODE_ID => Some(ScalarValue::UInt64(node_id)),
                _ => None,
            }) {
                continue;
            }
            let (active, pending) = CompactionState::current(&status.shared_state);
            for task in active.into_iter().chain(pending) {
                let entry = task.snapshot(node_id);
                if entry.matches(filters) {
                    entries.push(entry);
                }
            }
        }
        entries
    }
}

#[cfg(test)]
#[path = "scheduler_test.rs"]
mod tests;