reth_tasks/
cancel.rs

1//! Cooperative cancellation of spawned work.
2
3use std::{
4    cell::RefCell,
5    sync::{
6        atomic::{AtomicBool, AtomicU8, Ordering},
7        Arc,
8    },
9};
10
11/// The work is still in progress.
12const RUNNING: u8 = 0;
13/// The work should wrap up and keep what it has produced so far.
14const FINALIZATION_REQUESTED: u8 = 1;
15/// The work should stop and discard what it has produced so far.
16const CANCELLED: u8 = 2;
17
18/// Cancels execution on drop and supports cooperative finalization.
19///
20/// If dropped, it will set the `cancelled` flag to true.
21///
22/// This is most useful when a spawned job should stop once its owner goes away, e.g. a payload
23/// job or a blocking RPC call whose caller disconnected.
24#[derive(Default, Clone, Debug)]
25pub struct CancelOnDrop(Arc<AtomicU8>);
26
27// === impl CancelOnDrop ===
28
29impl CancelOnDrop {
30    /// Returns true if the current work should be interrupted.
31    pub fn is_interrupted(&self) -> bool {
32        self.0.load(Ordering::Relaxed) != RUNNING
33    }
34
35    /// Returns true if the job was cancelled.
36    pub fn is_cancelled(&self) -> bool {
37        self.0.load(Ordering::Relaxed) == CANCELLED
38    }
39
40    /// Requests that the current work be finalized without cancelling it.
41    pub fn request_finalization(&self) {
42        let _ = self.0.compare_exchange(
43            RUNNING,
44            FINALIZATION_REQUESTED,
45            Ordering::Relaxed,
46            Ordering::Relaxed,
47        );
48    }
49
50    /// Returns true if finalization was requested.
51    pub fn is_finalization_requested(&self) -> bool {
52        self.0.load(Ordering::Relaxed) == FINALIZATION_REQUESTED
53    }
54
55    /// Runs `f` with this as the cancellation state of the current thread, see [`is_cancelled`].
56    ///
57    /// This lets code deep inside `f`, such as a loop over transactions, observe cancellation
58    /// without threading the [`CancelOnDrop`] through every call. The previous state of the
59    /// thread is restored once `f` returns or unwinds. Leaving the scope does not cancel.
60    pub fn scope<R>(&self, f: impl FnOnce() -> R) -> R {
61        struct Restore(Option<Arc<AtomicU8>>);
62
63        impl Drop for Restore {
64            fn drop(&mut self) {
65                CURRENT.set(self.0.take());
66            }
67        }
68
69        let _restore = Restore(CURRENT.replace(Some(self.0.clone())));
70        f()
71    }
72}
73
74impl Drop for CancelOnDrop {
75    fn drop(&mut self) {
76        self.0.store(CANCELLED, Ordering::Relaxed);
77    }
78}
79
80/// A marker that can be used to cancel execution.
81///
82/// If dropped, it will NOT set the `cancelled` flag to true.
83/// If `cancel` is called, the `cancelled` flag will be set to true.
84///
85/// This is useful when an external signal should cancel many tasks at once.
86#[derive(Default, Clone, Debug)]
87pub struct ManualCancel(Arc<AtomicBool>);
88
89// === impl ManualCancel ===
90
91impl ManualCancel {
92    /// Returns true if the job was cancelled.
93    pub fn is_cancelled(&self) -> bool {
94        self.0.load(Ordering::Relaxed)
95    }
96
97    /// Drops the [`ManualCancel`], setting the cancelled flag to true.
98    pub fn cancel(self) {
99        self.0.store(true, Ordering::Relaxed);
100    }
101}
102
103/// Returns `true` if the work running on the current thread was cancelled.
104///
105/// This reads the state of the innermost [`CancelOnDrop::scope`] on this thread, for example a
106/// blocking RPC call whose request was dropped because the client disconnected. Long running work
107/// should check this between units of work, such as transactions, and stop early since nobody
108/// waits for the result.
109///
110/// Always returns `false` outside of a [`CancelOnDrop::scope`].
111pub fn is_cancelled() -> bool {
112    CURRENT.with_borrow(|state| {
113        state.as_ref().is_some_and(|state| state.load(Ordering::Relaxed) == CANCELLED)
114    })
115}
116
117thread_local! {
118    /// Cancellation state of the innermost [`CancelOnDrop::scope`] on this thread.
119    static CURRENT: RefCell<Option<Arc<AtomicU8>>> = const { RefCell::new(None) };
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125
126    #[test]
127    fn test_default_cancelled() {
128        let c = CancelOnDrop::default();
129        assert!(!c.is_interrupted());
130        assert!(!c.is_cancelled());
131    }
132
133    #[test]
134    fn test_default_cancel_task() {
135        let c = ManualCancel::default();
136        assert!(!c.is_cancelled());
137    }
138
139    #[test]
140    fn test_set_cancel_task() {
141        let c = ManualCancel::default();
142        assert!(!c.is_cancelled());
143        let c2 = c.clone();
144        let c3 = c.clone();
145        c.cancel();
146        assert!(c3.is_cancelled());
147        assert!(c2.is_cancelled());
148    }
149
150    #[test]
151    fn test_cancel_task_multiple_threads() {
152        let c = ManualCancel::default();
153        let cloned_cancel = c.clone();
154
155        // we want to make sure that:
156        // * we can spawn tasks that do things
157        // * those tasks can run to completion and the flag remains unset unless we call cancel
158        let mut handles = vec![];
159        for _ in 0..10 {
160            let c = c.clone();
161            let handle = std::thread::spawn(move || {
162                for _ in 0..1000 {
163                    if c.is_cancelled() {
164                        return;
165                    }
166                }
167            });
168            handles.push(handle);
169        }
170
171        // wait for all the threads to finish
172        for handle in handles {
173            handle.join().unwrap();
174        }
175
176        // check that the flag is still unset
177        assert!(!c.is_cancelled());
178
179        // cancel and check that the flag is set
180        c.cancel();
181        assert!(cloned_cancel.is_cancelled());
182    }
183
184    #[test]
185    fn test_cancelondrop_clone_behavior() {
186        let cancel = CancelOnDrop::default();
187        assert!(!cancel.is_cancelled());
188
189        // Clone the CancelOnDrop
190        let cloned_cancel = cancel.clone();
191        assert!(!cloned_cancel.is_cancelled());
192
193        // Drop the original - this should set the cancelled flag
194        drop(cancel);
195
196        // The cloned instance should now see the cancelled flag as true
197        assert!(cloned_cancel.is_interrupted());
198        assert!(cloned_cancel.is_cancelled());
199    }
200
201    #[test]
202    fn test_cancelondrop_multiple_clones() {
203        let cancel = CancelOnDrop::default();
204        let clone1 = cancel.clone();
205        let clone2 = cancel.clone();
206        let clone3 = cancel.clone();
207
208        assert!(!cancel.is_cancelled());
209        assert!(!clone1.is_cancelled());
210        assert!(!clone2.is_cancelled());
211        assert!(!clone3.is_cancelled());
212
213        // Drop one clone - this should cancel all instances
214        drop(clone1);
215
216        assert!(cancel.is_interrupted());
217        assert!(cancel.is_cancelled());
218        assert!(clone2.is_cancelled());
219        assert!(clone3.is_cancelled());
220    }
221
222    #[test]
223    fn test_cancel_on_drop_finalization_request() {
224        let cancel = CancelOnDrop::default();
225        let clone = cancel.clone();
226
227        cancel.request_finalization();
228
229        assert!(clone.is_interrupted());
230        assert!(clone.is_finalization_requested());
231        assert!(!clone.is_cancelled());
232
233        drop(cancel);
234
235        assert!(clone.is_cancelled());
236        assert!(!clone.is_finalization_requested());
237
238        clone.request_finalization();
239
240        assert!(clone.is_cancelled());
241        assert!(!clone.is_finalization_requested());
242    }
243
244    #[test]
245    fn test_scope_is_cancelled() {
246        assert!(!is_cancelled());
247
248        let outer = CancelOnDrop::default();
249        let inner = CancelOnDrop::default();
250        outer.scope(|| {
251            assert!(!is_cancelled());
252            inner.scope(|| {
253                drop(inner.clone());
254                assert!(is_cancelled());
255            });
256            // the outer state is restored
257            assert!(!is_cancelled());
258            outer.request_finalization();
259            assert!(!is_cancelled());
260        });
261
262        // leaving the scope does not cancel
263        assert!(!is_cancelled());
264        assert!(!outer.is_cancelled());
265    }
266
267    #[test]
268    fn test_scope_restores_on_panic() {
269        let cancel = CancelOnDrop::default();
270        drop(cancel.clone());
271
272        let res = std::panic::catch_unwind(|| cancel.scope(|| panic!("scope panicked")));
273        assert!(res.is_err());
274        assert!(!is_cancelled());
275    }
276}
reth_tasks/cancel.rs

reth_tasks/
cancel.rs
