diff --git a/tokio-util/src/task/join_map.rs b/tokio-util/src/task/join_map.rs index 13e27bb67..d5c5d2a42 100644 --- a/tokio-util/src/task/join_map.rs +++ b/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)> { - 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. diff --git a/tokio-util/tests/task_join_map.rs b/tokio-util/tests/task_join_map.rs index 1ab5f9ba8..79b47154c 100644 --- a/tokio-util/tests/task_join_map.rs +++ b/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()); +}