macros: join! start by polling a different future each time poll_fn i…

…s polled

Fixes: tokio-rs#4612

tokio-macros/src/join.rs

@@ -0,0 +1,109 @@
+use proc_macro::TokenStream;
+use quote::quote;
+use syn::parse::{Parse, ParseStream};
+use syn::{Expr, Token};
+
+#[doc(hidden)]
+struct Join {
+    fut_exprs: Vec<Expr>,
+}
+
+impl Parse for Join {
+    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
+        let mut exprs = Vec::new();
+
+        while !input.is_empty() {
+            exprs.push(input.parse::<Expr>()?);
+
+            if !input.is_empty() {
+                input.parse::<Token![,]>()?;
+            }
+        }
+
+        Ok(Join { fut_exprs: exprs })
+    }
+}
+
+pub(crate) fn join(input: TokenStream) -> TokenStream {
+    let parsed = syn::parse_macro_input!(input as Join);
+
+    let futures_count = parsed.fut_exprs.len() as u32;
+
+    let match_statement_branches = (0..futures_count).map(|i| {
+        let pos = syn::Index::from(i as usize);
+
+        quote! {
+            #pos => {
+                let fut = &mut futures.#pos;
+
+                // Safety: future is stored on the stack above
+                // and never moved.
+                let mut fut = unsafe { Pin::new_unchecked(fut) };
+
+                // Try polling
+                if fut.poll(cx).is_pending() {
+                    is_pending = true;
+                }
+            }
+        }
+    });
+
+    let ready_output = (0..futures_count).map(|i| {
+        let pos = syn::Index::from(i as usize);
+
+        quote! {{
+          let fut = &mut futures.#pos;
+
+            // Safety: future is stored on the stack above
+            // and never moved.
+            let mut fut = unsafe { Pin::new_unchecked(fut) };
+
+            fut.take_output().expect("expected completed future")
+        }}
+    });
+
+    let futures = parsed.fut_exprs.into_iter();
+
+    TokenStream::from(quote! {{
+      use tokio::macros::support::{maybe_done, poll_fn, Future, Pin};
+      use tokio::macros::support::Poll::{Ready, Pending};
+
+      // Safety: nothing must be moved out of `futures`. This is to satisfy
+      // the requirement of `Pin::new_unchecked` called below.
+      // let mut futures = #futures;
+      let mut futures = ( #( maybe_done(#futures), )* );
+
+      // When poll_fn is polled, start polling the future at this index.
+      let mut start_index = 0;
+
+      poll_fn(move |cx| {
+          let mut is_pending = false;
+
+          for i in 0..#futures_count {
+              let turn;
+
+              #[allow(clippy::modulo_one)]
+              {
+                  turn = (start_index + i) % #futures_count
+              };
+
+              match turn {
+                #( #match_statement_branches, )*
+                _ => unreachable!("reaching this means there probably is an off by one bug")
+              }
+          }
+
+          if is_pending {
+              // Start by polling the next future first the next time poll_fn is polled
+              #[allow(clippy::modulo_one)]
+              {
+                  start_index = (start_index + 1) % #futures_count;
+              }
+
+              Pending
+          } else {
+             Ready( ( #( #ready_output, )* ) )
+          }
+      }).await
+    }})
+}

tokio-macros/src/lib.rs

@@ -18,6 +18,7 @@
 extern crate proc_macro;
 
 mod entry;
+mod join;
 mod select;
 
 use proc_macro::TokenStream;
@@ -336,3 +337,62 @@ pub fn select_priv_declare_output_enum(input: TokenStream) -> TokenStream {
 pub fn select_priv_clean_pattern(input: TokenStream) -> TokenStream {
     select::clean_pattern_macro(input)
 }
+
+/// Waits on multiple concurrent branches, returning when **all** branches
+/// complete.
+///
+/// The `join!` macro must be used inside of async functions, closures, and
+/// blocks.
+///
+/// The `join!` macro takes a list of async expressions and evaluates them
+/// concurrently on the same task. Each async expression evaluates to a future
+/// and the futures from each expression are multiplexed on the current task.
+///
+/// When working with async expressions returning `Result`, `join!` will wait
+/// for **all** branches complete regardless if any complete with `Err`. Use
+/// [`try_join!`] to return early when `Err` is encountered.
+///
+/// [`try_join!`]: macro.try_join.html
+///
+/// # Notes
+///
+/// The supplied futures are stored inline and does not require allocating a
+/// `Vec`.
+///
+/// ### Runtime characteristics
+///
+/// By running all async expressions on the current task, the expressions are
+/// able to run **concurrently** but not in **parallel**. This means all
+/// expressions are run on the same thread and if one branch blocks the thread,
+/// all other expressions will be unable to continue. If parallelism is
+/// required, spawn each async expression using [`tokio::spawn`] and pass the
+/// join handle to `join!`.
+///
+/// [`tokio::spawn`]: fn.spawn.html
+///
+/// # Examples
+///
+/// Basic join with two branches
+///
+/// ```
+/// async fn do_stuff_async() {
+///     // async work
+/// }
+///
+/// async fn more_async_work() {
+///     // more here
+/// }
+///
+/// #[tokio::main]
+/// async fn main() {
+///     let (first, second) = tokio::join!(
+///         do_stuff_async(),
+///         more_async_work());
+///
+///     // do something with the values
+/// }
+/// ```
+#[proc_macro]
+pub fn join(input: TokenStream) -> TokenStream {
+    join::join(input)
+}

tokio/src/lib.rs

@@ -528,6 +528,8 @@ cfg_macros! {
     #[doc(hidden)]
     pub use tokio_macros::select_priv_clean_pattern;
 
+    pub use tokio_macros::join;
+
     cfg_rt! {
         #[cfg(feature = "rt-multi-thread")]
         #[cfg(not(test))] // Work around for rust-lang/rust#62127

tokio/src/macros/mod.rs

@@ -28,9 +28,6 @@ cfg_macros! {
     #[macro_use]
     mod select;
 
-    #[macro_use]
-    mod join;
-
     #[macro_use]
     mod try_join;
 }

tokio/tests/macros_join.rs

@@ -1,6 +1,8 @@
 #![cfg(feature = "macros")]
 #![allow(clippy::blacklisted_name)]
 
+use std::{sync::Arc, time::Duration};
+
 #[cfg(target_arch = "wasm32")]
 use wasm_bindgen_test::wasm_bindgen_test as test;
 #[cfg(target_arch = "wasm32")]
@@ -9,7 +11,7 @@ use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test;
 #[cfg(not(target_arch = "wasm32"))]
 use tokio::test as maybe_tokio_test;
 
-use tokio::sync::oneshot;
+use tokio::sync::{oneshot, Semaphore};
 use tokio_test::{assert_pending, assert_ready, task};
 
 #[maybe_tokio_test]
@@ -71,12 +73,50 @@ fn join_size() {
         let ready = future::ready(0i32);
         tokio::join!(ready)
     };
-    assert_eq!(mem::size_of_val(&fut), 16);
+    assert_eq!(mem::size_of_val(&fut), 20);
 
     let fut = async {
         let ready1 = future::ready(0i32);
         let ready2 = future::ready(0i32);
         tokio::join!(ready1, ready2)
     };
-    assert_eq!(mem::size_of_val(&fut), 28);
+    assert_eq!(mem::size_of_val(&fut), 32);
+}
+
+async fn non_cooperative_task(permits: Arc<Semaphore>) -> usize {
+    let mut exceeded_budget = 0;
+
+    for _ in 0..5 {
+        // Another task should run after after this task uses its whole budget
+        for _ in 0..128 {
+            let _permit = permits.clone().acquire_owned().await.unwrap();
+        }
+
+        exceeded_budget += 1;
+    }
+
+    exceeded_budget
+}
+
+async fn poor_little_task() -> usize {
+    let mut how_many_times_i_got_to_run = 0;
+
+    for _ in 0..5 {
+        tokio::time::sleep(Duration::from_millis(100)).await;
+        how_many_times_i_got_to_run += 1;
+    }
+
+    how_many_times_i_got_to_run
+}
+
+#[tokio::test]
+async fn join_does_not_allow_tasks_to_starve() {
+    let permits = Arc::new(Semaphore::new(10));
+
+    // non_cooperative_task should yield after its budget is exceeded and then poor_little_task should run.
+    let (non_cooperative_result, little_task_result) =
+        tokio::join!(non_cooperative_task(permits), poor_little_task());
+
+    assert_eq!(5, non_cooperative_result);
+    assert_eq!(5, little_task_result);
 }