rskit_stream/

task.rs

//! Cancellable owned tasks.
//!
//! [`SpawnedTask`] is the canonical owner for a background task that must be
//! cooperatively stopped: it bundles a `CancellationToken` with the
//! `JoinHandle`, so every task has explicit ownership, cancellation, and a
//! bounded, drain-then-abort shutdown. [`TaskGroup`] owns a set of such tasks
//! and shuts them all down together. This replaces the hand-rolled
//! `cancel + handle` pairs duplicated across consumers, watchers, and servers.

use std::time::Duration;

use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;

/// A single background task with cooperative cancellation.
///
/// Created by [`SpawnedTask::spawn`], which hands the body a child
/// `CancellationToken` to `select!` on. Stop it with [`shutdown`](Self::shutdown)
/// (cancel → await up to a grace period → abort) so a wedged task can never
/// block shutdown indefinitely. Extra teardown (closing connections, flushing
/// buffers) stays with the caller.
#[derive(Debug)]
pub struct SpawnedTask {
    cancel: CancellationToken,
    handle: JoinHandle<()>,
}

impl SpawnedTask {
    /// Spawn `body` with a fresh cancellation token passed to the closure.
    ///
    /// The body should `select!` on `token.cancelled()` to exit promptly.
    pub fn spawn<F, Fut>(body: F) -> Self
    where
        F: FnOnce(CancellationToken) -> Fut,
        Fut: Future<Output = ()> + Send + 'static,
    {
        let cancel = CancellationToken::new();
        let handle = tokio::spawn(body(cancel.clone()));
        Self { cancel, handle }
    }

    /// Spawn `body` with the supplied cancellation token (e.g. a shared parent).
    pub fn spawn_with<Fut>(cancel: CancellationToken, body: Fut) -> Self
    where
        Fut: Future<Output = ()> + Send + 'static,
    {
        let handle = tokio::spawn(body);
        Self { cancel, handle }
    }

    /// Adopt an already-spawned `JoinHandle` paired with the token that
    /// stops it. For callers that spawn the task themselves but still want
    /// the bounded drain-then-abort shutdown.
    #[must_use]
    pub const fn from_parts(cancel: CancellationToken, handle: JoinHandle<()>) -> Self {
        Self { cancel, handle }
    }

    /// Signal the task to stop without waiting for it to finish.
    pub fn cancel(&self) {
        self.cancel.cancel();
    }

    /// Return `true` once the task has finished.
    #[must_use]
    pub fn is_finished(&self) -> bool {
        self.handle.is_finished()
    }

    /// A clone of the task's cancellation token, for cooperative shutdown.
    #[must_use]
    pub fn cancellation(&self) -> CancellationToken {
        self.cancel.clone()
    }

    /// Cancel the task and wait up to `grace` for it to drain, aborting if it
    /// overruns. Always resolves; never blocks shutdown indefinitely.
    pub async fn shutdown(self, grace: Duration) {
        self.cancel.cancel();
        let mut handle = self.handle;
        if tokio::time::timeout(grace, &mut handle).await.is_err() {
            handle.abort();
            let _ = handle.await;
        }
    }

    /// Abort the task immediately without waiting.
    pub fn abort(self) {
        self.handle.abort();
    }

    /// Wait for the task to finish without cancelling it.
    pub async fn join(self) {
        let _ = self.handle.await;
    }
}

/// A set of [`SpawnedTask`]s shut down together.
///
/// Each task owns its own cancellation token; [`shutdown`](Self::shutdown)
/// cancels all of them, then drains each within a per-task grace period,
/// aborting stragglers so a single wedged task cannot stall teardown.
#[derive(Debug, Default)]
pub struct TaskGroup {
    tasks: Vec<SpawnedTask>,
}

impl TaskGroup {
    /// Create an empty group.
    #[must_use]
    pub const fn new() -> Self {
        Self { tasks: Vec::new() }
    }

    /// Add a task to the group.
    pub fn push(&mut self, task: SpawnedTask) {
        self.tasks.push(task);
    }

    /// Number of tasks currently owned.
    #[must_use]
    pub const fn len(&self) -> usize {
        self.tasks.len()
    }

    /// Whether the group owns no tasks.
    #[must_use]
    pub const fn is_empty(&self) -> bool {
        self.tasks.is_empty()
    }

    /// Cancel every task without waiting.
    pub fn cancel_all(&self) {
        for task in &self.tasks {
            task.cancel();
        }
    }

    /// Cancel and drain every task, each within `grace`, aborting stragglers.
    pub async fn shutdown(&mut self, grace: Duration) {
        let tasks = std::mem::take(&mut self.tasks);
        for task in &tasks {
            task.cancel();
        }
        for task in tasks {
            task.shutdown(grace).await;
        }
    }
}

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

    #[tokio::test]
    async fn shutdown_cancels_cooperative_task() {
        let ran = Arc::new(AtomicBool::new(false));
        let r = ran.clone();
        let task = SpawnedTask::spawn(move |cancel| async move {
            cancel.cancelled().await;
            r.store(true, Ordering::SeqCst);
        });
        task.shutdown(Duration::from_secs(1)).await;
        assert!(ran.load(Ordering::SeqCst));
    }

    #[tokio::test(start_paused = true)]
    async fn shutdown_aborts_wedged_task() {
        let task = SpawnedTask::spawn(|_cancel| async move {
            std::future::pending::<()>().await;
        });
        // Returns even though the body ignores cancellation.
        task.shutdown(Duration::from_millis(50)).await;
    }

    #[tokio::test]
    async fn group_shuts_down_all() {
        let mut group = TaskGroup::new();
        for _ in 0..3 {
            group.push(SpawnedTask::spawn(|cancel| async move {
                cancel.cancelled().await;
            }));
        }
        assert_eq!(group.len(), 3);
        group.shutdown(Duration::from_secs(1)).await;
        assert!(group.is_empty());
    }
}