runifold_core/
cancellation.rs

1use std::sync::{
2    Arc, Mutex, Weak,
3    atomic::{AtomicBool, Ordering},
4};
5
6use event_listener::Event;
7
8/// A hierarchical cancellation token.
9///
10/// Cancelling a parent is visible to all descendants. Cancelling a child does
11/// not affect its parent or siblings.
12#[derive(Clone, Debug)]
13pub struct CancellationToken {
14    state: Arc<CancellationState>,
15}
16
17#[derive(Debug)]
18struct CancellationState {
19    cancelled: AtomicBool,
20    event: Event,
21    children: Mutex<Vec<Weak<CancellationState>>>,
22}
23
24impl CancellationToken {
25    /// Creates an uncancelled root token.
26    pub fn new() -> Self {
27        Self {
28            state: Arc::new(CancellationState {
29                cancelled: AtomicBool::new(false),
30                event: Event::new(),
31                children: Mutex::new(Vec::new()),
32            }),
33        }
34    }
35
36    /// Creates a child token linked to this token.
37    #[must_use]
38    pub fn child_token(&self) -> Self {
39        let mut children = self
40            .state
41            .children
42            .lock()
43            .unwrap_or_else(std::sync::PoisonError::into_inner);
44        let child = Self {
45            state: Arc::new(CancellationState {
46                cancelled: AtomicBool::new(self.is_cancelled()),
47                event: Event::new(),
48                children: Mutex::new(Vec::new()),
49            }),
50        };
51        if !child.is_cancelled() {
52            children.push(Arc::downgrade(&child.state));
53        }
54        child
55    }
56
57    /// Cancels this token and, transitively, its descendants.
58    pub fn cancel(&self) {
59        cancel_state(&self.state);
60    }
61
62    /// Returns whether this token has been cancelled.
63    pub fn is_cancelled(&self) -> bool {
64        self.state.cancelled.load(Ordering::Acquire)
65    }
66
67    /// Waits until this token or one of its ancestors is cancelled.
68    pub async fn cancelled(&self) {
69        loop {
70            if self.is_cancelled() {
71                return;
72            }
73            let listener = self.state.event.listen();
74            if self.is_cancelled() {
75                return;
76            }
77            listener.await;
78        }
79    }
80}
81
82fn cancel_state(state: &Arc<CancellationState>) {
83    if state.cancelled.swap(true, Ordering::AcqRel) {
84        return;
85    }
86    state.event.notify(usize::MAX);
87
88    let children = state
89        .children
90        .lock()
91        .unwrap_or_else(std::sync::PoisonError::into_inner)
92        .iter()
93        .filter_map(Weak::upgrade)
94        .collect::<Vec<_>>();
95    for child in children {
96        cancel_state(&child);
97    }
98}
99
100impl Default for CancellationToken {
101    fn default() -> Self {
102        Self::new()
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::CancellationToken;
109
110    #[test]
111    fn parent_cancellation_reaches_descendants() {
112        let parent = CancellationToken::new();
113        let child = parent.child_token();
114        let grandchild = child.child_token();
115
116        parent.cancel();
117
118        assert!(parent.is_cancelled());
119        assert!(child.is_cancelled());
120        assert!(grandchild.is_cancelled());
121    }
122
123    #[test]
124    fn child_cancellation_is_isolated() {
125        let parent = CancellationToken::new();
126        let first_child = parent.child_token();
127        let second_child = parent.child_token();
128
129        first_child.cancel();
130
131        assert!(first_child.is_cancelled());
132        assert!(!parent.is_cancelled());
133        assert!(!second_child.is_cancelled());
134    }
135
136    #[test]
137    fn asynchronous_waiters_are_woken_by_ancestor_cancellation() {
138        let parent = CancellationToken::new();
139        let child = parent.child_token();
140        let waiter = child.clone();
141        let canceller = std::thread::spawn(move || parent.cancel());
142
143        futures_executor::block_on(waiter.cancelled());
144        canceller.join().unwrap();
145
146        assert!(child.is_cancelled());
147    }
148
149    #[test]
150    fn children_created_after_cancellation_start_cancelled() {
151        let parent = CancellationToken::new();
152        parent.cancel();
153
154        let child = parent.child_token();
155
156        assert!(child.is_cancelled());
157        futures_executor::block_on(child.cancelled());
158    }
159}
runifold_core/cancellation.rs

runifold_core/
cancellation.rs
