Skip to main content

made_core/ports/
execution_cancellation.rs

1use std::future::{poll_fn, Future};
2use std::sync::{Arc, Mutex};
3use std::task::{Poll, Waker};
4
5/// Cooperative loss-of-authority signal. Process adapters must kill and reap
6/// their owned execution before returning from cancellation.
7#[derive(Debug, Clone)]
8pub struct ExecutionCancellation {
9    state: Arc<Mutex<Option<Vec<Waker>>>>,
10}
11
12impl ExecutionCancellation {
13    #[must_use]
14    pub fn new() -> Self {
15        Self {
16            state: Arc::new(Mutex::new(Some(Vec::new()))),
17        }
18    }
19
20    pub fn cancel(&self) {
21        let waiters = self
22            .state
23            .lock()
24            .unwrap_or_else(std::sync::PoisonError::into_inner)
25            .take();
26        for waker in waiters.into_iter().flatten() {
27            waker.wake();
28        }
29    }
30
31    #[must_use]
32    pub fn is_cancelled(&self) -> bool {
33        self.state
34            .lock()
35            .unwrap_or_else(std::sync::PoisonError::into_inner)
36            .is_none()
37    }
38
39    pub async fn cancelled(&self) {
40        poll_fn(|cx| {
41            let mut state = self
42                .state
43                .lock()
44                .unwrap_or_else(std::sync::PoisonError::into_inner);
45            match state.as_mut() {
46                None => Poll::Ready(()),
47                Some(waiters) => {
48                    if !waiters.iter().any(|w| w.will_wake(cx.waker())) {
49                        waiters.push(cx.waker().clone());
50                    }
51                    Poll::Pending
52                }
53            }
54        })
55        .await;
56    }
57
58    /// Cancels a future without asserting that dropping it kills external work.
59    pub async fn run<F: Future>(&self, future: F) -> Option<F::Output> {
60        let mut future = std::pin::pin!(future);
61        let mut cancelled = std::pin::pin!(self.cancelled());
62        poll_fn(|cx| {
63            if cancelled.as_mut().poll(cx).is_ready() {
64                return Poll::Ready(None);
65            }
66            future.as_mut().poll(cx).map(Some)
67        })
68        .await
69    }
70}
71
72impl Default for ExecutionCancellation {
73    fn default() -> Self {
74        Self::new()
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::ExecutionCancellation;
81
82    #[tokio::test]
83    async fn cancellation_wakes_every_waiter() {
84        let cancellation = ExecutionCancellation::new();
85        let left = cancellation.clone();
86        let right = cancellation.clone();
87        let left = tokio::spawn(async move { left.cancelled().await });
88        let right = tokio::spawn(async move { right.cancelled().await });
89
90        tokio::task::yield_now().await;
91        cancellation.cancel();
92        left.await.unwrap();
93        right.await.unwrap();
94        assert!(cancellation.is_cancelled());
95    }
96
97    #[tokio::test]
98    async fn run_returns_none_after_authority_is_cancelled() {
99        let cancellation = ExecutionCancellation::new();
100        let trigger = cancellation.clone();
101        let result = cancellation
102            .run(async move {
103                trigger.cancel();
104                tokio::task::yield_now().await;
105                42
106            })
107            .await;
108        assert_eq!(result, None);
109    }
110}