vv_agent/runtime/
cancellation.rs

1use std::sync::atomic::{AtomicBool, Ordering};
2use std::sync::{Arc, Mutex};
3
4#[derive(Debug, Clone, PartialEq, Eq)]
5pub struct CancelledError {
6    message: String,
7}
8
9impl CancelledError {
10    pub fn new(message: impl Into<String>) -> Self {
11        Self {
12            message: message.into(),
13        }
14    }
15
16    pub fn message(&self) -> &str {
17        &self.message
18    }
19}
20
21impl std::fmt::Display for CancelledError {
22    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
23        formatter.write_str(&self.message)
24    }
25}
26
27impl std::error::Error for CancelledError {}
28
29#[derive(Clone, Default)]
30pub struct CancellationToken {
31    inner: Arc<CancellationState>,
32}
33
34#[derive(Default)]
35struct CancellationState {
36    cancelled: AtomicBool,
37    callbacks: Mutex<Vec<Arc<dyn Fn() + Send + Sync + 'static>>>,
38}
39
40impl std::fmt::Debug for CancellationToken {
41    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42        formatter
43            .debug_struct("CancellationToken")
44            .field("cancelled", &self.is_cancelled())
45            .finish()
46    }
47}
48
49impl CancellationToken {
50    pub fn cancel(&self) {
51        if self.inner.cancelled.swap(true, Ordering::SeqCst) {
52            return;
53        }
54        let callbacks = std::mem::take(
55            &mut *self
56                .inner
57                .callbacks
58                .lock()
59                .expect("cancellation callbacks lock"),
60        );
61        for callback in callbacks {
62            callback();
63        }
64    }
65
66    pub fn is_cancelled(&self) -> bool {
67        self.inner.cancelled.load(Ordering::SeqCst)
68    }
69
70    pub fn cancelled(&self) -> bool {
71        self.is_cancelled()
72    }
73
74    pub fn check(&self) -> Result<(), String> {
75        if self.is_cancelled() {
76            Err("Operation was cancelled".to_string())
77        } else {
78            Ok(())
79        }
80    }
81
82    pub fn on_cancel(&self, callback: impl Fn() + Send + Sync + 'static) {
83        let callback: Arc<dyn Fn() + Send + Sync + 'static> = Arc::new(callback);
84        let call_immediately = {
85            let mut callbacks = self
86                .inner
87                .callbacks
88                .lock()
89                .expect("cancellation callbacks lock");
90            if self.is_cancelled() {
91                true
92            } else {
93                callbacks.push(callback.clone());
94                false
95            }
96        };
97        if call_immediately {
98            callback();
99        }
100    }
101
102    pub fn child(&self) -> Self {
103        let child = Self::default();
104        let child_to_cancel = child.clone();
105        self.on_cancel(move || child_to_cancel.cancel());
106        child
107    }
108}
vv_agent/runtime/cancellation.rs

vv_agent/runtime/
cancellation.rs
