ntex_util/future/
mod.rs

1//! Future and stream utilities.
2use std::future::{Future, poll_fn};
3use std::task::{Context, Poll, ready};
4use std::{mem, pin::Pin};
5
6pub use futures_core::{Stream, TryFuture};
7
8mod either;
9mod join;
10mod lazy;
11mod on_drop;
12mod select;
13
14pub use self::either::Either;
15pub use self::join::{join, join_all};
16pub use self::lazy::{Lazy, lazy};
17pub use self::on_drop::{OnDropFn, OnDropFuture, OnDropFutureExt};
18pub use self::select::select;
19
20/// An owned, dynamically dispatched future.
21///
22/// This is useful when a concrete future type cannot be named or when type
23/// erasure is required.
24pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + 'a>>;
25
26/// Waits for and returns the next item from a stream.
27pub async fn stream_recv<S>(stream: &mut S) -> Option<S::Item>
28where
29    S: Stream + Unpin,
30{
31    poll_fn(|cx| Pin::new(&mut *stream).poll_next(cx)).await
32}
33
34enum MaybeDone<F>
35where
36    F: Future,
37{
38    Pending(F),
39    Done(F::Output),
40    Gone,
41}
42
43impl<F: Future> MaybeDone<F> {
44    fn take_output(self: Pin<&mut Self>) -> Option<F::Output> {
45        match &*self {
46            Self::Done(_) => {}
47            Self::Pending(_) | Self::Gone => return None,
48        }
49        unsafe {
50            match mem::replace(self.get_unchecked_mut(), Self::Gone) {
51                MaybeDone::Done(output) => Some(output),
52                _ => unreachable!(),
53            }
54        }
55    }
56}
57
58impl<F: Future> Future for MaybeDone<F> {
59    type Output = ();
60
61    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
62        unsafe {
63            match self.as_mut().get_unchecked_mut() {
64                MaybeDone::Pending(f) => {
65                    let res = ready!(Pin::new_unchecked(f).poll(cx));
66                    self.set(Self::Done(res));
67                }
68                MaybeDone::Done(_) => {}
69                MaybeDone::Gone => panic!("MaybeDone polled after value taken"),
70            }
71        }
72        Poll::Ready(())
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use std::{future::pending, future::ready, pin::pin};
79
80    use super::*;
81
82    /// STATEMENT: `MaybeDone` exposes an output only after completion and only once.
83    ///
84    /// HOW:
85    /// - Try taking output from a pending future.
86    /// - Complete a ready future, take its output, and verify a second take returns `None`.
87    #[ntex::test]
88    async fn maybe_done_take_output() {
89        let mut fut = pin!(MaybeDone::Pending(pending::<()>()));
90        assert!(fut.as_mut().take_output().is_none());
91
92        let mut fut = pin!(MaybeDone::Pending(ready(1)));
93        fut.as_mut().await;
94        assert_eq!(fut.as_mut().take_output(), Some(1));
95        // the output is taken once
96        assert!(fut.as_mut().take_output().is_none());
97    }
98
99    /// STATEMENT: Polling `MaybeDone` after taking its output panics with the expected message.
100    ///
101    /// HOW:
102    /// - Complete the future, extract and check its output, then await it again.
103    #[ntex::test]
104    #[should_panic(expected = "MaybeDone polled after value taken")]
105    async fn maybe_done_poll_after_take() {
106        let mut fut = pin!(MaybeDone::Pending(ready(1)));
107        fut.as_mut().await;
108        assert_eq!(fut.as_mut().take_output(), Some(1));
109        fut.as_mut().await;
110    }
111}
ntex_util/future/mod.rs

ntex_util/future/
mod.rs
