swansong 0.3.6

Graceful Shutdown
Documentation
#![forbid(unsafe_code, future_incompatible)]
#![deny(
    missing_debug_implementations,
    nonstandard_style,
    missing_copy_implementations,
    unused_qualifications,
    missing_docs,
    rustdoc::missing_crate_level_docs
)]
#![warn(clippy::pedantic)]

/*!
# 🦢 A graceful shutdown

## How this crate works

Each [`Swansong`] represents a functional unit of behavior that needs to coordinate termination. The
primary type [`Swansong`] provides the control interface.

To allow Swansong to interrupt/cancel/seal various types, use
[`Swansong::interrupt`]. [`Interrupt<T>`] has custom implementations of
[`Future`][std::future::Future], [`Stream`][futures_core::Stream],
[`AsyncRead`][futures_io::AsyncRead], [`AsyncWrite`][futures_io::AsyncWrite],
[`AsyncBufRead`][futures_io::AsyncBufRead], and [`Iterator`]. See further documentation for the
behaviors for each of these types at [`Interrupt`].

Each of the trait implementations will behave normally with a small overhead to check shutdown state
until shutdown has been initiated with [`Swansong::shut_down`]. When shutdown has been initiated,
any type inside of an associated Interrupt will be interrupted/canceled/sealed at an appropriate
time.

Swansong graceful shutdown is considered in progress while there are outstanding [`Guard`]s. All
guards must be dropped for shutdown to be considered complete. Guards can be created in three ways:

* [`Swansong::guard`] returns a standalone [`Guard`] that can be moved into a closure or future,
  stored in a struct, or otherwise retained temporarily during a block of code that should not be
  abruptly stopped. For example, a http server might move a [`Guard`] into a closure or async task
  that handle an individual request.
* [`Swansong::guarded`] returns a [`Guarded`] wrapper type provides transparent implementations of
  various traits like [`Future`][std::future::Future] and [`Stream`][futures_core::Stream], as well
  as [`Deref`][std::ops::Deref]ing to the wrapped type. This is identical to moving a [`Guard`] into the wrapped
  type, but sometimes it's easier to compose the guard around a named future than to move the guard
  into the future.
* [`Interrupt::guarded`] allows an Interrupt wrapper to also act as a guard.

## Debugging stragglers

Each [`Guard`] records the source location it was created at, captured with `#[track_caller]`. If
a shutdown stalls because of straggling guards, [`Swansong::guard_report`] returns a
[`GuardReport`] describing the outstanding guards, aggregated by creation site and including the
age of the oldest guard at each site.

Guards dropped after shutdown has been initiated also log their creation site at
[`log::Level::Debug`], so enabling debug logging shows each straggler as it resolves.

## Async Example
```rust
# #![cfg(not(miri))]
# use async_global_executor as executor;
# use async_io::Timer;
# use std::{future::pending, time::Duration};
# use swansong::Swansong;
# futures_lite::future::block_on(async {
let swansong = Swansong::new();
executor::spawn(swansong.interrupt(pending::<()>()).guarded()).detach();
executor::spawn(swansong.guarded(Timer::after(Duration::from_secs(1)))).detach();
swansong.shut_down().await;
# });
```

## Sync example

```rust
# #![cfg(not(miri))]
# use std::{iter, thread, time::Duration};
# use swansong::Swansong;
let swansong = Swansong::new();

thread::spawn({
   let guard = swansong.guard();
   move || {
       let _guard = guard;
       thread::sleep(Duration::from_secs(1));
    }
});

thread::spawn({
    let swansong = swansong.clone();
    move || {
        for n in swansong.interrupt(iter::repeat_with(|| fastrand::u8(..)))
        {
            thread::sleep(Duration::from_millis(100));
            println!("{n}");
        }
    }
});

thread::sleep(Duration::from_millis(500));

swansong.shut_down().block();
```
*/

use std::{future::IntoFuture, sync::Arc};

mod implementation;
use implementation::Inner;
pub use implementation::{
    Guard, GuardReport, GuardReportEntry, Guarded, Interrupt, ShutdownCompletion, ShuttingDown,
};

/// # 🦢 Shutdown manager
///
/// See crate level docs for overview and example.
///
/// Two Swansongs are [`Eq`] if and only if they are clones of each other.
#[derive(Debug)]
pub struct Swansong {
    inner: Arc<Inner>,
}

impl Clone for Swansong {
    fn clone(&self) -> Self {
        self.inner.swansong_clone();
        Self {
            inner: Arc::clone(&self.inner),
        }
    }
}

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

impl Eq for Swansong {}

