ferroui_base/threading/
cancellation_token.rs

1//! A minimal cooperative cancellation primitive.
2//!
3//! The dispatcher API takes cancellation tokens in the same places the
4//! reference implementation does (run loops, queued operations, frames).
5//! Tokens are cheap handles that can be shared across threads.
6
7use std::sync::atomic::{AtomicBool, Ordering};
8use std::sync::{Arc, Mutex, PoisonError};
9
10type Callback = Box<dyn FnOnce() + Send>;
11
12struct Inner {
13    canceled: AtomicBool,
14    callbacks: Mutex<Callbacks>,
15}
16
17#[derive(Default)]
18struct Callbacks {
19    next_id: u64,
20    entries: Vec<(u64, Callback)>,
21}
22
23/// Signals to one or more [`CancellationToken`]s that they should be
24/// canceled.
25#[derive(Clone)]
26pub struct CancellationTokenSource {
27    inner: Arc<Inner>,
28}
29
30impl Default for CancellationTokenSource {
31    fn default() -> Self {
32        Self::new()
33    }
34}
35
36impl CancellationTokenSource {
37    pub fn new() -> Self {
38        Self {
39            inner: Arc::new(Inner { canceled: AtomicBool::new(false), callbacks: Mutex::new(Callbacks::default()) }),
40        }
41    }
42
43    /// The token observing this source.
44    pub fn token(&self) -> CancellationToken {
45        CancellationToken { inner: Some(self.inner.clone()) }
46    }
47
48    pub fn is_cancellation_requested(&self) -> bool {
49        self.inner.canceled.load(Ordering::SeqCst)
50    }
51
52    /// Requests cancellation and runs the registered callbacks on the
53    /// calling thread. Subsequent calls do nothing.
54    pub fn cancel(&self) {
55        if self.inner.canceled.swap(true, Ordering::SeqCst) {
56            return;
57        }
58        let entries = {
59            let mut callbacks = self.inner.callbacks.lock().unwrap_or_else(PoisonError::into_inner);
60            std::mem::take(&mut callbacks.entries)
61        };
62        for (_, callback) in entries {
63            callback();
64        }
65    }
66}
67
68/// Propagates notification that operations should be canceled.
69#[derive(Clone, Default)]
70pub struct CancellationToken {
71    inner: Option<Arc<Inner>>,
72}
73
74impl CancellationToken {
75    /// A token that can never be canceled.
76    pub const fn none() -> Self {
77        Self { inner: None }
78    }
79
80    /// Whether this token is capable of being in the canceled state.
81    pub fn can_be_canceled(&self) -> bool {
82        self.inner.is_some()
83    }
84
85    pub fn is_cancellation_requested(&self) -> bool {
86        self.inner.as_ref().is_some_and(|inner| inner.canceled.load(Ordering::SeqCst))
87    }
88
89    /// Registers a callback that runs when the token is canceled. When the
90    /// token is already canceled the callback runs immediately.
91    pub fn register(&self, callback: impl FnOnce() + Send + 'static) -> CancellationTokenRegistration {
92        let Some(inner) = &self.inner else {
93            return CancellationTokenRegistration { inner: None, id: 0 };
94        };
95        {
96            let mut callbacks = inner.callbacks.lock().unwrap_or_else(PoisonError::into_inner);
97            // Checked under the lock so that a concurrent `cancel` either
98            // sees the new entry or has already flipped the flag.
99            if !inner.canceled.load(Ordering::SeqCst) {
100                callbacks.next_id += 1;
101                let id = callbacks.next_id;
102                callbacks.entries.push((id, Box::new(callback)));
103                return CancellationTokenRegistration { inner: Some(inner.clone()), id };
104            }
105        }
106        callback();
107        CancellationTokenRegistration { inner: None, id: 0 }
108    }
109}
110
111/// A callback registered with a [`CancellationToken`].
112pub struct CancellationTokenRegistration {
113    inner: Option<Arc<Inner>>,
114    id: u64,
115}
116
117impl CancellationTokenRegistration {
118    /// Removes the callback so that it no longer runs on cancellation.
119    pub fn dispose(&self) {
120        if let Some(inner) = &self.inner {
121            let removed = {
122                let mut callbacks = inner.callbacks.lock().unwrap_or_else(PoisonError::into_inner);
123                callbacks.entries.iter().position(|(id, _)| *id == self.id).map(|i| callbacks.entries.remove(i))
124            };
125            drop(removed);
126        }
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133    use std::sync::atomic::AtomicUsize;
134
135    #[test]
136    fn none_cannot_be_canceled() {
137        let token = CancellationToken::none();
138        assert!(!token.can_be_canceled());
139        assert!(!token.is_cancellation_requested());
140    }
141
142    #[test]
143    fn cancel_runs_callbacks_once() {
144        let count = Arc::new(AtomicUsize::new(0));
145        let source = CancellationTokenSource::new();
146        let token = source.token();
147        let c = count.clone();
148        token.register(move || {
149            c.fetch_add(1, Ordering::SeqCst);
150        });
151        let c = count.clone();
152        let removed = token.register(move || {
153            c.fetch_add(10, Ordering::SeqCst);
154        });
155        removed.dispose();
156        source.cancel();
157        source.cancel();
158        assert!(token.is_cancellation_requested());
159        assert_eq!(count.load(Ordering::SeqCst), 1);
160
161        let c = count.clone();
162        token.register(move || {
163            c.fetch_add(100, Ordering::SeqCst);
164        });
165        assert_eq!(count.load(Ordering::SeqCst), 101);
166    }
167}
ferroui_base/threading/cancellation_token.rs

ferroui_base/threading/
cancellation_token.rs
