Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

implement try_join_next #6280

Merged
merged 9 commits into from
Jan 14, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 54 additions & 1 deletion tokio/src/task/join_set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ use std::task::{Context, Poll};
use crate::runtime::Handle;
#[cfg(tokio_unstable)]
use crate::task::Id;
use crate::task::{AbortHandle, JoinError, JoinHandle, LocalSet};
use crate::task::{unconstrained, AbortHandle, JoinError, JoinHandle, LocalSet};
use crate::util::IdleNotifiedSet;

/// A collection of tasks spawned on a Tokio runtime.
Expand Down Expand Up @@ -306,6 +306,59 @@ impl<T: 'static> JoinSet<T> {
crate::future::poll_fn(|cx| self.poll_join_next_with_id(cx)).await
}

/// Tries to join one of the tasks in the set that has completed and return its output.
///
/// Returns `None` if the set is empty.
pub fn try_join_next(&mut self) -> Option<Result<T, JoinError>> {
// Loop over all notified `JoinHandle`s to find one that's ready, or until none are left.
loop {
let mut entry = self.inner.try_pop_notified()?;

let res = entry.with_value_and_context(|jh, ctx| {
// Since this function is not async and cannot be forced to yield, we should
// disable budgeting when we want to check for the `JoinHandle` readiness.
Pin::new(&mut unconstrained(jh)).poll(ctx)
});

if let Poll::Ready(res) = res {
let _entry = entry.remove();

return Some(res);
}
Comment on lines +323 to +327
Copy link
Contributor

Choose a reason for hiding this comment

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

Please add a comment about why you go around the loop when it returns Pending.

Copy link
Member Author

Choose a reason for hiding this comment

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

// Loop over all notified `JoinHandle`s to find one that's ready, or until none are left.

I've added this comment at the top of the function to explain why the loop exists. Is it too concise?

}
}

/// Tries to join one of the tasks in the set that has completed and return its output,
/// along with the [task ID] of the completed task.
///
/// Returns `None` if the set is empty.
///
/// When this method returns an error, then the id of the task that failed can be accessed
/// using the [`JoinError::id`] method.
///
/// [task ID]: crate::task::Id
/// [`JoinError::id`]: fn@crate::task::JoinError::id
#[cfg(tokio_unstable)]
#[cfg_attr(docsrs, doc(cfg(tokio_unstable)))]
pub fn try_join_next_with_id(&mut self) -> Option<Result<(Id, T), JoinError>> {
// Loop over all notified `JoinHandle`s to find one that's ready, or until none are left.
loop {
let mut entry = self.inner.try_pop_notified()?;

let res = entry.with_value_and_context(|jh, ctx| {
// Since this function is not async and cannot be forced to yield, we should
// disable budgeting when we want to check for the `JoinHandle` readiness.
Pin::new(&mut unconstrained(jh)).poll(ctx)
});

if let Poll::Ready(res) = res {
let entry = entry.remove();

return Some(res.map(|output| (entry.id(), output)));
}
}
}

/// Aborts all tasks and waits for them to finish shutting down.
///
/// Calling this method is equivalent to calling [`abort_all`] and then calling [`join_next`] in
Expand Down
28 changes: 28 additions & 0 deletions tokio/src/util/idle_notified_set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -203,6 +203,34 @@ impl<T> IdleNotifiedSet<T> {
Some(EntryInOneOfTheLists { entry, set: self })
}

/// Tries to pop an entry from the notified list to poll it. The entry is moved to
/// the idle list atomically.
pub(crate) fn try_pop_notified(&mut self) -> Option<EntryInOneOfTheLists<'_, T>> {
// We don't decrement the length because this call moves the entry to
// the idle list rather than removing it.
if self.length == 0 {
// Fast path.
return None;
}

let mut lock = self.lists.lock();

// Pop the entry, returning None if empty.
let entry = lock.notified.pop_back()?;

lock.idle.push_front(entry.clone());

// Safety: We are holding the lock.
entry.my_list.with_mut(|ptr| unsafe {
*ptr = List::Idle;
});

drop(lock);

// Safety: We just put the entry in the idle list, so it is in one of the lists.
Some(EntryInOneOfTheLists { entry, set: self })
}

/// Call a function on every element in this list.
pub(crate) fn for_each<F: FnMut(&mut T)>(&mut self, mut func: F) {
fn get_ptrs<T>(list: &mut LinkedList<T>, ptrs: &mut Vec<*mut T>) {
Expand Down
77 changes: 77 additions & 0 deletions tokio/tests/task_join_set.rs
Copy link
Contributor

Choose a reason for hiding this comment

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

The test is nice. Especially the fact that you cover the case with coop budgeting. You don't have any tests that handle the case where it returns None even though there are tasks, because all the tasks are still running. Could you do that?

Copy link
Member Author

Choose a reason for hiding this comment

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

Sure!
I've updated the tests so they also test the case where there are no finished task.

Original file line number Diff line number Diff line change
Expand Up @@ -227,3 +227,80 @@ async fn join_set_coop() {
assert!(coop_count >= 1);
assert_eq!(count, TASK_NUM);
}

#[tokio::test(flavor = "current_thread")]
async fn try_join_next() {
const TASK_NUM: u32 = 1000;

let (send, recv) = tokio::sync::watch::channel(());

let mut set = JoinSet::new();

for _ in 0..TASK_NUM {
let mut recv = recv.clone();
set.spawn(async move { recv.changed().await.unwrap() });
}
drop(recv);

assert!(set.try_join_next().is_none());

send.send_replace(());
send.closed().await;
Comment on lines +247 to +248
Copy link
Contributor

Choose a reason for hiding this comment

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

It's not actually guaranteed that the tasks are considered finished by the JoinSet when closed().await returns. This is because tasks are marked finished after the destructor runs, but the channels are dropped while the destructor runs.

If you make this into a current-thread runtime, then this race goes away.

Copy link
Member Author

Choose a reason for hiding this comment

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

Ah you're right. Fixed it.


let mut count = 0;
loop {
match set.try_join_next() {
Some(Ok(())) => {
count += 1;
}
Some(Err(err)) => panic!("failed: {}", err),
None => {
break;
}
}
}

assert_eq!(count, TASK_NUM);
}

#[cfg(tokio_unstable)]
#[tokio::test(flavor = "current_thread")]
async fn try_join_next_with_id() {
const TASK_NUM: u32 = 1000;

let (send, recv) = tokio::sync::watch::channel(());

let mut set = JoinSet::new();
let mut spawned = std::collections::HashSet::with_capacity(TASK_NUM as usize);

for _ in 0..TASK_NUM {
let mut recv = recv.clone();
let handle = set.spawn(async move { recv.changed().await.unwrap() });

spawned.insert(handle.id());
}
drop(recv);

assert!(set.try_join_next_with_id().is_none());

send.send_replace(());
send.closed().await;

let mut count = 0;
let mut joined = std::collections::HashSet::with_capacity(TASK_NUM as usize);
loop {
match set.try_join_next_with_id() {
Some(Ok((id, ()))) => {
count += 1;
joined.insert(id);
}
Some(Err(err)) => panic!("failed: {}", err),
None => {
break;
}
}
}

assert_eq!(count, TASK_NUM);
assert_eq!(joined, spawned);
}
Loading