taskvisor 0.8.0

//! Lets one task attempt to observe cancellation.
//!
//! Taskvisor passes a [`TaskContext`] to every [`Task::spawn`](crate::Task::spawn) call.
//! Await [`TaskContext::cancelled`] in a `tokio::select!`, check [`TaskContext::is_cancelled`]
//! between units of work, or wrap a drop-safe future with [`TaskContext::run_until_cancelled`].
//!
//! Task removal and runtime shutdown cancel the active attempt. A configured timeout cancels only that attempt,
//! then drops its future. Cancellation is cooperative: it cannot interrupt synchronous code, and timeout
//! does not poll the future again to run application cleanup after setting the signal.
//!
//! ```text
//! registry removal or runtime shutdown
//!                 │ actor token
//!                 ▼
//!             TaskActor
//!                 ▼
//!              run_once
//!                 ├── start attempt ──► child token ──► TaskContext
//!                 │                                      ▼
//!                 │                                  Task::spawn
//!                 └── timeout ──► cancel attempt token ──► drop future
//! ```
//!
//! A long-running attempt must observe the signal and return.
//! Use ordinary Rust guards for cleanup that must run when the future is dropped.

use std::future::Future;

use tokio_util::sync::CancellationToken;

use crate::error::TaskError;

/// A cloneable cancellation view for one [`Task`](crate::Task) attempt.
///
/// Clones share the same signal.
/// Await [`cancelled`](Self::cancelled), check [`is_cancelled`](Self::is_cancelled), or use
/// [`run_until_cancelled`](Self::run_until_cancelled) around a future that is safe to stop by dropping.
#[derive(Clone, Debug)]
pub struct TaskContext {
    cancel: CancellationToken,
}

impl TaskContext {
    /// Wraps the cancellation token created for one attempt.
    pub(crate) fn from_token(cancel: CancellationToken) -> Self {
        Self { cancel }
    }

    /// Creates an active context with no supervisor connection.
    ///
    /// Use it to test code that accepts a context directly.
    ///
    /// ```rust
    /// use taskvisor::TaskContext;
    ///
    /// let ctx = TaskContext::detached();
    /// assert!(!ctx.is_cancelled());
    /// ```
    #[cfg(feature = "test-util")]
    #[cfg_attr(docsrs, doc(cfg(feature = "test-util")))]
    #[must_use]
    pub fn detached() -> Self {
        Self::from_token(CancellationToken::new())
    }

    /// Creates an already-cancelled context for cancellation-path tests.
    ///
    /// ```rust
    /// use taskvisor::TaskContext;
    ///
    /// let ctx = TaskContext::detached_cancelled();
    /// assert!(ctx.is_cancelled());
    /// ```
    #[cfg(feature = "test-util")]
    #[cfg_attr(docsrs, doc(cfg(feature = "test-util")))]
    #[must_use]
    pub fn detached_cancelled() -> Self {
        let token = CancellationToken::new();
        token.cancel();
        Self::from_token(token)
    }

    /// Waits until cancellation is requested.
    ///
    /// Returns immediately after the signal has been set.
    /// The same context may be awaited more than once or used in `tokio::select!`.
    pub async fn cancelled(&self) {
        self.cancel.cancelled().await;
    }

    /// Returns whether cancellation has been requested.
    #[must_use]
    pub fn is_cancelled(&self) -> bool {
        self.cancel.is_cancelled()
    }

    /// Polls `fut` until it completes or cancellation wins.
    ///
    /// Returns `Ok(output)` when `fut` finishes first.
    /// Cancellation wins a tie and returns [`TaskError::Canceled`].
    /// An already-cancelled context does not poll `fut`.
    /// In both cancellation cases, `fut` is dropped.
    ///
    /// Use this method only with futures that are safe to cancel by dropping.
    ///
    /// ```rust,no_run
    /// use std::time::Duration;
    /// use taskvisor::{TaskFn, TaskRef};
    ///
    /// let poller: TaskRef = TaskFn::arc(|ctx| async move {
    ///     loop {
    ///         ctx.run_until_cancelled(tokio::time::sleep(Duration::from_secs(5)))
    ///             .await?;
    ///         // Do one unit of work.
    ///     }
    /// });
    /// ```
    pub async fn run_until_cancelled<F: Future>(&self, fut: F) -> Result<F::Output, TaskError> {
        tokio::select! {
            biased;
            _ = self.cancelled() => Err(TaskError::Canceled),
            output = fut => Ok(output),
        }
    }

    /// Creates a child cancellation scope.
    ///
    /// Parent cancellation reaches the child.
    /// Cancelling the child through an interop API does not cancel this context.
    #[must_use]
    pub fn child(&self) -> TaskContext {
        TaskContext {
            cancel: self.cancel.child_token(),
        }
    }

    /// Returns a [`CancellationToken`] that shares this context's state.
    ///
    /// This interop method requires the `tokio-util-interop` feature.
    #[cfg(feature = "tokio-util-interop")]
    #[cfg_attr(docsrs, doc(cfg(feature = "tokio-util-interop")))]
    #[must_use]
    pub fn cancellation_token(&self) -> CancellationToken {
        self.cancel.clone()
    }
}

