task: add JoinMap::try_join_next (#8099)

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<V, JoinError>)> {
+        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

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<usize, ()> = 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();