impl PartialEq for Swansong {
    fn eq(&self, other: &Self) -> bool {
        Arc::ptr_eq(&self.inner, &other.inner)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
/// enum that represents the current status of this `Swansong`.
pub enum ShutdownState {
    /// Shutdown has not been initiated.
    Running,
    /// Shutdown has been initiated, and we are waiting for any `Guard`s to drop.
    ShuttingDown,
    /// Shutdown has been initiated and all guards have been Dropped.
    Complete,
}

impl ShutdownState {
    /// Is the shutdown state [`ShutdownState::Running`]?
    #[must_use]
    pub fn is_running(&self) -> bool {
        matches!(self, ShutdownState::Running)
    }

    /// Is the shutdown state [`ShutdownState::ShuttingDown`]?
    #[must_use]
    pub fn is_shutting_down(&self) -> bool {
        matches!(self, ShutdownState::ShuttingDown)
    }

    /// Is the shutdown state [`ShutdownState::Complete`]?
    #[must_use]
    pub fn is_complete(&self) -> bool {
        matches!(self, ShutdownState::Complete)
    }
}

impl Swansong {
    /// Construct a new `Swansong`.
    #[must_use]
    pub fn new() -> Self {
        Self {
            inner: Inner::new_root(),
        }
    }

    /// Initiate graceful shutdown.
    ///
    /// This will gracefully stop any associated [`Interrupt`]s.
    ///
    /// This returns a [`ShutdownCompletion`] type, which can be blocked on with
    /// [`ShutdownCompletion::block`] or awaited as a [`Future`][std::future::Future].
    #[allow(clippy::must_use_candidate)] // It's fine not to use the returned ShutdownCompletion.
    pub fn shut_down(&self) -> ShutdownCompletion {
        self.inner.stop();
        ShutdownCompletion::new(&self.inner)
    }

    /// Blocks the current thread until graceful shutdown is complete
    ///
    /// Do not use this in async contexts. Instead, await the [`Swansong`]
    pub fn block_on_shutdown_completion(self) {
        ShutdownCompletion::new(&self.inner).block();
    }

    /// Determine if this `Swansong` is running, shutting down, or complete.
    #[must_use]
    pub fn state(&self) -> ShutdownState {
        let stopped = self.inner.is_stopped();
        let complete = stopped && self.inner.is_zero();
        if complete {
            ShutdownState::Complete
        } else if stopped {
            ShutdownState::ShuttingDown
        } else {
            ShutdownState::Running
        }
    }

    /// Returns a [`ShuttingDown`] future that resolves when shutdown has been
    /// initiated on this `Swansong`.
    ///
    /// This does not initiate shutdown — it only observes the transition from
    /// [`ShutdownState::Running`] to [`ShutdownState::ShuttingDown`]. The future
    /// resolves as soon as [`Swansong::shut_down`] is called (on this handle or
    /// any parent), or when the last root handle is dropped.
    ///
    /// The returned future does not wait for outstanding [`Guard`]s to drop. To
    /// wait for full shutdown (stopped _and_ all guards dropped), use
    /// [`Swansong::shut_down`] or [`Swansong::into_future`].
    ///
    /// [`ShuttingDown`] can also be used in blocking contexts via
    /// [`ShuttingDown::block`].
    #[must_use]
    pub fn shutting_down(&self) -> ShuttingDown {
        ShuttingDown::new(&self.inner)
    }

    /// Wrap any type with an [`Interrupt`], allowing it to cancel or seal on shutdown.
    ///
    /// [`Interrupt<T>`] has custom implementations of [`Future`][std::future::Future],
    /// [`Stream`][futures_core::Stream], [`AsyncRead`][futures_io::AsyncRead],
    /// [`AsyncWrite`][futures_io::AsyncWrite], [`AsyncBufRead`][futures_io::AsyncBufRead], and
    /// [`Iterator`]. See further documentation for the behaviors for each of these types at
    /// [`Interrupt`].
    ///
    /// Each of the trait implementations will behave normally with a small overhead to check
    /// shutdown state until shutdown has been initiated with [`Swansong::shut_down`]. When shutdown
    /// has been initiated, any type inside of an associated Interrupt will be
    /// interrupted/canceled/sealed at an appropriate time.
    #[must_use]
    pub fn interrupt<T>(&self, wrapped_type: T) -> Interrupt<T> {
        Interrupt::new(&self.inner, wrapped_type)
    }

    /// Returns a new [`Guard`], which forstalls shutdown until it is dropped.
    ///
    /// The guard records the source location of this call for diagnostic
    /// purposes; see [`Swansong::guard_report`].
    #[must_use]
    #[track_caller]
    pub fn guard(&self) -> Guard {
        Guard::new(&self.inner)
    }

    /// The total number of outstanding `Guard`s in this Swansong's subtree.
    ///
    /// For a root or leaf Swansong with no children, this is just the Guards created
    /// via this Swansong (and its clones). For a Swansong with children, it is the
    /// sum of Guards held in this node plus all Guards in all descendant subtrees.
    /// Child Swansong handles themselves are not counted as Guards; only real
    /// `Guard` instances are.
    ///
    /// This performs an `O(subtree)` walk; prefer not to call it in a hot path.
    #[must_use]
    pub fn guard_count(&self) -> usize {
        self.inner.guard_count_subtree()
    }

    /// Report the provenance of the outstanding [`Guard`]s in this Swansong's
    /// subtree.
    ///
    /// Every guard records the source location it was created at (captured
    /// with `#[track_caller]`, so it points at the `swansong.guard()`,
    /// `swansong.guarded(..)`, or `interrupt.guarded()` call in your code;
    /// clones inherit the location of the guard they were cloned from). This
    /// method aggregates the currently-live guards by creation site. It is
    /// useful for debugging a shutdown that is stalled on straggling guards —
    /// for example, by logging the report (via its [`Display`][std::fmt::Display]
    /// implementation) when [`Swansong::shut_down`] has not completed within a
    /// deadline.
    ///
    /// Like [`Swansong::guard_count`], this performs an `O(subtree)` walk, and
    /// additionally takes a brief per-node lock; it is intended for
    /// shutdown-time diagnostics rather than hot paths.
    ///
    /// # Example
    ///
    /// ```rust
    /// # use swansong::Swansong;
    /// let swansong = Swansong::new();
    /// let guard = swansong.guard();
    /// let report = swansong.guard_report();
    /// assert_eq!(report.guard_count(), 1);
    /// println!("{report}");
    /// drop(guard);
    /// assert!(swansong.guard_report().is_empty());
    /// ```
    #[must_use]
    pub fn guard_report(&self) -> GuardReport {
        let mut infos = Vec::new();
        self.inner.guard_infos_subtree(&mut infos);
        GuardReport::from_infos(infos)
    }

    /// Create a child [`Swansong`] linked to this parent.
    ///
    /// The child can be independently shut down without affecting the parent. However, when
    /// the parent is shut down, the child will also be shut down. The parent's shutdown will
    /// not be considered complete until the child [`Swansong`] and any associated
    /// [`ShutdownCompletion`]s are dropped.
    ///
    /// This supports multi-level hierarchies: a child can itself have children, forming a
    /// tree. Shutdown propagates downward through the entire tree.
    ///
    /// # Performance
    ///
    /// Creating a child has O(n) cost where n is the number of live children, due to
    /// cleanup of stale internal references. This is suitable for most use cases but may
    /// not be ideal for very high-frequency child creation with many concurrent children.
    ///
    /// # Example: HTTP/3 server
    ///
    /// ```rust
    /// # use swansong::Swansong;
    /// let server = Swansong::new();
    ///
    /// // Each connection gets its own child swansong
    /// let connection = server.child();
    ///
    /// // Connection can be shut down independently (e.g., on GOAWAY)
    /// connection.shut_down();
    ///
    /// // Or the entire server can shut down, which also stops all connections
    /// // server.shut_down().await;
    /// ```
    #[must_use]
    pub fn child(&self) -> Swansong {
        Swansong {
            inner: Inner::new_child(&self.inner),
        }
    }

    /// Attach a guard to the provided type, delaying shutdown until it drops.
    ///
    /// This function returns a [`Guarded`] wrapper type provides transparent implementations of
    /// various traits like [`Future`][std::future::Future] and [`Stream`][futures_core::Stream], as
    /// well as [`Deref`][std::ops::Deref]ing to the wrapped type. This is identical to moving a [`Guard`] into the
    /// wrapped type, but sometimes it's easier to compose the guard around a named future than to
    /// move the guard into the future.  See [`Guarded`] for more information about trait
    /// implementations on [`Guarded`]
    #[must_use]
    #[track_caller]
    pub fn guarded<T>(&self, wrapped_type: T) -> Guarded<T> {
        Guarded::new(&self.inner, wrapped_type)
    }
}

/// If we are dropping the last `Swansong` handle to a **root** node, stop the
/// associated futures. Dropping the last handle to a **child** (non-root) node
/// does not signal shutdown: the parent remains the governing authority, and
/// the child stays in its current state until the parent propagates stop or
/// until all Guards in the child's subset naturally drain.
impl Drop for Swansong {
    fn drop(&mut self) {
        if self.inner.swansong_drop() && self.inner.is_root() {
            self.inner.stop();
        }
    }
}

impl IntoFuture for Swansong {
    type Output = ();
    type IntoFuture = ShutdownCompletion;

    fn into_future(self) -> Self::IntoFuture {
        ShutdownCompletion::new(&self.inner)
    }
}