tokio/future/
try_join.rs

1use crate::future::maybe_done::{maybe_done, MaybeDone};
2
3use pin_project_lite::pin_project;
4use std::future::Future;
5use std::pin::Pin;
6use std::task::{Context, Poll};
7
8pub(crate) fn try_join3<T1, F1, T2, F2, T3, F3, E>(
9    future1: F1,
10    future2: F2,
11    future3: F3,
12) -> TryJoin3<F1, F2, F3>
13where
14    F1: Future<Output = Result<T1, E>>,
15    F2: Future<Output = Result<T2, E>>,
16    F3: Future<Output = Result<T3, E>>,
17{
18    TryJoin3 {
19        future1: maybe_done(future1),
20        future2: maybe_done(future2),
21        future3: maybe_done(future3),
22    }
23}
24
25pin_project! {
26    pub(crate) struct TryJoin3<F1, F2, F3>
27    where
28        F1: Future,
29        F2: Future,
30        F3: Future,
31    {
32        #[pin]
33        future1: MaybeDone<F1>,
34        #[pin]
35        future2: MaybeDone<F2>,
36        #[pin]
37        future3: MaybeDone<F3>,
38    }
39}
40
41impl<T1, F1, T2, F2, T3, F3, E> Future for TryJoin3<F1, F2, F3>
42where
43    F1: Future<Output = Result<T1, E>>,
44    F2: Future<Output = Result<T2, E>>,
45    F3: Future<Output = Result<T3, E>>,
46{
47    type Output = Result<(T1, T2, T3), E>;
48
49    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
50        let mut all_done = true;
51
52        let mut me = self.project();
53
54        if me.future1.as_mut().poll(cx).is_pending() {
55            all_done = false;
56        } else if me.future1.as_mut().output_mut().unwrap().is_err() {
57            return Poll::Ready(Err(me.future1.take_output().unwrap().err().unwrap()));
58        }
59
60        if me.future2.as_mut().poll(cx).is_pending() {
61            all_done = false;
62        } else if me.future2.as_mut().output_mut().unwrap().is_err() {
63            return Poll::Ready(Err(me.future2.take_output().unwrap().err().unwrap()));
64        }
65
66        if me.future3.as_mut().poll(cx).is_pending() {
67            all_done = false;
68        } else if me.future3.as_mut().output_mut().unwrap().is_err() {
69            return Poll::Ready(Err(me.future3.take_output().unwrap().err().unwrap()));
70        }
71
72        if all_done {
73            Poll::Ready(Ok((
74                me.future1.take_output().unwrap().ok().unwrap(),
75                me.future2.take_output().unwrap().ok().unwrap(),
76                me.future3.take_output().unwrap().ok().unwrap(),
77            )))
78        } else {
79            Poll::Pending
80        }
81    }
82}