veilid_tools/
stop_token.rs

1use super::*;
2
3use event_listener::{Event, EventListener};
4
5struct Inner {
6    stopped: AtomicBool,
7    event: Event,
8}
9
10/// Cancellation source; dropping it stops every associated [`StopToken`].
11#[must_use]
12pub struct StopSource {
13    inner: Arc<Inner>,
14}
15
16impl StopSource {
17    /// Create a new stop source.
18    pub fn new() -> Self {
19        Self {
20            inner: Arc::new(Inner {
21                stopped: AtomicBool::new(false),
22                event: Event::new(),
23            }),
24        }
25    }
26
27    /// A token that resolves when this source is dropped.
28    pub fn token(&self) -> StopToken {
29        StopToken {
30            inner: self.inner.clone(),
31            listener: None,
32        }
33    }
34}
35
36impl Default for StopSource {
37    fn default() -> Self {
38        Self::new()
39    }
40}
41
42impl Drop for StopSource {
43    fn drop(&mut self) {
44        self.inner.stopped.store(true, Ordering::SeqCst);
45        self.inner.event.notify(usize::MAX);
46    }
47}
48
49impl fmt::Debug for StopSource {
50    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
51        write!(f, "StopSource")
52    }
53}
54
55/// Future that resolves when its [`StopSource`] is dropped.
56#[must_use]
57pub struct StopToken {
58    inner: Arc<Inner>,
59    listener: Option<Pin<Box<EventListener>>>,
60}
61
62impl Clone for StopToken {
63    fn clone(&self) -> Self {
64        Self {
65            inner: self.inner.clone(),
66            listener: None,
67        }
68    }
69}
70
71impl fmt::Debug for StopToken {
72    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
73        write!(f, "StopToken")
74    }
75}
76
77impl Future for StopToken {
78    type Output = ();
79
80    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
81        let this = self.get_mut();
82        loop {
83            if this.inner.stopped.load(Ordering::SeqCst) {
84                this.listener = None;
85                return task::Poll::Ready(());
86            }
87            match &mut this.listener {
88                Some(listener) => match listener.as_mut().poll(cx) {
89                    task::Poll::Ready(()) => {
90                        this.listener = None;
91                    }
92                    task::Poll::Pending => return task::Poll::Pending,
93                },
94                None => {
95                    this.listener = Some(Box::pin(this.inner.event.listen()));
96                }
97            }
98        }
99    }
100}
101
102/// Error returned by [`future::FutureExt::timeout_at`] when the deadline fires first.
103#[derive(Copy, Clone, Debug, PartialEq, Eq)]
104pub struct TimedOutError;
105
106impl fmt::Display for TimedOutError {
107    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
108        write!(f, "timed out")
109    }
110}
111
112impl std::error::Error for TimedOutError {}
113
114/// Future combinators bounded by a deadline future.
115pub mod future {
116    use super::*;
117
118    pin_project_lite::pin_project! {
119        /// Future that resolves to `Err(TimedOutError)` if the deadline fires before the inner future.
120        #[must_use]
121        pub struct TimeoutAt<F, D> {
122            #[pin]
123            future: F,
124            #[pin]
125            deadline: D,
126        }
127    }
128
129    impl<F: Future, D: Future> Future for TimeoutAt<F, D> {
130        type Output = Result<F::Output, TimedOutError>;
131
132        fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
133            let this = self.project();
134            if let task::Poll::Ready(out) = this.future.poll(cx) {
135                return task::Poll::Ready(Ok(out));
136            }
137            match this.deadline.poll(cx) {
138                task::Poll::Ready(_) => task::Poll::Ready(Err(TimedOutError)),
139                task::Poll::Pending => task::Poll::Pending,
140            }
141        }
142    }
143
144    /// Deadline-bounding extension for futures.
145    pub trait FutureExt: Future + Sized {
146        /// Run this future until `deadline` resolves, then return `Err(TimedOutError)`.
147        fn timeout_at<D: Future>(self, deadline: D) -> TimeoutAt<Self, D> {
148            TimeoutAt {
149                future: self,
150                deadline,
151            }
152        }
153    }
154
155    impl<F: Future + Sized> FutureExt for F {}
156}
157
158/// Common imports for stop token users.
159pub mod prelude {
160    #[doc(no_inline)]
161    pub use super::future::FutureExt as _;
162    #[doc(no_inline)]
163    pub use super::{StopSource, StopToken, TimedOutError};
164}
veilid_tools/stop_token.rs

veilid_tools/
stop_token.rs
