aria_primitives/
cancel.rs

1//! Cooperative cancellation and operation deadlines.
2
3use std::sync::Arc;
4use std::sync::atomic::{AtomicBool, Ordering};
5use std::time::{Duration, Instant};
6
7/// A time budget, fixed when the operation is admitted.
8#[derive(Debug, Clone, Copy)]
9pub struct Deadline {
10    timeout: Duration,
11    expires_at: Instant,
12}
13
14impl Deadline {
15    pub fn after(timeout: Duration) -> Self {
16        Self {
17            timeout,
18            expires_at: Instant::now() + timeout,
19        }
20    }
21
22    pub fn expired(&self) -> bool {
23        Instant::now() >= self.expires_at
24    }
25
26    /// The budget as granted, for error reporting.
27    pub fn timeout(&self) -> Duration {
28        self.timeout
29    }
30}
31
32/// A cancellation signal polled by long-running work.
33#[derive(Clone)]
34pub struct CancelToken(Inner);
35
36#[derive(Clone)]
37enum Inner {
38    Never,
39    Flag(Arc<AtomicBool>),
40    Fn(Arc<dyn Fn() -> bool + Send + Sync>),
41}
42
43impl CancelToken {
44    pub fn new() -> Self {
45        Self(Inner::Flag(Arc::new(AtomicBool::new(false))))
46    }
47
48    /// A token that never cancels, for callers without an abort path.
49    pub const fn never() -> Self {
50        Self(Inner::Never)
51    }
52
53    /// Reads cancellation from an existing condition.
54    pub fn from_fn(should_cancel: impl Fn() -> bool + Send + Sync + 'static) -> Self {
55        Self(Inner::Fn(Arc::new(should_cancel)))
56    }
57
58    /// Cancels every clone of this token. A `never()` token ignores it, and a
59    /// `from_fn` token answers from its closure.
60    pub fn cancel(&self) {
61        if let Inner::Flag(flag) = &self.0 {
62            flag.store(true, Ordering::Release);
63        }
64    }
65
66    pub fn is_cancelled(&self) -> bool {
67        match &self.0 {
68            Inner::Never => false,
69            Inner::Flag(flag) => flag.load(Ordering::Acquire),
70            Inner::Fn(should_cancel) => should_cancel(),
71        }
72    }
73}
74
75impl Default for CancelToken {
76    fn default() -> Self {
77        Self::new()
78    }
79}
80
81impl std::fmt::Debug for CancelToken {
82    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83        match &self.0 {
84            Inner::Never => formatter.write_str("CancelToken::Never"),
85            Inner::Flag(flag) => formatter
86                .debug_tuple("CancelToken::Flag")
87                .field(&flag.load(Ordering::Acquire))
88                .finish(),
89            Inner::Fn(_) => formatter.write_str("CancelToken::Fn"),
90        }
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97
98    #[test]
99    fn a_fresh_token_is_not_cancelled() {
100        assert!(!CancelToken::new().is_cancelled());
101        assert!(!CancelToken::never().is_cancelled());
102    }
103
104    #[test]
105    fn cancellation_reaches_every_clone() {
106        let token = CancelToken::new();
107        let observer = token.clone();
108
109        token.cancel();
110
111        assert!(observer.is_cancelled());
112    }
113
114    #[test]
115    fn never_ignores_cancel() {
116        let token = CancelToken::never();
117
118        token.cancel();
119
120        assert!(!token.is_cancelled());
121    }
122
123    #[test]
124    fn a_deadline_expires_once_its_budget_is_spent() {
125        let deadline = Deadline::after(Duration::from_millis(10));
126
127        assert!(!deadline.expired());
128        assert_eq!(deadline.timeout(), Duration::from_millis(10));
129        std::thread::sleep(Duration::from_millis(15));
130        assert!(deadline.expired());
131    }
132
133    #[test]
134    fn from_fn_answers_from_the_closure() {
135        let flag = Arc::new(AtomicBool::new(false));
136        let observed = Arc::clone(&flag);
137        let token = CancelToken::from_fn(move || observed.load(Ordering::Acquire));
138
139        assert!(!token.is_cancelled());
140        flag.store(true, Ordering::Release);
141        assert!(token.is_cancelled());
142    }
143}
aria_primitives/cancel.rs

aria_primitives/
cancel.rs
