From 8e999e380643160f0e944fd5708de2bd3d1b0d51 Mon Sep 17 00:00:00 2001 From: Jess Izen <44884346+jlizen@users.noreply.github.com> Date: Mon, 9 Jun 2025 00:31:05 -0700 Subject: [PATCH] macros: add biased mode to join! and try_join! (#7307) --- tokio/src/macros/join.rs | 113 +++++++++++++++++++++++++++------ tokio/src/macros/support.rs | 2 + tokio/src/macros/try_join.rs | 91 ++++++++++++++++++++------ tokio/tests/macros_join.rs | 60 +++++++++++++++++ tokio/tests/macros_try_join.rs | 72 +++++++++++++++++++-- 5 files changed, 297 insertions(+), 41 deletions(-) diff --git a/tokio/src/macros/join.rs b/tokio/src/macros/join.rs index 163053b25..cd89da1cc 100644 --- a/tokio/src/macros/join.rs +++ b/tokio/src/macros/join.rs @@ -21,7 +21,7 @@ macro_rules! doc { /// The supplied futures are stored inline and do not require allocating a /// `Vec`. /// - /// ### Runtime characteristics + /// ## 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 @@ -32,6 +32,25 @@ macro_rules! doc { /// /// [`tokio::spawn`]: crate::spawn /// + /// ## Fairness + /// + /// By default, `join!`'s generated future rotates which contained + /// future is polled first whenever it is woken. + /// + /// This behavior can be overridden by adding `biased;` to the beginning of the + /// macro usage. See the examples for details. This will cause `join` to poll + /// the futures in the order they appear from top to bottom. + /// + /// You may want this if your futures may interact in a way where known polling order is significant. + /// + /// But there is an important caveat to this mode. It becomes your responsibility + /// to ensure that the polling order of your futures is fair. If for example you + /// are joining a stream and a shutdown future, and the stream has a + /// huge volume of messages that takes a long time to finish processing per poll, you should + /// place the shutdown future earlier in the `join!` list to ensure that it is + /// always polled, and will not be delayed due to the stream future taking a long time to return + /// `Poll::Pending`. + /// /// # Examples /// /// Basic join with two branches @@ -54,6 +73,30 @@ macro_rules! doc { /// // do something with the values /// } /// ``` + /// + /// Using the `biased;` mode to control polling order. + /// + /// ``` + /// async fn do_stuff_async() { + /// // async work + /// } + /// + /// async fn more_async_work() { + /// // more here + /// } + /// + /// #[tokio::main] + /// async fn main() { + /// let (first, second) = tokio::join!( + /// biased; + /// do_stuff_async(), + /// more_async_work() + /// ); + /// + /// // do something with the values + /// } + /// ``` + #[macro_export] #[cfg_attr(docsrs, doc(cfg(feature = "macros")))] $join @@ -62,12 +105,16 @@ macro_rules! doc { #[cfg(doc)] doc! {macro_rules! join { - ($($future:expr),*) => { unimplemented!() } + ($(biased;)? $($future:expr),*) => { unimplemented!() } }} #[cfg(not(doc))] doc! {macro_rules! join { (@ { + // Type of rotator that controls which inner future to start with + // when polling our output future. + rotator=$rotator:ty; + // One `_` for each branch in the `join!` macro. This is not used once // normalization is complete. ( $($count:tt)* ) @@ -96,25 +143,19 @@ doc! {macro_rules! join { // let mut futures = &mut futures; - // 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; + const COUNT: u32 = $($total)*; + + // Each time the future created by poll_fn is polled, if not using biased mode, + // a different future is polled first to ensure every future passed to join! + // can make progress even if one of the futures consumes the whole budget. + let mut rotator = <$rotator>::default(); 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 }; + let mut skip = rotator.num_skip(); // This loop runs twice and the first `skip` futures // are not polled in the first iteration. @@ -164,15 +205,51 @@ doc! {macro_rules! join { // ===== Normalize ===== - (@ { ( $($s:tt)* ) ( $($n:tt)* ) $($t:tt)* } $e:expr, $($r:tt)* ) => { - $crate::join!(@{ ($($s)* _) ($($n)* + 1) $($t)* ($($s)*) $e, } $($r)*) + (@ { rotator=$rotator:ty; ( $($s:tt)* ) ( $($n:tt)* ) $($t:tt)* } $e:expr, $($r:tt)* ) => { + $crate::join!(@{ rotator=$rotator; ($($s)* _) ($($n)* + 1) $($t)* ($($s)*) $e, } $($r)*) }; // ===== Entry point ===== + ( biased; $($e:expr),+ $(,)?) => { + $crate::join!(@{ rotator=$crate::macros::support::BiasedRotator; () (0) } $($e,)*) + }; ( $($e:expr),+ $(,)?) => { - $crate::join!(@{ () (0) } $($e,)*) + $crate::join!(@{ rotator=$crate::macros::support::Rotator; () (0) } $($e,)*) }; + (biased;) => { async {}.await }; + () => { async {}.await } }} + +/// Rotates by one each [`Self::num_skip`] call up to COUNT - 1. +#[derive(Default, Debug)] +pub struct Rotator { + next: u32, +} + +impl Rotator { + /// Rotates by one each [`Self::num_skip`] call up to COUNT - 1 + #[inline] + pub fn num_skip(&mut self) -> u32 { + let num_skip = self.next; + self.next += 1; + if self.next == COUNT { + self.next = 0; + } + num_skip + } +} + +/// [`Self::num_skip`] always returns 0. +#[derive(Default, Debug)] +pub struct BiasedRotator {} + +impl BiasedRotator { + /// Always returns 0. + #[inline] + pub fn num_skip(&mut self) -> u32 { + 0 + } +} diff --git a/tokio/src/macros/support.rs b/tokio/src/macros/support.rs index ff3ccb406..213d85cde 100644 --- a/tokio/src/macros/support.rs +++ b/tokio/src/macros/support.rs @@ -3,6 +3,8 @@ cfg_macros! { pub use std::future::poll_fn; + pub use crate::macros::join::{BiasedRotator, Rotator}; + #[doc(hidden)] pub fn thread_rng_n(n: u32) -> u32 { crate::runtime::context::thread_rng_n(n) diff --git a/tokio/src/macros/try_join.rs b/tokio/src/macros/try_join.rs index fd862f072..471b5ebb4 100644 --- a/tokio/src/macros/try_join.rs +++ b/tokio/src/macros/try_join.rs @@ -19,7 +19,7 @@ macro_rules! doc { /// The supplied futures are stored inline and do not require allocating a /// `Vec`. /// - /// ### Runtime characteristics + /// ## 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 @@ -30,6 +30,25 @@ macro_rules! doc { /// /// [`tokio::spawn`]: crate::spawn /// + /// ## Fairness + /// + /// By default, `try_join!`'s generated future rotates which + /// contained future is polled first whenever it is woken. + /// + /// This behavior can be overridden by adding `biased;` to the beginning of the + /// macro usage. See the examples for details. This will cause `try_join` to poll + /// the futures in the order they appear from top to bottom. + /// + /// You may want this if your futures may interact in a way where known polling order is significant. + /// + /// But there is an important caveat to this mode. It becomes your responsibility + /// to ensure that the polling order of your futures is fair. If for example you + /// are joining a stream and a shutdown future, and the stream has a + /// huge volume of messages that takes a long time to finish processing per poll, you should + /// place the shutdown future earlier in the `try_join!` list to ensure that it is + /// always polled, and will not be delayed due to the stream future taking a long time to return + /// `Poll::Pending`. + /// /// # Examples /// /// Basic `try_join` with two branches. @@ -100,6 +119,37 @@ macro_rules! doc { /// } /// } /// ``` + /// Using the `biased;` mode to control polling order. + /// + /// ``` + /// async fn do_stuff_async() -> Result<(), &'static str> { + /// // async work + /// # Ok(()) + /// } + /// + /// async fn more_async_work() -> Result<(), &'static str> { + /// // more here + /// # Ok(()) + /// } + /// + /// #[tokio::main] + /// async fn main() { + /// let res = tokio::try_join!( + /// biased; + /// do_stuff_async(), + /// more_async_work() + /// ); + /// + /// match res { + /// Ok((first, second)) => { + /// // do something with the values + /// } + /// Err(err) => { + /// println!("processing failed; error = {}", err); + /// } + /// } + /// } + /// ``` #[macro_export] #[cfg_attr(docsrs, doc(cfg(feature = "macros")))] $try_join @@ -108,12 +158,16 @@ macro_rules! doc { #[cfg(doc)] doc! {macro_rules! try_join { - ($($future:expr),*) => { unimplemented!() } + ($(biased;)? $($future:expr),*) => { unimplemented!() } }} #[cfg(not(doc))] doc! {macro_rules! try_join { (@ { + // Type of rotator that controls which inner future to start with + // when polling our output future. + rotator=$rotator:ty; + // One `_` for each branch in the `try_join!` macro. This is not used once // normalization is complete. ( $($count:tt)* ) @@ -142,25 +196,19 @@ doc! {macro_rules! try_join { // let mut futures = &mut futures; - // 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; + const COUNT: u32 = $($total)*; + + // Each time the future created by poll_fn is polled, if not using biased mode, + // a different future is polled first to ensure every future passed to try_join! + // can make progress even if one of the futures consumes the whole budget. + let mut rotator = <$rotator>::default(); 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 }; + // The number of futures that will be skipped in the first loop iteration. + let mut skip = rotator.num_skip(); // This loop runs twice and the first `skip` futures // are not polled in the first iteration. @@ -216,15 +264,20 @@ doc! {macro_rules! try_join { // ===== Normalize ===== - (@ { ( $($s:tt)* ) ( $($n:tt)* ) $($t:tt)* } $e:expr, $($r:tt)* ) => { - $crate::try_join!(@{ ($($s)* _) ($($n)* + 1) $($t)* ($($s)*) $e, } $($r)*) + (@ { rotator=$rotator:ty; ( $($s:tt)* ) ( $($n:tt)* ) $($t:tt)* } $e:expr, $($r:tt)* ) => { + $crate::try_join!(@{ rotator=$rotator; ($($s)* _) ($($n)* + 1) $($t)* ($($s)*) $e, } $($r)*) }; // ===== Entry point ===== + ( biased; $($e:expr),+ $(,)?) => { + $crate::try_join!(@{ rotator=$crate::macros::support::BiasedRotator; () (0) } $($e,)*) + }; ( $($e:expr),+ $(,)?) => { - $crate::try_join!(@{ () (0) } $($e,)*) + $crate::try_join!(@{ rotator=$crate::macros::support::Rotator; () (0) } $($e,)*) }; + (biased;) => { async { Ok(()) }.await }; + () => { async { Ok(()) }.await } }} diff --git a/tokio/tests/macros_join.rs b/tokio/tests/macros_join.rs index 083eecf29..fd4fdae3a 100644 --- a/tokio/tests/macros_join.rs +++ b/tokio/tests/macros_join.rs @@ -17,28 +17,36 @@ use tokio_test::{assert_pending, assert_ready, task}; #[maybe_tokio_test] async fn sync_one_lit_expr_comma() { let foo = tokio::join!(async { 1 },); + assert_eq!(foo, (1,)); + let foo = tokio::join!(biased; async { 1 },); assert_eq!(foo, (1,)); } #[maybe_tokio_test] async fn sync_one_lit_expr_no_comma() { let foo = tokio::join!(async { 1 }); + assert_eq!(foo, (1,)); + let foo = tokio::join!(biased; async { 1 }); assert_eq!(foo, (1,)); } #[maybe_tokio_test] async fn sync_two_lit_expr_comma() { let foo = tokio::join!(async { 1 }, async { 2 },); + assert_eq!(foo, (1, 2)); + let foo = tokio::join!(biased; async { 1 }, async { 2 },); assert_eq!(foo, (1, 2)); } #[maybe_tokio_test] async fn sync_two_lit_expr_no_comma() { let foo = tokio::join!(async { 1 }, async { 2 }); + assert_eq!(foo, (1, 2)); + let foo = tokio::join!(biased; async { 1 }, async { 2 }); assert_eq!(foo, (1, 2)); } @@ -154,10 +162,62 @@ async fn a_different_future_is_polled_first_every_time_poll_fn_is_polled() { ); } +#[tokio::test] +async fn futures_are_polled_in_order_in_biased_mode() { + let poll_order = Arc::new(std::sync::Mutex::new(vec![])); + + let fut = |x, poll_order: Arc>>| async move { + for _ in 0..4 { + { + let mut guard = poll_order.lock().unwrap(); + + guard.push(x); + } + + tokio::task::yield_now().await; + } + }; + + tokio::join!( + biased; + fut(1, Arc::clone(&poll_order)), + fut(2, Arc::clone(&poll_order)), + fut(3, Arc::clone(&poll_order)), + ); + + // Each time the future created by join! is polled, it should start + // by polling in the order as declared in the macro inputs. + assert_eq!( + vec![1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3], + *poll_order.lock().unwrap() + ); +} + +#[test] +#[cfg(target_pointer_width = "64")] +fn join_size_biased() { + use futures::future; + use std::mem; + + let fut = async { + let ready = future::ready(0i32); + tokio::join!(biased; ready) + }; + assert_eq!(mem::size_of_val(&fut), 24); + + let fut = async { + let ready1 = future::ready(0i32); + let ready2 = future::ready(0i32); + tokio::join!(biased; ready1, ready2) + }; + assert_eq!(mem::size_of_val(&fut), 40); +} + #[tokio::test] #[allow(clippy::unit_cmp)] async fn empty_join() { assert_eq!(tokio::join!(), ()); + assert_eq!(tokio::join!(biased;), ()); } #[tokio::test] diff --git a/tokio/tests/macros_try_join.rs b/tokio/tests/macros_try_join.rs index 76958f167..e68c3400f 100644 --- a/tokio/tests/macros_try_join.rs +++ b/tokio/tests/macros_try_join.rs @@ -1,7 +1,7 @@ #![cfg(feature = "macros")] #![allow(clippy::disallowed_names)] -use std::sync::Arc; +use std::{convert::Infallible, sync::Arc}; use tokio::sync::{oneshot, Semaphore}; use tokio_test::{assert_pending, assert_ready, task}; @@ -15,28 +15,36 @@ use tokio::test as maybe_tokio_test; #[maybe_tokio_test] async fn sync_one_lit_expr_comma() { let foo = tokio::try_join!(async { ok(1) },); + assert_eq!(foo, Ok((1,))); + let foo = tokio::try_join!(biased; async { ok(1) },); assert_eq!(foo, Ok((1,))); } #[maybe_tokio_test] async fn sync_one_lit_expr_no_comma() { let foo = tokio::try_join!(async { ok(1) }); + assert_eq!(foo, Ok((1,))); + let foo = tokio::try_join!(biased; async { ok(1) }); assert_eq!(foo, Ok((1,))); } #[maybe_tokio_test] async fn sync_two_lit_expr_comma() { let foo = tokio::try_join!(async { ok(1) }, async { ok(2) },); + assert_eq!(foo, Ok((1, 2))); + let foo = tokio::try_join!(biased;async { ok(1) }, async { ok(2) },); assert_eq!(foo, Ok((1, 2))); } #[maybe_tokio_test] async fn sync_two_lit_expr_no_comma() { let foo = tokio::try_join!(async { ok(1) }, async { ok(2) }); + assert_eq!(foo, Ok((1, 2))); + let foo = tokio::try_join!(biased; async { ok(1) }, async { ok(2) }); assert_eq!(foo, Ok((1, 2))); } @@ -84,7 +92,7 @@ async fn err_abort_early() { #[test] #[cfg(target_pointer_width = "64")] -fn join_size() { +fn try_join_size() { use futures::future; use std::mem; @@ -163,13 +171,15 @@ async fn a_different_future_is_polled_first_every_time_poll_fn_is_polled() { tokio::task::yield_now().await; } + Ok::<(), Infallible>(()) }; - tokio::join!( + tokio::try_join!( fut(1, Arc::clone(&poll_order)), fut(2, Arc::clone(&poll_order)), fut(3, Arc::clone(&poll_order)), - ); + ) + .unwrap(); // Each time the future created by join! is polled, it should start // by polling a different future first. @@ -179,7 +189,61 @@ async fn a_different_future_is_polled_first_every_time_poll_fn_is_polled() { ); } +#[tokio::test] +async fn futures_are_polled_in_order_in_biased_mode() { + let poll_order = Arc::new(std::sync::Mutex::new(vec![])); + + let fut = |x, poll_order: Arc>>| async move { + for _ in 0..4 { + { + let mut guard = poll_order.lock().unwrap(); + + guard.push(x); + } + + tokio::task::yield_now().await; + } + Ok::<(), Infallible>(()) + }; + + tokio::try_join!( + biased; + fut(1, Arc::clone(&poll_order)), + fut(2, Arc::clone(&poll_order)), + fut(3, Arc::clone(&poll_order)), + ) + .unwrap(); + + // Each time the future created by join! is polled, it should start + // by polling in the order as declared in the macro inputs. + assert_eq!( + vec![1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3], + *poll_order.lock().unwrap() + ); +} + +#[test] +#[cfg(target_pointer_width = "64")] +fn try_join_size_biased() { + use futures::future; + use std::mem; + + let fut = async { + let ready = future::ready(ok(0i32)); + tokio::try_join!(biased; ready) + }; + assert_eq!(mem::size_of_val(&fut), 24); + + let fut = async { + let ready1 = future::ready(ok(0i32)); + let ready2 = future::ready(ok(0i32)); + tokio::try_join!(biased; ready1, ready2) + }; + assert_eq!(mem::size_of_val(&fut), 40); +} + #[tokio::test] async fn empty_try_join() { assert_eq!(tokio::try_join!() as Result<_, ()>, Ok(())); + assert_eq!(tokio::try_join!(biased;) as Result<_, ()>, Ok(())); }