bloq_graph/
cancellation.rs

1//! Request-scoped cooperative cancellation for graph analysis.
2
3use std::cell::RefCell;
4use std::sync::{
5    Arc,
6    atomic::{AtomicBool, Ordering},
7};
8
9/// A computation was cancelled before producing its result.
10/// Cancellation is neither a resource failure nor a semantic diagnosis.
11#[derive(Debug, Clone, Copy, thiserror::Error)]
12#[error("computation cancelled")]
13pub struct ComputationCancelled;
14
15/// Shared cancellation state for one request.
16///
17/// Run synchronous analysis inside [`run`](Self::run). Existing normalization
18/// checkpoints observe this token, and [`crate::map_jobs`] carries it into its
19/// scoped workers. Other threads and later requests keep independent state.
20#[derive(Debug, Clone, Default)]
21pub struct CancellationToken(Arc<AtomicBool>);
22
23thread_local! {
24    static ACTIVE: RefCell<Option<CancellationToken>> = const { RefCell::new(None) };
25}
26
27impl CancellationToken {
28    /// Creates a token for a new request.
29    pub fn new() -> Self {
30        Self::default()
31    }
32
33    /// Requests cancellation. Repeated calls are harmless.
34    pub fn cancel(&self) {
35        self.0.store(true, Ordering::Relaxed);
36    }
37
38    /// Whether this request has been cancelled.
39    pub fn is_cancelled(&self) -> bool {
40        self.0.load(Ordering::Relaxed)
41    }
42
43    /// Checks this request at a safe checkpoint.
44    ///
45    /// # Errors
46    ///
47    /// Returns [`ComputationCancelled`] after [`cancel`](Self::cancel).
48    pub fn check(&self) -> Result<(), ComputationCancelled> {
49        if self.is_cancelled() {
50            Err(ComputationCancelled)
51        } else {
52            Ok(())
53        }
54    }
55
56    /// Runs analysis with this token, checking before entry and before returning
57    /// a result. Work stops at cooperative checkpoints; threads are not killed.
58    ///
59    /// # Errors
60    ///
61    /// Returns cancellation or the computation's own error. No partial result
62    /// is returned after cancellation. The preceding scope is restored on exit.
63    pub fn run<T, E: From<ComputationCancelled>>(
64        &self,
65        f: impl FnOnce() -> Result<T, E>,
66    ) -> Result<T, E> {
67        self.check()?;
68        self.scope(|| {
69            let result = f();
70            self.check()?;
71            result
72        })
73    }
74
75    pub(crate) fn current() -> Option<Self> {
76        ACTIVE.with(|active| active.borrow().clone())
77    }
78
79    pub(crate) fn scope<T>(&self, f: impl FnOnce() -> T) -> T {
80        struct Restore(Option<CancellationToken>);
81        impl Drop for Restore {
82            fn drop(&mut self) {
83                ACTIVE.with(|active| {
84                    active.replace(self.0.take());
85                });
86            }
87        }
88        let previous = ACTIVE.with(|active| active.replace(Some(self.clone())));
89        let _restore = Restore(previous);
90        f()
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97
98    #[test]
99    fn scoped_workers_observe_cancellation_without_leaking_into_the_next_request() {
100        for _ in 0..2 {
101            let token = CancellationToken::new();
102            token.scope(|| {
103                let workers =
104                    crate::map_jobs(&[1, 2], std::num::NonZeroUsize::new(2).unwrap(), |_| {
105                        let active = CancellationToken::current().expect("worker inherits request");
106                        active.cancel();
107                        active.check()
108                    });
109                assert!(workers.iter().all(Result::is_err));
110            });
111            assert!(CancellationToken::current().is_none());
112        }
113        CancellationToken::new()
114            .run(|| Ok::<_, ComputationCancelled>(()))
115            .unwrap();
116    }
117}
bloq_graph/cancellation.rs

bloq_graph/
cancellation.rs
