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    /// Starts the non-cancellable commit phase.
46    ///
47    /// Returns false if cancellation was requested first.
48    pub(crate) fn mark_commit_started(&self) -> bool {
49        let mut commit_started = self.commit_started.lock().unwrap();
50        if self.cancel_handle.is_cancelled() {
51            return false;
52        }
53        *commit_started = true;
54        true
55    }
56
57    pub(crate) fn request_cancel(&self) -> RequestCancelResult {
58        // Hold the commit lock while cancelling to serialize cancellation with commit startup.
59        let commit_started = self.commit_started.lock().unwrap();
60        if *commit_started {
61            return RequestCancelResult::TooLateToCancel;
62        }
63        if self.cancel_handle.is_cancelled() {
64            return RequestCancelResult::AlreadyCancelling;
65        }
66
67        self.cancel_handle.cancel();
68        RequestCancelResult::CancelIssued
69    }
70}
71
72#[derive(Debug, Clone, Copy, PartialEq, Eq)]
73pub(crate) enum RequestCancelResult {
74    CancelIssued,
75    AlreadyCancelling,
76    TooLateToCancel,
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    #[test]
84    fn test_cancellable_task_state_transitions() {
85        let state = CancellableTaskState::new();
86        assert_eq!(RequestCancelResult::CancelIssued, state.request_cancel());
87        assert!(state.is_cancelled());
88        assert_eq!(
89            RequestCancelResult::AlreadyCancelling,
90            state.request_cancel()
91        );
92        assert!(!state.mark_commit_started());
93
94        let state = CancellableTaskState::new();
95        assert!(state.mark_commit_started());
96        assert_eq!(RequestCancelResult::TooLateToCancel, state.request_cancel());
97    }
98}
mito2/schedule.rs

mito2/
schedule.rs
