Skip to main content

mito2/compaction/
scheduler.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15mod local;
16mod planning;
17mod state;
18
19use std::collections::HashMap;
20use std::sync::{Arc, LazyLock};
21
22use api::v1::region::compact_request;
23use common_base::Plugins;
24use common_memory_manager::OnExhaustedPolicy;
25use common_meta::key::SchemaMetadataManagerRef;
26use common_telemetry::{debug, error, info, warn};
27use common_time::range::TimestampRange;
28pub(crate) use planning::CompactionPickFinished;
29pub use planning::CompactionRequest;
30pub(crate) use state::CompactionExecution;
31use state::{CompactionStatus, CompactionTrigger, DdlExecution, PendingCompaction};
32use store_api::storage::RegionId;
33use tokio::sync::mpsc::Sender;
34
35use crate::access_layer::AccessLayerRef;
36use crate::cache::CacheManagerRef;
37use crate::compaction::memory_manager::CompactionMemoryManager;
38use crate::compaction::task::MAX_PARALLEL_COMPACTION;
39use crate::config::MitoConfig;
40use crate::error::{
41    CompactionCancelledSnafu, Error, ManualCompactionAlreadyRunningSnafu, RegionClosedSnafu,
42    RegionDroppedSnafu, RegionTruncatedSnafu, Result, UnsupportedOperationSnafu,
43};
44use crate::region::version::VersionControlRef;
45use crate::region::{ManifestContextRef, RegionLeaderState, RegionRoleState};
46use crate::request::{DdlRequest, OptionOutputTx, SenderDdlRequest, WorkerRequestWithTime};
47#[cfg(test)]
48use crate::schedule::RequestCancelResult;
49use crate::schedule::scheduler::SchedulerRef;
50use crate::worker::WorkerListener;
51
52/// Internal testing switch to prevent SST compaction from changing file layouts
53/// during query regression tests. This is not a public configuration option.
54const DISABLE_COMPACTION_ENV: &str = "GREPTIME_DISABLE_COMPACTION";
55
56/// Internal query-regression switch, read once when the first scheduler starts.
57/// Applies to SST compaction only; ingestion and flushing remain enabled.
58static DISABLE_COMPACTION: LazyLock<bool> = LazyLock::new(|| {
59    let value = std::env::var(DISABLE_COMPACTION_ENV).ok();
60    let disabled = parse_disable_compaction(value.as_deref());
61    if disabled {
62        warn!(
63            "SST compaction is disabled by {DISABLE_COMPACTION_ENV}; restart without it to resume"
64        );
65    }
66    disabled
67});
68
69fn parse_disable_compaction(value: Option<&str>) -> bool {
70    match value {
71        None | Some("false") => false,
72        Some("true") => true,
73        Some(_) => {
74            warn!(
75                "Invalid {DISABLE_COMPACTION_ENV}: expected true or false; compaction remains enabled"
76            );
77            false
78        }
79    }
80}
81
82/// Compaction scheduler tracks and manages compaction tasks.
83pub(crate) struct CompactionScheduler {
84    /// Disables SST compaction for internal query regression tests.
85    /// Initialized from [`DISABLE_COMPACTION`] and fixed for the scheduler's lifetime
86    /// in production; unit tests can override it without changing the process environment.
87    disable_compaction: bool,
88    scheduler: SchedulerRef,
89    /// All loaded regions, including idle regions and regions executing DDLs.
90    region_status: HashMap<RegionId, CompactionStatus>,
91    /// Request sender of the worker that this scheduler belongs to.
92    request_sender: Sender<WorkerRequestWithTime>,
93    cache_manager: CacheManagerRef,
94    engine_config: Arc<MitoConfig>,
95    memory_manager: Arc<CompactionMemoryManager>,
96    memory_policy: OnExhaustedPolicy,
97    listener: WorkerListener,
98    /// Plugins for the compaction scheduler.
99    plugins: Plugins,
100    /// Scheduler-wide generation counter for compaction plans and executions.
101    /// It outlives region statuses so close/reopen cannot reuse an old identity.
102    next_plan_id: u64,
103}
104
105/// Describes the immediate action produced by a compaction terminal transition.
106#[derive(Debug)]
107pub(crate) enum CompactionTransition {
108    /// No follow-up work was dispatched and no DDL became ready.
109    NoAction,
110    /// An automatic follow-up planning cycle was dispatched.
111    AutomaticFollowupScheduled,
112    /// DDL requests became ready after the compaction lifecycle terminated.
113    DdlReady(Vec<SenderDdlRequest>),
114}
115
116impl CompactionTransition {
117    fn from_pending_ddls(pending_ddls: Vec<SenderDdlRequest>) -> Self {
118        if pending_ddls.is_empty() {
119            Self::NoAction
120        } else {
121            Self::DdlReady(pending_ddls)
122        }
123    }
124
125    #[cfg(test)]
126    fn is_empty(&self) -> bool {
127        !matches!(self, Self::DdlReady(pending_ddls) if !pending_ddls.is_empty())
128    }
129
130    #[cfg(test)]
131    fn len(&self) -> usize {
132        match self {
133            Self::DdlReady(pending_ddls) => pending_ddls.len(),
134            Self::NoAction | Self::AutomaticFollowupScheduled => 0,
135        }
136    }
137}
138
139// API used by callers outside the compaction scheduler module.
140impl CompactionScheduler {
141    /// Creates an empty scheduler bound to one region worker.
142    ///
143    /// # Effects
144    ///
145    /// Stores the worker request sender and shared scheduling resources. No
146    /// compaction is dispatched until a scheduling method is called.
147    ///
148    /// # Constraints
149    ///
150    /// All lifecycle methods on the returned scheduler must be driven by the
151    /// same worker so region state transitions remain serialized.
152    #[allow(clippy::too_many_arguments)]
153    pub(crate) fn new(
154        scheduler: SchedulerRef,
155        request_sender: Sender<WorkerRequestWithTime>,
156        cache_manager: CacheManagerRef,
157        engine_config: Arc<MitoConfig>,
158        listener: WorkerListener,
159        plugins: Plugins,
160        memory_manager: Arc<CompactionMemoryManager>,
161        memory_policy: OnExhaustedPolicy,
162    ) -> Self {
163        Self {
164            disable_compaction: *DISABLE_COMPACTION,
165            scheduler,
166            region_status: HashMap::new(),
167            request_sender,
168            cache_manager,
169            engine_config,
170            memory_manager,
171            memory_policy,
172            listener,
173            plugins,
174            next_plan_id: 0,
175        }
176    }
177
178    /// Registers a loaded region; replacing its VersionControl invalidates the old incarnation.
179    pub(crate) fn register_region(
180        &mut self,
181        version_control: &VersionControlRef,
182        access_layer: &AccessLayerRef,
183    ) {
184        let region_id = version_control.region_id();
185        if let Some(status) = self.region_status.get_mut(&region_id) {
186            if Arc::ptr_eq(&status.version_control, version_control) {
187                return;
188            }
189            status.on_failure(Arc::new(RegionClosedSnafu { region_id }.build()));
190        }
191        self.region_status.insert(
192            region_id,
193            CompactionStatus::new(region_id, version_control.clone(), access_layer.clone()),
194        );
195    }
196
197    /// Reports active work rather than the existence of a resident region record.
198    #[cfg(test)]
199    fn is_compacting(&self, region_id: RegionId) -> bool {
200        self.region_status
201            .get(&region_id)
202            .is_some_and(CompactionStatus::is_compacting)
203    }
204
205    /// Accepts an automatic compaction trigger for a region.
206    ///
207    /// # Effects
208    ///
209    /// Starts planning when the registered region is idle.
210    /// Otherwise, an active status coalesces the trigger into one unrestricted
211    /// follow-up cycle. Returns `true` only when this call dispatches planning.
212    ///
213    /// # Constraints
214    ///
215    /// The owning region worker must call this method serially. Automatic
216    /// triggers have no waiter and may be ignored while a DDL fence is active.
217    /// The region identity comes from `version_control`; `access_layer` and
218    /// `manifest_ctx` must belong to the same region.
219    pub(crate) fn schedule_automatic_compaction(
220        &mut self,
221        compact_options: compact_request::Options,
222        version_control: &VersionControlRef,
223        access_layer: &AccessLayerRef,
224        manifest_ctx: &ManifestContextRef,
225        schema_metadata_manager: SchemaMetadataManagerRef,
226    ) -> Result<bool> {
227        self.schedule_compaction(
228            CompactionTrigger::Automatic,
229            compact_options,
230            version_control,
231            access_layer,
232            OptionOutputTx::none(),
233            manifest_ctx,
234            schema_metadata_manager,
235            1, // Default for automatic compaction
236            None,
237        )
238    }
239
240    /// Accepts a manual compaction request for a region.
241    ///
242    /// # Effects
243    ///
244    /// Starts planning when the registered region is idle. During
245    /// an automatic compaction, queues the request and replaces any older
246    /// pending manual request; during a manual compaction, rejects the new
247    /// request through `waiter`. Returns `true` only when this call dispatches
248    /// planning.
249    ///
250    /// # Constraints
251    ///
252    /// The owning region worker must call this method serially. The supplied
253    /// waiter is completed exactly once by cycle completion, replacement,
254    /// cancellation, failure, or region teardown. The region identity comes
255    /// from `version_control`; `access_layer` and `manifest_ctx` must belong to
256    /// the same region.
257    #[allow(clippy::too_many_arguments)]
258    pub(crate) fn schedule_manual_compaction(
259        &mut self,
260        compact_options: compact_request::Options,
261        version_control: &VersionControlRef,
262        access_layer: &AccessLayerRef,
263        waiter: OptionOutputTx,
264        manifest_ctx: &ManifestContextRef,
265        schema_metadata_manager: SchemaMetadataManagerRef,
266        max_parallelism: usize,
267        time_range: Option<TimestampRange>,
268    ) -> Result<bool> {
269        self.schedule_compaction(
270            CompactionTrigger::Manual,
271            compact_options,
272            version_control,
273            access_layer,
274            waiter,
275            manifest_ctx,
276            schema_metadata_manager,
277            max_parallelism,
278            time_range,
279        )
280    }
281
282    /// Applies a background planning result to the current region state.
283    ///
284    /// # Effects
285    ///
286    /// Rejects stale plans, submits a prepared local or remote execution, or
287    /// completes the cycle when planning produced no executable plan. The
288    /// returned transition reports a dispatched automatic follow-up or DDL
289    /// requests released by the resulting terminal transition.
290    ///
291    /// # Constraints
292    ///
293    /// The owning worker must call this method for planning notifications and
294    /// must execute every returned DDL request. The notification may be stale;
295    /// stale notifications intentionally have no effect.
296    pub(crate) async fn handle_compaction_pick_finished(
297        &mut self,
298        finished: CompactionPickFinished,
299        manifest_ctx: &ManifestContextRef,
300        schema_metadata_manager: SchemaMetadataManagerRef,
301    ) -> CompactionTransition {
302        self.handle_compaction_pick_finished_inner(finished, manifest_ctx, schema_metadata_manager)
303            .await
304    }
305
306    /// Returns whether a terminal notification belongs to the installed execution.
307    ///
308    /// # Effects
309    ///
310    /// Performs a read-only identity check against the region's current plan.
311    ///
312    /// # Constraints
313    ///
314    /// Callers must check this before applying a compaction edit. Matching only
315    /// the region id is insufficient because close/reopen can replace a status.
316    pub(crate) fn is_current_execution(
317        &self,
318        region_id: RegionId,
319        execution: &CompactionExecution,
320    ) -> bool {
321        self.region_status
322            .get(&region_id)
323            .is_some_and(|status| status.matches_execution(execution))
324    }
325
326    /// Completes the installed execution after its edit has been applied.
327    ///
328    /// # Effects
329    ///
330    /// Notifies waiters, schedules a pending manual or automatic follow-up, or
331    /// returns the region to idle. A follow-up is scheduled whenever the cycle
332    /// latched an automatic trigger, or when `made_progress` is true (the
333    /// execution produced output or reduced the file count): successful compaction
334    /// keeps draining the region until the picker returns no plan. The returned
335    /// transition reports a dispatched automatic follow-up or DDLs that are now
336    /// safe to execute.
337    ///
338    /// # Constraints
339    ///
340    /// The owning worker must call this only after accepting and applying the
341    /// execution's result, and must execute every returned DDL request. Stale
342    /// executions are ignored.
343    pub(crate) async fn on_execution_finished(
344        &mut self,
345        region_id: RegionId,
346        execution: &CompactionExecution,
347        manifest_ctx: &ManifestContextRef,
348        schema_metadata_manager: SchemaMetadataManagerRef,
349        made_progress: bool,
350    ) -> CompactionTransition {
351        if !self.is_current_execution(region_id, execution) {
352            return CompactionTransition::NoAction;
353        }
354        self.on_compaction_finished(
355            region_id,
356            manifest_ctx,
357            schema_metadata_manager,
358            made_progress,
359        )
360        .await
361    }
362
363    /// Records failure of the installed execution.
364    ///
365    /// # Effects
366    ///
367    /// Resets the matching region to idle and fails all compaction waiters and
368    /// dependent DDL requests with the supplied error.
369    ///
370    /// # Constraints
371    ///
372    /// A stale execution failure is ignored so it cannot tear down replacement
373    /// state.
374    pub(crate) fn on_execution_failed(
375        &mut self,
376        region_id: RegionId,
377        execution: &CompactionExecution,
378        err: Arc<Error>,
379    ) {
380        if !self.is_current_execution(region_id, execution) {
381            return;
382        }
383        self.on_compaction_failed(region_id, err);
384    }
385
386    /// Removes compaction state because the region was dropped.
387    ///
388    /// # Effects
389    ///
390    /// Fails all compaction and DDL waiters owned by the status.
391    ///
392    /// # Constraints
393    ///
394    /// The owning worker must invoke this as part of serialized region teardown.
395    pub(crate) fn on_region_dropped(&mut self, region_id: RegionId) {
396        self.remove_region(
397            region_id,
398            Arc::new(RegionDroppedSnafu { region_id }.build()),
399        );
400    }
401
402    /// Removes compaction state because the region was closed.
403    ///
404    /// # Effects
405    ///
406    /// Fails all compaction and DDL waiters owned by the status.
407    ///
408    /// # Constraints
409    ///
410    /// The owning worker must invoke this as part of serialized region teardown.
411    pub(crate) fn on_region_closed(&mut self, region_id: RegionId) {
412        self.remove_region(region_id, Arc::new(RegionClosedSnafu { region_id }.build()));
413    }
414
415    /// Resets the cycle after truncate, retaining registration and outstanding DDL replies.
416    ///
417    /// # Effects
418    ///
419    /// Fails all compaction and DDL waiters owned by the status.
420    ///
421    /// # Constraints
422    ///
423    /// The owning worker must invoke this after truncate completes.
424    pub(crate) fn on_region_truncated(&mut self, region_id: RegionId) {
425        self.fail_compaction(
426            region_id,
427            Arc::new(RegionTruncatedSnafu { region_id }.build()),
428        );
429    }
430
431    /// Cancels a running compaction and queues its dependent DDL atomically.
432    ///
433    /// # Effects
434    ///
435    /// Requests cancellation when the active phase can still stop, then queues
436    /// the DDL. The queued DDL forms a scheduling fence for subsequent
437    /// compaction triggers. If no compaction is running, returns the sender and
438    /// typed request unchanged.
439    ///
440    /// # Constraints
441    ///
442    /// Production callers use this only for [`DdlRequest::Truncate`] and
443    /// [`DdlRequest::EnterStaging`]. `Ok(())` means the caller must wait for a
444    /// terminal callback to receive and execute the queued DDL; `Err` means the
445    /// caller must execute the returned request directly.
446    pub(crate) fn try_cancel_and_add_ddl<T>(
447        &mut self,
448        region_id: RegionId,
449        sender: OptionOutputTx,
450        request: T,
451        into_ddl_request: impl FnOnce(T) -> DdlRequest,
452    ) -> std::result::Result<(), (OptionOutputTx, T)> {
453        let Some(status) = self.region_status.get_mut(&region_id) else {
454            return Err((sender, request));
455        };
456        if !status.is_compacting() {
457            return Err((sender, request));
458        }
459        status.request_cancel();
460
461        let request = SenderDdlRequest {
462            region_id,
463            sender,
464            request: into_ddl_request(request),
465        };
466        debug!(
467            "Added pending DDL request for region: {}, ddl: {:?}",
468            request.region_id, request.request
469        );
470        status.pending_ddl_requests.push(request);
471        Ok(())
472    }
473}
474
475// Internal state-machine helpers.
476impl CompactionScheduler {
477    /// Returns the current plan id and advances the counter.
478    ///
479    /// Takes the counter instead of `&mut self` so callers can bump it while
480    /// holding a mutable borrow of a region status.
481    fn next_plan_id(counter: &mut u64) -> u64 {
482        let plan_id = *counter;
483        *counter = counter.wrapping_add(1);
484        plan_id
485    }
486
487    #[allow(clippy::too_many_arguments)]
488    fn schedule_compaction(
489        &mut self,
490        trigger: CompactionTrigger,
491        compact_options: compact_request::Options,
492        version_control: &VersionControlRef,
493        access_layer: &AccessLayerRef,
494        waiter: OptionOutputTx,
495        manifest_ctx: &ManifestContextRef,
496        schema_metadata_manager: SchemaMetadataManagerRef,
497        max_parallelism: usize,
498        time_range: Option<TimestampRange>,
499    ) -> Result<bool> {
500        if self.disable_compaction {
501            if matches!(trigger, CompactionTrigger::Manual) {
502                waiter.send(
503                    UnsupportedOperationSnafu {
504                        err_msg: format!("SST compaction is disabled by {DISABLE_COMPACTION_ENV}"),
505                    }
506                    .fail(),
507                );
508            }
509            return Ok(false);
510        }
511        let region_id = version_control.region_id();
512        let status = self.region_status.entry(region_id).or_insert_with(|| {
513            CompactionStatus::new(region_id, version_control.clone(), access_layer.clone())
514        });
515        if status.has_ddl() {
516            info!(
517                "Region {} has pending DDL requests, ignoring compaction: queued={}, executing={}",
518                region_id,
519                status.pending_ddl_requests.len(),
520                status.executing_ddl.is_some(),
521            );
522            waiter.send(CompactionCancelledSnafu.fail());
523            return Ok(false);
524        }
525        let current_state = manifest_ctx.current_state();
526        if current_state == RegionRoleState::Leader(RegionLeaderState::Staging) {
527            info!(
528                "Skipping compaction for region {} in staging mode, options: {:?}",
529                region_id, compact_options
530            );
531            waiter.send(Ok(0));
532            return Ok(false);
533        }
534
535        if status.is_compacting() {
536            match trigger {
537                CompactionTrigger::Automatic => status.mark_automatic_trigger(),
538                CompactionTrigger::Manual if status.is_manual_compaction() => {
539                    waiter.send(ManualCompactionAlreadyRunningSnafu { region_id }.fail());
540                    info!(
541                        "Region {} already has a manually triggered compaction running",
542                        region_id
543                    );
544                }
545                CompactionTrigger::Manual => {
546                    status.set_pending_request(PendingCompaction {
547                        options: compact_options,
548                        waiter,
549                        max_parallelism,
550                        time_range,
551                    });
552                    info!(
553                        "Region {} is running an automatic compaction; manual compaction will be re-scheduled",
554                        region_id
555                    );
556                }
557            }
558            return Ok(false);
559        }
560
561        // Publish the picking phase before dispatching background planning.
562        let plan_id = Self::next_plan_id(&mut self.next_plan_id);
563        status.start_picking_with_trigger(plan_id, trigger);
564        let request = status.new_compaction_request(
565            self.request_sender.clone(),
566            self.engine_config.clone(),
567            self.cache_manager.clone(),
568            manifest_ctx,
569            self.listener.clone(),
570            schema_metadata_manager,
571            max_parallelism,
572        );
573        status.merge_waiter(waiter);
574        self.dispatch_compaction_planning(plan_id, request, compact_options, time_range, trigger);
575        self.listener.on_compaction_scheduled(region_id);
576        Ok(true)
577    }
578
579    // Handle pending manual compaction request for the region.
580    //
581    // Returns true if should early return, false otherwise.
582    fn handle_pending_compaction_request(
583        &mut self,
584        region_id: RegionId,
585        manifest_ctx: &ManifestContextRef,
586        schema_metadata_manager: SchemaMetadataManagerRef,
587    ) -> bool {
588        let Some(status) = self.region_status.get_mut(&region_id) else {
589            return true;
590        };
591
592        // If there is a pending manual compaction request, schedule it.
593        // and defer returning the pending DDL requests to the caller.
594        let Some(pending_request) = std::mem::take(&mut status.pending_request) else {
595            return false;
596        };
597
598        let PendingCompaction {
599            options,
600            waiter,
601            max_parallelism,
602            time_range,
603        } = pending_request;
604
605        let request = status.new_compaction_request(
606            self.request_sender.clone(),
607            self.engine_config.clone(),
608            self.cache_manager.clone(),
609            manifest_ctx,
610            self.listener.clone(),
611            schema_metadata_manager,
612            max_parallelism,
613        );
614        status.merge_waiter(waiter);
615        // Bump the counter through a disjoint field borrow so the `status`
616        // borrow stays alive; nothing could have removed the status since it
617        // was fetched above.
618        let plan_id = Self::next_plan_id(&mut self.next_plan_id);
619        status.start_picking_with_trigger(plan_id, CompactionTrigger::Manual);
620        self.dispatch_compaction_planning(
621            plan_id,
622            request,
623            options,
624            time_range,
625            CompactionTrigger::Manual,
626        );
627        debug!(
628            "Successfully scheduled manual compaction planning for region id: {}",
629            region_id
630        );
631        true
632    }
633
634    /// Notifies the scheduler that the compaction job is finished successfully.
635    async fn on_compaction_finished(
636        &mut self,
637        region_id: RegionId,
638        manifest_ctx: &ManifestContextRef,
639        schema_metadata_manager: SchemaMetadataManagerRef,
640        made_progress: bool,
641    ) -> CompactionTransition {
642        if !self.region_status.contains_key(&region_id) {
643            return CompactionTransition::NoAction;
644        }
645
646        if let Some(ddls) = self.take_ready_ddls(region_id) {
647            return CompactionTransition::DdlReady(ddls);
648        }
649
650        if self.handle_pending_compaction_request(
651            region_id,
652            manifest_ctx,
653            schema_metadata_manager.clone(),
654        ) {
655            return CompactionTransition::NoAction;
656        }
657
658        let Some(status) = self.region_status.get_mut(&region_id) else {
659            return CompactionTransition::NoAction;
660        };
661        for waiter in status.take_waiters() {
662            waiter.send(Ok(0));
663        }
664
665        // Keep draining when the cycle latched an automatic trigger, or when the
666        // execution produced output or reduced the file count. An empty edit stops
667        // here; the next flush trigger resumes.
668        let should_continue = status.active.reset_automatic_followup() || made_progress;
669        if should_continue
670            && self.schedule_automatic_followup(region_id, manifest_ctx, schema_metadata_manager)
671        {
672            return CompactionTransition::AutomaticFollowupScheduled;
673        }
674        if let Some(status) = self.region_status.get_mut(&region_id) {
675            status.become_idle();
676        }
677        CompactionTransition::NoAction
678    }
679
680    fn schedule_automatic_followup(
681        &mut self,
682        region_id: RegionId,
683        manifest_ctx: &ManifestContextRef,
684        schema_metadata_manager: SchemaMetadataManagerRef,
685    ) -> bool {
686        if self.disable_compaction {
687            return false;
688        }
689        let Some(status) = self.region_status.get_mut(&region_id) else {
690            return false;
691        };
692        // An external automatic trigger requires one unrestricted follow-up pick.
693        let request = status.new_compaction_request(
694            self.request_sender.clone(),
695            self.engine_config.clone(),
696            self.cache_manager.clone(),
697            manifest_ctx,
698            self.listener.clone(),
699            schema_metadata_manager,
700            MAX_PARALLEL_COMPACTION,
701        );
702        // Bump the counter through a disjoint field borrow so the `status`
703        // borrow stays alive; nothing could have removed the status since it
704        // was fetched above.
705        let plan_id = Self::next_plan_id(&mut self.next_plan_id);
706        status.start_regular_picking(plan_id);
707        self.dispatch_compaction_planning(
708            plan_id,
709            request,
710            compact_request::Options::Regular(Default::default()),
711            None,
712            CompactionTrigger::Automatic,
713        );
714        debug!(
715            "Successfully scheduled next compaction planning for region id: {}",
716            region_id
717        );
718        true
719    }
720
721    /// Notifies the scheduler that the compaction job is cancelled cooperatively.
722    #[cfg(test)]
723    async fn on_compaction_cancelled(&mut self, region_id: RegionId) -> Vec<SenderDdlRequest> {
724        self.finish_compaction_on_cancel(region_id)
725    }
726
727    /// Notifies the scheduler that the compaction job is failed.
728    fn on_compaction_failed(&mut self, region_id: RegionId, err: Arc<Error>) {
729        error!(err; "Region {} failed to compact, cancel all pending tasks", region_id);
730        self.fail_compaction(region_id, err);
731    }
732
733    #[cfg(test)]
734    fn add_ddl_request_to_pending(&mut self, request: SenderDdlRequest) {
735        self.region_status
736            .get_mut(&request.region_id)
737            .unwrap()
738            .pending_ddl_requests
739            .push(request);
740    }
741
742    #[cfg(test)]
743    fn has_pending_ddls(&self, region_id: RegionId) -> bool {
744        let has_pending = self
745            .region_status
746            .get(&region_id)
747            .map(|status| !status.pending_ddl_requests.is_empty())
748            .unwrap_or(false);
749        debug!(
750            "Checked pending DDL requests for region: {}, has_pending: {}",
751            region_id, has_pending
752        );
753        has_pending
754    }
755
756    #[cfg(test)]
757    fn request_cancel(&mut self, region_id: RegionId) -> RequestCancelResult {
758        self.region_status
759            .get_mut(&region_id)
760            .unwrap()
761            .request_cancel()
762    }
763
764    /// Unregisters a closed or dropped region and fails its queued requests.
765    fn remove_region(&mut self, region_id: RegionId, err: Arc<Error>) {
766        if let Some(mut status) = self.region_status.remove(&region_id) {
767            status.on_failure(err);
768        }
769    }
770
771    /// Fails a cycle without discarding the loaded region's registration.
772    fn fail_compaction(&mut self, region_id: RegionId, err: Arc<Error>) {
773        if let Some(status) = self.region_status.get_mut(&region_id) {
774            status.on_failure(err);
775        }
776    }
777
778    /// Returns the cycle to idle and records any DDL handoff in the same resident entry.
779    fn finish_compaction_on_cancel(&mut self, region_id: RegionId) -> Vec<SenderDdlRequest> {
780        let Some(status) = self.region_status.get_mut(&region_id) else {
781            return Vec::new();
782        };
783
784        let ddls = status.on_cancel();
785        self.fence_ddl_replies(region_id, ddls)
786    }
787
788    /// Hands off pending DDLs after compaction ends and cancels remaining compaction waiters.
789    fn take_ready_ddls(&mut self, region_id: RegionId) -> Option<Vec<SenderDdlRequest>> {
790        if self
791            .region_status
792            .get(&region_id)?
793            .pending_ddl_requests
794            .is_empty()
795        {
796            return None;
797        }
798        Some(self.finish_compaction_on_cancel(region_id))
799    }
800
801    /// Hands off a drained batch of DDLs and fences new compaction until every reply lands.
802    ///
803    /// Installs the fence, intercepts each DDL's reply, and returns the same batch so the caller
804    /// can execute it. The real result of each DDL is still forwarded to its original sender.
805    ///
806    /// # Constraints
807    /// Must only be called from a compaction terminal transition, and only by the worker that owns
808    /// `region_id`.
809    fn fence_ddl_replies(
810        &mut self,
811        region_id: RegionId,
812        mut ddls: Vec<SenderDdlRequest>,
813    ) -> Vec<SenderDdlRequest> {
814        if ddls.is_empty() {
815            return ddls;
816        }
817        let generation = Self::next_plan_id(&mut self.next_plan_id);
818        if let Some(status) = self.region_status.get_mut(&region_id) {
819            // Set region status to "executing DDLs" so future compaction won't be scheduled.
820            status.executing_ddl = Some(DdlExecution::new(generation, ddls.len()));
821        }
822
823        for (request_index, ddl) in ddls.iter_mut().enumerate() {
824            // Swap the DDL's real reply sender for a temporary one so we can observe it. The
825            // spawned task forwards the observed result to the original caller and, in between,
826            // notifies the scheduler that this slot of the batch has completed.
827            let (tx, rx) = tokio::sync::oneshot::channel();
828            let reply = std::mem::replace(&mut ddl.sender, tx.into());
829            let sender = self.request_sender.clone();
830            common_runtime::spawn_global(async move {
831                use snafu::ResultExt;
832                let result = rx
833                    .await
834                    .context(crate::error::RecvSnafu)
835                    .and_then(|result| result);
836                // Route the completion through the worker channel so `on_ddl_complete` sees it
837                // serialized with other region state transitions.
838                //
839                // At most one batch can be outstanding per region: [`CompactionStatus::has_ddl`] blocks
840                // compaction admission while `executing_ddl` is set, so `fence_ddl_replies` cannot run again
841                // for the same region before the current batch drains. A single `(generation, request_index)`
842                // pair therefore identifies every reply. `generation` comes from the scheduler-wide counter,
843                // which outlives region statuses, so a late reply from a previous incarnation cannot clear the
844                // current fence. Inside the batch `request_index` makes completion idempotent: only the last
845                // distinct reply releases the fence.
846                let _ = sender
847                    .send(WorkerRequestWithTime::new(
848                        crate::request::WorkerRequest::Background {
849                            region_id,
850                            notify: crate::request::BackgroundNotify::CompactionDdlComplete {
851                                generation,
852                                request_index,
853                            },
854                        },
855                    ))
856                    .await;
857                reply.send(result);
858            });
859        }
860        ddls
861    }
862
863    /// Releases the DDL fence after its last reply, ignoring replies from an older generation.
864    pub(crate) fn on_ddl_complete(
865        &mut self,
866        region_id: RegionId,
867        generation: u64,
868        request_index: usize,
869    ) {
870        let Some(status) = self.region_status.get_mut(&region_id) else {
871            return;
872        };
873        // `complete` returns false for a stale generation or a duplicate `request_index`, so a late or
874        // repeated notify cannot end the batch early. It only clears the fence on the last distinct
875        // reply, which is what lets `has_ddl` unblock compaction admission again.
876        if let Some(ddl) = &mut status.executing_ddl
877            && ddl.complete(generation, request_index)
878        {
879            status.executing_ddl = None;
880        }
881    }
882}
883
884impl Drop for CompactionScheduler {
885    fn drop(&mut self) {
886        for (region_id, mut status) in self.region_status.drain() {
887            // We are shutting down so notify all pending tasks.
888            status.on_failure(Arc::new(RegionClosedSnafu { region_id }.build()));
889        }
890    }
891}
892
893#[cfg(test)]
894#[path = "scheduler_test.rs"]
895mod tests;