made_core/ports/
execution_cancellation.rs1use std::future::{poll_fn, Future};
2use std::sync::{Arc, Mutex};
3use std::task::{Poll, Waker};
4
5#[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 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}