tokio_tasker/
future_ext.rs

1use std::fmt::{self, Debug};
2use std::future::Future;
3use std::pin::Pin;
4use std::task::{Context, Poll};
5
6use futures_util::future::FusedFuture;
7use pin_project_lite::pin_project;
8
9pin_project! {
10    /// A [`Future`][std::future::Future] for the [`unless()`][FutureExt::unless()] function.
11    #[project = UnlessProj]
12    #[project_replace = UnlessProjReplace]
13    #[must_use = "futures do nothing unless you `.await` or poll them"]
14    pub enum Unless<F, U> {
15        Awaiting {
16            #[pin] future: F,
17            #[pin] unless: U,
18        },
19        Done,
20    }
21}
22
23impl<F, U> Unless<F, U> {
24    pub(crate) fn new(future: F, unless: U) -> Self {
25        Self::Awaiting { future, unless }
26    }
27}
28
29impl<F, U> Future for Unless<F, U>
30where
31    F: Future,
32    U: Future,
33{
34    type Output = Result<F::Output, U::Output>;
35
36    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
37        use UnlessProj::*;
38
39        let output = match self.as_mut().project() {
40            Awaiting { future, unless } => {
41                if let Poll::Ready(output) = unless.poll(cx) {
42                    Poll::Ready(Err(output))
43                } else if let Poll::Ready(output) = future.poll(cx) {
44                    Poll::Ready(Ok(output))
45                } else {
46                    return Poll::Pending;
47                }
48            }
49            Done => panic!("Unless polled after it returned `Poll::Ready`"),
50        };
51
52        // If we get here, either of the futures became ready and we have the output.
53        // Update state and return result:
54        self.as_mut().project_replace(Unless::Done);
55        output
56    }
57}
58
59impl<F, U> FusedFuture for Unless<F, U>
60where
61    F: Future,
62    U: Future,
63{
64    fn is_terminated(&self) -> bool {
65        matches!(self, Unless::Done)
66    }
67}
68
69impl<F, U> Debug for Unless<F, U>
70where
71    F: Debug,
72    U: Debug,
73{
74    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
75        match self {
76            Unless::Awaiting { future, unless } => f
77                .debug_struct("Unless::Awaiting")
78                .field("future", future)
79                .field("unless", unless)
80                .finish(),
81            Unless::Done => write!(f, "Unless::Done"),
82        }
83    }
84}
85
86/// A [`Future`][std::future::Future] extension adding the [`.unless()`][FutureExt::unless()]
87/// method that is desinged to be used with a [`Stopper`][crate::Stopper].
88pub trait FutureExt {
89    /// Wrap a future to be stopped from resolving if the `unless` future resolves first.
90    ///
91    /// This is exactly the same operation as future's `.select()`, but better designed
92    /// for stopping futures with the [`Stopper`] and without the `Unpin` requirement.
93    ///
94    /// Unlike `Select`, this future yields a `Result`
95    /// where in the `Ok()` variant the result of the original future is returned,
96    /// and in the `Err()` variant the result of the `unless` future
97    /// is returned in case the `unless` future resolved first.
98    ///
99    /// When used with [`Stopper`], the future will yield [`Err(Stopped)`]
100    /// if it is stopped by the associated [`Tasker`].
101    ///
102    /// [`Tasker`]: crate::Tasker
103    /// [`Stopper`]: crate::Stopper
104    /// [`Err(Stopped)`]: crate::Stopped
105    fn unless<U>(self, unless: U) -> Unless<Self, U>
106    where
107        Self: Sized,
108        U: Future;
109}
110
111impl<F> FutureExt for F
112where
113    F: Future,
114{
115    fn unless<U>(self, unless: U) -> Unless<Self, U>
116    where
117        Self: Sized,
118        U: Future,
119    {
120        Unless::new(self, unless)
121    }
122}