Repository: tokio-rs/tokio
File: tokio/src/macros/try_join.rs (Rust)

Code before:

        // Normalized try_join! branches
        $( ( $($skip:tt)* ) ( $($branch_index:tt)* ) $e:expr, )*

    }) => {{
        use $crate::macros::support::{maybe_done, poll_fn, Future, Pin};
        use $crate::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 = ( $( maybe_done($e), )* );

        // When poll_fn is polled, start polling the future at this index.
        let mut start_index = 0;

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

            let mut is_pending = false;

            let mut turn = start_index;

            for _ in 0..FUTURE_COUNT {
                $(
                  {
                      const INDEX: u32 = $($branch_index)*;
                      if turn == INDEX {
                          // 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.as_mut().poll(cx).is_pending() {
                              is_pending = true;
                          } else if fut.as_mut().output_mut().expect("expected completed future").is_err() {
                              return Ready(Err(fut.take_output().expect("expected completed future").err().unwrap()))
                          }

                          turn = if turn + 1 == FUTURE_COUNT {
                              0
                          } else {
                             turn + 1
                          };

                          continue;
                      }
                  }
                )*

Review comment (suggestion):
This doesn't actually poll every future every time as far as I can tell, like the `loop` approach I suggested does.

Code after:

        // Normalized try_join! branches
        $( ( $($skip:tt)* ) $e:expr, )*

    }) => {{
        use $crate::macros::support::{maybe_done, poll_fn, Future, Pin};
        use $crate::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 = ( $( maybe_done($e), )* );

        // Each time the future created by poll_fn is polled, a different future will be polled first
        // to ensure every future passed to join! gets a chance to make progress even if
        // one of the futures consumes the whole budget.
        //
        // This is number of futures that will be skipped in the first loop
        // iteration the next time.
        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;

            // The number of futures that will be skipped in the first loop iteration
            let mut skip = skip_next_time;

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

            // This loop runs twice and the first `skip` futures
            // are not polled in the first iteration.
            loop {
            $(
                if skip == 0 {
                    if to_run == 0 {
                        // Every future has been polled
                        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