Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

util: add JoinMap::keys #6046

Merged
merged 3 commits into from
Oct 15, 2023
Merged

util: add JoinMap::keys #6046

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1a5bc08
util: add `JoinMap::keys`
andreastedile Oct 3, 2023
bf31e91
Use dedicated type instead of `impl Iterator`
Darksonn Oct 14, 2023
9c64a7d
Update tokio-util/src/task/join_map.rs
Darksonn Oct 14, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions tokio-util/src/task/join_map.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -626,6 +626,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
Darksonn marked this conversation as resolved.
Show resolved Hide resolved
/// call to [`join_next`], this method will still return its key.
///
/// [`join_next`]: fn@Self::join_next
pub fn keys<Q: ?Sized>(&self) -> impl Iterator<Item = &K>
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's definitely much simpler to return an impl Iterator here. However, I will note that the hash_map::Keys iterator that this wraps also implements some other traits: in addition to being an Iterator, it also implements Clone, ExactSizeIterator, and FusedIterator. It might be better if we returned an iterator that also implements those traits.

In particular, ExactSizeIterator is used by the standard library to implement optimizations when collect()ing an iterator into particular collections, like Vecs and HashMaps, since the size of the iterator can be reserved up front by with_capacity, instead of allocating multiple times while collecting.

Therefore, we may want to change this to return a named Keys iterator, like this:

use std::iterator::{ExactSizeIterator, FusedIterator};

pub struct Keys<'map, K, S = RandomState> {
    inner: hashbrown::hash_map::Keys<'map, Key<K>, AbortHandle, S>>,
}

impl<'map, K, S> Iterator for Keys<'map, K, S> {
    type Item = &'map K;

    fn next(&mut self) -> Option<Self::Item> {
         self.inner.next.map(|key| &key.key)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.inner.size_hint()    
    }
}

impl<'map, K, S> ExactSizeIterator for Keys<'map, K, S> {
    fn len(&self) -> usize {
         self.inner.len()
    }

    fn is_empty(&self) -> bool {
         self.inner.is_empty()
    }
}

impl<'map, K, S> FusedIterator for Keys<'map, K, S> {}

Or, alternatively, I think we could write

Suggested change
pub fn keys<Q: ?Sized>(&self) -> impl Iterator<Item = &K>
pub fn keys<Q: ?Sized>(&self) -> impl Iterator<Item = &K> + ExactSizeIterator + FusedIterator

but I think the named type might be a nicer API surface for users?

where
Q: Hash,
K: Borrow<Q>,
{
self.tasks_by_key.keys().map(|key| &key.key)
}

/// 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
Expand Down
24 changes: 24 additions & 0 deletions tokio-util/tests/task_join_map.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,30 @@ async fn alternating() {
}
}

#[tokio::test]
async fn test_keys() {
use std::collections::HashSet;

let mut map = JoinMap::new();

assert_eq!(map.len(), 0);
map.spawn(1, async {});
assert_eq!(map.len(), 1);
map.spawn(2, async {});
assert_eq!(map.len(), 2);

let keys = map.keys().collect::<HashSet<&u32>>();
assert!(keys.contains(&1));
assert!(keys.contains(&2));

let _ = map.join_next().await.unwrap();
let _ = map.join_next().await.unwrap();

assert_eq!(map.len(), 0);
let keys = map.keys().collect::<HashSet<&u32>>();
assert!(keys.is_empty());
}

#[tokio::test(start_paused = true)]
async fn abort_by_key() {
let mut map = JoinMap::new();
Expand Down
Loading