task: add try_join_next and try_join_next_with_id on JoinQueue (#7636)

tokio-util/tests/task_join_queue.rs

@@ -192,12 +192,12 @@ async fn test_join_queue_join_next_with_id() {
 
     let (send, recv) = tokio::sync::watch::channel(());
 
-    let mut set = JoinQueue::new();
+    let mut queue = JoinQueue::new();
     let mut spawned = Vec::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() });
+        let handle = queue.spawn(async move { recv.changed().await.unwrap() });
 
         spawned.push(handle.id());
     }
@@ -208,7 +208,7 @@ async fn test_join_queue_join_next_with_id() {
 
     let mut count = 0;
     let mut joined = Vec::with_capacity(TASK_NUM as usize);
-    while let Some(res) = set.join_next_with_id().await {
+    while let Some(res) = queue.join_next_with_id().await {
         match res {
             Ok((id, ())) => {
                 count += 1;
@@ -221,3 +221,141 @@ async fn test_join_queue_join_next_with_id() {
     assert_eq!(count, TASK_NUM);
     assert_eq!(joined, spawned);
 }
+
+#[tokio::test]
+async fn test_join_queue_try_join_next() {
+    let mut queue = JoinQueue::new();
+    let (tx1, rx1) = oneshot::channel::<()>();
+    queue.spawn(async {
+        let _ = rx1.await;
+    });
+    let (tx2, rx2) = oneshot::channel::<()>();
+    queue.spawn(async {
+        let _ = rx2.await;
+    });
+    let (tx3, rx3) = oneshot::channel::<()>();
+    queue.spawn(async {
+        let _ = rx3.await;
+    });
+
+    // This function also checks that calling `queue.try_join_next()` repeatedly when
+    // no task is ready is idempotent, i.e. that it does not change the queue state.
+    fn check_try_join_next_is_noop(queue: &mut JoinQueue<()>) {
+        let len = queue.len();
+        for _ in 0..5 {
+            assert!(queue.try_join_next().is_none());
+            assert_eq!(queue.len(), len);
+        }
+    }
+
+    assert_eq!(queue.len(), 3);
+    check_try_join_next_is_noop(&mut queue);
+
+    tx1.send(()).unwrap();
+    tokio::task::yield_now().await;
+
+    assert_eq!(queue.len(), 3);
+    assert!(queue.try_join_next().is_some());
+    assert_eq!(queue.len(), 2);
+    check_try_join_next_is_noop(&mut queue);
+
+    tx3.send(()).unwrap();
+    tokio::task::yield_now().await;
+
+    assert_eq!(queue.len(), 2);
+    check_try_join_next_is_noop(&mut queue);
+
+    tx2.send(()).unwrap();
+    tokio::task::yield_now().await;
+
+    assert_eq!(queue.len(), 2);
+    assert!(queue.try_join_next().is_some());
+    assert_eq!(queue.len(), 1);
+    assert!(queue.try_join_next().is_some());
+    assert!(queue.is_empty());
+    check_try_join_next_is_noop(&mut queue);
+}
+
+#[tokio::test]
+async fn test_join_queue_try_join_next_disabled_coop() {
+    // This number is large enough to trigger coop. Without using `tokio::task::coop::unconstrained`
+    // inside `try_join_next` this test fails on `assert!(coop_count == 0)`.
+    const TASK_NUM: u32 = 1000;
+
+    let sem: std::sync::Arc<tokio::sync::Semaphore> =
+        std::sync::Arc::new(tokio::sync::Semaphore::new(0));
+
+    let mut queue = JoinQueue::new();
+
+    for _ in 0..TASK_NUM {
+        let sem = sem.clone();
+        queue.spawn(async move {
+            sem.add_permits(1);
+        });
+    }
+
+    let _ = sem.acquire_many(TASK_NUM).await.unwrap();
+
+    let mut count = 0;
+    let mut coop_count = 0;
+    while !queue.is_empty() {
+        match queue.try_join_next() {
+            Some(Ok(())) => count += 1,
+            Some(Err(err)) => panic!("failed: {err}"),
+            None => {
+                coop_count += 1;
+                tokio::task::yield_now().await;
+            }
+        }
+    }
+    assert_eq!(coop_count, 0);
+    assert_eq!(count, TASK_NUM);
+}
+
+#[tokio::test]
+async fn test_join_queue_try_join_next_with_id_disabled_coop() {
+    // Note that this number is large enough to trigger coop as in
+    // `test_join_queue_try_join_next_coop` test. Without using
+    // `tokio::task::coop::unconstrained` inside `try_join_next_with_id`
+    // this test fails on `assert_eq!(count, TASK_NUM)`.
+    const TASK_NUM: u32 = 1000;
+
+    let (send, recv) = tokio::sync::watch::channel(());
+
+    let mut queue = JoinQueue::new();
+    let mut spawned = Vec::with_capacity(TASK_NUM as usize);
+
+    for _ in 0..TASK_NUM {
+        let mut recv = recv.clone();
+        let handle = queue.spawn(async move { recv.changed().await.unwrap() });
+
+        spawned.push(handle.id());
+    }
+    drop(recv);
+
+    assert!(queue.try_join_next_with_id().is_none());
+
+    send.send_replace(());
+    send.closed().await;
+
+    let mut count = 0;
+    let mut coop_count = 0;
+    let mut joined = Vec::with_capacity(TASK_NUM as usize);
+    while !queue.is_empty() {
+        match queue.try_join_next_with_id() {
+            Some(Ok((id, ()))) => {
+                count += 1;
+                joined.push(id);
+            }
+            Some(Err(err)) => panic!("failed: {err}"),
+            None => {
+                coop_count += 1;
+                tokio::task::yield_now().await;
+            }
+        }
+    }
+
+    assert_eq!(coop_count, 0);
+    assert_eq!(count, TASK_NUM);
+    assert_eq!(joined, spawned);
+}