reth_cli_util/cancellation.rs
1//! Thread-safe cancellation primitives for cooperative task cancellation.
2
3use std::sync::{
4 atomic::{AtomicBool, Ordering},
5 Arc,
6};
7
8/// A thread-safe cancellation token that can be shared across threads.
9///
10/// This token allows cooperative cancellation by providing a way to signal
11/// cancellation and check cancellation status. The token can be cloned and
12/// shared across multiple threads, with all clones sharing the same cancellation state.
13///
14/// # Example
15///
16/// ```
17/// use reth_cli_util::cancellation::CancellationToken;
18/// use std::{thread, time::Duration};
19///
20/// let token = CancellationToken::new();
21/// let worker_token = token.clone();
22///
23/// let handle = thread::spawn(move || {
24/// while !worker_token.is_cancelled() {
25/// // Do work...
26/// thread::sleep(Duration::from_millis(100));
27/// }
28/// });
29///
30/// // Cancel from main thread
31/// token.cancel();
32/// handle.join().unwrap();
33/// ```
34#[derive(Clone, Debug)]
35pub struct CancellationToken {
36 cancelled: Arc<AtomicBool>,
37}
38
39impl CancellationToken {
40 /// Creates a new cancellation token in the non-cancelled state.
41 pub fn new() -> Self {
42 Self { cancelled: Arc::new(AtomicBool::new(false)) }
43 }
44
45 /// Signals cancellation to all holders of this token and its clones.
46 ///
47 /// Once cancelled, the token cannot be reset. This operation is thread-safe
48 /// and can be called multiple times without issue.
49 pub fn cancel(&self) {
50 self.cancelled.store(true, Ordering::Release);
51 }
52
53 /// Checks whether cancellation has been requested.
54 ///
55 /// Returns `true` if [`cancel`](Self::cancel) has been called on this token
56 /// or any of its clones.
57 pub fn is_cancelled(&self) -> bool {
58 self.cancelled.load(Ordering::Relaxed)
59 }
60
61 /// Creates a guard that automatically cancels this token when dropped.
62 ///
63 /// This is useful for ensuring cancellation happens when a scope exits,
64 /// either normally or via panic.
65 ///
66 /// # Example
67 ///
68 /// ```
69 /// use reth_cli_util::cancellation::CancellationToken;
70 ///
71 /// let token = CancellationToken::new();
72 /// {
73 /// let _guard = token.drop_guard();
74 /// assert!(!token.is_cancelled());
75 /// // Guard dropped here, triggering cancellation
76 /// }
77 /// assert!(token.is_cancelled());
78 /// ```
79 pub fn drop_guard(&self) -> CancellationGuard {
80 CancellationGuard { token: self.clone() }
81 }
82}
83
84impl Default for CancellationToken {
85 fn default() -> Self {
86 Self::new()
87 }
88}
89
90/// A guard that cancels its associated [`CancellationToken`] when dropped.
91///
92/// Created by calling [`CancellationToken::drop_guard`]. When this guard is dropped,
93/// it automatically calls [`cancel`](CancellationToken::cancel) on the token.
94#[derive(Debug)]
95pub struct CancellationGuard {
96 token: CancellationToken,
97}
98
99impl Drop for CancellationGuard {
100 fn drop(&mut self) {
101 self.token.cancel();
102 }
103}