tokio_task_supervisor/

lib.rs

use std::{future::Future, time::Duration};

use tokio::runtime::Handle;
use tokio::task::{JoinHandle, LocalSet};
use tokio_util::sync::{CancellationToken, DropGuardRef};

pub use tokio_util::task::task_tracker::{
    TaskTracker, TaskTrackerToken, TaskTrackerWaitFuture, TrackedFuture,
};

/// The outcome of a task that races against cancellation.
///
/// This enum is returned by [`spawn_with_cancel`](TaskManager::spawn_with_cancel) variants
/// to indicate whether the task completed normally or was cancelled.
#[derive(Debug, PartialEq, Eq)]
pub enum CancelOutcome<T> {
    /// The task future completed before cancellation was requested.
    Completed(T),
    /// Cancellation won the race; the task future was dropped.
    Cancelled,
}

impl<T> CancelOutcome<T> {
    /// Creates a `CancelOutcome` from an optional result.
    ///
    /// # Arguments
    ///
    /// * `result` - `Some(value)` indicates completion, `None` indicates cancellation.
    #[inline]
    pub fn outcome(result: Option<T>) -> CancelOutcome<T> {
        match result {
            Some(value) => CancelOutcome::Completed(value),
            None => CancelOutcome::Cancelled,
        }
    }
}

/// Manages a collection of asynchronous tasks and coordinates their shutdown.
///
/// `TaskSupervisor` wraps [`TaskTracker`] to keep count of outstanding tasks while also exposing a
/// process-wide [`CancellationToken`] that can be used to request cooperative shutdown.
///
/// # Examples
///
/// ```rust
/// use tokio_task_supervisor::TaskSupervisor;
/// use tokio::time::{sleep, Duration};
///
/// #[tokio::main]
/// async fn main() {
///     let supervisor = TaskSupervisor::new();
///     
///     // Spawn a task that cooperatively handles cancellation
///     let handle = supervisor.spawn_with_token(|token| async move {
///         loop {
///             if token.is_cancelled() {
///                 break;
///             }
///             // Do work...
///             sleep(Duration::from_millis(100)).await;
///         }
///     });
///     
///     // Later, request shutdown
///     supervisor.shutdown().await;
/// }
/// ```
#[derive(Clone)]
pub struct TaskSupervisor {
    tracker: TaskTracker,
    shutdown: CancellationToken,
}

impl TaskSupervisor {
    // === Construction ===

    /// Creates a new task manager.
    #[must_use]
    pub fn new() -> Self {
        Self {
            tracker: TaskTracker::new(),
            shutdown: CancellationToken::new(),
        }
    }

    // === Accessors ===

    /// Returns a reference to the underlying [`TaskTracker`].
    #[inline]
    pub fn tracker(&self) -> &TaskTracker {
        &self.tracker
    }

    /// Returns a clone of the shared cancellation token.
    #[inline]
    pub fn token(&self) -> CancellationToken {
        self.shutdown.clone()
    }

