reth_revm/
cancelled.rs

1use alloc::sync::Arc;
2use core::sync::atomic::{AtomicBool, AtomicU8, Ordering};
3
4/// Payload building is still in progress.
5const RUNNING: u8 = 0;
6/// Payload building should stop accepting transactions and seal the accumulated work.
7const FINALIZATION_REQUESTED: u8 = 1;
8/// Payload building should stop and discard the accumulated work.
9const CANCELLED: u8 = 2;
10
11/// Cancels execution on drop and supports cooperative finalization.
12///
13/// If dropped, it will set the `cancelled` flag to true.
14///
15/// This is most useful when a payload job needs to be cancelled.
16#[derive(Default, Clone, Debug)]
17pub struct CancelOnDrop(Arc<AtomicU8>);
18
19// === impl CancelOnDrop ===
20
21impl CancelOnDrop {
22    /// Returns true if the current work should be interrupted.
23    pub fn is_interrupted(&self) -> bool {
24        self.0.load(Ordering::Relaxed) != RUNNING
25    }
26
27    /// Returns true if the job was cancelled.
28    pub fn is_cancelled(&self) -> bool {
29        self.0.load(Ordering::Relaxed) == CANCELLED
30    }
31
32    /// Requests that the current work be finalized without cancelling it.
33    pub fn request_finalization(&self) {
34        let _ = self.0.compare_exchange(
35            RUNNING,
36            FINALIZATION_REQUESTED,
37            Ordering::Relaxed,
38            Ordering::Relaxed,
39        );
40    }
41
42    /// Returns true if finalization was requested.
43    pub fn is_finalization_requested(&self) -> bool {
44        self.0.load(Ordering::Relaxed) == FINALIZATION_REQUESTED
45    }
46}
47
48impl Drop for CancelOnDrop {
49    fn drop(&mut self) {
50        self.0.store(CANCELLED, Ordering::Relaxed);
51    }
52}
53
54/// A marker that can be used to cancel execution.
55///
56/// If dropped, it will NOT set the `cancelled` flag to true.
57/// If `cancel` is called, the `cancelled` flag will be set to true.
58///
59/// This is useful in prewarming, when an external signal is received to cancel many prewarming
60/// tasks.
61#[derive(Default, Clone, Debug)]
62pub struct ManualCancel(Arc<AtomicBool>);
63
64// === impl ManualCancel ===
65
66impl ManualCancel {
67    /// Returns true if the job was cancelled.
68    pub fn is_cancelled(&self) -> bool {
69        self.0.load(core::sync::atomic::Ordering::Relaxed)
70    }
71
72    /// Drops the [`ManualCancel`], setting the cancelled flag to true.
73    pub fn cancel(self) {
74        self.0.store(true, core::sync::atomic::Ordering::Relaxed);
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81
82    #[test]
83    fn test_default_cancelled() {
84        let c = CancelOnDrop::default();
85        assert!(!c.is_interrupted());
86        assert!(!c.is_cancelled());
87    }
88
89    #[test]
90    fn test_default_cancel_task() {
91        let c = ManualCancel::default();
92        assert!(!c.is_cancelled());
93    }
94
95    #[test]
96    fn test_set_cancel_task() {
97        let c = ManualCancel::default();
98        assert!(!c.is_cancelled());
99        let c2 = c.clone();
100        let c3 = c.clone();
101        c.cancel();
102        assert!(c3.is_cancelled());
103        assert!(c2.is_cancelled());
104    }
105
106    #[test]
107    fn test_cancel_task_multiple_threads() {
108        let c = ManualCancel::default();
109        let cloned_cancel = c.clone();
110
111        // we want to make sure that:
112        // * we can spawn tasks that do things
113        // * those tasks can run to completion and the flag remains unset unless we call cancel
114        let mut handles = vec![];
115        for _ in 0..10 {
116            let c = c.clone();
117            let handle = std::thread::spawn(move || {
118                for _ in 0..1000 {
119                    if c.is_cancelled() {
120                        return;
121                    }
122                }
123            });
124            handles.push(handle);
125        }
126
127        // wait for all the threads to finish
128        for handle in handles {
129            handle.join().unwrap();
130        }
131
132        // check that the flag is still unset
133        assert!(!c.is_cancelled());
134
135        // cancel and check that the flag is set
136        c.cancel();
137        assert!(cloned_cancel.is_cancelled());
138    }
139
140    #[test]
141    fn test_cancelondrop_clone_behavior() {
142        let cancel = CancelOnDrop::default();
143        assert!(!cancel.is_cancelled());
144
145        // Clone the CancelOnDrop
146        let cloned_cancel = cancel.clone();
147        assert!(!cloned_cancel.is_cancelled());
148
149        // Drop the original - this should set the cancelled flag
150        drop(cancel);
151
152        // The cloned instance should now see the cancelled flag as true
153        assert!(cloned_cancel.is_interrupted());
154        assert!(cloned_cancel.is_cancelled());
155    }
156
157    #[test]
158    fn test_cancelondrop_multiple_clones() {
159        let cancel = CancelOnDrop::default();
160        let clone1 = cancel.clone();
161        let clone2 = cancel.clone();
162        let clone3 = cancel.clone();
163
164        assert!(!cancel.is_cancelled());
165        assert!(!clone1.is_cancelled());
166        assert!(!clone2.is_cancelled());
167        assert!(!clone3.is_cancelled());
168
169        // Drop one clone - this should cancel all instances
170        drop(clone1);
171
172        assert!(cancel.is_interrupted());
173        assert!(cancel.is_cancelled());
174        assert!(clone2.is_cancelled());
175        assert!(clone3.is_cancelled());
176    }
177
178    #[test]
179    fn test_cancel_on_drop_finalization_request() {
180        let cancel = CancelOnDrop::default();
181        let clone = cancel.clone();
182
183        cancel.request_finalization();
184
185        assert!(clone.is_interrupted());
186        assert!(clone.is_finalization_requested());
187        assert!(!clone.is_cancelled());
188
189        drop(cancel);
190
191        assert!(clone.is_cancelled());
192        assert!(!clone.is_finalization_requested());
193
194        clone.request_finalization();
195
196        assert!(clone.is_cancelled());
197        assert!(!clone.is_finalization_requested());
198    }
199}
reth_revm/cancelled.rs

reth_revm/
cancelled.rs
