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

jlizen · 2025-05-05T18:59:21Z

Closes: #7304

Adding a biased; mode to the join! and try_join! macros that always polls the wrapped futures in their order of declaration, rather than rotating each poll.

This parallels the select! API (though is simpler under the hood, since these macros rotate rather than randomize which branch is polled first).

Motivation

This is useful when the contained futures interact in a way where polling order is significant.

My specific case is that I want to join! or try_join! two separate listener futures in an AWS Lambda. Lambdas receive one request at a time. One of my listeners is my primary function handler, which is higher priority. The second is an extension that does post-request 'telemetry flushing'. So, there is never a need for the second future to be polled first.

Here is a rough example: ref.

Solution

I added an optional biased; token to each macro's entrypoint. This resolves to a literal variable, $rotate_poll_order, as either true or false (with no biased; declaration resulting in true).

Then, inside the generated code, we only increment the skip counter used to rotate polling, if $rotate_poll_order is true.

I initially changed the skip order counter from u32 to Option<u32>, but that would have resulted in a bit muddier code. Adding an if clause around the single line incrementing the skip counter, was more readable and kept the footprint of this change more contained.

The performance impact of this change on cases that do NOT use the biased mode is:

one additional boolean stored on the stack
one extra branch every poll (checking the boolean), which branch prediction should handle cleanly

This perf impact seems very marginal, but we could always make a more complicated macro that generates distinct futures based on bias mode. Probably the external API should stay the same as select! though (ie, just add the biased; token to the existing macros).

jlizen · 2025-05-05T19:02:09Z

Here's the PR related to our discussion yesterday @maminrayej - thanks for the guidance!

tokio/tests/macros_join.rs

Darksonn · 2025-05-06T08:33:12Z

tokio/src/macros/try_join.rs

        let mut skip_next_time: u32 = 0;
+        let rotate_skip = $rotate_poll_order;


I would like to avoid having this variable at all when using biased. I don't think it can be optimized out because it gets stored as a field in the closure type.

My newest commit removes this variable, along with skip_next_time, from the biased branch's poll_fn() (so it will now be smaller than the unbiased equivalent).

It's a bit verbose... since each branch of match $rotate_poll_order {} needs to generate its own unnameable poll function. I initially was pulling out common logic into a closure, but I felt it was ultimately less readable.

Glad to take another stab at simplifying if you have any feedback.

Yeah, I don't think we want the match. It duplicates the entire logic, which isn't nice. One option could be to define two helper types:

#[derive(Default)] pub struct Rotater<const COUNT: usize> { next: usize, } impl<const COUNT: usize> Rotater<COUNT> { #[inline] pub fn num_skip(&mut self) -> usize { let mut num_skip = self.next; self.next += 1; if self.next == COUNT { self.next = 0; } num_skip } } #[derive(Default)] pub struct BiasedRotater {} impl BiasedRotater { #[inline] pub fn num_skip(&mut self) -> usize { 0 } }

and then have the macro select which type to use:

( biased; $($e:expr),+ $(,)?) => { $crate::join!(@{ rotator=$crate::macros::support::BiasedRotator; () (0) } $($e,)*) }; ( $($e:expr),+ $(,)?) => { $crate::join!(@{ rotator=$crate::macros::support::Rotator; () (0) } $($e,)*) };

Ah, that's clever. Thanks for the advice! Will take a crack at it.

Got it working, thanks!

tokio/src/macros/join.rs

tokio/tests/macros_try_join.rs

tokio/tests/macros_join.rs

jlizen · 2025-05-06T14:24:05Z

Appreciate the eyes, Darksonn@ - I'll take a stab at implementing the requested feedback.

jlizen · 2025-05-08T20:02:17Z

@Darksonn latest refactor is up.

I noticed our existing try_join! poll sequencing test was actually using join!, I assume by accident. Updated that as well.

Darksonn · 2025-05-09T09:47:20Z

tokio/src/macros/support.rs

