task: properly handle removed entries in JoinMap (#7264)

tokio-util/src/task/join_map.rs

@@ -469,16 +469,19 @@ where
     ///
     /// [`tokio::select!`]: tokio::select
     pub async fn join_next(&mut self) -> Option<(K, Result<V, JoinError>)> {
-        let (res, id) = match self.tasks.join_next_with_id().await {
-            Some(Ok((id, output))) => (Ok(output), id),
-            Some(Err(e)) => {
-                let id = e.id();
-                (Err(e), id)
+        loop {
+            let (res, id) = match self.tasks.join_next_with_id().await {
+                Some(Ok((id, output))) => (Ok(output), id),
+                Some(Err(e)) => {
+                    let id = e.id();
+                    (Err(e), id)
+                }
+                None => return None,
+            };
+            if let Some(key) = self.remove_by_id(id) {
+                break Some((key, res));
             }
-            None => return None,
-        };
-        let key = self.remove_by_id(id)?;
-        Some((key, res))
+        }
     }
 
     /// Aborts all tasks and waits for them to finish shutting down.

tokio-util/tests/task_join_map.rs

@@ -297,3 +297,49 @@ async fn abort_all() {
         assert!(was_seen);
     }
 }
+
+#[tokio::test]
+async fn duplicate_keys() {
+    let mut map = JoinMap::new();
+    map.spawn(1, async { 1 });
+    map.spawn(1, async { 2 });
+
+    assert_eq!(map.len(), 1);
+
+    let (key, res) = map.join_next().await.unwrap();
+    assert_eq!(key, 1);
+    assert_eq!(res.unwrap(), 2);
+
+    assert!(map.join_next().await.is_none());
+}
+
+#[tokio::test]
+async fn duplicate_keys2() {
+    let (send, recv) = oneshot::channel::<()>();
+
+    let mut map = JoinMap::new();
+    map.spawn(1, async { 1 });
+    map.spawn(1, async {
+        recv.await.unwrap();
+        2
+    });
+
+    assert_eq!(map.len(), 1);
+
+    tokio::select! {
+        biased;
+        res = map.join_next() => match res {
+            Some((_key, res)) => panic!("Task {res:?} exited."),
+            None => panic!("Phantom task completeion."),
+        },
+        () = tokio::task::yield_now() => {},
+    }
+
+    send.send(()).unwrap();
+
+    let (key, res) = map.join_next().await.unwrap();
+    assert_eq!(key, 1);
+    assert_eq!(res.unwrap(), 2);
+
+    assert!(map.join_next().await.is_none());
+}