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, RwLock};
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::compaction_task::{self, CompactionTaskEntry};
25use store_api::storage::RegionId;
26use tokio::sync::mpsc::Sender;
27
28use crate::access_layer::AccessLayerRef;
29use crate::cache::CacheManagerRef;
30use crate::compaction::CompactionOutput;
31use crate::compaction::compactor::{CompactionRegion, CompactionVersion};
32use crate::compaction::picker::PickerOutput;
33use crate::compaction::scheduler::history::PickHistory;
34use crate::compaction::scheduler::local::{LocalCompaction, LocalCompactionState};
35use crate::compaction::scheduler::planning::CompactionRequest;
36use crate::config::MitoConfig;
37use crate::error::{
38    CompactRegionSnafu, CompactionCancelledSnafu, Error, ManualCompactionOverrideSnafu,
39};
40use crate::region::ManifestContextRef;
41use crate::region::version::VersionControlRef;
42use crate::request::{OptionOutputTx, OutputTx, SenderDdlRequest, WorkerRequestWithTime};
43#[cfg(test)]
44use crate::schedule::CancellableTaskState;
45use crate::schedule::RequestCancelResult;
46use crate::sst::file::{FileHandle, RegionFileId};
47use crate::worker::WorkerListener;
48
49/// Identifies an accepted compaction attempt and keeps its SST reservations alive.
50/// The plan id fences terminal notifications from superseded attempts.
51#[derive(Debug, Clone)]
52pub(crate) struct CompactionExecution {
53    plan_id: u64,
54    _files: CompactingFiles,
55}
56
57impl CompactionExecution {
58    /// Returns the attempt ID used to reject stale execution notifications.
59    pub(crate) fn plan_id(&self) -> u64 {
60        self.plan_id
61    }
62    pub(super) fn new(plan_id: u64, files: CompactingFiles) -> Self {
63        Self {
64            plan_id,
65            _files: files,
66        }
67    }
68
69    pub(crate) fn matches(&self, other: &Self) -> bool {
70        self.plan_id == other.plan_id
71    }
72
73    #[cfg(test)]
74    pub(crate) fn for_test(plan_id: u64) -> Self {
75        Self::new(plan_id, CompactingFiles::empty())
76    }
77}
78
79#[derive(Debug)]
80pub(crate) enum CompactionPhase {
81    /// The region is registered but has no compaction in progress.
82    Idle,
83    Picking {
84        plan_id: u64,
85        cancelled: bool,
86    },
87    #[cfg(test)]
88    Local {
89        state: CancellableTaskState,
90        execution: CompactionExecution,
91    },
92    Remote {
93        execution: CompactionExecution,
94    },
95    Units(LocalCompaction),
96}
97
98/// Describes how the current compaction cycle was triggered.
99#[derive(Debug, Clone, Copy, PartialEq, Eq)]
100pub(crate) enum CompactionTrigger {
101    Automatic,
102    Manual,
103}
104
105#[derive(Debug)]
106pub(super) struct ActiveCompaction {
107    /// Query metadata owned by this live cycle, dropped when it becomes idle.
108    pub(crate) task_state: Option<CompactionTaskState>,
109    pub(super) phase: CompactionPhase,
110    trigger: CompactionTrigger,
111    /// Waiters satisfied by the current planning or execution cycle. Picking waiters move into
112    /// the submitted task; regular triggers coalesced during execution accumulate here.
113    pub(super) waiters: Vec<OutputTx>,
114    /// An automatic trigger arrived during this cycle and requires one unrestricted regular
115    /// picking cycle after the current cycle finishes. Recorded in every phase so an external
116    /// trigger always resets the continuation scope, even during local/remote execution.
117    pub(super) automatic_followup_required: bool,
118}
119
120impl ActiveCompaction {
121    /// Creates an empty cycle while retaining the region's scheduling record.
122    fn idle() -> Self {
123        Self {
124            task_state: None,
125            phase: CompactionPhase::Idle,
126            trigger: CompactionTrigger::Automatic,
127            waiters: Vec::new(),
128            automatic_followup_required: false,
129        }
130    }
131
132    #[cfg(test)]
133    pub(super) fn picking(
134        plan_id: u64,
135        waiters: Vec<OutputTx>,
136        trigger: CompactionTrigger,
137    ) -> Self {
138        Self {
139            task_state: None,
140            phase: CompactionPhase::Picking {
141                plan_id,
142                cancelled: false,
143            },
144            trigger,
145            waiters,
146            automatic_followup_required: false,
147        }
148    }
149
150    pub(super) fn start_picking(&mut self, plan_id: u64, trigger: CompactionTrigger) {
151        self.task_state = None;
152        self.phase = CompactionPhase::Picking {
153            plan_id,
154            cancelled: false,
155        };
156        self.trigger = trigger;
157    }
158
159    pub(super) fn start_regular_picking(&mut self, plan_id: u64) {
160        self.start_picking(plan_id, CompactionTrigger::Automatic);
161    }
162
163    /// Marks an automatic trigger. The trigger is coalesced into a single unrestricted
164    /// follow-up cycle regardless of the current phase.
165    pub(super) fn mark_automatic_trigger(&mut self) {
166        self.automatic_followup_required = true;
167    }
168
169    /// Resets whether an unrestricted follow-up cycle is required and
170    /// return the previous value.
171    pub(super) fn reset_automatic_followup(&mut self) -> bool {
172        std::mem::take(&mut self.automatic_followup_required)
173    }
174
175    pub(super) fn is_manual(&self) -> bool {
176        self.trigger == CompactionTrigger::Manual
177    }
178
179    pub(super) fn is_picking(&self, expected_plan_id: u64) -> bool {
180        matches!(
181            self.phase,
182            CompactionPhase::Picking { plan_id, .. } if plan_id == expected_plan_id
183        )
184    }
185
186    pub(super) fn accept_plan(&self, expected_plan_id: u64) -> bool {
187        matches!(
188            self.phase,
189            CompactionPhase::Picking {
190                plan_id,
191                cancelled: false,
192            } if plan_id == expected_plan_id
193        )
194    }
195
196    pub(super) fn matches_execution(&self, execution: &CompactionExecution) -> bool {
197        match &self.phase {
198            CompactionPhase::Idle | CompactionPhase::Picking { .. } => None,
199            #[cfg(test)]
200            CompactionPhase::Local { execution, .. } => Some(execution),
201            CompactionPhase::Units(_) => None,
202            CompactionPhase::Remote { execution } => Some(execution),
203        }
204        .is_some_and(|current| current.matches(execution))
205    }
206
207    pub(super) fn request_cancel(&mut self) -> RequestCancelResult {
208        match &mut self.phase {
209            CompactionPhase::Idle => RequestCancelResult::TooLateToCancel,
210            CompactionPhase::Picking { cancelled, .. } => {
211                if *cancelled {
212                    RequestCancelResult::AlreadyCancelling
213                } else {
214                    *cancelled = true;
215                    RequestCancelResult::CancelIssued
216                }
217            }
218            #[cfg(test)]
219            CompactionPhase::Local { state, .. } => state.request_cancel(),
220            CompactionPhase::Units(units) => units.request_cancel(),
221            CompactionPhase::Remote { .. } => RequestCancelResult::TooLateToCancel,
222        }
223    }
224
225    pub(super) fn merge_waiter(&mut self, mut waiter: OptionOutputTx) {
226        if let Some(waiter) = waiter.take_inner() {
227            self.waiters.push(waiter);
228        }
229    }
230    /// Captures worker-owned metadata while retaining live background task handles.
231    fn current_task_state(&self) -> Option<CompactionTaskState> {
232        let mut task = self.task_state.clone()?;
233        let entry = &mut task.entry;
234        match &self.phase {
235            CompactionPhase::Idle => return None,
236            CompactionPhase::Picking { cancelled, .. } => {
237                entry.state = if *cancelled {
238                    compaction_task::STATE_CANCELLING
239                } else {
240                    compaction_task::STATE_RUNNING
241                }
242                .into();
243                entry.phase = compaction_task::PHASE_PLANNING.into();
244            }
245            CompactionPhase::Remote { .. } => {
246                entry.state = compaction_task::STATE_RUNNING.into();
247                entry.phase = compaction_task::PHASE_REMOTE.into();
248            }
249            CompactionPhase::Units(local) => task.local_state = Some(local.current_state()),
250            #[cfg(test)]
251            CompactionPhase::Local { .. } => {
252                entry.state = compaction_task::STATE_RUNNING.into();
253                entry.phase = compaction_task::PHASE_EXECUTING.into();
254            }
255        }
256        Some(task)
257    }
258}
259
260/// Owns atomic reservations for every SST selected by a compaction plan.
261#[derive(Debug, Clone)]
262pub(crate) struct CompactingFiles {
263    files: HashMap<RegionFileId, Arc<CompactingFile>>,
264}
265
266/// Keeps an SST reserved until the last shared lease is dropped.
267#[derive(Debug)]
268struct CompactingFile {
269    file: FileHandle,
270}
271
272impl CompactingFiles {
273    /// Reserves both merge inputs and expired files from an accepted picker result.
274    pub(super) fn try_new(output: &PickerOutput) -> Option<Self> {
275        Self::try_reserve(
276            output
277                .outputs
278                .iter()
279                .flat_map(|o| &o.inputs)
280                .chain(&output.expired_ssts),
281        )
282    }
283
284    /// Reserves distinct inputs, releasing acquired leases if any input is already busy.
285    fn try_reserve<'a>(selected_files: impl IntoIterator<Item = &'a FileHandle>) -> Option<Self> {
286        let mut files = HashMap::new();
287
288        for file in selected_files {
289            let hash_map::Entry::Vacant(entry) = files.entry(file.file_id()) else {
290                continue;
291            };
292            if !file.try_set_compacting() {
293                return None;
294            }
295            entry.insert(Arc::new(CompactingFile { file: file.clone() }));
296        }
297
298        Some(Self { files })
299    }
300
301    /// Clones the shared leases of a dependency group's inputs without briefly releasing any input.
302    /// `self` retains its leases, so an input is released only after every owner, including
303    /// the original collection, drops its lease.
304    /// A delayed remote notifier may retain the same leases across local fallback.
305    /// Indexing keeps splitting proportional to the group's input count.
306    pub(crate) fn for_inputs(&self, inputs: &[FileHandle]) -> Self {
307        Self {
308            files: inputs
309                .iter()
310                .filter_map(|file| {
311                    let id = file.file_id();
312                    self.files.get(&id).map(|lease| (id, lease.clone()))
313                })
314                .collect(),
315        }
316    }
317
318    #[cfg(test)]
319    pub(super) fn empty() -> Self {
320        Self {
321            files: HashMap::new(),
322        }
323    }
324}
325
326impl Drop for CompactingFile {
327    fn drop(&mut self) {
328        self.file.set_compacting(false);
329    }
330}
331
332/// A handed-off DDL batch, retained until each reply has been observed once.
333#[derive(Debug)]
334pub(crate) struct DdlExecution {
335    generation: u64,
336    pending_replies: HashSet<usize>,
337}
338
339impl DdlExecution {
340    /// Tracks a batch's individual replies so duplicate notifications are harmless.
341    pub(crate) fn new(generation: u64, count: usize) -> Self {
342        Self {
343            generation,
344            pending_replies: (0..count).collect(),
345        }
346    }
347
348    /// Returns true only for the last distinct reply of this batch.
349    pub(crate) fn complete(&mut self, generation: u64, request_index: usize) -> bool {
350        self.generation == generation
351            && self.pending_replies.remove(&request_index)
352            && self.pending_replies.is_empty()
353    }
354}
355
356/// Scheduling state retained for the lifetime of a loaded region.
357pub(super) struct CompactionStatus {
358    /// Current cycles shared with queries without entering the region worker.
359    pub(crate) shared_state: CompactionStateRef,
360    /// Resident automatic-pick history, retained across batches until region unload.
361    pub(crate) pick_history: PickHistory,
362    /// Id of the region.
363    pub(super) region_id: RegionId,
364    /// Version control of the region.
365    pub(super) version_control: VersionControlRef,
366    /// Access layer of the region.
367    pub(super) access_layer: AccessLayerRef,
368    /// Current compaction lifecycle.
369    pub(super) active: ActiveCompaction,
370    /// A manual compaction waiting for the current automatic compaction to finish.
371    ///
372    /// A manual request is rejected if the current compaction is also manual. Automatic requests
373    /// are merged into the active compaction instead of using this slot.
374    pub(super) pending_request: Option<PendingCompaction>,
375    /// Metadata for the pending manual request, moved to the active cycle on promotion.
376    pub(crate) pending_task_state: Option<CompactionTaskState>,
377    /// Pending DDL requests that should run when compaction is done.
378    ///
379    /// Although [`SenderDdlRequest`] can wrap any DDL variant, production code only queues
380    /// [`crate::request::DdlRequest::Truncate`] and [`crate::request::DdlRequest::EnterStaging`] here. Both must serialize with
381    /// compaction so they observe the version after compaction terminates.
382    pub(super) pending_ddl_requests: Vec<SenderDdlRequest>,
383    /// DDLs already handed to the worker keep new compaction paused until their replies.
384    pub(crate) executing_ddl: Option<DdlExecution>,
385}
386
387impl CompactionStatus {
388    /// Publishes both slots together without holding the lock while building their metadata.
389    pub(crate) fn publish_state(&self) {
390        let active = self.active.current_task_state().map(Arc::new);
391        let pending = self
392            .pending_request
393            .as_ref()
394            .and(self.pending_task_state.as_ref())
395            .cloned()
396            .map(Arc::new);
397        let previous = {
398            let mut shared = self.shared_state.write().unwrap();
399            std::mem::replace(&mut *shared, CompactionState { active, pending })
400        };
401        drop(previous);
402    }
403
404    /// Registers an idle region without starting a compaction cycle.
405    pub(super) fn new(
406        region_id: RegionId,
407        version_control: VersionControlRef,
408        access_layer: AccessLayerRef,
409    ) -> CompactionStatus {
410        CompactionStatus {
411            shared_state: Arc::new(RwLock::new(CompactionState::default())),
412            pick_history: PickHistory::default(),
413            region_id,
414            version_control,
415            access_layer,
416            active: ActiveCompaction::idle(),
417            pending_request: None,
418            pending_task_state: None,
419            pending_ddl_requests: Vec::new(),
420            executing_ddl: None,
421        }
422    }
423
424    #[cfg(test)]
425    pub(super) fn for_test(
426        region_id: RegionId,
427        version_control: VersionControlRef,
428        access_layer: AccessLayerRef,
429    ) -> Self {
430        let mut status = Self::new(region_id, version_control, access_layer);
431        status.active = ActiveCompaction::picking(0, Vec::new(), CompactionTrigger::Automatic);
432        status
433    }
434
435    /// Registration alone does not indicate active work.
436    pub(crate) fn is_compacting(&self) -> bool {
437        !matches!(self.active.phase, CompactionPhase::Idle)
438    }
439
440    /// Both queued and handed-off DDLs prevent new compaction admission.
441    pub(crate) fn has_ddl(&self) -> bool {
442        !self.pending_ddl_requests.is_empty() || self.executing_ddl.is_some()
443    }
444
445    /// Ends the current cycle without losing region or DDL execution state.
446    pub(crate) fn become_idle(&mut self) {
447        self.active = ActiveCompaction::idle();
448        self.publish_state();
449    }
450
451    #[cfg(test)]
452    pub(super) fn start_picking(&mut self, plan_id: u64) {
453        self.start_picking_with_trigger(plan_id, CompactionTrigger::Automatic);
454    }
455
456    pub(super) fn start_picking_with_trigger(&mut self, plan_id: u64, trigger: CompactionTrigger) {
457        self.active.start_picking(plan_id, trigger);
458    }
459
460    pub(super) fn start_regular_picking(&mut self, plan_id: u64) {
461        self.active.start_regular_picking(plan_id);
462    }
463
464    pub(super) fn is_picking(&self, expected_plan_id: u64) -> bool {
465        self.active.is_picking(expected_plan_id)
466    }
467
468    pub(super) fn accept_plan(&self, expected_plan_id: u64) -> bool {
469        self.active.accept_plan(expected_plan_id)
470    }
471
472    pub(super) fn is_manual_compaction(&self) -> bool {
473        self.active.is_manual()
474    }
475
476    pub(super) fn matches_execution(&self, execution: &CompactionExecution) -> bool {
477        self.active.matches_execution(execution)
478    }
479
480    #[cfg(test)]
481    pub(super) fn start_local_task(&mut self) -> CancellableTaskState {
482        let state = CancellableTaskState::new();
483        let execution = CompactionExecution::new(0, CompactingFiles::empty());
484        let phase = CompactionPhase::Local {
485            state: state.clone(),
486            execution,
487        };
488        self.active.phase = phase;
489        state
490    }
491
492    #[cfg(test)]
493    pub(super) fn start_remote_task(&mut self) {
494        let execution = CompactionExecution::new(0, CompactingFiles::empty());
495        let phase = CompactionPhase::Remote { execution };
496        self.active.phase = phase;
497    }
498
499    pub(super) fn request_cancel(&mut self) -> RequestCancelResult {
500        let result = self.active.request_cancel();
501        self.publish_state();
502        result
503    }
504
505    pub(super) fn mark_automatic_trigger(&mut self) {
506        self.active.mark_automatic_trigger();
507    }
508
509    /// Merge the waiter to the pending compaction.
510    pub(super) fn merge_waiter(&mut self, waiter: OptionOutputTx) {
511        self.active.merge_waiter(waiter);
512    }
513
514    pub(super) fn take_waiters(&mut self) -> Vec<OutputTx> {
515        std::mem::take(&mut self.active.waiters)
516    }
517
518    pub(super) fn extend_waiters(&mut self, waiters: Vec<OutputTx>) {
519        self.active.waiters.extend(waiters);
520    }
521
522    pub(super) fn set_phase(&mut self, phase: CompactionPhase) {
523        self.active.phase = phase;
524        self.publish_state();
525    }
526
527    /// Sets a pending manual compaction request, replacing an older pending request.
528    pub(super) fn set_pending_request(&mut self, pending: PendingCompaction) {
529        self.pending_task_state = Some(CompactionTaskState::new(
530            self.region_id,
531            &pending.options,
532            pending.max_parallelism,
533            pending.time_range,
534            CompactionTrigger::Manual,
535        ));
536        if let Some(prev) = self.pending_request.replace(pending) {
537            debug!(
538                "Replace pending compaction options with new request {:?} for region: {}",
539                prev.options, self.region_id
540            );
541            prev.waiter.send(ManualCompactionOverrideSnafu.fail());
542        }
543        self.publish_state();
544    }
545
546    pub(super) fn on_failure(&mut self, err: Arc<Error>) {
547        let mut active = std::mem::replace(&mut self.active, ActiveCompaction::idle());
548        for waiter in active.waiters.drain(..) {
549            waiter.send(Err(err.clone()).context(CompactRegionSnafu {
550                region_id: self.region_id,
551            }));
552        }
553
554        self.pending_task_state = None;
555        if let Some(pending_compaction) = self.pending_request.take() {
556            pending_compaction
557                .waiter
558                .send(Err(err.clone()).context(CompactRegionSnafu {
559                    region_id: self.region_id,
560                }));
561        }
562
563        for pending_ddl in self.pending_ddl_requests.drain(..) {
564            pending_ddl
565                .sender
566                .send(Err(err.clone()).context(CompactRegionSnafu {
567                    region_id: self.region_id,
568                }));
569        }
570        self.publish_state();
571    }
572
573    #[must_use]
574    pub(super) fn on_cancel(&mut self) -> Vec<SenderDdlRequest> {
575        let mut active = std::mem::replace(&mut self.active, ActiveCompaction::idle());
576        for waiter in active.waiters.drain(..) {
577            waiter.send(CompactionCancelledSnafu.fail());
578        }
579
580        self.pending_task_state = None;
581        if let Some(pending_compaction) = self.pending_request.take() {
582            pending_compaction.waiter.send(
583                Err(Arc::new(CompactionCancelledSnafu.build())).context(CompactRegionSnafu {
584                    region_id: self.region_id,
585                }),
586            );
587        }
588
589        self.publish_state();
590        std::mem::take(&mut self.pending_ddl_requests)
591    }
592
593    /// Creates an immutable request for background compaction planning.
594    #[allow(clippy::too_many_arguments)]
595    pub(super) fn new_compaction_request(
596        &self,
597        request_sender: Sender<WorkerRequestWithTime>,
598        engine_config: Arc<MitoConfig>,
599        cache_manager: CacheManagerRef,
600        manifest_ctx: &ManifestContextRef,
601        listener: WorkerListener,
602        schema_metadata_manager: SchemaMetadataManagerRef,
603        max_parallelism: usize,
604        pick_history: Option<PickHistory>,
605    ) -> CompactionRequest {
606        let current_version = CompactionVersion::from(self.version_control.current().version);
607        let start_time = Instant::now();
608
609        CompactionRequest {
610            pick_history,
611            engine_config,
612            current_version,
613            access_layer: self.access_layer.clone(),
614            request_sender: request_sender.clone(),
615            start_time,
616            cache_manager,
617            manifest_ctx: manifest_ctx.clone(),
618            listener,
619            schema_metadata_manager,
620            max_parallelism,
621        }
622    }
623}
624
625/// A manual compaction request waiting for an automatic compaction to finish.
626pub(super) struct PendingCompaction {
627    /// Compaction options.
628    pub(crate) options: compact_request::Options,
629    /// Waiters of pending requests.
630    pub(crate) waiter: OptionOutputTx,
631    /// Max parallelism for pending compaction.
632    pub(crate) max_parallelism: usize,
633    /// Optional time range that constrains candidate compaction windows.
634    pub(crate) time_range: Option<TimestampRange>,
635}
636
637/// Current cycles of one loaded region, without completed-task history.
638#[derive(Debug, Default)]
639pub(crate) struct CompactionState {
640    active: Option<Arc<CompactionTaskState>>,
641    pending: Option<Arc<CompactionTaskState>>,
642}
643
644/// Shared current cycles of a loaded region.
645pub(crate) type CompactionStateRef = Arc<RwLock<CompactionState>>;
646
647impl CompactionState {
648    /// Clones current cycle handles under a short read lock.
649    pub(crate) fn current(
650        state: &CompactionStateRef,
651    ) -> (
652        Option<Arc<CompactionTaskState>>,
653        Option<Arc<CompactionTaskState>>,
654    ) {
655        let state = state.read().unwrap();
656        (state.active.clone(), state.pending.clone())
657    }
658}
659
660/// Per-cycle request metadata and live execution handles, without retaining SSTs.
661#[derive(Debug, Clone)]
662pub(crate) struct CompactionTaskState {
663    pub(crate) entry: CompactionTaskEntry,
664    pub(crate) submitted: Instant,
665    local_state: Option<LocalCompactionState>,
666}
667
668impl CompactionTaskState {
669    /// Captures request parameters and assigns the identity of a new compaction cycle.
670    pub(crate) fn new(
671        region_id: RegionId,
672        options: &compact_request::Options,
673        max_parallelism: usize,
674        time_range: Option<TimestampRange>,
675        trigger: CompactionTrigger,
676    ) -> Self {
677        let (kind, window) = match options {
678            compact_request::Options::Regular(_) => ("regular", None),
679            compact_request::Options::StrictWindow(window) => {
680                ("strict_window", Some(window.window_seconds))
681            }
682        };
683        let request_options = serde_json::json!({
684            "type": kind,
685            "window_seconds": window,
686            "max_parallelism": max_parallelism,
687            "time_range": time_range,
688        })
689        .to_string();
690        Self {
691            submitted: Instant::now(),
692            local_state: None,
693            entry: CompactionTaskEntry {
694                task_id: uuid::Uuid::new_v4().to_string().into(),
695                region_id,
696                table_id: region_id.table_id(),
697                node_id: None,
698                trigger_type: match trigger {
699                    CompactionTrigger::Manual => compaction_task::TRIGGER_TYPE_MANUAL,
700                    CompactionTrigger::Automatic => compaction_task::TRIGGER_TYPE_AUTOMATIC,
701                }
702                .into(),
703                state: compaction_task::STATE_QUEUED.into(),
704                phase: compaction_task::PHASE_WAITING.into(),
705                submitted_at: common_time::util::current_time_millis(),
706                elapsed_millis: 0,
707                input_files: None,
708                input_bytes: None,
709                input_time_start: None,
710                input_time_end: None,
711                output_time_start: None,
712                output_time_end: None,
713                total_units: None,
714                completed_units: None,
715                request_options: request_options.into(),
716                effective_options: None,
717            },
718        }
719    }
720
721    /// Captures the resolved parameters and work selected by an accepted plan.
722    pub(crate) fn accept_plan(&mut self, region: &CompactionRegion, output: &PickerOutput) {
723        let max_parallelism = region.max_parallelism.max(1);
724        let max_picker_outputs = region
725            .engine_config
726            .max_background_compactions
727            .min(max_parallelism);
728        self.entry.effective_options = Some(
729            serde_json::json!({
730                "compaction": region.region_options.compaction,
731                "ttl": region.ttl,
732                "merge_mode": region.region_options.merge_mode(),
733                "max_picker_outputs": max_picker_outputs,
734                "max_parallelism": max_parallelism,
735                "time_window_seconds": output.time_window_size,
736                "max_output_file_size": output.max_file_size,
737            })
738            .to_string()
739            .into(),
740        );
741        self.set_plan_stats(
742            output
743                .outputs
744                .iter()
745                .flat_map(|output| &output.inputs)
746                .chain(&output.expired_ssts),
747            output.outputs.iter(),
748        );
749    }
750
751    /// Records the cycle's accepted inputs and planned output envelope, without retaining SSTs.
752    /// Called again after local memory admission to exclude rejected units. These aggregates
753    /// remain stable as individual units finish; inputs include expired SSTs and output ends
754    /// retain half-open semantics. If any output lacks a planned range, both output bounds
755    /// are None (SQL NULL). An unbounded side propagates independently across outputs.
756    pub(crate) fn set_plan_stats<'a>(
757        &mut self,
758        inputs: impl Iterator<Item = &'a FileHandle>,
759        outputs: impl Iterator<Item = &'a CompactionOutput>,
760    ) {
761        let mut seen = HashSet::new();
762        let mut bytes = 0;
763        let mut input_start = None;
764        let mut input_end = None;
765        for file in inputs {
766            if !seen.insert(file.file_id()) {
767                continue;
768            }
769            bytes += file.meta_ref().file_size;
770            let (start, end) = file.time_range();
771            input_start =
772                Some(input_start.map_or(start, |previous| std::cmp::min(previous, start)));
773            input_end = Some(input_end.map_or(end, |previous| std::cmp::max(previous, end)));
774        }
775        self.entry.input_files = Some(seen.len() as u64);
776        self.entry.input_bytes = Some(bytes);
777        self.entry.input_time_start = input_start;
778        self.entry.input_time_end = input_end;
779
780        let envelope = outputs
781            .map(|output| {
782                output
783                    .output_time_range
784                    .map(|range| (*range.start(), *range.end()))
785                    .unwrap_or((None, None))
786            })
787            .reduce(|(start, end), (next_start, next_end)| {
788                (
789                    start.zip(next_start).map(|(a, b)| a.min(b)),
790                    end.zip(next_end).map(|(a, b)| a.max(b)),
791                )
792            });
793        let (start, end) = envelope.unwrap_or((None, None));
794        self.entry.output_time_start = start;
795        self.entry.output_time_end = end;
796    }
797
798    /// Copies cycle metadata with the reporting node and current monotonic elapsed time.
799    pub(crate) fn snapshot(&self, node_id: Option<u64>) -> CompactionTaskEntry {
800        let mut entry = self.entry.clone();
801        if let Some(local) = &self.local_state {
802            local.inspect(&mut entry);
803        }
804        entry.node_id = node_id;
805        entry.elapsed_millis = self.submitted.elapsed().as_millis().min(u64::MAX as u128) as u64;
806        entry
807    }
808}
809
810#[cfg(test)]
811mod tests {
812    use common_time::Timestamp;
813    use store_api::storage::FileId;
814
815    use super::*;
816    use crate::compaction::test_util::new_file_handle;
817    use crate::test_util::new_noop_file_purger;
818
819    fn plan_time_range(start: Option<Timestamp>, end: Option<Timestamp>) -> TimestampRange {
820        TimestampRange::from((
821            start.map_or(std::ops::Bound::Unbounded, std::ops::Bound::Included),
822            end.map_or(std::ops::Bound::Unbounded, std::ops::Bound::Excluded),
823        ))
824    }
825
826    fn plan_stats_fixture() -> (CompactionTaskState, [FileHandle; 3], [CompactionOutput; 2]) {
827        let file = |start, end, size| {
828            let original = new_file_handle(FileId::random(), 0, 0, 0);
829            let mut meta = original.meta_ref().clone();
830            meta.time_range = (start, end);
831            meta.file_size = size;
832            FileHandle::new(meta, new_noop_file_purger())
833        };
834        // Compare timestamps in their original units: numeric values alone have the wrong order.
835        let early = file(
836            Timestamp::new_second(1),
837            Timestamp::new_millisecond(1500),
838            10,
839        );
840        let late = file(
841            Timestamp::new_microsecond(3_000_001),
842            Timestamp::new_second(4),
843            20,
844        );
845        let expired = file(Timestamp::new_second(-2), Timestamp::new_second(-1), 30);
846        let output = |inputs, start, end| CompactionOutput {
847            output_level: 1,
848            inputs,
849            filter_deleted: false,
850            output_time_range: Some(plan_time_range(start, end)),
851        };
852        let outputs = [
853            output(
854                vec![early.clone()],
855                Some(Timestamp::new_second(0)),
856                Some(Timestamp::new_second(2)),
857            ),
858            output(
859                vec![late.clone()],
860                Some(Timestamp::new_second(3)),
861                Some(Timestamp::new_second(5)),
862            ),
863        ];
864        let info = CompactionTaskState::new(
865            RegionId::new(1, 1),
866            &compact_request::Options::Regular(Default::default()),
867            1,
868            None,
869            CompactionTrigger::Manual,
870        );
871        (info, [early, late, expired], outputs)
872    }
873
874    type PlanStats = (
875        Option<u64>,
876        Option<u64>,
877        Option<Timestamp>,
878        Option<Timestamp>,
879        Option<Timestamp>,
880        Option<Timestamp>,
881    );
882
883    fn plan_stats(info: &CompactionTaskState) -> PlanStats {
884        let entry = &info.entry;
885        (
886            entry.input_files,
887            entry.input_bytes,
888            entry.input_time_start,
889            entry.input_time_end,
890            entry.output_time_start,
891            entry.output_time_end,
892        )
893    }
894
895    #[test]
896    fn test_compaction_plan_stats_deduplication_and_admission() {
897        let (mut info, [early, late, expired], outputs) = plan_stats_fixture();
898        // Repeated inputs across logical windows count once; expired SSTs extend only inputs.
899        info.set_plan_stats(
900            [&early, &late, &early, &expired].into_iter(),
901            outputs.iter(),
902        );
903        assert_eq!(
904            plan_stats(&info),
905            (
906                Some(3),
907                Some(60),
908                Some(Timestamp::new_second(-2)),
909                Some(Timestamp::new_second(4)),
910                Some(Timestamp::new_second(0)),
911                Some(Timestamp::new_second(5))
912            )
913        );
914        // Recompute after admission excludes a unit, including its output window.
915        info.set_plan_stats([&early, &expired].into_iter(), outputs[..1].iter());
916        assert_eq!(
917            plan_stats(&info),
918            (
919                Some(2),
920                Some(40),
921                Some(Timestamp::new_second(-2)),
922                Some(Timestamp::new_millisecond(1500)),
923                Some(Timestamp::new_second(0)),
924                Some(Timestamp::new_second(2))
925            )
926        );
927    }
928
929    #[test]
930    fn test_compaction_plan_stats_expiration_and_empty_reset() {
931        let (mut info, [early, late, expired], outputs) = plan_stats_fixture();
932        info.set_plan_stats([&early, &late, &expired].into_iter(), outputs.iter());
933        // Expiration-only cycles have input coverage but no output envelope.
934        info.set_plan_stats([&expired].into_iter(), std::iter::empty());
935        assert_eq!(
936            plan_stats(&info),
937            (
938                Some(1),
939                Some(30),
940                Some(Timestamp::new_second(-2)),
941                Some(Timestamp::new_second(-1)),
942                None,
943                None
944            )
945        );
946        info.set_plan_stats(std::iter::empty(), std::iter::empty());
947        assert_eq!(
948            plan_stats(&info),
949            (Some(0), Some(0), None, None, None, None)
950        );
951    }
952
953    #[test]
954    fn test_compaction_plan_stats_missing_and_unbounded_ranges() {
955        let (mut info, [early, late, _], outputs) = plan_stats_fixture();
956        // An unbounded side of any output remains unbounded in the cycle envelope.
957        for (range, expected_start, expected_end) in [
958            (None, None, None),
959            (
960                Some(plan_time_range(None, Some(Timestamp::new_second(5)))),
961                None,
962                Some(Timestamp::new_second(5)),
963            ),
964            (
965                Some(plan_time_range(Some(Timestamp::new_second(3)), None)),
966                Some(Timestamp::new_second(0)),
967                None,
968            ),
969        ] {
970            let mut mixed = outputs.clone();
971            mixed[1].output_time_range = range;
972            // Order must not affect propagation of unbounded ranges.
973            for reverse in [false, true] {
974                if reverse {
975                    mixed.reverse();
976                }
977                info.set_plan_stats([&early, &late].into_iter(), mixed.iter());
978                assert_eq!(info.entry.output_time_start, expected_start);
979                assert_eq!(info.entry.output_time_end, expected_end);
980            }
981        }
982    }
983
984    #[test]
985    fn test_compaction_unit_input_leases_release_independently() {
986        let x = new_file_handle(FileId::random(), 0, 100, 0);
987        let y = new_file_handle(FileId::random(), 100, 200, 0);
988        let unused = new_file_handle(FileId::random(), 200, 300, 0);
989        let reserved = CompactingFiles::try_reserve([&x, &y, &x, &unused]).unwrap();
990        let first = reserved.for_inputs(&[x.clone(), x.clone()]);
991        let second = reserved.for_inputs(std::slice::from_ref(&y));
992        drop(reserved);
993        assert!(x.compacting() && y.compacting());
994        assert!(
995            !unused.compacting(),
996            "unit leases must not retain unrelated inputs"
997        );
998        drop(first);
999        assert!(!x.compacting() && y.compacting());
1000        let notification = second.clone();
1001        drop(second);
1002        assert!(y.compacting());
1003        drop(notification);
1004        assert!(!y.compacting());
1005    }
1006}
mito2/compaction/scheduler/state.rs

mito2/compaction/scheduler/
state.rs
