mito2/compaction/scheduler/
state.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
15use std::collections::{HashMap, HashSet, hash_map};
16use std::sync::Arc;
17use std::time::Instant;
18
19use api::v1::region::compact_request;
20use common_meta::key::SchemaMetadataManagerRef;
21use common_telemetry::debug;
22use common_time::range::TimestampRange;
23use snafu::ResultExt;
24use store_api::storage::RegionId;
25use tokio::sync::mpsc::Sender;
26
27use crate::access_layer::AccessLayerRef;
28use crate::cache::CacheManagerRef;
29use crate::compaction::compactor::CompactionVersion;
30use crate::compaction::picker::PickerOutput;
31use crate::compaction::scheduler::local::LocalCompaction;
32use crate::compaction::scheduler::planning::CompactionRequest;
33use crate::config::MitoConfig;
34use crate::error::{
35    CompactRegionSnafu, CompactionCancelledSnafu, Error, ManualCompactionOverrideSnafu,
36};
37use crate::region::ManifestContextRef;
38use crate::region::version::VersionControlRef;
39use crate::request::{OptionOutputTx, OutputTx, SenderDdlRequest, WorkerRequestWithTime};
40#[cfg(test)]
41use crate::schedule::CancellableTaskState;
42use crate::schedule::RequestCancelResult;
43use crate::sst::file::{FileHandle, RegionFileId};
44use crate::worker::WorkerListener;
45
46/// Identifies an accepted compaction attempt and keeps its SST reservations alive.
47/// The plan id fences terminal notifications from superseded attempts.
48#[derive(Debug, Clone)]
49pub(crate) struct CompactionExecution {
50    plan_id: u64,
51    _files: CompactingFiles,
52}
53
54impl CompactionExecution {
55    /// Returns the attempt ID used to reject stale execution notifications.
56    pub(crate) fn plan_id(&self) -> u64 {
57        self.plan_id
58    }
59    pub(super) fn new(plan_id: u64, files: CompactingFiles) -> Self {
60        Self {
61            plan_id,
62            _files: files,
63        }
64    }
65
66    pub(crate) fn matches(&self, other: &Self) -> bool {
67        self.plan_id == other.plan_id
68    }
69
70    #[cfg(test)]
71    pub(crate) fn for_test(plan_id: u64) -> Self {
72        Self::new(plan_id, CompactingFiles::empty())
73    }
74}
75
76#[derive(Debug)]
77pub(crate) enum CompactionPhase {
78    /// The region is registered but has no compaction in progress.
79    Idle,
80    Picking {
81        plan_id: u64,
82        cancelled: bool,
83    },
84    #[cfg(test)]
85    Local {
86        state: CancellableTaskState,
87        execution: CompactionExecution,
88    },
89    Remote {
90        execution: CompactionExecution,
91    },
92    Units(LocalCompaction),
93}
94
95/// Describes how the current compaction cycle was triggered.
96#[derive(Debug, Clone, Copy, PartialEq, Eq)]
97pub(super) enum CompactionTrigger {
98    Automatic,
99    Manual,
100}
101
102#[derive(Debug)]
103pub(super) struct ActiveCompaction {
104    pub(super) phase: CompactionPhase,
105    trigger: CompactionTrigger,
106    /// Waiters satisfied by the current planning or execution cycle. Picking waiters move into
107    /// the submitted task; regular triggers coalesced during execution accumulate here.
108    pub(super) waiters: Vec<OutputTx>,
109    /// An automatic trigger arrived during this cycle and requires one unrestricted regular
110    /// picking cycle after the current cycle finishes. Recorded in every phase so an external
111    /// trigger always resets the continuation scope, even during local/remote execution.
112    pub(super) automatic_followup_required: bool,
113}
114
115impl ActiveCompaction {
116    /// Creates an empty cycle while retaining the region's scheduling record.
117    fn idle() -> Self {
118        Self {
119            phase: CompactionPhase::Idle,
120            trigger: CompactionTrigger::Automatic,
121            waiters: Vec::new(),
122            automatic_followup_required: false,
123        }
124    }
125
126    #[cfg(test)]
127    pub(super) fn picking(
128        plan_id: u64,
129        waiters: Vec<OutputTx>,
130        trigger: CompactionTrigger,
131    ) -> Self {
132        Self {
133            phase: CompactionPhase::Picking {
134                plan_id,
135                cancelled: false,
136            },
137            trigger,
138            waiters,
139            automatic_followup_required: false,
140        }
141    }
142
143    pub(super) fn start_picking(&mut self, plan_id: u64, trigger: CompactionTrigger) {
144        self.phase = CompactionPhase::Picking {
145            plan_id,
146            cancelled: false,
147        };
148        self.trigger = trigger;
149    }
150
151    pub(super) fn start_regular_picking(&mut self, plan_id: u64) {
152        self.start_picking(plan_id, CompactionTrigger::Automatic);
153    }
154
155    /// Marks an automatic trigger. The trigger is coalesced into a single unrestricted
156    /// follow-up cycle regardless of the current phase.
157    pub(super) fn mark_automatic_trigger(&mut self) {
158        self.automatic_followup_required = true;
159    }
160
161    /// Resets whether an unrestricted follow-up cycle is required and
162    /// return the previous value.
163    pub(super) fn reset_automatic_followup(&mut self) -> bool {
164        std::mem::take(&mut self.automatic_followup_required)
165    }
166
167    pub(super) fn is_manual(&self) -> bool {
168        self.trigger == CompactionTrigger::Manual
169    }
170
171    pub(super) fn is_picking(&self, expected_plan_id: u64) -> bool {
172        matches!(
173            self.phase,
174            CompactionPhase::Picking { plan_id, .. } if plan_id == expected_plan_id
175        )
176    }
177
178    pub(super) fn accept_plan(&self, expected_plan_id: u64) -> bool {
179        matches!(
180            self.phase,
181            CompactionPhase::Picking {
182                plan_id,
183                cancelled: false,
184            } if plan_id == expected_plan_id
185        )
186    }
187
188    pub(super) fn matches_execution(&self, execution: &CompactionExecution) -> bool {
189        match &self.phase {
190            CompactionPhase::Idle | CompactionPhase::Picking { .. } => None,
191            #[cfg(test)]
192            CompactionPhase::Local { execution, .. } => Some(execution),
193            CompactionPhase::Units(_) => None,
194            CompactionPhase::Remote { execution } => Some(execution),
195        }
196        .is_some_and(|current| current.matches(execution))
197    }
198
199    pub(super) fn request_cancel(&mut self) -> RequestCancelResult {
200        match &mut self.phase {
201            CompactionPhase::Idle => RequestCancelResult::TooLateToCancel,
202            CompactionPhase::Picking { cancelled, .. } => {
203                if *cancelled {
204                    RequestCancelResult::AlreadyCancelling
205                } else {
206                    *cancelled = true;
207                    RequestCancelResult::CancelIssued
208                }
209            }
210            #[cfg(test)]
211            CompactionPhase::Local { state, .. } => state.request_cancel(),
212            CompactionPhase::Units(units) => units.request_cancel(),
213            CompactionPhase::Remote { .. } => RequestCancelResult::TooLateToCancel,
214        }
215    }
216
217    pub(super) fn merge_waiter(&mut self, mut waiter: OptionOutputTx) {
218        if let Some(waiter) = waiter.take_inner() {
219            self.waiters.push(waiter);
220        }
221    }
222}
223
224/// Owns atomic reservations for every SST selected by a compaction plan.
225#[derive(Debug, Clone)]
226pub(crate) struct CompactingFiles {
227    files: HashMap<RegionFileId, Arc<CompactingFile>>,
228}
229
230/// Keeps an SST reserved until the last shared lease is dropped.
231#[derive(Debug)]
232struct CompactingFile {
233    file: FileHandle,
234}
235
236impl CompactingFiles {
237    /// Reserves both merge inputs and expired files from an accepted picker result.
238    pub(super) fn try_new(output: &PickerOutput) -> Option<Self> {
239        Self::try_reserve(
240            output
241                .outputs
242                .iter()
243                .flat_map(|o| &o.inputs)
244                .chain(&output.expired_ssts),
245        )
246    }
247
248    /// Reserves distinct inputs, releasing acquired leases if any input is already busy.
249    fn try_reserve<'a>(selected_files: impl IntoIterator<Item = &'a FileHandle>) -> Option<Self> {
250        let mut files = HashMap::new();
251
252        for file in selected_files {
253            let hash_map::Entry::Vacant(entry) = files.entry(file.file_id()) else {
254                continue;
255            };
256            if !file.try_set_compacting() {
257                return None;
258            }
259            entry.insert(Arc::new(CompactingFile { file: file.clone() }));
260        }
261
262        Some(Self { files })
263    }
264
265    /// Clones the shared leases of a dependency group's inputs without briefly releasing any input.
266    /// `self` retains its leases, so an input is released only after every owner, including
267    /// the original collection, drops its lease.
268    /// A delayed remote notifier may retain the same leases across local fallback.
269    /// Indexing keeps splitting proportional to the group's input count.
270    pub(crate) fn for_inputs(&self, inputs: &[FileHandle]) -> Self {
271        Self {
272            files: inputs
273                .iter()
274                .filter_map(|file| {
275                    let id = file.file_id();
276                    self.files.get(&id).map(|lease| (id, lease.clone()))
277                })
278                .collect(),
279        }
280    }
281
282    #[cfg(test)]
283    pub(super) fn empty() -> Self {
284        Self {
285            files: HashMap::new(),
286        }
287    }
288}
289
290impl Drop for CompactingFile {
291    fn drop(&mut self) {
292        self.file.set_compacting(false);
293    }
294}
295
296/// A handed-off DDL batch, retained until each reply has been observed once.
297#[derive(Debug)]
298pub(crate) struct DdlExecution {
299    generation: u64,
300    pending_replies: HashSet<usize>,
301}
302
303impl DdlExecution {
304    /// Tracks a batch's individual replies so duplicate notifications are harmless.
305    pub(crate) fn new(generation: u64, count: usize) -> Self {
306        Self {
307            generation,
308            pending_replies: (0..count).collect(),
309        }
310    }
311
312    /// Returns true only for the last distinct reply of this batch.
313    pub(crate) fn complete(&mut self, generation: u64, request_index: usize) -> bool {
314        self.generation == generation
315            && self.pending_replies.remove(&request_index)
316            && self.pending_replies.is_empty()
317    }
318}
319
320/// Scheduling state retained for the lifetime of a loaded region.
321pub(super) struct CompactionStatus {
322    /// Id of the region.
323    pub(super) region_id: RegionId,
324    /// Version control of the region.
325    pub(super) version_control: VersionControlRef,
326    /// Access layer of the region.
327    pub(super) access_layer: AccessLayerRef,
328    /// Current compaction lifecycle.
329    pub(super) active: ActiveCompaction,
330    /// A manual compaction waiting for the current automatic compaction to finish.
331    ///
332    /// A manual request is rejected if the current compaction is also manual. Automatic requests
333    /// are merged into the active compaction instead of using this slot.
334    pub(super) pending_request: Option<PendingCompaction>,
335    /// Pending DDL requests that should run when compaction is done.
336    ///
337    /// Although [`SenderDdlRequest`] can wrap any DDL variant, production code only queues
338    /// [`crate::request::DdlRequest::Truncate`] and [`crate::request::DdlRequest::EnterStaging`] here. Both must serialize with
339    /// compaction so they observe the version after compaction terminates.
340    pub(super) pending_ddl_requests: Vec<SenderDdlRequest>,
341    /// DDLs already handed to the worker keep new compaction paused until their replies.
342    pub(crate) executing_ddl: Option<DdlExecution>,
343}
344
345impl CompactionStatus {
346    /// Registers an idle region without starting a compaction cycle.
347    pub(super) fn new(
348        region_id: RegionId,
349        version_control: VersionControlRef,
350        access_layer: AccessLayerRef,
351    ) -> CompactionStatus {
352        CompactionStatus {
353            region_id,
354            version_control,
355            access_layer,
356            active: ActiveCompaction::idle(),
357            pending_request: None,
358            pending_ddl_requests: Vec::new(),
359            executing_ddl: None,
360        }
361    }
362
363    #[cfg(test)]
364    pub(super) fn for_test(
365        region_id: RegionId,
366        version_control: VersionControlRef,
367        access_layer: AccessLayerRef,
368    ) -> Self {
369        let mut status = Self::new(region_id, version_control, access_layer);
370        status.active = ActiveCompaction::picking(0, Vec::new(), CompactionTrigger::Automatic);
371        status
372    }
373
374    /// Registration alone does not indicate active work.
375    pub(crate) fn is_compacting(&self) -> bool {
376        !matches!(self.active.phase, CompactionPhase::Idle)
377    }
378
379    /// Both queued and handed-off DDLs prevent new compaction admission.
380    pub(crate) fn has_ddl(&self) -> bool {
381        !self.pending_ddl_requests.is_empty() || self.executing_ddl.is_some()
382    }
383
384    /// Ends the current cycle without losing region or DDL execution state.
385    pub(crate) fn become_idle(&mut self) {
386        self.active = ActiveCompaction::idle();
387    }
388
389    #[cfg(test)]
390    pub(super) fn start_picking(&mut self, plan_id: u64) {
391        self.start_picking_with_trigger(plan_id, CompactionTrigger::Automatic);
392    }
393
394    pub(super) fn start_picking_with_trigger(&mut self, plan_id: u64, trigger: CompactionTrigger) {
395        self.active.start_picking(plan_id, trigger);
396    }
397
398    pub(super) fn start_regular_picking(&mut self, plan_id: u64) {
399        self.active.start_regular_picking(plan_id);
400    }
401
402    pub(super) fn is_picking(&self, expected_plan_id: u64) -> bool {
403        self.active.is_picking(expected_plan_id)
404    }
405
406    pub(super) fn accept_plan(&self, expected_plan_id: u64) -> bool {
407        self.active.accept_plan(expected_plan_id)
408    }
409
410    pub(super) fn is_manual_compaction(&self) -> bool {
411        self.active.is_manual()
412    }
413
414    pub(super) fn matches_execution(&self, execution: &CompactionExecution) -> bool {
415        self.active.matches_execution(execution)
416    }
417
418    #[cfg(test)]
419    pub(super) fn start_local_task(&mut self) -> CancellableTaskState {
420        let state = CancellableTaskState::new();
421        let execution = CompactionExecution::new(0, CompactingFiles::empty());
422        let phase = CompactionPhase::Local {
423            state: state.clone(),
424            execution,
425        };
426        self.active.phase = phase;
427        state
428    }
429
430    #[cfg(test)]
431    pub(super) fn start_remote_task(&mut self) {
432        let execution = CompactionExecution::new(0, CompactingFiles::empty());
433        let phase = CompactionPhase::Remote { execution };
434        self.active.phase = phase;
435    }
436
437    pub(super) fn request_cancel(&mut self) -> RequestCancelResult {
438        self.active.request_cancel()
439    }
440
441    pub(super) fn mark_automatic_trigger(&mut self) {
442        self.active.mark_automatic_trigger();
443    }
444
445    /// Merge the waiter to the pending compaction.
446    pub(super) fn merge_waiter(&mut self, waiter: OptionOutputTx) {
447        self.active.merge_waiter(waiter);
448    }
449
450    pub(super) fn take_waiters(&mut self) -> Vec<OutputTx> {
451        std::mem::take(&mut self.active.waiters)
452    }
453
454    pub(super) fn extend_waiters(&mut self, waiters: Vec<OutputTx>) {
455        self.active.waiters.extend(waiters);
456    }
457
458    pub(super) fn set_phase(&mut self, phase: CompactionPhase) {
459        self.active.phase = phase;
460    }
461
462    /// Sets a pending manual compaction request, replacing an older pending request.
463    pub(super) fn set_pending_request(&mut self, pending: PendingCompaction) {
464        if let Some(prev) = self.pending_request.replace(pending) {
465            debug!(
466                "Replace pending compaction options with new request {:?} for region: {}",
467                prev.options, self.region_id
468            );
469            prev.waiter.send(ManualCompactionOverrideSnafu.fail());
470        }
471    }
472
473    pub(super) fn on_failure(&mut self, err: Arc<Error>) {
474        let mut active = std::mem::replace(&mut self.active, ActiveCompaction::idle());
475        for waiter in active.waiters.drain(..) {
476            waiter.send(Err(err.clone()).context(CompactRegionSnafu {
477                region_id: self.region_id,
478            }));
479        }
480
481        if let Some(pending_compaction) = self.pending_request.take() {
482            pending_compaction
483                .waiter
484                .send(Err(err.clone()).context(CompactRegionSnafu {
485                    region_id: self.region_id,
486                }));
487        }
488
489        for pending_ddl in self.pending_ddl_requests.drain(..) {
490            pending_ddl
491                .sender
492                .send(Err(err.clone()).context(CompactRegionSnafu {
493                    region_id: self.region_id,
494                }));
495        }
496    }
497
498    #[must_use]
499    pub(super) fn on_cancel(&mut self) -> Vec<SenderDdlRequest> {
500        let mut active = std::mem::replace(&mut self.active, ActiveCompaction::idle());
501        for waiter in active.waiters.drain(..) {
502            waiter.send(CompactionCancelledSnafu.fail());
503        }
504
505        if let Some(pending_compaction) = self.pending_request.take() {
506            pending_compaction.waiter.send(
507                Err(Arc::new(CompactionCancelledSnafu.build())).context(CompactRegionSnafu {
508                    region_id: self.region_id,
509                }),
510            );
511        }
512
513        std::mem::take(&mut self.pending_ddl_requests)
514    }
515
516    /// Creates an immutable request for background compaction planning.
517    #[allow(clippy::too_many_arguments)]
518    pub(super) fn new_compaction_request(
519        &self,
520        request_sender: Sender<WorkerRequestWithTime>,
521        engine_config: Arc<MitoConfig>,
522        cache_manager: CacheManagerRef,
523        manifest_ctx: &ManifestContextRef,
524        listener: WorkerListener,
525        schema_metadata_manager: SchemaMetadataManagerRef,
526        max_parallelism: usize,
527    ) -> CompactionRequest {
528        let current_version = CompactionVersion::from(self.version_control.current().version);
529        let start_time = Instant::now();
530
531        CompactionRequest {
532            engine_config,
533            current_version,
534            access_layer: self.access_layer.clone(),
535            request_sender: request_sender.clone(),
536            start_time,
537            cache_manager,
538            manifest_ctx: manifest_ctx.clone(),
539            listener,
540            schema_metadata_manager,
541            max_parallelism,
542        }
543    }
544}
545
546/// A manual compaction request waiting for an automatic compaction to finish.
547pub(super) struct PendingCompaction {
548    /// Compaction options.
549    pub(crate) options: compact_request::Options,
550    /// Waiters of pending requests.
551    pub(crate) waiter: OptionOutputTx,
552    /// Max parallelism for pending compaction.
553    pub(crate) max_parallelism: usize,
554    /// Optional time range that constrains candidate compaction windows.
555    pub(crate) time_range: Option<TimestampRange>,
556}
557
558#[cfg(test)]
559mod tests {
560    use store_api::storage::FileId;
561
562    use super::*;
563    use crate::compaction::test_util::new_file_handle;
564
565    #[test]
566    fn test_compaction_unit_input_leases_release_independently() {
567        let x = new_file_handle(FileId::random(), 0, 100, 0);
568        let y = new_file_handle(FileId::random(), 100, 200, 0);
569        let unused = new_file_handle(FileId::random(), 200, 300, 0);
570        let reserved = CompactingFiles::try_reserve([&x, &y, &x, &unused]).unwrap();
571        let first = reserved.for_inputs(&[x.clone(), x.clone()]);
572        let second = reserved.for_inputs(std::slice::from_ref(&y));
573        drop(reserved);
574        assert!(x.compacting() && y.compacting());
575        assert!(
576            !unused.compacting(),
577            "unit leases must not retain unrelated inputs"
578        );
579        drop(first);
580        assert!(!x.compacting() && y.compacting());
581        let notification = second.clone();
582        drop(second);
583        assert!(y.compacting());
584        drop(notification);
585        assert!(!y.compacting());
586    }
587}
mito2/compaction/scheduler/state.rs

mito2/compaction/scheduler/
state.rs
