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

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());
+}