macros: join! start by polling a different future each time poll_fn is polled

[PoorlyDefinedBehaviour]

Motivation

Fixes: #4612

Since tokio::join! runs the futures in the same task, the futures were are the same budget which makes it possible for one future to consume the whole budget leaving nothing for the other futures to consume.

Starvation example taken from #4612

    use std::{sync::Arc, time::Duration};
    use tokio::sync::Semaphore;

    #[tokio::main]
    async fn main() {
        let permits = Arc::new(Semaphore::new(10));

        tokio::join!(non_cooperative_task(permits), poor_little_task());
    }

    async fn non_cooperative_task(permits: Arc<Semaphore>) {
        loop {
            let permit = permits.clone().acquire_owned().await.unwrap();
            // uncommenting the following makes it work
            // tokio::time::sleep(Duration::from_millis(1)).await;
        }
    }
    
    async fn poor_little_task() {
        loop {
            tokio::time::sleep(Duration::from_secs(1)).await;
            // This println! never gets to run.
            println!("Hello!")
        }
    }

Solution

Start by polling a different future each time the future generated by poll_ fn is polled. This way, each future will get the chance to make progress even if there are futures that consume the whole budget.

Given futures A, B and C passsed to join!

tokio::join!(A, B, C);

The polling will look like this:

The first time the future generated by poll_fn is polled:

poll A
poll B
poll C

The second time:

poll B
poll C
poll A

The third time:

poll C
poll A
poll B

The fourth time (we are back at the start):

poll A 
poll B
poll C

Cargo expand example

#![feature(prelude_import)]
#[prelude_import]
use std::prelude::rust_2021::*;
#[macro_use]
extern crate std;
use std::{sync::Arc, time::Duration};
use tokio::sync::Semaphore;
fn main() {
    let body = async {
        let permits = Arc::new(Semaphore::new(10));
        {
            use ::tokio::macros::support::{maybe_done, poll_fn, Future, Pin};
            use ::tokio::macros::support::Poll::{Ready, Pending};
            let mut futures = (
                maybe_done(non_cooperative_task(permits)),
                maybe_done(poor_little_task()),
            );
            let mut skip_next_time: u32 = 0;
            poll_fn(move |cx| {
                const COUNT: u32 = 0 + 1 + 1;
                let mut is_pending = false;
                let mut to_run = COUNT;
                let mut skip = skip_next_time;
                skip_next_time = if skip + 1 == COUNT { 0 } else { skip + 1 };
                loop {
                    if skip == 0 {
                        if to_run == 0 {
                            break;
                        }
                        to_run -= 1;
                        let (fut, ..) = &mut futures;
                        let mut fut = unsafe { Pin::new_unchecked(fut) };
                        if fut.poll(cx).is_pending() {
                            is_pending = true;
                        }
                    } else {
                        skip -= 1;
                    }
                    if skip == 0 {
                        if to_run == 0 {
                            break;
                        }
                        to_run -= 1;
                        let (_, fut, ..) = &mut futures;
                        let mut fut = unsafe { Pin::new_unchecked(fut) };
                        if fut.poll(cx).is_pending() {
                            is_pending = true;
                        }
                    } else {
                        skip -= 1;
                    }
                }
                if is_pending {
                    Pending
                } else {
                    Ready((
                        {
                            let (fut, ..) = &mut futures;
                            let mut fut = unsafe { Pin::new_unchecked(fut) };
                            fut.take_output().expect("expected completed future")
                        },
                        {
                            let (_, fut, ..) = &mut futures;
                            let mut fut = unsafe { Pin::new_unchecked(fut) };
                            fut.take_output().expect("expected completed future")
                        },
                    ))
                }
            })
            .await
        };
    };
    #[allow(clippy::expect_used)]
    tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .expect("Failed building the Runtime")
        .block_on(body);
}
async fn non_cooperative_task(permits: Arc<Semaphore>) {
    loop {
        let _permit = permits.clone().acquire_owned().await.unwrap();
    }
}
async fn poor_little_task() {
    loop {
        tokio::time::sleep(Duration::from_secs(1)).await;
        ::std::io::_print(::core::fmt::Arguments::new_v1(&["Hello!\n"], &[]))
    }
}

[Darksonn]

If the join! macro previously worked with more than 64 branches, then this will break any such code. Please double check whether the old version had a branch limit.

[PoorlyDefinedBehaviour]

We can pass 125 branches to the old version without changing the recursion limit 🤔

[PoorlyDefinedBehaviour]

