Skip to main content

mito2/compaction/scheduler/
history.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::cmp::Reverse;
16use std::collections::{HashMap, HashSet};
17use std::time::{Duration, Instant};
18
19use prometheus::IntCounterVec;
20
21use crate::compaction::twcs::{PickClass, WindowCandidate};
22use crate::metrics::TWCS_SELECTED_CANDIDATES;
23
24/// Waiting compensation starts after one minute. This is a promotion threshold,
25/// not a deadline: an executing batch is never preempted to run an overdue pick.
26/// It is checked at pick time and does not wake the scheduler on its own.
27const PROMOTION_DELAY: Duration = Duration::from_secs(60);
28
29/// A work identity survives file replacement and picking epochs, but not a region
30/// incarnation. Different classes in the same window retain independent ages.
31#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
32pub(crate) struct PickKey {
33    /// Window identity, independent of the current input files.
34    pub(crate) window: i64,
35    /// Work category whose waiting age is tracked independently in this window.
36    pub(crate) class: PickClass,
37}
38
39impl From<&WindowCandidate> for PickKey {
40    fn from(candidate: &WindowCandidate) -> Self {
41        Self {
42            window: candidate.window,
43            class: candidate.pick.class,
44        }
45    }
46}
47
48/// Region-resident scheduling metadata. It owns no file handles or reservations.
49/// Planning works on a private snapshot; the accepted plan consumes one fairness
50/// turn per selected key. Execution success or failure does not rewrite ages.
51#[derive(Debug, Clone)]
52pub(crate) struct PickHistory {
53    /// First observation or last fairness turn for each eligible key.
54    /// Repeated picks preserve this timestamp until the key is selected.
55    waiting: HashMap<PickKey, Instant>,
56    selected_metrics: IntCounterVec,
57}
58
59impl Default for PickHistory {
60    fn default() -> Self {
61        Self {
62            waiting: HashMap::new(),
63            selected_metrics: TWCS_SELECTED_CANDIDATES.clone(),
64        }
65    }
66}
67
68/// Ascending fairness order: promoted work, oldest age, class, then newest window.
69#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
70struct PickPriority {
71    not_promoted: bool,
72    since: Option<Instant>,
73    class: PickClass,
74    window: Reverse<i64>,
75}
76
77impl PickHistory {
78    /// Refreshes all eligible identities before selection limits hide any work.
79    /// Preserves surviving ages and drops absent keys without changing selection.
80    fn observe(&mut self, keys: impl Iterator<Item = PickKey>, now: Instant) {
81        let live: HashSet<_> = keys.collect();
82        self.waiting.retain(|key, _| live.contains(key));
83        for key in live {
84            self.waiting.entry(key).or_insert(now);
85        }
86    }
87
88    /// Returns an ascending sort key: promoted work first, oldest waiting first.
89    /// Below [`PROMOTION_DELAY`], age does not affect class/newest-window ordering.
90    /// One shared `now` keeps the comparator consistent across the entire sort.
91    fn priority(&self, key: PickKey, now: Instant) -> PickPriority {
92        let overdue = self
93            .waiting
94            .get(&key)
95            .copied()
96            .filter(|since| now.saturating_duration_since(*since) >= PROMOTION_DELAY);
97        PickPriority {
98            not_promoted: overdue.is_none(),
99            since: overdue,
100            class: key.class,
101            window: Reverse(key.window),
102        }
103    }
104
105    /// Consumes the fairness turn for `key`. Only an accepted plan installs the
106    /// private snapshot, so rejected or stale plans never advance resident ages.
107    fn serve(&mut self, key: PickKey, now: Instant) {
108        self.waiting.insert(key, now);
109    }
110
111    /// Observes all eligible work and ranks one alternative per window without
112    /// consuming a fairness turn. The caller acknowledges only attempted work
113    /// after applying its output budget and semantic closure policy.
114    pub(crate) fn rank_candidates(
115        &mut self,
116        mut candidates: Vec<WindowCandidate>,
117        now: Instant,
118    ) -> Vec<WindowCandidate> {
119        self.observe(candidates.iter().map(PickKey::from), now);
120        candidates.sort_by_key(|candidate| self.priority(PickKey::from(candidate), now));
121        // Alternatives may share files or semantic dependencies within a window.
122        let mut windows = HashSet::new();
123        candidates.retain(|candidate| windows.insert(candidate.window));
124        candidates
125    }
126
127    /// Consumes turns for attempted seeds, including absorbed or rejected closure
128    /// seeds, but not work left behind by the output budget. Only an accepted
129    /// plan installs this private snapshot; execution never updates waiting ages.
130    pub(crate) fn acknowledge(&mut self, keys: impl Iterator<Item = PickKey>, now: Instant) {
131        for key in keys {
132            let priority = if self.priority(key, now).since.is_some() {
133                "promoted"
134            } else {
135                "normal"
136            };
137            self.selected_metrics
138                .with_label_values(&[key.class.as_str(), priority])
139                .inc();
140            self.serve(key, now);
141        }
142    }
143}
144
145#[cfg(test)]
146impl PickHistory {
147    /// Uses isolated metric handles while retaining the production acknowledgement path.
148    pub(crate) fn with_selected_metrics_for_test(selected_metrics: IntCounterVec) -> Self {
149        Self {
150            waiting: HashMap::new(),
151            selected_metrics,
152        }
153    }
154
155    /// Seeds waiting ages for cross-module scheduling tests.
156    pub(crate) fn observe_for_test(&mut self, keys: impl Iterator<Item = PickKey>, now: Instant) {
157        self.observe(keys, now);
158    }
159
160    /// Seeds one waiting age without pruning unrelated keys.
161    pub(crate) fn seed_for_test(&mut self, key: PickKey, now: Instant) {
162        self.serve(key, now);
163    }
164
165    /// Inspects promoted ages without exposing the private priority representation.
166    pub(crate) fn promoted_since_for_test(&self, key: PickKey, now: Instant) -> Option<Instant> {
167        self.priority(key, now).since
168    }
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174
175    fn key(window: i64, class: PickClass) -> PickKey {
176        PickKey { window, class }
177    }
178
179    #[test]
180    fn test_aging_survives_repicks_and_other_service() {
181        let start = Instant::now();
182        let l0 = key(100, PickClass::L0Merge);
183        let old = key(0, PickClass::FileReduction);
184        let overlap = key(0, PickClass::OverlapResolution);
185        let mut history = PickHistory::default();
186        history.observe([l0, old, overlap].into_iter(), start);
187        assert!(history.priority(l0, start) < history.priority(old, start));
188        for seconds in 1..=60 {
189            let now = start + Duration::from_secs(seconds);
190            history.observe([l0, old, overlap].into_iter(), now);
191            history.serve(l0, now);
192        }
193        let now = start + PROMOTION_DELAY;
194        assert!(history.priority(old, now) < history.priority(l0, now));
195        history.serve(old, now);
196        assert!(history.priority(overlap, now) < history.priority(old, now));
197    }
198
199    #[test]
200    fn test_old_windows_progress_and_obsolete_work_is_pruned() {
201        let start = Instant::now();
202        let old = key(0, PickClass::L0Merge);
203        let new = key(1, PickClass::L0Merge);
204        let mut history = PickHistory::default();
205        history.observe([old].into_iter(), start);
206        let now = start + PROMOTION_DELAY;
207        history.observe([old, new].into_iter(), now);
208        assert!(history.priority(old, now) < history.priority(new, now));
209        history.observe([new].into_iter(), now);
210        assert_eq!(1, history.waiting.len());
211        assert!(!history.waiting.contains_key(&old));
212        history.observe(std::iter::empty(), now);
213        assert!(history.waiting.is_empty());
214    }
215}