@@ -28,3 +28,34 @@ cfg_macros! {
 pub use std::future::{Future, IntoFuture};
 pub use std::pin::Pin;
 pub use std::task::{Context, Poll};
+
+#[doc(hidden)]


The entire module is #[doc(hidden)] already so you don't need this.

Thanks, will update.

Darksonn · 2025-05-09T09:48:09Z

tokio/src/macros/support.rs

+
+#[doc(hidden)]
+#[derive(Default, Debug)]
+pub struct Rotator<const COUNT: u32> {


Can we make it clear that these types are here for use with the join! macro? In fact, maybe move them to src/macros/join.rs and re-export them here?

Sure, makes sense. Will do.

Darksonn · 2025-05-09T09:49:05Z

tokio/src/macros/try_join.rs

+        ///     let res = tokio::try_join!(
+        ///         //do_stuff_async() will always be polled first when woken
+        ///         biased;
+        ///         do_stuff_async(),
+        ///         more_async_work());


It's confusing that you have a comment inside the macro before biased;. Also, please format this code properly by moving the ); to its own line.

Will do on both counts. Sorry, was overly reliant on rustfmt here, will manually clean up.

Darksonn · 2025-05-09T09:49:58Z

tokio/src/macros/try_join.rs

            loop {
-            $(
-                if skip == 0 {
-                    if to_run == 0 {
-                        // Every future has been polled
-                        break;
-                    }
-                    to_run -= 1;
+                $(
+                    if skip == 0 {
+                        if to_run == 0 {
+                            // Every future has been polled
+                            break;
+                        }
+                        to_run -= 1;



Why is the indentation changed throughout the entire macro?

Will tweak, was overly reliant on rustfmt which has trouble with macros I guess.

…expand unit tests, add new token to macro docs

jlizen · 2025-05-09T16:27:02Z

tokio/tests/macros_try_join.rs

    };

-    tokio::join!(
+    tokio::try_join!(


this was a mistake in existing test, i think

…emove comment inside macro invocation in example

jlizen · 2025-05-09T18:33:56Z

@Darksonn latest changes are up per your comments, thanks for your time on this PR review.

Darksonn · 2025-05-10T15:25:33Z

tokio/src/macros/support.rs

+    // Used with `join!` and `try_join!`
+    pub use crate::macros::join::Rotator;
+    // Used with `join!` and `try_join!`
+    pub use crate::macros::join::BiasedRotator;


Suggested change

// Used with `join!` and `try_join!`

pub use crate::macros::join::Rotator;

// Used with `join!` and `try_join!`

pub use crate::macros::join::BiasedRotator;

// Used with `join!` and `try_join!`

pub use crate::macros::join::{BiasedRotator, Rotator};

Darksonn · 2025-05-10T15:25:49Z

tokio/src/macros/support.rs

+
+
 }


Unnecessary empty lines.

Darksonn · 2025-05-19T21:55:47Z

tokio/src/macros/join.rs

+#[derive(Default, Debug)]
+/// Rotates by one every [`Self::num_skip`] call up to COUNT - 1.
+pub struct Rotator<const COUNT: u32> {


Suggested change

#[derive(Default, Debug)]

/// Rotates by one every [`Self::num_skip`] call up to COUNT - 1.

pub struct Rotator<const COUNT: u32> {

/// Rotates by one every [`Self::num_skip`] call up to COUNT - 1.

#[derive(Default, Debug)]

pub struct Rotator<const COUNT: u32> {

Darksonn · 2025-05-19T22:02:55Z

tokio/src/macros/join.rs

+        /// 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`.


I don't think that a shutdown future is an amazing example here. That makes sense for select!, but you wouldn't really have a shutdown future in join!. Not sure what a better example would be, though.

Here is a real-world example usage involving a shutdown signal: https://github.com/awslabs/aws-lambda-rust-runtime/blob/main/lambda-runtime/src/lib.rs#L227

The distinction here is that the other task is a server future, not a stream future, meaning it will generally keep running and generating new request futures. So we need to continue driving it. Even if the shutdown future will only ever come up once and then resolve, we still need to check it first every time.

I guess technically the server future will generally be dispatching work to read from sockets rather than processing the messages directly (ie not long polls as a concern), but anyway you want to shut down before dispatching that work.

Anyway I'll tweak wording from stream -> server and finesse it a bit more.

Also renaming the section Fairness -> Poll Ordering, it's not really about fairness anymore in the example given.

Actually, I changed it back. My initial rework was basically restating that deterministic ordering is useful. But it's important for readers to consider the case that one future might take a long time to poll and starve the other, that was the point of the fairness discussion.

Based on the above lambda runtime link, I do think there are cases where you would be joining a shutdown future, so it's probably still relevant? Anyway I think it does a good job of illustrating the fairness risk even if the specific use case isn't the most universal?

Glad to keep thinking about it though if not sold and try to come up with another example.

…ng/import organization

jlizen force-pushed the master branch from 97522c0 to 2c27422 Compare May 5, 2025 19:00

jlizen force-pushed the master branch from 2c27422 to f8a83e4 Compare May 5, 2025 19:04

Darksonn added A-tokio Area: The main tokio crate M-macros Module: macros in the main Tokio crate labels May 6, 2025

Darksonn reviewed May 6, 2025

View reviewed changes

tokio/tests/macros_join.rs Outdated Show resolved Hide resolved

jlizen requested a review from Darksonn May 7, 2025 21:04

Darksonn reviewed May 9, 2025

View reviewed changes

jlizen added 4 commits May 9, 2025 16:16

feat(tokio/macros): add biased mode to join! and try_join!

1b574ea

pr feedback: avoid allocating unused vars in biased-derived poll_fn, …

eccc5bb

…expand unit tests, add new token to macro docs

rustfmt

39380de

refactor rotation into standalone type, fix try_join sequencing test

3cc5e26

jlizen force-pushed the master branch from 071b4a4 to d138825 Compare May 9, 2025 16:26

jlizen commented May 9, 2025

View reviewed changes

tokio/tests/macros_try_join.rs

};

tokio::join!(

tokio::try_join!(

Copy link

Member Author

jlizen May 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this was a mistake in existing test, i think

jlizen force-pushed the master branch from 641b69d to 2944509 Compare May 9, 2025 16:50

move rotator types into /macros/join and re-export, fix formatting, r…

367d7b7

…emove comment inside macro invocation in example

jlizen force-pushed the master branch from 2944509 to 367d7b7 Compare May 9, 2025 17:02

jlizen requested a review from Darksonn May 9, 2025 18:33

Darksonn reviewed May 19, 2025

View reviewed changes

Merge branch 'master' of https://github.com/tokio-rs/tokio

bc38c87

jlizen force-pushed the master branch from d55ddee to 78fbd85 Compare May 20, 2025 20:06

tweak verbiage of try_join/join fairness examples, tweaks to formatti…

4865079

…ng/import organization

jlizen force-pushed the master branch from 78fbd85 to 4865079 Compare May 20, 2025 20:10

jlizen requested a review from Darksonn May 20, 2025 20:31

change try_join/join fairness section back to original verbiage

8acfe2d

		let mut skip_next_time: u32 = 0;
		let rotate_skip = $rotate_poll_order;

Uh oh!

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

Are you sure you want to change the base?

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

Uh oh!

Conversation

jlizen commented May 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Solution

Uh oh!

jlizen commented May 5, 2025

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jlizen commented May 6, 2025

Uh oh!

jlizen commented May 8, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jlizen commented May 9, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jlizen May 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jlizen May 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jlizen May 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jlizen commented May 5, 2025 •

edited

Loading

jlizen May 20, 2025 •

edited

Loading

jlizen May 20, 2025 •

edited

Loading

jlizen May 21, 2025 •

edited

Loading