reth_revm/
cancelled.rs

1use alloc::sync::Arc;
2use core::sync::atomic::AtomicBool;
3
4/// A marker that can be used to cancel execution.
5///
6/// If dropped, it will set the `cancelled` flag to true.
7///
8/// This is most useful when a payload job needs to be cancelled.
9#[derive(Default, Clone, Debug)]
10pub struct CancelOnDrop(Arc<AtomicBool>);
11
12// === impl CancelOnDrop ===
13
14impl CancelOnDrop {
15    /// Returns true if the job was cancelled.
16    pub fn is_cancelled(&self) -> bool {
17        self.0.load(core::sync::atomic::Ordering::Relaxed)
18    }
19}
20
21impl Drop for CancelOnDrop {
22    fn drop(&mut self) {
23        self.0.store(true, core::sync::atomic::Ordering::Relaxed);
24    }
25}
26
27/// A marker that can be used to cancel execution.
28///
29/// If dropped, it will NOT set the `cancelled` flag to true.
30/// If `cancel` is called, the `cancelled` flag will be set to true.
31///
32/// This is useful in prewarming, when an external signal is received to cancel many prewarming
33/// tasks.
34#[derive(Default, Clone, Debug)]
35pub struct ManualCancel(Arc<AtomicBool>);
36
37// === impl ManualCancel ===
38
39impl ManualCancel {
40    /// Returns true if the job was cancelled.
41    pub fn is_cancelled(&self) -> bool {
42        self.0.load(core::sync::atomic::Ordering::Relaxed)
43    }
44
45    /// Drops the [`ManualCancel`], setting the cancelled flag to true.
46    pub fn cancel(self) {
47        self.0.store(true, core::sync::atomic::Ordering::Relaxed);
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54
55    #[test]
56    fn test_default_cancelled() {
57        let c = CancelOnDrop::default();
58        assert!(!c.is_cancelled());
59    }
60
61    #[test]
62    fn test_default_cancel_task() {
63        let c = ManualCancel::default();
64        assert!(!c.is_cancelled());
65    }
66
67    #[test]
68    fn test_set_cancel_task() {
69        let c = ManualCancel::default();
70        assert!(!c.is_cancelled());
71        let c2 = c.clone();
72        let c3 = c.clone();
73        c.cancel();
74        assert!(c3.is_cancelled());
75        assert!(c2.is_cancelled());
76    }
77
78    #[test]
79    fn test_cancel_task_multiple_threads() {
80        let c = ManualCancel::default();
81        let cloned_cancel = c.clone();
82
83        // we want to make sure that:
84        // * we can spawn tasks that do things
85        // * those tasks can run to completion and the flag remains unset unless we call cancel
86        let mut handles = vec![];
87        for _ in 0..10 {
88            let c = c.clone();
89            let handle = std::thread::spawn(move || {
90                for _ in 0..1000 {
91                    if c.is_cancelled() {
92                        return;
93                    }
94                }
95            });
96            handles.push(handle);
97        }
98
99        // wait for all the threads to finish
100        for handle in handles {
101            handle.join().unwrap();
102        }
103
104        // check that the flag is still unset
105        assert!(!c.is_cancelled());
106
107        // cancel and check that the flag is set
108        c.cancel();
109        assert!(cloned_cancel.is_cancelled());
110    }
111
112    #[test]
113    fn test_cancelondrop_clone_behavior() {
114        let cancel = CancelOnDrop::default();
115        assert!(!cancel.is_cancelled());
116
117        // Clone the CancelOnDrop
118        let cloned_cancel = cancel.clone();
119        assert!(!cloned_cancel.is_cancelled());
120
121        // Drop the original - this should set the cancelled flag
122        drop(cancel);
123
124        // The cloned instance should now see the cancelled flag as true
125        assert!(cloned_cancel.is_cancelled());
126    }
127
128    #[test]
129    fn test_cancelondrop_multiple_clones() {
130        let cancel = CancelOnDrop::default();
131        let clone1 = cancel.clone();
132        let clone2 = cancel.clone();
133        let clone3 = cancel.clone();
134
135        assert!(!cancel.is_cancelled());
136        assert!(!clone1.is_cancelled());
137        assert!(!clone2.is_cancelled());
138        assert!(!clone3.is_cancelled());
139
140        // Drop one clone - this should cancel all instances
141        drop(clone1);
142
143        assert!(cancel.is_cancelled());
144        assert!(clone2.is_cancelled());
145        assert!(clone3.is_cancelled());
146    }
147}
reth_revm/cancelled.rs

reth_revm/
cancelled.rs
