tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-08-24 00:00:11 +02:00
Code Issues Packages Projects Releases Wiki Activity

macros: add biased mode to join! and try_join! (#7307)

Browse Source
This commit is contained in:
Jess Izen
2025-06-09 09:31:05 +02:00
committed by GitHub
parent 1d980145cb
commit 8e999e3806
5 changed files with 297 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio/src/macros/join.rs
+95 -18
View File
@@ -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 {
// <https://internals.rust-lang.org/t/surprising-soundness-trouble-around-pollfn/17484>
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<COUNT>; () (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<const COUNT: u32> {
next: u32,
}
impl<const COUNT: u32> Rotator<COUNT> {
/// 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
}
}
tokio/src/macros/support.rs
+2
View File
@@ -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)
tokio/src/macros/try_join.rs
+72 -19
View File
@@ -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 {
// <https://internals.rust-lang.org/t/surprising-soundness-trouble-around-pollfn/17484>
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<COUNT>; () (0) } $($e,)*)
};
(biased;) => { async { Ok(()) }.await };
() => { async { Ok(()) }.await }
}}
tokio/tests/macros_join.rs
+60
View File
@@ -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<std::sync::Mutex<Vec<i32>>>| 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]
tokio/tests/macros_try_join.rs
+68 -4
View File
@@ -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<std::sync::Mutex<Vec<i32>>>| 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(()));
}