util: add JoinDeque structure

Author: nickkuk

tokio-util/src/task/join_deque.rs

@@ -0,0 +1,87 @@
+use std::{collections::VecDeque, future::Future};
+use tokio::task::{AbortHandle, JoinError, JoinHandle};
+
+/// A collection of tasks spawned on a Tokio runtime.
+///
+/// A `JoinDeque` can be used to await the completion of the tasks in FIFO
+/// order. That is, if tasks are spawned in the order A, B, C, then
+/// awaiting the next completed task will always return A first, then B,
+/// then C, regardless of the order in which the tasks actually complete.
+///
+/// All of the tasks must have the same return type `T`.
+///
+/// When the `JoinDeque` is dropped, all tasks in the `JoinDeque` are
+/// immediately aborted.
+#[derive(Debug)]
+pub struct JoinDeque<T>(VecDeque<JoinHandle<T>>);
+
+impl<T> Default for JoinDeque<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<T> Drop for JoinDeque<T> {
+    fn drop(&mut self) {
+        for join_handle in self.0.drain(..) {
+            join_handle.abort();
+        }
+    }
+}
+
+impl<T> JoinDeque<T> {
+    /// Create a new empty `JoinDeque`.
+    pub const fn new() -> Self {
+        Self(VecDeque::new())
+    }
+    /// Creates an empty `JoinDeque` with space for at least `capacity` tasks.
+    pub fn with_capacity(capacity: usize) -> Self {
+        Self(VecDeque::with_capacity(capacity))
+    }
+    /// Returns the number of tasks currently in the `JoinDeque`.
+    ///
+    /// This includes both tasks that are currently running and tasks that have
+    /// completed but not yet been removed from the queue because outputting of
+    /// them waits for FIFO order.
+    pub fn len(&self) -> usize {
+        self.0.len()
+    }
+    /// Returns whether the `JoinDeque` is empty.
+    pub fn is_empty(&self) -> bool {
+        self.0.is_empty()
+    }
+    /// Spawn the provided task on the `JoinDeque`, returning an [`AbortHandle`]
+    /// that can be used to remotely cancel the task.
+    ///
+    /// The provided future will start running in the background immediately
+    /// when this method is called, even if you don't await anything on this
+    /// `JoinDeque`.
+    ///
+    /// # Panics
+    ///
+    /// This method panics if called outside of a Tokio runtime.
+    pub fn spawn<F>(&mut self, task: F) -> AbortHandle
+    where
+        F: Future<Output = T> + Send + 'static,
+        T: Send + 'static,
+    {
+        let join_handle = tokio::spawn(task);
+        let abort_handle = join_handle.abort_handle();
+        self.0.push_back(join_handle);
+        abort_handle
+    }
+    /// Waits until the next task in FIFO order completes and returns its output.
+    ///
+    /// Returns `None` if the queue 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 `JoinDeque`.
+    pub async fn join_next(&mut self) -> Option<Result<T, JoinError>> {
+        let result = self.0.front_mut()?.await;
+        self.0.pop_front();
+        Some(result)
+    }
+}

tokio-util/src/task/mod.rs

@@ -13,6 +13,9 @@ cfg_rt! {
 
     mod abort_on_drop;
     pub use abort_on_drop::AbortOnDropHandle;
+
+    mod join_deque;
+    pub use join_deque::JoinDeque;
 }
 
 #[cfg(feature = "join-map")]