klukai_types/tripwire/
preempt.rs

1use pin_project_lite::pin_project;
2use std::{
3    future::Future,
4    pin::Pin,
5    task::{Context, Poll},
6    time::Duration,
7};
8use tokio::time::Sleep;
9
10/// Whether a preemptible future was preempted or completed
11#[derive(Debug)]
12pub enum Outcome<P, C> {
13    /// The future was preempted by another future
14    Preempted(P),
15    /// The future completed uninterrupted
16    Completed(C),
17}
18
19impl<P, C> Outcome<P, C> {
20    /// True if the variant is Preempted
21    pub fn is_preempted(&self) -> bool {
22        matches!(self, Outcome::Preempted(_))
23    }
24
25    /// True if the variant is Completed
26    pub fn is_completed(&self) -> bool {
27        matches!(self, Outcome::Completed(_))
28    }
29}
30
31pin_project! {
32    /// A future that can be preempted by another one
33    pub struct PreemptibleFuture<P, C> {
34        #[pin]
35        preempt: P,
36
37        #[pin]
38        complete: C,
39    }
40}
41
42impl<P, C> Future for PreemptibleFuture<P, C>
43where
44    P: Future,
45    C: Future,
46{
47    type Output = Outcome<P::Output, C::Output>;
48
49    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
50        let this = self.project();
51
52        if let Poll::Ready(v) = this.preempt.poll(cx) {
53            return Poll::Ready(Outcome::Preempted(v));
54        }
55
56        this.complete.poll(cx).map(Outcome::Completed)
57    }
58}
59
60impl<P, C, T, E> PreemptibleFuture<P, C>
61where
62    P: Future,
63    C: Future<Output = Result<T, E>>,
64{
65    /// Map a "preempted" outcome to an error via a closure
66    pub async fn complete_or_else(self, f: impl FnOnce(P::Output) -> E) -> Result<T, E> {
67        match self.await {
68            Outcome::Preempted(v) => Err(f(v)),
69            Outcome::Completed(res) => res,
70        }
71    }
72
73    /// Map a "preempted" outcome to an error
74    pub async fn complete_or(self, err: E) -> Result<T, E> {
75        match self.await {
76            Outcome::Preempted(_) => Err(err),
77            Outcome::Completed(res) => res,
78        }
79    }
80}
81
82/// Extension trait to allow wrapping futures such that they can be preempted
83pub trait PreemptibleFutureExt<P>: Sized {
84    /// Wrap a future such that it can be preempted by another one
85    fn preemptible(self, done: P) -> PreemptibleFuture<P, Self>;
86}
87
88/// Extension trait to allow wrapping futures such that they can time out
89/// Kinda like `tokio::time::timeout` but postfix
90pub trait TimeoutFutureExt: Sized {
91    /// Wrap a future so that it can be preempted by a timeout
92    fn with_timeout(self, delay: Duration) -> PreemptibleFuture<Sleep, Self>;
93}
94
95impl<P, C> PreemptibleFutureExt<P> for C
96where
97    P: Future,
98    C: Future,
99{
100    fn preemptible(self, preempt: P) -> PreemptibleFuture<P, Self> {
101        PreemptibleFuture {
102            preempt,
103            complete: self,
104        }
105    }
106}
107
108impl<C> TimeoutFutureExt for C
109where
110    C: Future,
111{
112    fn with_timeout(self, delay: Duration) -> PreemptibleFuture<Sleep, Self> {
113        self.preemptible(tokio::time::sleep(delay))
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use std::{
120        sync::{
121            Arc,
122            atomic::{AtomicBool, Ordering},
123        },
124        time::Duration,
125    };
126
127    use crate::tripwire::{Outcome, PreemptibleFutureExt, Tripwire};
128    use assert2::let_assert;
129
130    #[tokio::test]
131    async fn preemptible_tripwire() {
132        let (tripwire, tripwire_worker, tripwire_tx) = Tripwire::new_simple();
133
134        {
135            let signal: Arc<AtomicBool> = Default::default();
136            let f = {
137                let signal = signal.clone();
138                async move {
139                    tokio::time::sleep(Duration::from_millis(300)).await;
140                    signal.store(true, Ordering::SeqCst);
141                }
142                .preemptible(tripwire.clone())
143            };
144            let_assert!(Outcome::Completed(_) = f.await);
145            assert!(signal.load(Ordering::SeqCst))
146        }
147
148        {
149            let signal: Arc<AtomicBool> = Default::default();
150            let f = {
151                let signal = signal.clone();
152                async move {
153                    tokio::time::sleep(Duration::from_millis(300)).await;
154                    signal.store(true, Ordering::SeqCst);
155                }
156                .preemptible(tripwire.clone())
157            };
158            tripwire_tx.send(()).await.ok();
159            tripwire_worker.await;
160            let_assert!(Outcome::Preempted(()) = f.await);
161            assert!(!signal.load(Ordering::SeqCst))
162        }
163    }
164}
klukai_types/tripwire/preempt.rs

klukai_types/tripwire/
preempt.rs
