monoio/macros/
join.rs

1/// Wait on multiple concurrent branches, returning when **all** branches
2/// complete.
3///
4/// The `join!` macro must be used inside of async functions, closures, and
5/// blocks.
6///
7/// The `join!` macro takes a list of async expressions and evaluates them
8/// concurrently on the same task. Each async expression evaluates to a future
9/// and the futures from each expression are multiplexed on the current task.
10///
11/// When working with async expressions returning `Result`, `join!` will wait
12/// for **all** branches complete regardless if any complete with `Err`. Use
13/// [`try_join!`] to return early when `Err` is encountered.
14///
15/// [`try_join!`]: macro@try_join
16///
17/// # Notes
18///
19/// The supplied futures are stored inline and does not require allocating a
20/// `Vec`.
21///
22/// ### Runtime characteristics
23///
24/// By running all async expressions on the current task, the expressions are
25/// able to run **concurrently** but not in **parallel**. This means all
26/// expressions are run on the same thread and if one branch blocks the thread,
27/// all other expressions will be unable to continue. If parallelism is
28/// required, spawn each async expression using [`monoio::spawn`] and pass the
29/// join handle to `join!`.
30///
31/// [`monoio::spawn`]: crate::spawn
32///
33/// # Examples
34///
35/// Basic join with two branches
36///
37/// ```
38/// async fn do_stuff_async() {
39///     // async work
40/// }
41///
42/// async fn more_async_work() {
43///     // more here
44/// }
45///
46/// #[monoio::main]
47/// async fn main() {
48///     let (first, second) = monoio::join!(do_stuff_async(), more_async_work());
49///
50///     // do something with the values
51/// }
52/// ```
53#[macro_export]
54#[cfg_attr(docsrs, doc(cfg(feature = "macros")))]
55macro_rules! join {
56    (@ {
57        // One `_` for each branch in the `join!` macro. This is not used once
58        // normalization is complete.
59        ( $($count:tt)* )
60
61        // Normalized join! branches
62        $( ( $($skip:tt)* ) $e:expr, )*
63
64    }) => {{
65        use $crate::macros::support::{maybe_done, poll_fn, Future, Pin};
66        use $crate::macros::support::Poll::{Ready, Pending};
67
68        // Safety: nothing must be moved out of `futures`. This is to satisfy
69        // the requirement of `Pin::new_unchecked` called below.
70        let mut futures = ( $( maybe_done($e), )* );
71
72        poll_fn(move |cx| {
73            let mut is_pending = false;
74
75            $(
76                // Extract the future for this branch from the tuple.
77                let ( $($skip,)* fut, .. ) = &mut futures;
78
79                // Safety: future is stored on the stack above
80                // and never moved.
81                let fut = unsafe { Pin::new_unchecked(fut) };
82
83                // Try polling
84                if fut.poll(cx).is_pending() {
85                    is_pending = true;
86                }
87            )*
88
89            if is_pending {
90                Pending
91            } else {
92                Ready(($({
93                    // Extract the future for this branch from the tuple.
94                    let ( $($skip,)* fut, .. ) = &mut futures;
95
96                    // Safety: future is stored on the stack above
97                    // and never moved.
98                    let fut = unsafe { Pin::new_unchecked(fut) };
99
100                    fut.take_output().expect("expected completed future")
101                },)*))
102            }
103        }).await
104    }};
105
106    // ===== Normalize =====
107
108    (@ { ( $($s:tt)* ) $($t:tt)* } $e:expr, $($r:tt)* ) => {
109        $crate::join!(@{ ($($s)* _) $($t)* ($($s)*) $e, } $($r)*)
110    };
111
112    // ===== Entry point =====
113
114    ( $($e:expr),* $(,)?) => {
115        $crate::join!(@{ () } $($e,)*)
116    };
117}