    /// Returns a guard that cancels the shutdown token when dropped.
    #[must_use]
    #[inline]
    pub fn cancel_on_drop(&self) -> DropGuardRef<'_> {
        self.shutdown.drop_guard_ref()
    }

    // === State Queries ===

    /// Returns `true` if the shutdown token has been cancelled.
    #[inline]
    pub fn is_cancelled(&self) -> bool {
        self.shutdown.is_cancelled()
    }

    /// Returns `true` if the task tracker is closed.
    #[inline]
    pub fn is_closed(&self) -> bool {
        self.tracker.is_closed()
    }

    /// Returns the number of outstanding tasks.
    #[inline]
    pub fn len(&self) -> usize {
        self.tracker.len()
    }

    /// Returns a future that completes when all tasks finish.
    #[inline]
    pub fn wait(&self) -> TaskTrackerWaitFuture<'_> {
        self.tracker.wait()
    }

    // === Control Operations ===

    /// Cancels the shared shutdown token.
    ///
    /// Tasks spawned through the managed API can observe this and exit cooperatively.
    #[inline]
    pub fn cancel(&self) {
        self.shutdown.cancel();
    }

    /// Initiates graceful shutdown by closing the tracker and cancelling all tasks.
    ///
    /// This method will:
    /// 1. Close the task tracker to prevent new tasks from being spawned
    /// 2. Cancel the shutdown token to signal all existing tasks
    /// 3. Wait for all tasks to complete
    pub async fn shutdown(&self) {
        self.tracker.close();
        self.shutdown.cancel();
        self.tracker.wait().await;
    }

    /// Initiates graceful shutdown with a timeout.
    ///
    /// # Arguments
    ///
    /// * `timeout` - Maximum time to wait for shutdown to complete
    ///
    /// # Returns
    ///
    /// * `Ok(())` if shutdown completed within the timeout
    /// * `Err(Elapsed)` if the timeout was exceeded
    #[inline]
    pub async fn shutdown_with_timeout(
        &self,
        timeout: Duration,
    ) -> Result<(), tokio::time::error::Elapsed> {
        tokio::time::timeout(timeout, self.shutdown()).await
    }

    // === Spawn Methods with Cancellation Race ===

    /// Spawns a task that races against the shared cancellation token.
    ///
    /// The returned future resolves with [`CancelOutcome`], indicating whether the task finished
    /// normally or was cancelled. When cancellation wins the race, the task future is dropped, so it
    /// should not rely on `Drop` for cleanup.
    ///
    /// # Arguments
    ///
    /// * `task` - A closure that returns the future to execute
    ///
    /// # Returns
    ///
    /// A `JoinHandle` that resolves to the task's outcome
    #[must_use]
    pub fn spawn_with_cancel<F, Fut>(&self, task: F) -> JoinHandle<CancelOutcome<Fut::Output>>
    where
        F: FnOnce() -> Fut + Send + 'static,
        Fut: Future + Send + 'static,
        Fut::Output: Send + 'static,
    {
        let token = self.token();
        self.tracker
            .spawn(async move { CancelOutcome::outcome(token.run_until_cancelled(task()).await) })
    }

    /// Spawns a task with cancellation handling on a specific runtime handle.
    ///
    /// # Arguments
    ///
    /// * `task` - A closure that returns the future to execute
    /// * `handle` - The runtime handle to spawn the task on
    ///
    /// # Returns
    ///
    /// A `JoinHandle` that resolves to the task's outcome
    #[must_use]
    pub fn spawn_on_with_cancel<F, Fut>(
        &self,
        task: F,
        handle: &Handle,
    ) -> JoinHandle<CancelOutcome<Fut::Output>>
    where
        F: FnOnce() -> Fut + Send + 'static,
        Fut: Future + Send + 'static,
        Fut::Output: Send + 'static,
    {
        let token = self.token();
        self.tracker.spawn_on(
            async move { CancelOutcome::outcome(token.run_until_cancelled(task()).await) },
            handle,
        )
    }

    /// Spawns a !Send task that races against the shared cancellation token.
    ///
    /// # Arguments
    ///
    /// * `task` - A closure that returns the future to execute
    ///
    /// # Returns
    ///
    /// A `JoinHandle` that resolves to the task's outcome
    #[must_use]
    pub fn spawn_local_with_cancel<F, Fut>(&self, task: F) -> JoinHandle<CancelOutcome<Fut::Output>>
    where
        F: FnOnce() -> Fut + 'static,
        Fut: Future + 'static,
        Fut::Output: 'static,
    {
        let token = self.token();
        self.tracker.spawn_local(async move {
            CancelOutcome::outcome(token.run_until_cancelled(task()).await)
        })
    }

    /// Spawns a !Send task on a [`LocalSet`] with cancellation handling.
    ///
    /// # Arguments
    ///
    /// * `task` - A closure that returns the future to execute
    /// * `local_set` - The local set to spawn the task on
    ///
    /// # Returns
    ///
    /// A `JoinHandle` that resolves to the task's outcome
    #[must_use]
    pub fn spawn_local_on_with_cancel<F, Fut>(
        &self,
        task: F,
        local_set: &LocalSet,
    ) -> JoinHandle<CancelOutcome<Fut::Output>>
    where
        F: FnOnce() -> Fut + 'static,
        Fut: Future + 'static,
        Fut::Output: 'static,
    {
        let token = self.token();
        self.tracker.spawn_local_on(
            async move { CancelOutcome::outcome(token.run_until_cancelled(task()).await) },
            local_set,
        )
    }

    // === Spawn Methods with Token ===

    /// Spawns a task that receives the shared cancellation token.
    ///
    /// # Arguments
    ///
    /// * `task` - A closure that takes a cancellation token and returns a future
    ///
    /// # Returns
    ///
    /// A `JoinHandle` that resolves to the task's output
    #[must_use]
    pub fn spawn_with_token<F, Fut>(&self, task: F) -> JoinHandle<Fut::Output>
    where
        F: FnOnce(CancellationToken) -> Fut + Send + 'static,
        Fut: Future + Send + 'static,
        Fut::Output: Send + 'static,
    {
        let token = self.shutdown.child_token();
        self.tracker.spawn(async move { task(token).await })
    }

    /// Spawns a task with the shared cancellation token on a specific runtime handle.
    ///
    /// # Arguments
    ///
    /// * `task` - A closure that takes a cancellation token and returns a future
    /// * `handle` - The runtime handle to spawn the task on
    ///
    /// # Returns
    ///
    /// A `JoinHandle` that resolves to the task's output
    #[must_use]
    pub fn spawn_on_with_token<F, Fut>(&self, task: F, handle: &Handle) -> JoinHandle<Fut::Output>
    where
        F: FnOnce(CancellationToken) -> Fut + Send + 'static,
        Fut: Future + Send + 'static,
        Fut::Output: Send + 'static,
    {
        let token = self.shutdown.child_token();
        self.tracker
            .spawn_on(async move { task(token).await }, handle)
    }

    /// Spawns a local task that receives the shared cancellation token.
    ///
    /// # Arguments
    ///
    /// * `task` - A closure that takes a cancellation token and returns a future
    ///
    /// # Returns
    ///
    /// A `JoinHandle` that resolves to the task's output
    #[must_use]
    pub fn spawn_local_with_token<F, Fut>(&self, task: F) -> JoinHandle<Fut::Output>
    where
        F: FnOnce(CancellationToken) -> Fut + 'static,
        Fut: Future + 'static,
        Fut::Output: 'static,
    {
        let token = self.shutdown.child_token();
        self.tracker.spawn_local(async move { task(token).await })
    }

    /// Spawns a local task with a cancellation token on a specific local set.
    ///
    /// # Arguments
    ///
    /// * `task` - A closure that takes a cancellation token and returns a future
    /// * `local_set` - The local set to spawn the task on
    ///
    /// # Returns
    ///
    /// A `JoinHandle` that resolves to the task's output
    #[must_use]
    pub fn spawn_local_on_with_token<F, Fut>(
        &self,
        task: F,
        local_set: &LocalSet,
    ) -> JoinHandle<Fut::Output>
    where
        F: FnOnce(CancellationToken) -> Fut + 'static,
        Fut: Future + 'static,
        Fut::Output: 'static,
    {
        let token = self.shutdown.child_token();
        self.tracker
            .spawn_local_on(async move { task(token).await }, local_set)
    }

    /// Spawns a blocking task that receives a cancellation context.
    ///
    /// # Arguments
    ///
    /// * `task` - A closure that takes a cancellation token and returns a value
    ///
    /// # Returns
    ///
    /// A `JoinHandle` that resolves to the task's output
    #[cfg(not(target_family = "wasm"))]
    #[must_use]
    pub fn spawn_blocking_with_token<F, T>(&self, task: F) -> JoinHandle<T>
    where
        F: FnOnce(CancellationToken) -> T + Send + 'static,
        T: Send + 'static,
    {
        let token = self.shutdown.child_token();
        self.tracker.spawn_blocking(move || task(token))
    }

    /// Spawns a blocking task with context on a specific runtime handle.
    ///
    /// # Arguments
    ///
    /// * `task` - A closure that takes a cancellation token and returns a value
    /// * `handle` - The runtime handle to spawn the task on
    ///
    /// # Returns
    ///
    /// A `JoinHandle` that resolves to the task's output
    #[cfg(not(target_family = "wasm"))]
    #[must_use]
    pub fn spawn_blocking_on_with_token<F, T>(&self, task: F, handle: &Handle) -> JoinHandle<T>
    where
        F: FnOnce(CancellationToken) -> T + Send + 'static,
        T: Send + 'static,
    {
        let token = self.shutdown.child_token();
        self.tracker.spawn_blocking_on(move || task(token), handle)
    }
}

