From 1fe1b0e727dd1c48ad82173e35e125a3ba43b955 Mon Sep 17 00:00:00 2001 From: Joe Grund Date: Thu, 21 May 2026 09:41:55 -0400 Subject: [PATCH] task: add `JoinMap::try_join_next` (#8099) --- tokio-util/src/task/join_map.rs | 56 ++++++++++++- tokio-util/tests/task_join_map.rs | 125 ++++++++++++++++++++++++++++++ 2 files changed, 180 insertions(+), 1 deletion(-) diff --git a/tokio-util/src/task/join_map.rs b/tokio-util/src/task/join_map.rs index 49c192537..027470257 100644 --- a/tokio-util/src/task/join_map.rs +++ b/tokio-util/src/task/join_map.rs @@ -446,7 +446,7 @@ where /// * `Some((key, Ok(value)))` if one of the tasks in this `JoinMap` has /// completed. The `value` is the return value of that ask, and `key` is /// the key associated with the task. - /// * `Some((key, Err(err))` if one of the tasks in this `JoinMap` has + /// * `Some((key, Err(err)))` if one of the tasks in this `JoinMap` has /// panicked or been aborted. `key` is the key associated with the task /// that panicked or was aborted. /// * `None` if the `JoinMap` is empty. @@ -468,6 +468,60 @@ where } } + /// Tries to join one of the tasks in the map that has completed and + /// returns its output, along with the key corresponding to that task. + /// + /// Returns `None` if there are no completed tasks, or if the map is empty. + /// + /// # Returns + /// + /// This function returns: + /// + /// * `Some((key, Ok(value)))` if one of the tasks in this `JoinMap` has + /// completed. The `value` is the return value of that task, and `key` + /// is the key associated with the task. + /// * `Some((key, Err(err)))` if one of the tasks in this `JoinMap` has + /// panicked or been aborted. `key` is the key associated with the task + /// that panicked or was aborted. + /// * `None` if there are no completed tasks ready to be joined, or the + /// `JoinMap` is empty. + /// + /// # Examples + /// + /// ``` + /// use tokio_util::task::JoinMap; + /// + /// # #[tokio::main(flavor = "current_thread")] + /// # async fn main() { + /// let mut map = JoinMap::new(); + /// map.spawn("answer", async { 42 }); + /// + /// let (key, res) = loop { + /// if let Some(joined) = map.try_join_next() { + /// break joined; + /// } + /// tokio::task::yield_now().await; + /// }; + /// + /// assert_eq!(key, "answer"); + /// assert_eq!(res.unwrap(), 42); + /// # } + /// ``` + pub fn try_join_next(&mut self) -> Option<(K, Result)> { + loop { + let (res, id) = match self.tasks.try_join_next_with_id()? { + Ok((id, output)) => (Ok(output), id), + Err(e) => { + let id = e.id(); + (Err(e), id) + } + }; + if let Some(key) = self.remove_by_id(id) { + break Some((key, res)); + } + } + } + /// 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 diff --git a/tokio-util/tests/task_join_map.rs b/tokio-util/tests/task_join_map.rs index 70c33d8a9..0748fb866 100644 --- a/tokio-util/tests/task_join_map.rs +++ b/tokio-util/tests/task_join_map.rs @@ -358,6 +358,131 @@ async fn abort_all() { } } +#[tokio::test] +async fn try_join_next_empty() { + let mut map: JoinMap = JoinMap::new(); + assert!(map.try_join_next().is_none()); +} + +#[tokio::test] +async fn try_join_next_no_ready_task() { + let mut map = JoinMap::new(); + let (_tx, rx) = oneshot::channel::<()>(); + map.spawn("pending", async move { + let _ = rx.await; + }); + + // Task is not yet ready. + assert!(map.try_join_next().is_none()); + assert_eq!(map.len(), 1); +} + +#[tokio::test] +async fn try_join_next_completed_task() { + let mut map = JoinMap::new(); + map.spawn("hello", async { 42 }); + + let mut got = None; + while got.is_none() { + got = map.try_join_next(); + if got.is_none() { + tokio::task::yield_now().await; + } + } + let (key, res) = got.unwrap(); + assert_eq!(key, "hello"); + assert_eq!(res.unwrap(), 42); + assert!(map.is_empty()); +} + +#[tokio::test] +async fn try_join_next_aborted_task() { + let mut map = JoinMap::new(); + map.spawn("forever", async { + futures::future::pending::<()>().await; + }); + + assert!(map.abort("forever")); + + let mut got = None; + while got.is_none() { + got = map.try_join_next(); + if got.is_none() { + tokio::task::yield_now().await; + } + } + let (key, res) = got.unwrap(); + assert_eq!(key, "forever"); + assert!(res.unwrap_err().is_cancelled()); + assert!(map.is_empty()); +} + +#[tokio::test(flavor = "current_thread")] +async fn try_join_next_advances_through_multiple() { + const N: u32 = 8; + + static SEM: tokio::sync::Semaphore = tokio::sync::Semaphore::const_new(0); + + let mut map = JoinMap::new(); + for i in 0..N { + map.spawn(i, async move { + SEM.add_permits(1); + i + }); + } + + // Wait until all tasks have signalled completion. On the current_thread + // runtime this means they have actually finished. + let _ = SEM.acquire_many(N).await.unwrap(); + + let mut seen = vec![false; N as usize]; + let mut count = 0; + loop { + match map.try_join_next() { + Some((key, res)) => { + let v = res.expect("task should have completed successfully"); + assert_eq!(key, v); + seen[v as usize] = true; + count += 1; + } + None if map.is_empty() => break, + None => tokio::task::yield_now().await, + } + } + + assert_eq!(count, N); + assert!(seen.into_iter().all(|b| b)); + assert!(map.try_join_next().is_none()); +} + +#[tokio::test] +async fn try_join_next_skips_replaced_task() { + let mut map = JoinMap::new(); + + let (tx1, rx1) = oneshot::channel::<()>(); + map.spawn(1, async { + let _ = rx1.await; + 11 + }); + tx1.send(()).unwrap(); + tokio::task::yield_now().await; + + let (tx2, rx2) = oneshot::channel::<()>(); + map.spawn(1, async { + let _ = rx2.await; + 22 + }); + tx2.send(()).unwrap(); + tokio::task::yield_now().await; + + let (key, res) = map.try_join_next().unwrap(); + assert_eq!(key, 1); + assert_eq!(res.unwrap(), 22); + + assert!(map.try_join_next().is_none()); + assert!(map.is_empty()); +} + #[tokio::test] async fn duplicate_keys() { let mut map = JoinMap::new();