mito2/compaction/scheduler/
local.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, VecDeque};
16use std::sync::Arc;
17
18use common_meta::key::SchemaMetadataManagerRef;
19use store_api::storage::RegionId;
20
21use crate::compaction::compactor::{CompactionContext, CompactionRegion};
22use crate::compaction::scheduler::state::{CompactingFiles, CompactionExecution, CompactionPhase};
23use crate::compaction::scheduler::{CompactionScheduler, CompactionTransition};
24use crate::compaction::task::CompactionTask;
25use crate::compaction::unit::CompactionUnit;
26use crate::error::{self, Error};
27use crate::metrics::INFLIGHT_COMPACTION_COUNT;
28use crate::region::ManifestContextRef;
29use crate::request::{
30    BackgroundNotify, CompactionUnitNotification, WorkerRequest, WorkerRequestWithTime,
31};
32use crate::schedule::{CancellableTaskState, RequestCancelResult};
33use crate::sst::file::UncommittedSsts;
34
35/// Tracks one unit through its terminal result and subsequent task exit.
36#[derive(Debug)]
37struct UnitExecution {
38    state: CancellableTaskState,
39    finished: bool,
40}
41
42/// A request's local units. All transitions run on the owning region worker.
43#[derive(Debug)]
44pub(crate) struct LocalCompaction {
45    /// All pending and submitted units, retained until task exit. Tasks own their input leases.
46    /// Keyed by each unit's execution plan ID, allocated from the scheduler's `next_plan_id`;
47    /// it equals `CompactionExecution::plan_id()` and the `plan_id` in unit notifications.
48    units: HashMap<u64, UnitExecution>,
49    pending: VecDeque<CompactionTask>,
50    parallelism: usize,
51    error: Option<Arc<Error>>,
52    made_progress: bool,
53}
54
55impl LocalCompaction {
56    /// Requests cancellation for all units that have not started committing.
57    pub(crate) fn request_cancel(&self) -> RequestCancelResult {
58        let mut result = RequestCancelResult::TooLateToCancel;
59        for unit in self.units.values() {
60            match unit.state.request_cancel() {
61                RequestCancelResult::CancelIssued => result = RequestCancelResult::CancelIssued,
62                RequestCancelResult::AlreadyCancelling
63                    if result != RequestCancelResult::CancelIssued =>
64                {
65                    result = RequestCancelResult::AlreadyCancelling;
66                }
67                _ => {}
68            }
69        }
70        result
71    }
72
73    /// Counts submitted units that have not exited, including queued jobs.
74    fn num_running(&self) -> usize {
75        self.units.len() - self.pending.len()
76    }
77
78    /// Checks this attempt's phase without depending on any sibling's apply status.
79    fn can_apply(&self, id: u64) -> bool {
80        self.units
81            .get(&id)
82            .is_some_and(|unit| !unit.finished && unit.state.commit_started())
83    }
84
85    /// Records a terminal result once, but still awaits task exit.
86    fn finish(&mut self, id: u64, result: std::result::Result<bool, Arc<Error>>) {
87        let Some(unit) = self.units.get_mut(&id) else {
88            return;
89        };
90        if unit.finished {
91            return;
92        }
93        unit.finished = true;
94        match result {
95            Ok(progress) => self.made_progress |= progress,
96            Err(err) => {
97                // A real failure must not be hidden by a sibling's cooperative cancellation.
98                if self.error.as_ref().is_none_or(|current| {
99                    matches!(current.as_ref(), Error::CompactionCancelled { .. })
100                }) {
101                    self.error = Some(err);
102                }
103            }
104        }
105    }
106}
107
108impl Drop for LocalCompaction {
109    /// Stops cancellable work when its region or compaction cycle is removed.
110    fn drop(&mut self) {
111        self.request_cancel();
112    }
113}
114
115impl CompactionScheduler {
116    /// Builds unit tasks with fresh execution IDs and shares their already-reserved inputs.
117    pub(crate) fn prepare_local_units(
118        &mut self,
119        region: CompactionRegion,
120        units: Vec<CompactionUnit>,
121        reserved: CompactingFiles,
122    ) -> CompactionPhase {
123        let mut local = LocalCompaction {
124            units: HashMap::new(),
125            pending: VecDeque::new(),
126            parallelism: region.max_parallelism.max(1),
127            error: None,
128            made_progress: false,
129        };
130        let context = CompactionContext::from(&region);
131        // Planning is complete; unit tasks must not pin the other units' inputs.
132        drop(region);
133        for unit in units {
134            let files = reserved.for_inputs(&unit.inputs);
135            let id = Self::next_plan_id(&mut self.next_plan_id);
136            let execution = CompactionExecution::new(id, files);
137            let state = CancellableTaskState::new();
138            local.units.insert(
139                id,
140                UnitExecution {
141                    state: state.clone(),
142                    finished: false,
143                },
144            );
145            local.pending.push_back(CompactionTask {
146                state,
147                execution,
148                compaction_context: context.clone(),
149                request_sender: self.request_sender.clone(),
150                listener: self.listener.clone(),
151                estimated_memory_bytes: unit.estimated_memory_bytes(),
152                unit,
153                memory_manager: self.memory_manager.clone(),
154                memory_policy: self.memory_policy,
155                uncommitted: UncommittedSsts::new(
156                    context.region_id,
157                    context.access_layer.clone(),
158                    Some(context.cache_manager.clone()),
159                ),
160            });
161        }
162        CompactionPhase::Units(local)
163    }
164
165    /// Returns local unit state only while the region is in a local execution phase.
166    fn local_units_mut(&mut self, region_id: RegionId) -> Option<&mut LocalCompaction> {
167        match &mut self.region_status.get_mut(&region_id)?.active.phase {
168            CompactionPhase::Units(local) => Some(local),
169            _ => None,
170        }
171    }
172
173    /// Submits pending units up to the request limit and arranges task-exit notifications.
174    pub(crate) fn dispatch_local_units(&mut self, region_id: RegionId) {
175        loop {
176            let Some(local) = self.local_units_mut(region_id) else {
177                return;
178            };
179            if local.num_running() >= local.parallelism {
180                return;
181            }
182            let Some(mut task) = local.pending.pop_front() else {
183                return;
184            };
185            let id = task.execution.plan_id();
186            let sender = self.request_sender.clone();
187            let completion = Box::pin(async move {
188                let _ = sender
189                    .send(WorkerRequestWithTime::new(WorkerRequest::Background {
190                        region_id,
191                        notify: BackgroundNotify::CompactionUnit(
192                            CompactionUnitNotification::Released { plan_id: id },
193                        ),
194                    }))
195                    .await;
196            });
197            let job = Box::pin(async move {
198                INFLIGHT_COMPACTION_COUNT.inc();
199                task.run().await;
200                drop(task);
201                INFLIGHT_COMPACTION_COUNT.dec();
202            });
203            if let Err(err) = self.scheduler.schedule_with_completion(job, completion)
204                && let Some(local) = self.local_units_mut(region_id)
205            {
206                local.finish(id, Err(Arc::new(err)));
207                local.units.remove(&id);
208            }
209        }
210    }
211
212    /// Accepts an edit from any current unit that has started committing and is not finished.
213    pub(crate) fn can_apply_unit(&self, region_id: RegionId, id: u64) -> bool {
214        self.region_status.get(&region_id).is_some_and(|status| {
215            matches!(&status.active.phase, CompactionPhase::Units(local)
216                if local.can_apply(id))
217        })
218    }
219
220    /// Records apply progress or failure without ending sibling executions.
221    pub(crate) fn on_unit_finished(
222        &mut self,
223        region_id: RegionId,
224        id: u64,
225        result: std::result::Result<bool, Arc<Error>>,
226    ) {
227        if let Some(local) = self.local_units_mut(region_id) {
228            local.finish(id, result);
229        }
230    }
231
232    /// Accounts for task exit, starts pending units, and checks request completion.
233    pub(crate) async fn on_unit_released(
234        &mut self,
235        region_id: RegionId,
236        id: u64,
237        manifest: &ManifestContextRef,
238        schemas: SchemaMetadataManagerRef,
239    ) -> CompactionTransition {
240        let Some(local) = self.local_units_mut(region_id) else {
241            return CompactionTransition::NoAction;
242        };
243        if !local.units.contains_key(&id) {
244            return CompactionTransition::NoAction;
245        }
246        local.finish(
247            id,
248            Err(Arc::new(
249                error::UnexpectedSnafu {
250                    reason: "Compaction unit exited without a terminal notification",
251                }
252                .build(),
253            )),
254        );
255        local.units.remove(&id);
256        self.dispatch_local_units(region_id);
257        self.finish_local_units(region_id, manifest, schemas).await
258    }
259
260    /// Completes the request only after all units exit, using the existing terminal paths.
261    pub(crate) async fn finish_local_units(
262        &mut self,
263        region_id: RegionId,
264        manifest: &ManifestContextRef,
265        schemas: SchemaMetadataManagerRef,
266    ) -> CompactionTransition {
267        let Some(local) = self.local_units_mut(region_id) else {
268            return CompactionTransition::NoAction;
269        };
270        if !local.units.is_empty() {
271            return CompactionTransition::NoAction;
272        }
273        if let Some(err) = local.error.take() {
274            if matches!(err.as_ref(), Error::CompactionCancelled { .. }) {
275                return CompactionTransition::from_pending_ddls(
276                    self.finish_compaction_on_cancel(region_id),
277                );
278            }
279            self.on_compaction_failed(region_id, err);
280            return CompactionTransition::NoAction;
281        }
282        let progress = local.made_progress;
283        if let Some(status) = self.region_status.get_mut(&region_id) {
284            for waiter in status.take_waiters() {
285                waiter.send(Ok(0));
286            }
287        }
288        self.on_compaction_finished(region_id, manifest, schemas, progress)
289            .await
290    }
291}
292
293#[cfg(test)]
294mod tests {
295    use super::*;
296
297    #[test]
298    fn test_compaction_unit_apply_accepts_independent_executions() {
299        let units = [1, 2]
300            .into_iter()
301            .map(|id| {
302                (
303                    id,
304                    UnitExecution {
305                        state: CancellableTaskState::new(),
306                        finished: false,
307                    },
308                )
309            })
310            .collect();
311        let mut local = LocalCompaction {
312            units,
313            pending: VecDeque::new(),
314            parallelism: 2,
315            error: None,
316            made_progress: false,
317        };
318        assert!(!local.can_apply(1));
319        assert!(!local.can_apply(2));
320        assert!(!local.can_apply(99));
321        for unit in local.units.values() {
322            assert!(unit.state.mark_commit_started());
323        }
324        assert!(local.can_apply(1) && local.can_apply(2));
325        local.finish(2, Ok(true));
326        assert!(local.can_apply(1));
327        assert!(!local.can_apply(2));
328        local.finish(2, Err(Arc::new(error::CompactionCancelledSnafu.build())));
329        local.finish(99, Err(Arc::new(error::CompactionCancelledSnafu.build())));
330        assert!(local.can_apply(1));
331        assert!(local.error.is_none());
332        assert!(local.made_progress);
333        local.finish(1, Ok(false));
334        assert!(!local.can_apply(1));
335        assert_eq!(
336            2,
337            local.units.len(),
338            "terminal units remain tracked until task exit"
339        );
340    }
341}
mito2/compaction/scheduler/local.rs

mito2/compaction/scheduler/
local.rs