impl Default for TaskSupervisor {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
    use std::sync::Arc;
    use tokio::time::{sleep, Duration};

    #[tokio::test]
    async fn test_spawn_with_token_provides_token() {
        let supervisor = TaskSupervisor::new();

        let handle = supervisor.spawn_with_token(|token| async move { token.is_cancelled() });

        let result = handle.await.unwrap();
        assert!(!result);
    }

    #[tokio::test]
    #[cfg(feature = "rt")]
    async fn test_spawn_on_with_token_provides_token() {
        let supervisor = TaskSupervisor::new();
        let handle = tokio::runtime::Handle::current();

        let result = supervisor
            .spawn_on_with_token(|token| async move { token.is_cancelled() }, &handle)
            .await
            .unwrap();

        assert!(!result);
    }

    #[tokio::test]
    #[cfg(all(feature = "rt", not(target_family = "wasm")))]
    async fn test_spawn_blocking_with_token_provides_token() {
        let supervisor = TaskSupervisor::new();

        let handle = supervisor.spawn_blocking_with_token(move |token| token.is_cancelled());

        let result = handle.await.unwrap();
        assert!(!result);
    }

    #[tokio::test]
    async fn test_cancel_sets_cancelled_state() {
        let supervisor = TaskSupervisor::new();
        assert!(!supervisor.is_cancelled());

        supervisor.cancel();
        assert!(supervisor.is_cancelled());
    }

