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}