tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-08-22 00:00:11 +02:00
Code Issues Packages Projects Releases Wiki Activity

task: remove raw-entry feature from hashbrown dep (#7252)

Browse Source
This commit is contained in:
Conrad Ludgate
2025-07-22 15:52:54 +02:00
committed by GitHub
parent 0d234c3cf9
commit 9e94fa7e15
3 changed files with 100 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio-util/Cargo.toml
+1 -1
View File
@@ -45,7 +45,7 @@ slab = { version = "0.4.4", optional = true } # Backs `DelayQueue`
tracing = { version = "0.1.29", default-features = false, features = ["std"], optional = true }
[target.'cfg(tokio_unstable)'.dependencies]
hashbrown = { version = "0.15.0", default-features = false, features = ["raw-entry"], optional = true }
hashbrown = { version = "0.15.0", default-features = false, optional = true }
[dev-dependencies]
tokio = { version = "1.0.0", path = "../tokio", features = ["full"] }
tokio-util/src/task/join_map.rs
+66 -104
View File
@@ -1,5 +1,5 @@
use hashbrown::hash_map::RawEntryMut;
use hashbrown::HashMap;
use hashbrown::hash_table::Entry;
use hashbrown::{HashMap, HashTable};
use std::borrow::Borrow;
use std::collections::hash_map::RandomState;
use std::fmt;
@@ -103,13 +103,8 @@ use tokio::task::{AbortHandle, Id, JoinError, JoinSet, LocalSet};
#[cfg_attr(docsrs, doc(cfg(all(feature = "rt", tokio_unstable))))]
pub struct JoinMap<K, V, S = RandomState> {
/// A map of the [`AbortHandle`]s of the tasks spawned on this `JoinMap`,
/// indexed by their keys and task IDs.
///
/// The [`Key`] type contains both the task's `K`-typed key provided when
/// spawning tasks, and the task's IDs. The IDs are stored here to resolve
/// hash collisions when looking up tasks based on their pre-computed hash
/// (as stored in the `hashes_by_task` map).
tasks_by_key: HashMap<Key<K>, AbortHandle, S>,
/// indexed by their keys.
tasks_by_key: HashTable<(K, AbortHandle)>,
/// A map from task IDs to the hash of the key associated with that task.
///
@@ -125,21 +120,6 @@ pub struct JoinMap<K, V, S = RandomState> {
tasks: JoinSet<V>,
}
/// A [`JoinMap`] key.
///
/// This holds both a `K`-typed key (the actual key as seen by the user), _and_
/// a task ID, so that hash collisions between `K`-typed keys can be resolved
/// using either `K`'s `Eq` impl *or* by checking the task IDs.
///
/// This allows looking up a task using either an actual key (such as when the
/// user queries the map with a key), *or* using a task ID and a hash (such as
/// when removing completed tasks from the map).
#[derive(Debug)]
struct Key<K> {
key: K,
id: Id,
}
impl<K, V> JoinMap<K, V> {
/// Creates a new empty `JoinMap`.
///
@@ -176,7 +156,7 @@ impl<K, V> JoinMap<K, V> {
}
}
impl<K, V, S: Clone> JoinMap<K, V, S> {
impl<K, V, S> JoinMap<K, V, S> {
/// Creates an empty `JoinMap` which will use the given hash builder to hash
/// keys.
///
@@ -226,7 +206,7 @@ impl<K, V, S: Clone> JoinMap<K, V, S> {
#[must_use]
pub fn with_capacity_and_hasher(capacity: usize, hash_builder: S) -> Self {
Self {
tasks_by_key: HashMap::with_capacity_and_hasher(capacity, hash_builder.clone()),
tasks_by_key: HashTable::with_capacity(capacity),
hashes_by_task: HashMap::with_capacity_and_hasher(capacity, hash_builder),
tasks: JoinSet::new(),
}
@@ -415,33 +395,42 @@ where
self.insert(key, task)
}
fn insert(&mut self, key: K, abort: AbortHandle) {
let hash = self.hash(&key);
fn insert(&mut self, mut key: K, mut abort: AbortHandle) {
let hash_builder = self.hashes_by_task.hasher();
let hash = hash_one(hash_builder, &key);
let id = abort.id();
let map_key = Key { id, key };
// Insert the new key into the map of tasks by keys.
let entry = self
.tasks_by_key
.raw_entry_mut()
.from_hash(hash, |k| k.key == map_key.key);
let entry =
self.tasks_by_key
.entry(hash, |(k, _)| *k == key, |(k, _)| hash_one(hash_builder, k));
match entry {
RawEntryMut::Occupied(mut occ) => {
Entry::Occupied(occ) => {
// There was a previous task spawned with the same key! Cancel
// that task, and remove its ID from the map of hashes by task IDs.
let Key { id: prev_id, .. } = occ.insert_key(map_key);
occ.insert(abort).abort();
let _prev_hash = self.hashes_by_task.remove(&prev_id);
(key, abort) = std::mem::replace(occ.into_mut(), (key, abort));
// Remove the old task ID.
let _prev_hash = self.hashes_by_task.remove(&abort.id());
debug_assert_eq!(Some(hash), _prev_hash);
// Associate the key's hash with the new task's ID, for looking up tasks by ID.
let _prev = self.hashes_by_task.insert(id, hash);
debug_assert!(_prev.is_none(), "no prior task should have had the same ID");
// Note: it's important to drop `key` and abort the task here.
// This defends against any panics during drop handling for causing inconsistent state.
abort.abort();
drop(key);
}
RawEntryMut::Vacant(vac) => {
vac.insert(map_key, abort);
Entry::Vacant(vac) => {
vac.insert((key, abort));
// Associate the key's hash with this task's ID, for looking up tasks by ID.
let _prev = self.hashes_by_task.insert(id, hash);
debug_assert!(_prev.is_none(), "no prior task should have had the same ID");
}
};
// Associate the key's hash with this task's ID, for looking up tasks by ID.
let _prev = self.hashes_by_task.insert(id, hash);
debug_assert!(_prev.is_none(), "no prior task should have had the same ID");
}
/// Waits until one of the tasks in the map completes and returns its
@@ -623,7 +612,7 @@ where
// Note: this method iterates over the tasks and keys *without* removing
// any entries, so that the keys from aborted tasks can still be
// returned when calling `join_next` in the future.
for (Key { ref key, .. }, task) in &self.tasks_by_key {
for (key, task) in &self.tasks_by_key {
if predicate(key) {
task.abort();
}
@@ -638,7 +627,7 @@ where
/// [`join_next`]: fn@Self::join_next
pub fn keys(&self) -> JoinMapKeys<'_, K, V> {
JoinMapKeys {
iter: self.tasks_by_key.keys(),
iter: self.tasks_by_key.iter(),
_value: PhantomData,
}
}
@@ -666,7 +655,7 @@ where
/// [`join_next`]: fn@Self::join_next
/// [task ID]: tokio::task::Id
pub fn contains_task(&self, task: &Id) -> bool {
self.get_by_id(task).is_some()
self.hashes_by_task.contains_key(task)
}
/// Reserves capacity for at least `additional` more tasks to be spawned
@@ -690,7 +679,9 @@ where
/// ```
#[inline]
pub fn reserve(&mut self, additional: usize) {
self.tasks_by_key.reserve(additional);
let hash_builder = self.hashes_by_task.hasher();
self.tasks_by_key
.reserve(additional, |(k, _)| hash_one(hash_builder, k));
self.hashes_by_task.reserve(additional);
}
@@ -716,7 +707,9 @@ where
#[inline]
pub fn shrink_to_fit(&mut self) {
self.hashes_by_task.shrink_to_fit();
self.tasks_by_key.shrink_to_fit();
let hash_builder = self.hashes_by_task.hasher();
self.tasks_by_key
.shrink_to_fit(|(k, _)| hash_one(hash_builder, k));
}
/// Shrinks the capacity of the map with a lower limit. It will drop
@@ -745,27 +738,20 @@ where
#[inline]
pub fn shrink_to(&mut self, min_capacity: usize) {
self.hashes_by_task.shrink_to(min_capacity);
self.tasks_by_key.shrink_to(min_capacity)
let hash_builder = self.hashes_by_task.hasher();
self.tasks_by_key
.shrink_to(min_capacity, |(k, _)| hash_one(hash_builder, k))
}
/// Look up a task in the map by its key, returning the key and abort handle.
fn get_by_key<'map, Q: ?Sized>(&'map self, key: &Q) -> Option<(&'map Key<K>, &'map AbortHandle)>
fn get_by_key<'map, Q: ?Sized>(&'map self, key: &Q) -> Option<&'map (K, AbortHandle)>
where
Q: Hash + Eq,
K: Borrow<Q>,
{
let hash = self.hash(key);
self.tasks_by_key
.raw_entry()
.from_hash(hash, |k| k.key.borrow() == key)
}
/// Look up a task in the map by its task ID, returning the key and abort handle.
fn get_by_id<'map>(&'map self, id: &Id) -> Option<(&'map Key<K>, &'map AbortHandle)> {
let hash = self.hashes_by_task.get(id)?;
self.tasks_by_key
.raw_entry()
.from_hash(*hash, |k| &k.id == id)
let hash_builder = self.hashes_by_task.hasher();
let hash = hash_one(hash_builder, key);
self.tasks_by_key.find(hash, |(k, _)| k.borrow() == key)
}
/// Remove a task from the map by ID, returning the key for that task.
@@ -776,28 +762,25 @@ where
// Remove the entry for that hash.
let entry = self
.tasks_by_key
.raw_entry_mut()
.from_hash(hash, |k| k.id == id);
let (Key { id: _key_id, key }, handle) = match entry {
RawEntryMut::Occupied(entry) => entry.remove_entry(),
.find_entry(hash, |(_, abort)| abort.id() == id);
let (key, _) = match entry {
Ok(entry) => entry.remove().0,
_ => return None,
};
debug_assert_eq!(_key_id, id);
debug_assert_eq!(id, handle.id());
self.hashes_by_task.remove(&id);
Some(key)
}
}
/// Returns the hash for a given key.
#[inline]
fn hash<Q: ?Sized>(&self, key: &Q) -> u64
where
Q: Hash,
{
let mut hasher = self.tasks_by_key.hasher().build_hasher();
key.hash(&mut hasher);
hasher.finish()
}
/// Returns the hash for a given key.
#[inline]
fn hash_one<S: BuildHasher, Q: ?Sized>(hash_builder: &S, key: &Q) -> u64
where
Q: Hash,
{
let mut hasher = hash_builder.build_hasher();
key.hash(&mut hasher);
hasher.finish()
}
impl<K, V, S> JoinMap<K, V, S>
@@ -831,11 +814,11 @@ impl<K: fmt::Debug, V, S> fmt::Debug for JoinMap<K, V, S> {
// printing the key and task ID pairs, without format the `Key` struct
// itself or the `AbortHandle`, which would just format the task's ID
// again.
struct KeySet<'a, K: fmt::Debug, S>(&'a HashMap<Key<K>, AbortHandle, S>);
impl<K: fmt::Debug, S> fmt::Debug for KeySet<'_, K, S> {
struct KeySet<'a, K: fmt::Debug>(&'a HashTable<(K, AbortHandle)>);
impl<K: fmt::Debug> fmt::Debug for KeySet<'_, K> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_map()
.entries(self.0.keys().map(|Key { key, id }| (key, id)))
.entries(self.0.iter().map(|(key, abort)| (key, abort.id())))
.finish()
}
}
@@ -856,31 +839,10 @@ impl<K, V> Default for JoinMap<K, V> {
}
}
// === impl Key ===
impl<K: Hash> Hash for Key<K> {
// Don't include the task ID in the hash.
#[inline]
fn hash<H: Hasher>(&self, hasher: &mut H) {
self.key.hash(hasher);
}
}
// Because we override `Hash` for this type, we must also override the
// `PartialEq` impl, so that all instances with the same hash are equal.
impl<K: PartialEq> PartialEq for Key<K> {
#[inline]
fn eq(&self, other: &Self) -> bool {
self.key == other.key
}
}
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>,
iter: hashbrown::hash_table::Iter<'a, (K, AbortHandle)>,
/// To make it easier to change `JoinMap` in the future, keep V as a generic
/// parameter.
_value: PhantomData<&'a V>,
@@ -890,7 +852,7 @@ 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)
self.iter.next().map(|(key, _)| key)
}
fn size_hint(&self) -> (usize, Option<usize>) {
tokio-util/tests/task_join_map.rs
+33
View File
@@ -1,6 +1,8 @@
#![warn(rust_2018_idioms)]
#![cfg(all(feature = "rt", tokio_unstable))]
use std::panic::AssertUnwindSafe;
use tokio::sync::oneshot;
use tokio::time::Duration;
use tokio_util::task::JoinMap;
@@ -343,3 +345,34 @@ async fn duplicate_keys2() {
assert!(map.join_next().await.is_none());
}
#[cfg_attr(not(panic = "unwind"), ignore)]
#[tokio::test]
async fn duplicate_keys_drop() {
#[derive(Hash, Debug, PartialEq, Eq)]
struct Key;
impl Drop for Key {
fn drop(&mut self) {
panic!("drop called for key");
}
}
let (send, recv) = oneshot::channel::<()>();
let mut map = JoinMap::new();
map.spawn(Key, async { recv.await.unwrap() });
// replace the task, force it to drop the key and abort the task
// we should expect it to panic when dropping the key.
let _ = std::panic::catch_unwind(AssertUnwindSafe(|| map.spawn(Key, async {}))).unwrap_err();
// don't panic when this key drops.
let (key, _) = map.join_next().await.unwrap();
std::mem::forget(key);
// original task should have been aborted, so the sender should be dangling.
assert!(send.is_closed());
assert!(map.join_next().await.is_none());
}