    #[tokio::test]
    async fn test_cancel_propagates_to_all_tasks() {
        let supervisor = TaskSupervisor::new();
        let count = Arc::new(AtomicUsize::new(0));

        for _ in 0..3 {
            let count_clone = count.clone();
            let _ = supervisor.spawn_with_token(|token| async move {
                token.cancelled().await;
                count_clone.fetch_add(1, Ordering::SeqCst);
            });
        }

        sleep(Duration::from_millis(50)).await;
        supervisor.cancel();
        sleep(Duration::from_millis(100)).await;

        assert_eq!(count.load(Ordering::SeqCst), 3);
    }

    #[tokio::test]
    async fn test_shutdown_cancels_and_waits() {
        let supervisor = TaskSupervisor::new();
        let task_finished = Arc::new(AtomicBool::new(false));
        let task_finished_clone = task_finished.clone();

        let _ = supervisor.spawn_with_token(|_token| async move {
            sleep(Duration::from_millis(100)).await;
            task_finished_clone.store(true, Ordering::SeqCst);
        });

        assert!(!supervisor.is_cancelled());
        assert!(!supervisor.is_closed());

        supervisor.shutdown().await;

        assert!(supervisor.is_cancelled());
        assert!(supervisor.is_closed());
        assert!(task_finished.load(Ordering::SeqCst));
        assert_eq!(supervisor.len(), 0);
    }

    #[tokio::test]
    async fn test_shutdown_with_timeout_completes_in_time() {
        let supervisor = TaskSupervisor::new();

        let _ = supervisor.spawn_with_token(|_token| async move {
            sleep(Duration::from_millis(50)).await;
        });

        let result = supervisor
            .shutdown_with_timeout(Duration::from_secs(1))
            .await;
        assert!(result.is_ok());
        assert!(supervisor.is_cancelled());
        assert!(supervisor.is_closed());
    }

    #[tokio::test]
    async fn test_shutdown_with_timeout_times_out() {
        let supervisor = TaskSupervisor::new();

        let _ = supervisor.tracker().spawn(async {
            sleep(Duration::from_secs(10)).await;
        });

        let result = supervisor
            .shutdown_with_timeout(Duration::from_millis(50))
            .await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_cooperative_cancellation_in_loop() {
        let supervisor = TaskSupervisor::new();
        let iterations = Arc::new(AtomicUsize::new(0));
        let iterations_clone = iterations.clone();

        let _ = supervisor.spawn_with_token(|token| async move {
            loop {
                if token.is_cancelled() {
                    break;
                }
                iterations_clone.fetch_add(1, Ordering::SeqCst);
                sleep(Duration::from_millis(10)).await;
            }
        });

        sleep(Duration::from_millis(55)).await;
        supervisor.cancel();
        sleep(Duration::from_millis(50)).await;

        let count = iterations.load(Ordering::SeqCst);
        assert!(count >= 3 && count < 20);
    }

    #[tokio::test]
    async fn test_spawn_with_cancel_completes() {
        let supervisor = TaskSupervisor::new();

        let handle = supervisor.spawn_with_cancel(|| async move {
            sleep(Duration::from_millis(30)).await;
            42
        });

        match handle.await.unwrap() {
            CancelOutcome::Completed(value) => assert_eq!(value, 42),
            CancelOutcome::Cancelled => panic!("task should have completed"),
        }
    }

    #[tokio::test]
    async fn test_spawn_with_cancel_reports_cancellation() {
        let supervisor = TaskSupervisor::new();

        let handle = supervisor.spawn_with_cancel(|| async move {
            loop {
                sleep(Duration::from_millis(10)).await;
            }
        });

        sleep(Duration::from_millis(35)).await;
        supervisor.cancel();

        match handle.await.unwrap() {
            CancelOutcome::Completed(_) => panic!("task should have been cancelled"),
            CancelOutcome::Cancelled => {}
        }
    }
}