task: add JoinMap::keys (#6046)

Co-authored-by: Alice Ryhl <[email protected]>

tokio-util/src/task/join_map.rs

@@ -5,6 +5,7 @@ use std::collections::hash_map::RandomState;
 use std::fmt;
 use std::future::Future;
 use std::hash::{BuildHasher, Hash, Hasher};
+use std::marker::PhantomData;
 use tokio::runtime::Handle;
 use tokio::task::{AbortHandle, Id, JoinError, JoinSet, LocalSet};
 
@@ -626,6 +627,19 @@ where
         }
     }
 
+    /// Returns an iterator visiting all keys in this `JoinMap` in arbitrary order.
+    ///
+    /// If a task has completed, but its output hasn't yet been consumed by a
+    /// call to [`join_next`], this method will still return its key.
+    ///
+    /// [`join_next`]: fn@Self::join_next
+    pub fn keys(&self) -> JoinMapKeys<'_, K, V> {
+        JoinMapKeys {
+            iter: self.tasks_by_key.keys(),
+            _value: PhantomData,
+        }
+    }
+
     /// Returns `true` if this `JoinMap` contains a task for the provided key.
     ///
     /// If the task has completed, but its output hasn't yet been consumed by a
@@ -859,3 +873,32 @@ impl<K: PartialEq> PartialEq for Key<K> {
 }
 
 impl<K: Eq> Eq for Key<K> {}
+
+/// An iterator over the keys of a [`JoinMap`].
+#[derive(Debug, Clone)]
+pub struct JoinMapKeys<'a, K, V> {
+    iter: hashbrown::hash_map::Keys<'a, Key<K>, AbortHandle>,
+    /// To make it easier to change JoinMap in the future, keep V as a generic
+    /// parameter.
+    _value: PhantomData<&'a V>,
+}
+
+impl<'a, K, V> Iterator for JoinMapKeys<'a, K, V> {
+    type Item = &'a K;
+
+    fn next(&mut self) -> Option<&'a K> {
+        self.iter.next().map(|key| &key.key)
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.iter.size_hint()
+    }
+}
+
+impl<'a, K, V> ExactSizeIterator for JoinMapKeys<'a, K, V> {
+    fn len(&self) -> usize {
+        self.iter.len()
+    }
+}
+
+impl<'a, K, V> std::iter::FusedIterator for JoinMapKeys<'a, K, V> {}