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