async_macros/
try_select.rs

1#![allow(non_snake_case)]
2
3/// Waits for either one of several similarly-typed futures to complete.
4///
5/// Awaits multiple futures simultaneously, returning all results once complete.
6///
7/// `try_select!` is similar to [`select!`], but keeps going if a future
8/// resolved to an error until all futures have been resolved. In which case
9/// the error of the last item in the list will be returned.
10///
11/// This macro is only usable inside of async functions, closures, and blocks.
12///
13/// # Examples
14///
15/// ```
16/// # futures::executor::block_on(async {
17/// # async fn main() -> Result<(), std::io::Error> {
18/// use async_macros::try_select;
19/// use futures::future;
20/// use std::io::{Error, ErrorKind};
21///
22/// let a = future::pending::<Result<_, Error>>();
23/// let b = future::ready(Err(Error::from(ErrorKind::Other)));
24/// let c = future::ready(Ok(1u8));
25///
26/// assert_eq!(try_select!(a, b, c).await?, 1u8);
27/// # Ok(())
28/// # }
29/// # main().await.unwrap();
30/// # });
31/// ```
32#[macro_export]
33macro_rules! try_select {
34    ($($fut:ident),+ $(,)?) => { {
35        async {
36            use $crate::utils::future::Future;
37            use $crate::utils::pin::Pin;
38            use $crate::utils::poll_fn;
39            use $crate::utils::result::Result;
40            use $crate::utils::task::Poll;
41
42            $(
43                // Move future into a local so that it is pinned in one place and
44                // is no longer accessible by the end user.
45                let mut $fut = $crate::MaybeDone::new($fut);
46            )*
47
48            let res: Result<_, _> = poll_fn(move |cx| {
49                let mut all_done = true;
50
51                $(
52                    let fut = unsafe { Pin::new_unchecked(&mut $fut) };
53                    if Future::poll(fut, cx).is_ready() {
54                        let fut = Pin::new(&$fut);
55                        if fut.output().unwrap().is_ok() {
56                            let fut = unsafe { Pin::new_unchecked(&mut $fut) };
57                            let res = fut.take().unwrap();
58                            return Poll::Ready(res);
59                        } else {
60                            all_done = false;
61                        }
62                    } else {
63                        all_done = false;
64                    }
65                )*
66
67                if all_done {
68                    // We need to iterate over all items to get the last error.
69                    let mut err = None;
70                    $(
71                        if err.is_none() {
72                            let fut = unsafe { Pin::new_unchecked(&mut $fut) };
73                            err = Some(fut.take().unwrap());
74                        }
75                    )*
76                    return Poll::Ready(err.unwrap());
77                } else {
78                    Poll::Pending
79                }
80            }).await;
81            res
82        }
83    } }
84}