mirror of
https://github.com/tokio-rs/tokio.git
synced 2025-09-28 12:10:37 +00:00
811 lines
28 KiB
Rust
811 lines
28 KiB
Rust
use hashbrown::hash_map::RawEntryMut;
|
|
use hashbrown::HashMap;
|
|
use std::borrow::Borrow;
|
|
use std::collections::hash_map::RandomState;
|
|
use std::fmt;
|
|
use std::future::Future;
|
|
use std::hash::{BuildHasher, Hash, Hasher};
|
|
use tokio::runtime::Handle;
|
|
use tokio::task::{AbortHandle, Id, JoinError, JoinSet, LocalSet};
|
|
|
|
/// A collection of tasks spawned on a Tokio runtime, associated with hash map
|
|
/// keys.
|
|
///
|
|
/// This type is very similar to the [`JoinSet`] type in `tokio::task`, with the
|
|
/// addition of a set of keys associated with each task. These keys allow
|
|
/// [cancelling a task][abort] or [multiple tasks][abort_matching] in the
|
|
/// `JoinMap` based on their keys, or [test whether a task corresponding to a
|
|
/// given key exists][contains] in the `JoinMap`.
|
|
///
|
|
/// In addition, when tasks in the `JoinMap` complete, they will return the
|
|
/// associated key along with the value returned by the task, if any.
|
|
///
|
|
/// A `JoinMap` can be used to await the completion of some or all of the tasks
|
|
/// in the map. The map is not ordered, and the tasks will be returned in the
|
|
/// order they complete.
|
|
///
|
|
/// All of the tasks must have the same return type `V`.
|
|
///
|
|
/// When the `JoinMap` is dropped, all tasks in the `JoinMap` are immediately aborted.
|
|
///
|
|
/// **Note**: This type depends on Tokio's [unstable API][unstable]. See [the
|
|
/// documentation on unstable features][unstable] for details on how to enable
|
|
/// Tokio's unstable features.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// Spawn multiple tasks and wait for them:
|
|
///
|
|
/// ```
|
|
/// use tokio_util::task::JoinMap;
|
|
///
|
|
/// #[tokio::main]
|
|
/// async fn main() {
|
|
/// let mut map = JoinMap::new();
|
|
///
|
|
/// for i in 0..10 {
|
|
/// // Spawn a task on the `JoinMap` with `i` as its key.
|
|
/// map.spawn(i, async move { /* ... */ });
|
|
/// }
|
|
///
|
|
/// let mut seen = [false; 10];
|
|
///
|
|
/// // When a task completes, `join_next` returns the task's key along
|
|
/// // with its output.
|
|
/// while let Some((key, res)) = map.join_next().await {
|
|
/// seen[key] = true;
|
|
/// assert!(res.is_ok(), "task {} completed successfully!", key);
|
|
/// }
|
|
///
|
|
/// for i in 0..10 {
|
|
/// assert!(seen[i]);
|
|
/// }
|
|
/// }
|
|
/// ```
|
|
///
|
|
/// Cancel tasks based on their keys:
|
|
///
|
|
/// ```
|
|
/// use tokio_util::task::JoinMap;
|
|
///
|
|
/// #[tokio::main]
|
|
/// async fn main() {
|
|
/// let mut map = JoinMap::new();
|
|
///
|
|
/// map.spawn("hello world", async move { /* ... */ });
|
|
/// map.spawn("goodbye world", async move { /* ... */});
|
|
///
|
|
/// // Look up the "goodbye world" task in the map and abort it.
|
|
/// let aborted = map.abort("goodbye world");
|
|
///
|
|
/// // `JoinMap::abort` returns `true` if a task existed for the
|
|
/// // provided key.
|
|
/// assert!(aborted);
|
|
///
|
|
/// while let Some((key, res)) = map.join_next().await {
|
|
/// if key == "goodbye world" {
|
|
/// // The aborted task should complete with a cancelled `JoinError`.
|
|
/// assert!(res.unwrap_err().is_cancelled());
|
|
/// } else {
|
|
/// // Other tasks should complete normally.
|
|
/// assert!(res.is_ok());
|
|
/// }
|
|
/// }
|
|
/// }
|
|
/// ```
|
|
///
|
|
/// [`JoinSet`]: tokio::task::JoinSet
|
|
/// [unstable]: tokio#unstable-features
|
|
/// [abort]: fn@Self::abort
|
|
/// [abort_matching]: fn@Self::abort_matching
|
|
/// [contains]: fn@Self::contains_key
|
|
#[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>,
|
|
|
|
/// A map from task IDs to the hash of the key associated with that task.
|
|
///
|
|
/// This map is used to perform reverse lookups of tasks in the
|
|
/// `tasks_by_key` map based on their task IDs. When a task terminates, the
|
|
/// ID is provided to us by the `JoinSet`, so we can look up the hash value
|
|
/// of that task's key, and then remove it from the `tasks_by_key` map using
|
|
/// the raw hash code, resolving collisions by comparing task IDs.
|
|
hashes_by_task: HashMap<Id, u64, S>,
|
|
|
|
/// The [`JoinSet`] that awaits the completion of tasks spawned on this
|
|
/// `JoinMap`.
|
|
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`.
|
|
///
|
|
/// The `JoinMap` is initially created with a capacity of 0, so it will not
|
|
/// allocate until a task is first spawned on it.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use tokio_util::task::JoinMap;
|
|
/// let map: JoinMap<&str, i32> = JoinMap::new();
|
|
/// ```
|
|
#[inline]
|
|
#[must_use]
|
|
pub fn new() -> Self {
|
|
Self::with_hasher(RandomState::new())
|
|
}
|
|
|
|
/// Creates an empty `JoinMap` with the specified capacity.
|
|
///
|
|
/// The `JoinMap` will be able to hold at least `capacity` tasks without
|
|
/// reallocating.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use tokio_util::task::JoinMap;
|
|
/// let map: JoinMap<&str, i32> = JoinMap::with_capacity(10);
|
|
/// ```
|
|
#[inline]
|
|
#[must_use]
|
|
pub fn with_capacity(capacity: usize) -> Self {
|
|
JoinMap::with_capacity_and_hasher(capacity, Default::default())
|
|
}
|
|
}
|
|
|
|
impl<K, V, S: Clone> JoinMap<K, V, S> {
|
|
/// Creates an empty `JoinMap` which will use the given hash builder to hash
|
|
/// keys.
|
|
///
|
|
/// The created map has the default initial capacity.
|
|
///
|
|
/// Warning: `hash_builder` is normally randomly generated, and
|
|
/// is designed to allow `JoinMap` to be resistant to attacks that
|
|
/// cause many collisions and very poor performance. Setting it
|
|
/// manually using this function can expose a DoS attack vector.
|
|
///
|
|
/// The `hash_builder` passed should implement the [`BuildHasher`] trait for
|
|
/// the `JoinMap` to be useful, see its documentation for details.
|
|
#[inline]
|
|
#[must_use]
|
|
pub fn with_hasher(hash_builder: S) -> Self {
|
|
Self::with_capacity_and_hasher(0, hash_builder)
|
|
}
|
|
|
|
/// Creates an empty `JoinMap` with the specified capacity, using `hash_builder`
|
|
/// to hash the keys.
|
|
///
|
|
/// The `JoinMap` will be able to hold at least `capacity` elements without
|
|
/// reallocating. If `capacity` is 0, the `JoinMap` will not allocate.
|
|
///
|
|
/// Warning: `hash_builder` is normally randomly generated, and
|
|
/// is designed to allow HashMaps to be resistant to attacks that
|
|
/// cause many collisions and very poor performance. Setting it
|
|
/// manually using this function can expose a DoS attack vector.
|
|
///
|
|
/// The `hash_builder` passed should implement the [`BuildHasher`] trait for
|
|
/// the `JoinMap`to be useful, see its documentation for details.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// # #[tokio::main]
|
|
/// # async fn main() {
|
|
/// use tokio_util::task::JoinMap;
|
|
/// use std::collections::hash_map::RandomState;
|
|
///
|
|
/// let s = RandomState::new();
|
|
/// let mut map = JoinMap::with_capacity_and_hasher(10, s);
|
|
/// map.spawn(1, async move { "hello world!" });
|
|
/// # }
|
|
/// ```
|
|
#[inline]
|
|
#[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()),
|
|
hashes_by_task: HashMap::with_capacity_and_hasher(capacity, hash_builder),
|
|
tasks: JoinSet::new(),
|
|
}
|
|
}
|
|
|
|
/// Returns the number of tasks currently in the `JoinMap`.
|
|
pub fn len(&self) -> usize {
|
|
let len = self.tasks_by_key.len();
|
|
debug_assert_eq!(len, self.hashes_by_task.len());
|
|
len
|
|
}
|
|
|
|
/// Returns whether the `JoinMap` is empty.
|
|
pub fn is_empty(&self) -> bool {
|
|
let empty = self.tasks_by_key.is_empty();
|
|
debug_assert_eq!(empty, self.hashes_by_task.is_empty());
|
|
empty
|
|
}
|
|
|
|
/// Returns the number of tasks the map can hold without reallocating.
|
|
///
|
|
/// This number is a lower bound; the `JoinMap` might be able to hold
|
|
/// more, but is guaranteed to be able to hold at least this many.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use tokio_util::task::JoinMap;
|
|
///
|
|
/// let map: JoinMap<i32, i32> = JoinMap::with_capacity(100);
|
|
/// assert!(map.capacity() >= 100);
|
|
/// ```
|
|
#[inline]
|
|
pub fn capacity(&self) -> usize {
|
|
let capacity = self.tasks_by_key.capacity();
|
|
debug_assert_eq!(capacity, self.hashes_by_task.capacity());
|
|
capacity
|
|
}
|
|
}
|
|
|
|
impl<K, V, S> JoinMap<K, V, S>
|
|
where
|
|
K: Hash + Eq,
|
|
V: 'static,
|
|
S: BuildHasher,
|
|
{
|
|
/// Spawn the provided task and store it in this `JoinMap` with the provided
|
|
/// key.
|
|
///
|
|
/// If a task previously existed in the `JoinMap` for this key, that task
|
|
/// will be cancelled and replaced with the new one. The previous task will
|
|
/// be removed from the `JoinMap`; a subsequent call to [`join_next`] will
|
|
/// *not* return a cancelled [`JoinError`] for that task.
|
|
///
|
|
/// # Panics
|
|
///
|
|
/// This method panics if called outside of a Tokio runtime.
|
|
///
|
|
/// [`join_next`]: Self::join_next
|
|
#[track_caller]
|
|
pub fn spawn<F>(&mut self, key: K, task: F)
|
|
where
|
|
F: Future<Output = V>,
|
|
F: Send + 'static,
|
|
V: Send,
|
|
{
|
|
let task = self.tasks.spawn(task);
|
|
self.insert(key, task)
|
|
}
|
|
|
|
/// Spawn the provided task on the provided runtime and store it in this
|
|
/// `JoinMap` with the provided key.
|
|
///
|
|
/// If a task previously existed in the `JoinMap` for this key, that task
|
|
/// will be cancelled and replaced with the new one. The previous task will
|
|
/// be removed from the `JoinMap`; a subsequent call to [`join_next`] will
|
|
/// *not* return a cancelled [`JoinError`] for that task.
|
|
///
|
|
/// [`join_next`]: Self::join_next
|
|
#[track_caller]
|
|
pub fn spawn_on<F>(&mut self, key: K, task: F, handle: &Handle)
|
|
where
|
|
F: Future<Output = V>,
|
|
F: Send + 'static,
|
|
V: Send,
|
|
{
|
|
let task = self.tasks.spawn_on(task, handle);
|
|
self.insert(key, task);
|
|
}
|
|
|
|
/// Spawn the provided task on the current [`LocalSet`] and store it in this
|
|
/// `JoinMap` with the provided key.
|
|
///
|
|
/// If a task previously existed in the `JoinMap` for this key, that task
|
|
/// will be cancelled and replaced with the new one. The previous task will
|
|
/// be removed from the `JoinMap`; a subsequent call to [`join_next`] will
|
|
/// *not* return a cancelled [`JoinError`] for that task.
|
|
///
|
|
/// # Panics
|
|
///
|
|
/// This method panics if it is called outside of a `LocalSet`.
|
|
///
|
|
/// [`LocalSet`]: tokio::task::LocalSet
|
|
/// [`join_next`]: Self::join_next
|
|
#[track_caller]
|
|
pub fn spawn_local<F>(&mut self, key: K, task: F)
|
|
where
|
|
F: Future<Output = V>,
|
|
F: 'static,
|
|
{
|
|
let task = self.tasks.spawn_local(task);
|
|
self.insert(key, task);
|
|
}
|
|
|
|
/// Spawn the provided task on the provided [`LocalSet`] and store it in
|
|
/// this `JoinMap` with the provided key.
|
|
///
|
|
/// If a task previously existed in the `JoinMap` for this key, that task
|
|
/// will be cancelled and replaced with the new one. The previous task will
|
|
/// be removed from the `JoinMap`; a subsequent call to [`join_next`] will
|
|
/// *not* return a cancelled [`JoinError`] for that task.
|
|
///
|
|
/// [`LocalSet`]: tokio::task::LocalSet
|
|
/// [`join_next`]: Self::join_next
|
|
#[track_caller]
|
|
pub fn spawn_local_on<F>(&mut self, key: K, task: F, local_set: &LocalSet)
|
|
where
|
|
F: Future<Output = V>,
|
|
F: 'static,
|
|
{
|
|
let task = self.tasks.spawn_local_on(task, local_set);
|
|
self.insert(key, task)
|
|
}
|
|
|
|
fn insert(&mut self, key: K, abort: AbortHandle) {
|
|
let hash = self.hash(&key);
|
|
let id = abort.id();
|
|
let map_key = Key {
|
|
id: id.clone(),
|
|
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);
|
|
match entry {
|
|
RawEntryMut::Occupied(mut 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);
|
|
debug_assert_eq!(Some(hash), _prev_hash);
|
|
}
|
|
RawEntryMut::Vacant(vac) => {
|
|
vac.insert(map_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");
|
|
}
|
|
|
|
/// Waits until one of the tasks in the map completes and returns its
|
|
/// output, along with the key corresponding to that task.
|
|
///
|
|
/// Returns `None` if the map is empty.
|
|
///
|
|
/// # Cancel Safety
|
|
///
|
|
/// This method is cancel safe. If `join_next` is used as the event in a [`tokio::select!`]
|
|
/// statement and some other branch completes first, it is guaranteed that no tasks were
|
|
/// removed from this `JoinMap`.
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// This function returns:
|
|
///
|
|
/// * `Some((key, Ok(value)))` if one of the tasks in this `JoinMap` has
|
|
/// completed. The `value` is the return value of that ask, and `key` is
|
|
/// the key associated with the task.
|
|
/// * `Some((key, Err(err))` if one of the tasks in this JoinMap` has
|
|
/// panicked or been aborted. `key` is the key associated with the task
|
|
/// that panicked or was aborted.
|
|
/// * `None` if the `JoinMap` is empty.
|
|
///
|
|
/// [`tokio::select!`]: tokio::select
|
|
pub async fn join_next(&mut self) -> Option<(K, Result<V, JoinError>)> {
|
|
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,
|
|
};
|
|
let key = self.remove_by_id(id)?;
|
|
Some((key, res))
|
|
}
|
|
|
|
/// Aborts all tasks and waits for them to finish shutting down.
|
|
///
|
|
/// Calling this method is equivalent to calling [`abort_all`] and then calling [`join_next`] in
|
|
/// a loop until it returns `None`.
|
|
///
|
|
/// This method ignores any panics in the tasks shutting down. When this call returns, the
|
|
/// `JoinMap` will be empty.
|
|
///
|
|
/// [`abort_all`]: fn@Self::abort_all
|
|
/// [`join_next`]: fn@Self::join_next
|
|
pub async fn shutdown(&mut self) {
|
|
self.abort_all();
|
|
while self.join_next().await.is_some() {}
|
|
}
|
|
|
|
/// Abort the task corresponding to the provided `key`.
|
|
///
|
|
/// If this `JoinMap` contains a task corresponding to `key`, this method
|
|
/// will abort that task and return `true`. Otherwise, if no task exists for
|
|
/// `key`, this method returns `false`.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// Aborting a task by key:
|
|
///
|
|
/// ```
|
|
/// use tokio_util::task::JoinMap;
|
|
///
|
|
/// # #[tokio::main]
|
|
/// # async fn main() {
|
|
/// let mut map = JoinMap::new();
|
|
///
|
|
/// map.spawn("hello world", async move { /* ... */ });
|
|
/// map.spawn("goodbye world", async move { /* ... */});
|
|
///
|
|
/// // Look up the "goodbye world" task in the map and abort it.
|
|
/// map.abort("goodbye world");
|
|
///
|
|
/// while let Some((key, res)) = map.join_next().await {
|
|
/// if key == "goodbye world" {
|
|
/// // The aborted task should complete with a cancelled `JoinError`.
|
|
/// assert!(res.unwrap_err().is_cancelled());
|
|
/// } else {
|
|
/// // Other tasks should complete normally.
|
|
/// assert!(res.is_ok());
|
|
/// }
|
|
/// }
|
|
/// # }
|
|
/// ```
|
|
///
|
|
/// `abort` returns `true` if a task was aborted:
|
|
/// ```
|
|
/// use tokio_util::task::JoinMap;
|
|
///
|
|
/// # #[tokio::main]
|
|
/// # async fn main() {
|
|
/// let mut map = JoinMap::new();
|
|
///
|
|
/// map.spawn("hello world", async move { /* ... */ });
|
|
/// map.spawn("goodbye world", async move { /* ... */});
|
|
///
|
|
/// // A task for the key "goodbye world" should exist in the map:
|
|
/// assert!(map.abort("goodbye world"));
|
|
///
|
|
/// // Aborting a key that does not exist will return `false`:
|
|
/// assert!(!map.abort("goodbye universe"));
|
|
/// # }
|
|
/// ```
|
|
pub fn abort<Q: ?Sized>(&mut self, key: &Q) -> bool
|
|
where
|
|
Q: Hash + Eq,
|
|
K: Borrow<Q>,
|
|
{
|
|
match self.get_by_key(key) {
|
|
Some((_, handle)) => {
|
|
handle.abort();
|
|
true
|
|
}
|
|
None => false,
|
|
}
|
|
}
|
|
|
|
/// Aborts all tasks with keys matching `predicate`.
|
|
///
|
|
/// `predicate` is a function called with a reference to each key in the
|
|
/// map. If it returns `true` for a given key, the corresponding task will
|
|
/// be cancelled.
|
|
///
|
|
/// # Examples
|
|
/// ```
|
|
/// use tokio_util::task::JoinMap;
|
|
///
|
|
/// # // use the current thread rt so that spawned tasks don't
|
|
/// # // complete in the background before they can be aborted.
|
|
/// # #[tokio::main(flavor = "current_thread")]
|
|
/// # async fn main() {
|
|
/// let mut map = JoinMap::new();
|
|
///
|
|
/// map.spawn("hello world", async move {
|
|
/// // ...
|
|
/// # tokio::task::yield_now().await; // don't complete immediately, get aborted!
|
|
/// });
|
|
/// map.spawn("goodbye world", async move {
|
|
/// // ...
|
|
/// # tokio::task::yield_now().await; // don't complete immediately, get aborted!
|
|
/// });
|
|
/// map.spawn("hello san francisco", async move {
|
|
/// // ...
|
|
/// # tokio::task::yield_now().await; // don't complete immediately, get aborted!
|
|
/// });
|
|
/// map.spawn("goodbye universe", async move {
|
|
/// // ...
|
|
/// # tokio::task::yield_now().await; // don't complete immediately, get aborted!
|
|
/// });
|
|
///
|
|
/// // Abort all tasks whose keys begin with "goodbye"
|
|
/// map.abort_matching(|key| key.starts_with("goodbye"));
|
|
///
|
|
/// let mut seen = 0;
|
|
/// while let Some((key, res)) = map.join_next().await {
|
|
/// seen += 1;
|
|
/// if key.starts_with("goodbye") {
|
|
/// // The aborted task should complete with a cancelled `JoinError`.
|
|
/// assert!(res.unwrap_err().is_cancelled());
|
|
/// } else {
|
|
/// // Other tasks should complete normally.
|
|
/// assert!(key.starts_with("hello"));
|
|
/// assert!(res.is_ok());
|
|
/// }
|
|
/// }
|
|
///
|
|
/// // All spawned tasks should have completed.
|
|
/// assert_eq!(seen, 4);
|
|
/// # }
|
|
/// ```
|
|
pub fn abort_matching(&mut self, mut predicate: impl FnMut(&K) -> bool) {
|
|
// 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 {
|
|
if predicate(key) {
|
|
task.abort();
|
|
}
|
|
}
|
|
}
|
|
|
|
/// 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
|
|
/// call to [`join_next`], this method will still return `true`.
|
|
///
|
|
/// [`join_next`]: fn@Self::join_next
|
|
pub fn contains_key<Q: ?Sized>(&self, key: &Q) -> bool
|
|
where
|
|
Q: Hash + Eq,
|
|
K: Borrow<Q>,
|
|
{
|
|
self.get_by_key(key).is_some()
|
|
}
|
|
|
|
/// Returns `true` if this `JoinMap` contains a task with the provided
|
|
/// [task ID].
|
|
///
|
|
/// If the task has completed, but its output hasn't yet been consumed by a
|
|
/// call to [`join_next`], this method will still return `true`.
|
|
///
|
|
/// [`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()
|
|
}
|
|
|
|
/// Reserves capacity for at least `additional` more tasks to be spawned
|
|
/// on this `JoinMap` without reallocating for the map of task keys. The
|
|
/// collection may reserve more space to avoid frequent reallocations.
|
|
///
|
|
/// Note that spawning a task will still cause an allocation for the task
|
|
/// itself.
|
|
///
|
|
/// # Panics
|
|
///
|
|
/// Panics if the new allocation size overflows [`usize`].
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use tokio_util::task::JoinMap;
|
|
///
|
|
/// let mut map: JoinMap<&str, i32> = JoinMap::new();
|
|
/// map.reserve(10);
|
|
/// ```
|
|
#[inline]
|
|
pub fn reserve(&mut self, additional: usize) {
|
|
self.tasks_by_key.reserve(additional);
|
|
self.hashes_by_task.reserve(additional);
|
|
}
|
|
|
|
/// Shrinks the capacity of the `JoinMap` as much as possible. It will drop
|
|
/// down as much as possible while maintaining the internal rules
|
|
/// and possibly leaving some space in accordance with the resize policy.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// # #[tokio::main]
|
|
/// # async fn main() {
|
|
/// use tokio_util::task::JoinMap;
|
|
///
|
|
/// let mut map: JoinMap<i32, i32> = JoinMap::with_capacity(100);
|
|
/// map.spawn(1, async move { 2 });
|
|
/// map.spawn(3, async move { 4 });
|
|
/// assert!(map.capacity() >= 100);
|
|
/// map.shrink_to_fit();
|
|
/// assert!(map.capacity() >= 2);
|
|
/// # }
|
|
/// ```
|
|
#[inline]
|
|
pub fn shrink_to_fit(&mut self) {
|
|
self.hashes_by_task.shrink_to_fit();
|
|
self.tasks_by_key.shrink_to_fit();
|
|
}
|
|
|
|
/// Shrinks the capacity of the map with a lower limit. It will drop
|
|
/// down no lower than the supplied limit while maintaining the internal rules
|
|
/// and possibly leaving some space in accordance with the resize policy.
|
|
///
|
|
/// If the current capacity is less than the lower limit, this is a no-op.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// # #[tokio::main]
|
|
/// # async fn main() {
|
|
/// use tokio_util::task::JoinMap;
|
|
///
|
|
/// let mut map: JoinMap<i32, i32> = JoinMap::with_capacity(100);
|
|
/// map.spawn(1, async move { 2 });
|
|
/// map.spawn(3, async move { 4 });
|
|
/// assert!(map.capacity() >= 100);
|
|
/// map.shrink_to(10);
|
|
/// assert!(map.capacity() >= 10);
|
|
/// map.shrink_to(0);
|
|
/// assert!(map.capacity() >= 2);
|
|
/// # }
|
|
/// ```
|
|
#[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)
|
|
}
|
|
|
|
/// 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)>
|
|
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)
|
|
}
|
|
|
|
/// Remove a task from the map by ID, returning the key for that task.
|
|
fn remove_by_id(&mut self, id: Id) -> Option<K> {
|
|
// Get the hash for the given ID.
|
|
let hash = self.hashes_by_task.remove(&id)?;
|
|
|
|
// 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(),
|
|
_ => 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()
|
|
}
|
|
}
|
|
|
|
impl<K, V, S> JoinMap<K, V, S>
|
|
where
|
|
V: 'static,
|
|
{
|
|
/// Aborts all tasks on this `JoinMap`.
|
|
///
|
|
/// This does not remove the tasks from the `JoinMap`. To wait for the tasks to complete
|
|
/// cancellation, you should call `join_next` in a loop until the `JoinMap` is empty.
|
|
pub fn abort_all(&mut self) {
|
|
self.tasks.abort_all()
|
|
}
|
|
|
|
/// Removes all tasks from this `JoinMap` without aborting them.
|
|
///
|
|
/// The tasks removed by this call will continue to run in the background even if the `JoinMap`
|
|
/// is dropped. They may still be aborted by key.
|
|
pub fn detach_all(&mut self) {
|
|
self.tasks.detach_all();
|
|
self.tasks_by_key.clear();
|
|
self.hashes_by_task.clear();
|
|
}
|
|
}
|
|
|
|
// Hand-written `fmt::Debug` implementation in order to avoid requiring `V:
|
|
// Debug`, since no value is ever actually stored in the map.
|
|
impl<K: fmt::Debug, V, S> fmt::Debug for JoinMap<K, V, S> {
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
// format the task keys and abort handles a little nicer by just
|
|
// 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> {
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
f.debug_map()
|
|
.entries(self.0.keys().map(|Key { key, id }| (key, id)))
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
f.debug_struct("JoinMap")
|
|
// The `tasks_by_key` map is the only one that contains information
|
|
// that's really worth formatting for the user, since it contains
|
|
// the tasks' keys and IDs. The other fields are basically
|
|
// implementation details.
|
|
.field("tasks", &KeySet(&self.tasks_by_key))
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
impl<K, V> Default for JoinMap<K, V> {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
// === 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> {}
|