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_set_cancel_task() {
135        let c = ManualCancel::default();
136        assert!(!c.is_cancelled());
137
138        // Dropping a clone must not cancel the original, unlike `CancelOnDrop`.
139        drop(c.clone());
140        assert!(!c.is_cancelled());
141
142        let c2 = c.clone();
143        let c3 = c.clone();
144        c.cancel();
145        assert!(c3.is_cancelled());
146        assert!(c2.is_cancelled());
147    }
148
149    #[test]
150    fn test_cancelondrop_clone_behavior() {
151        let cancel = CancelOnDrop::default();
152        assert!(!cancel.is_cancelled());
153
154        // Clone the CancelOnDrop
155        let cloned_cancel = cancel.clone();
156        assert!(!cloned_cancel.is_cancelled());
157
158        // Drop the original - this should set the cancelled flag
159        drop(cancel);
160
161        // The cloned instance should now see the cancelled flag as true
162        assert!(cloned_cancel.is_interrupted());
163        assert!(cloned_cancel.is_cancelled());
164    }
165
166    #[test]
167    fn test_cancelondrop_multiple_clones() {
168        let cancel = CancelOnDrop::default();
169        let clone1 = cancel.clone();
170        let clone2 = cancel.clone();
171        let clone3 = cancel.clone();
172
173        assert!(!cancel.is_cancelled());
174        assert!(!clone1.is_cancelled());
175        assert!(!clone2.is_cancelled());
176        assert!(!clone3.is_cancelled());
177
178        // Drop one clone - this should cancel all instances
179        drop(clone1);
180
181        assert!(cancel.is_interrupted());
182        assert!(cancel.is_cancelled());
183        assert!(clone2.is_cancelled());
184        assert!(clone3.is_cancelled());
185    }
186
187    #[test]
188    fn test_cancel_on_drop_finalization_request() {
189        let cancel = CancelOnDrop::default();
190        let clone = cancel.clone();
191
192        cancel.request_finalization();
193
194        assert!(clone.is_interrupted());
195        assert!(clone.is_finalization_requested());
196        assert!(!clone.is_cancelled());
197
198        drop(cancel);
199
200        assert!(clone.is_cancelled());
201        assert!(!clone.is_finalization_requested());
202
203        clone.request_finalization();
204
205        assert!(clone.is_cancelled());
206        assert!(!clone.is_finalization_requested());
207    }
208
209    #[test]
210    fn test_scope_is_cancelled() {
211        assert!(!is_cancelled());
212
213        let outer = CancelOnDrop::default();
214        let inner = CancelOnDrop::default();
215        outer.scope(|| {
216            assert!(!is_cancelled());
217            inner.scope(|| {
218                drop(inner.clone());
219                assert!(is_cancelled());
220            });
221            // the outer state is restored
222            assert!(!is_cancelled());
223            outer.request_finalization();
224            assert!(!is_cancelled());
225        });
226
227        // leaving the scope does not cancel
228        assert!(!is_cancelled());
229        assert!(!outer.is_cancelled());
230    }
231
232    #[test]
233    fn test_scope_restores_on_panic() {
234        let cancel = CancelOnDrop::default();
235        drop(cancel.clone());
236
237        let res = std::panic::catch_unwind(|| cancel.scope(|| panic!("scope panicked")));
238        assert!(res.is_err());
239        assert!(!is_cancelled());
240    }
241}
reth_tasks/cancel.rs

reth_tasks/
cancel.rs
