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

mito2/
schedule.rs