I was trying to make the new join! version accept 125 branches too and kept hitting the macro recursion limit. I think it its because of the way the join! input is normalized.

    // ===== Normalize =====

    (@ { ( $($s:tt)* ) $($t:tt)* } $e:expr, $($r:tt)* ) => {
        $crate::join!(@{ ($($s)* _) $($t)* ($($s)*) $e, } $($r)*)
    };

    // ===== Entry point =====

    ( $($e:expr),* $(,)?) => {
        $crate::join!(@{ () } $($e,)*)
    };

I implemented join! as a proc macro, it was easier to avoid the recursion limit this way. Would a proc macro be ok? (we are accepting way more branches at the moment)

[Darksonn]

I would prefer to avoid a procedural macro if possible. In general, I recommend you ask about an implementation strategy first before implementing it, especially if it's a larger change.

[PoorlyDefinedBehaviour]

I recommend you ask about an implementation strategy first before implementing it, especially if it's a larger change

I think i read that in the CONTRIBUTING.md file, i will make sure to ask first next time. I chose to implement a proc macro without asking because i can use the practice if it is not accepted.

I went back to the version using just a declarative macro and was playing around with it the last couple of days but i'm pretty much stuck at the moment.

I keep hitting the recursion limit when i pass 125 futures to join!(the maximum supported without changing the recursion limit at the moment) because after join! normalizes its input, it is 126 levels deep in macro recursion and i need to call a macro (count!) once again.

Any ideas i could try?

[Darksonn]

How about having the normalization count the number of branches?

         // One `_` for each branch in the `join!` macro. This is not used once
         // normalization is complete.
         ( $($count:tt)* )
+        // The expression `0+1+1+ ... +1` equal to the number of branches.
+        ( $($total:tt)* )
 
         // Normalized join! branches
         $( ( $($skip:tt)* ) $e:expr, )*

     // ===== Normalize =====
 
-    (@ { ( $($s:tt)* ) $($t:tt)* } $e:expr, $($r:tt)* ) => {
-        $crate::join!(@{ ($($s)* _) $($t)* ($($s)*) $e, } $($r)*)
+    (@ { ( $($s:tt)* ) ( $($n:tt)* ) $($t:tt)* } $e:expr, $($r:tt)* ) => {
+        $crate::join!(@{ ($($s)* _) ($($n)* + 1) $($t)* ($($s)*) $e, } $($r)*)
     };
 
     // ===== Entry point =====
 
     ( $($e:expr),* $(,)?) => {
-        $crate::join!(@{ () } $($e,)*)
+        $crate::join!(@{ () (0) } $($e,)*)
     };
 }

Then you could use the count in this way:

// Safety: nothing must be moved out of `futures`. This is to satisfy
// the requirement of `Pin::new_unchecked` called below.
let mut futures = ( $( maybe_done($e), )* );

let mut skip_next_time: u32 = 0;

poll_fn(move |cx| {
    const COUNT: u32 = $($total)*;

    let mut is_pending = false;

    let mut to_run = COUNT;
    let mut skip = skip_next_time;
    skip_next_time = if skip+1 == COUNT { 0 } else { skip+1 };

    loop {
    $(
        if skip == 0 {
            if to_run == 0 {
                break;
            }
            to_run -= 1;

            // Extract the future for this branch from the tuple.
            let ( $($skip,)* fut, .. ) = &mut futures;

            // 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;
            }
        } else {
            skip -= 1;
        }
    )*
    }

    ...
}).await

Also, I think the same issue exists for the try_join! macro, so if the above works, it would be good to fix that one too.

[PoorlyDefinedBehaviour]

It works 👍

I should have thought about that 🤦

[Darksonn]

This doesn't actually poll every future every time as far as I can tell, like the loop approach I suggested does.

[Darksonn]

Maybe you could also add a test for this?

[PoorlyDefinedBehaviour]

Added a test, i think the futures are getting polled in the expected order. Let me know if i am missing something.

I can switch to the approach you suggested if you prefer.

I added the cargo expand output to the PR description if you want to take a look.

[Darksonn]

Ah, I misread it. Anyway, the current loop is quite wasteful because you go through every future with a smaller index every time you poll something.

[PoorlyDefinedBehaviour]

Got it.

[Darksonn]

If you want to use time in tests, then you need to use Tokio's pause-time feature. Having a test that takes 100ms to run is way too long. Our test suite would take forever if we had many tests like this.

[PoorlyDefinedBehaviour]

Stopped using tokio::time

[Darksonn]

In general, I think this looks quite good! Are you able to add a comment on the skip variables that explain why they exist? (in both the join! and try_join! file)

[Darksonn]

Looks good to me.