#[cfg(test)]
mod tests {
    use super::TaskContext;
    use crate::error::TaskError;
    use std::sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
    };
    use std::time::Duration;
    use tokio_util::sync::CancellationToken;

    #[cfg(feature = "test-util")]
    #[tokio::test]
    async fn detached_context_starts_active() {
        let ctx = TaskContext::detached();

        assert!(!ctx.is_cancelled());
        let out = ctx.run_until_cancelled(async { 7 }).await;
        assert_eq!(out.expect("detached context starts active"), 7);
    }

    #[cfg(feature = "test-util")]
    #[tokio::test]
    async fn detached_cancelled_context_short_circuits() {
        let ctx = TaskContext::detached_cancelled();

        assert!(ctx.is_cancelled());
        let out = ctx.run_until_cancelled(async { 7 }).await;
        assert!(
            matches!(out, Err(TaskError::Canceled)),
            "an already-cancelled context must win the race"
        );
    }

    #[tokio::test]
    async fn run_until_cancelled_returns_output_when_future_completes_first() {
        let ctx = TaskContext::from_token(CancellationToken::new());

        let result = ctx.run_until_cancelled(async { 7 }).await;

        assert_eq!(
            result.expect("future completed without cancellation, must yield Ok"),
            7,
            "the future's output must pass through unchanged"
        );
    }

    #[tokio::test]
    async fn run_until_cancelled_returns_canceled_when_cancelled_mid_flight() {
        let token = CancellationToken::new();
        let ctx = TaskContext::from_token(token.clone());

        let running =
            tokio::spawn(
                async move { ctx.run_until_cancelled(std::future::pending::<()>()).await },
            );
        token.cancel();

        let result = tokio::time::timeout(Duration::from_secs(1), running)
            .await
            .expect("run_until_cancelled must resolve promptly after cancellation")
            .expect("the spawned task must not panic");
        assert!(
            matches!(result, Err(TaskError::Canceled)),
            "cancellation mid-flight must yield Err(TaskError::Canceled), got {result:?}"
        );
    }

    #[tokio::test]
    async fn run_until_cancelled_does_not_poll_future_when_already_cancelled() {
        let token = CancellationToken::new();
        let ctx = TaskContext::from_token(token.clone());
        token.cancel();

        let polled = Arc::new(AtomicBool::new(false));
        let flag = Arc::clone(&polled);
        let result = ctx
            .run_until_cancelled(async move {
                flag.store(true, Ordering::SeqCst);
                7
            })
            .await;

        assert!(
            matches!(result, Err(TaskError::Canceled)),
            "an already-cancelled context must yield Err(TaskError::Canceled), got {result:?}"
        );
        assert!(
            !polled.load(Ordering::SeqCst),
            "the future must not be polled when the context is already cancelled (cancellation wins ties)"
        );
    }

    #[test]
    fn context_and_clone_share_underlying_cancellation_state() {
        let token = CancellationToken::new();
        let ctx = TaskContext::from_token(token.clone());
        let clone = ctx.clone();

        assert!(!ctx.is_cancelled(), "fresh context must not be cancelled");
        assert!(!clone.is_cancelled(), "fresh clone must not be cancelled");
        token.cancel();
        assert!(
            ctx.is_cancelled(),
            "context must observe the underlying token's cancellation"
        );
        assert!(
            clone.is_cancelled(),
            "a cloned context must share cancellation state with the original"
        );
    }

    #[tokio::test]
    async fn cancelled_resolves_once_token_is_cancelled() {
        let token = CancellationToken::new();
        let ctx = TaskContext::from_token(token.clone());
        token.cancel();

        tokio::time::timeout(Duration::from_secs(1), ctx.cancelled())
            .await
            .expect("cancelled() must resolve promptly after the token is cancelled");
    }

    #[test]
    fn child_is_cancelled_when_parent_cancels() {
        let token = CancellationToken::new();
        let ctx = TaskContext::from_token(token.clone());
        let child = ctx.child();

        assert!(!child.is_cancelled(), "fresh child must not be cancelled");
        token.cancel();
        assert!(
            child.is_cancelled(),
            "cancelling the parent must propagate to a child context"
        );
    }

    #[cfg(feature = "tokio-util-interop")]
    #[test]
    fn cancellation_token_shares_state_with_context() {
        let token = CancellationToken::new();
        let ctx = TaskContext::from_token(token.clone());
        let raw = ctx.cancellation_token();

        assert!(!raw.is_cancelled());
        token.cancel();
        assert!(
            raw.is_cancelled(),
            "the raw token must share cancellation state with the context"
        );
    }

    #[cfg(feature = "tokio-util-interop")]
    #[test]
    fn child_cancellation_does_not_affect_parent() {
        let ctx = TaskContext::from_token(CancellationToken::new());
        let child = ctx.child();

        child.cancellation_token().cancel();
        assert!(
            child.is_cancelled(),
            "child must observe its own cancellation"
        );
        assert!(
            !ctx.is_cancelled(),
            "cancelling a child must not cancel the parent (child() must use a child token, not a clone)"
        );
    }
}