mito2/
schedule.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::sync::atomic::{AtomicBool, Ordering};
16use std::sync::{Arc, Mutex};
17
18use common_base::cancellation::CancellationHandle;
19
20pub mod remote_job_scheduler;
21pub mod scheduler;
22
23/// Shared state for cooperatively cancelling a task until it starts committing.
24#[derive(Debug, Clone)]
25pub(crate) struct CancellableTaskState {
26    cancel_handle: Arc<CancellationHandle>,
27    commit_started: Arc<Mutex<bool>>,
28    started: Arc<AtomicBool>,
29}
30
31impl CancellableTaskState {
32    pub(crate) fn new() -> Self {
33        Self {
34            cancel_handle: Arc::new(CancellationHandle::default()),
35            commit_started: Arc::new(Mutex::new(false)),
36            started: Arc::new(AtomicBool::new(false)),
37        }
38    }
39
40    /// Records that the executor has started this task.
41    pub(crate) fn mark_started(&self) {
42        self.started.store(true, Ordering::Relaxed);
43    }
44
45    /// Returns whether the executor has started this task.
46    pub(crate) fn started(&self) -> bool {
47        self.started.load(Ordering::Relaxed)
48    }
49
50    pub(crate) fn cancel_handle(&self) -> Arc<CancellationHandle> {
51        self.cancel_handle.clone()
52    }
53
54    pub(crate) fn is_cancelled(&self) -> bool {
55        self.cancel_handle.is_cancelled()
56    }
57
58    /// Returns whether the task has crossed the non-cancellable publication boundary.
59    pub(crate) fn commit_started(&self) -> bool {
60        *self.commit_started.lock().unwrap()
61    }
62
63    /// Starts the non-cancellable commit phase.
64    ///
65    /// Returns false if cancellation was requested first.
66    pub(crate) fn mark_commit_started(&self) -> bool {
67        let mut commit_started = self.commit_started.lock().unwrap();
68        if self.cancel_handle.is_cancelled() {
69            return false;
70        }
71        *commit_started = true;
72        true
73    }
74
75    pub(crate) fn request_cancel(&self) -> RequestCancelResult {
76        // Hold the commit lock while cancelling to serialize cancellation with commit startup.
77        let commit_started = self.commit_started.lock().unwrap();
78        if *commit_started {
79            return RequestCancelResult::TooLateToCancel;
80        }
81        if self.cancel_handle.is_cancelled() {
82            return RequestCancelResult::AlreadyCancelling;
83        }
84
85        self.cancel_handle.cancel();
86        RequestCancelResult::CancelIssued
87    }
88}
89
90#[derive(Debug, Clone, Copy, PartialEq, Eq)]
91pub(crate) enum RequestCancelResult {
92    CancelIssued,
93    AlreadyCancelling,
94    TooLateToCancel,
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[test]
102    fn test_cancellable_task_state_transitions() {
103        let state = CancellableTaskState::new();
104        assert_eq!(RequestCancelResult::CancelIssued, state.request_cancel());
105        assert!(state.is_cancelled());
106        assert_eq!(
107            RequestCancelResult::AlreadyCancelling,
108            state.request_cancel()
109        );
110        assert!(!state.mark_commit_started());
111
112        let state = CancellableTaskState::new();
113        assert!(state.mark_commit_started());
114        assert_eq!(RequestCancelResult::TooLateToCancel, state.request_cancel());
115    }
116}
mito2/schedule.rs

mito2/
schedule.rs
