reth_revm/
cancelled.rs

Help
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}
reth_revm/cancelled.rs

reth_revm/
cancelled.rs
