ntex_util/future/
join.rs

1use std::{future::Future, mem, pin::Pin, task::Context, task::Poll};
2
3use crate::future::MaybeDone;
4
5/// Future for the `join` combinator, waiting for two futures to complete.
6pub async fn join<A, B>(fut_a: A, fut_b: B) -> (A::Output, B::Output)
7where
8    A: Future,
9    B: Future,
10{
11    Join {
12        fut_a: MaybeDone::Pending(fut_a),
13        fut_b: MaybeDone::Pending(fut_b),
14    }
15    .await
16}
17
18pin_project_lite::pin_project! {
19    pub(crate) struct Join<A: Future, B: Future> {
20        #[pin]
21        fut_a: MaybeDone<A>,
22        #[pin]
23        fut_b: MaybeDone<B>,
24    }
25}
26
27impl<A, B> Future for Join<A, B>
28where
29    A: Future,
30    B: Future,
31{
32    type Output = (A::Output, B::Output);
33
34    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
35        let mut this = self.project();
36        let mut all_done = true;
37        all_done &= this.fut_a.as_mut().poll(cx).is_ready();
38        all_done &= this.fut_b.as_mut().poll(cx).is_ready();
39        if all_done {
40            Poll::Ready((
41                this.fut_a.take_output().unwrap(),
42                this.fut_b.take_output().unwrap(),
43            ))
44        } else {
45            Poll::Pending
46        }
47    }
48}
49
50/// Creates a future which represents a collection of the outputs of the futures
51/// given.
52pub async fn join_all<I>(i: I) -> Vec<<I::Item as Future>::Output>
53where
54    I: IntoIterator,
55    I::Item: Future,
56{
57    let elems: Box<[_]> = i.into_iter().map(MaybeDone::Pending).collect();
58    JoinAll {
59        elems: elems.into(),
60    }
61    .await
62}
63
64pub(crate) struct JoinAll<T: Future> {
65    elems: Pin<Box<[MaybeDone<T>]>>,
66}
67
68fn iter_pin_mut<T>(slice: Pin<&mut [T]>) -> impl Iterator<Item = Pin<&mut T>> {
69    // Safety: `std` _could_ make this unsound if it were to decide Pin's
70    // invariants aren't required to transmit through slices. Otherwise this has
71    // the same safety as a normal field pin projection.
72    unsafe { slice.get_unchecked_mut() }
73        .iter_mut()
74        .map(|t| unsafe { Pin::new_unchecked(t) })
75}
76
77impl<T> Future for JoinAll<T>
78where
79    T: Future,
80{
81    type Output = Vec<T::Output>;
82
83    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
84        let mut all_done = true;
85
86        for elem in iter_pin_mut(self.elems.as_mut()) {
87            if elem.poll(cx).is_pending() {
88                all_done = false;
89            }
90        }
91        if all_done {
92            let mut elems = mem::replace(&mut self.elems, Box::pin([]));
93            let result = iter_pin_mut(elems.as_mut())
94                .map(|e| e.take_output().unwrap())
95                .collect();
96            Poll::Ready(result)
97        } else {
98            Poll::Pending
99        }
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106    use crate::{future::Ready, time};
107
108    #[ntex_macros::rt_test2]
109    async fn join_tests() {
110        let res = join(Ready::<_, ()>::Ok("test"), time::sleep(time::Millis(50))).await;
111        assert_eq!(res, (Ok("test"), ()));
112
113        let res =
114            join_all([time::sleep(time::Millis(50)), time::sleep(time::Millis(60))]).await;
115        assert_eq!(res, vec![(), ()]);
116    }
117}
ntex_util/future/join.rs

ntex_util/future/
join.rs
