/home/runner/work/lyquor/lyquor/platform/api/src/subsys.rs

Source
//! Supervised lifecycle and dependency tracking for long-lived node subsystems.
//!
//! A subsystem is brought up by [`Subsystem::start`], returns a cloneable capability
//! [`Subsystem::Handle`] for its dependents, and winds down when its [`SubsysCtx`] is cancelled.
//! A [`Supervisor`] records the dependency graph as subsystems are spawned and derives shutdown
//! from it: walking the graph in reverse, each subsystem is cancelled and drained as soon as all
//! of its dependents have drained, with independent chains stopping in parallel.
//!
//! Startup is concurrent. [`Supervisor::spawn`] launches each subsystem's `start` in its own
//! lifecycle task and returns a [`Running`] promise of the handle, backed by a
//! [`tokio::sync::SetOnce`]. A dependent holds `Running` clones of its dependencies and parks on
//! [`Running::handle`] inside its own `start`, so it begins the moment the handles it needs exist
//! — never earlier, and never serialized behind unrelated subsystems. A failed (or panicked)
//! start fails the boot: [`Supervisor::wait_all_ready`] reports the cause, and
//! [`Supervisor::shutdown`] aborts the starts still parked on the missing handle.
//!
//! Dependency edges are declared with [`SubsysId`]s at spawn time. A supervisor is the single
//! composition root for one node process, and its ids are plain graph indices scoped to that
//! supervisor. Registration is sequential (spawning takes `&mut Supervisor`) even though starting
//! is concurrent, so an id only ever names an earlier-registered subsystem and the edges are
//! acyclic by construction; only the explicit [`Supervisor::shutdown_order`] overrides need a
//! runtime cycle check.
//!
//! Passive values with no tasks of their own (stores, registries, resolvers) can be recorded with
//! [`Supervisor::resource`] so the graph stays complete; they are never "stopped" but ordering
//! constraints propagate through them.

use std::{
    future::Future,
    sync::Arc,
    time::{Duration, Instant},
};

use tokio::sync::SetOnce;
use tokio_util::{sync::CancellationToken, task::TaskTracker};

/// Handle cell shared between a subsystem's lifecycle task and its [`Running`] clones. Set only
/// on success; a failed start leaves it forever pending.
type HandleCell<H> = Arc<SetOnce<H>>;

/// Readiness of one subsystem, observed by [`Supervisor::wait_all_ready`]: unset while starting,
/// `Ok` once ready, `Err` if the start failed.
type ReadyCell = Arc<SetOnce<anyhow::Result<()>>>;

/// Interval at which the supervisor warns about a subsystem that is slow to start or stop. The
/// supervisor keeps waiting: this is observability, not a deadline.
const SLOW_WARN_INTERVAL: Duration = Duration::from_secs(10);

/// A long-lived node component with a supervised lifecycle.
///
/// The implementing type is the subsystem's *configuration*: it carries the settings and the
/// [`Running`] dependencies the subsystem needs and is consumed by [`Subsystem::start`]. "Ready"
/// is defined as `start` returning; acquiring dependency handles ([`Running::handle`]), binding
/// sockets, replaying state, and other slow readiness all belong inside `start`.
///
/// Contract for implementations:
/// - Acquire dependency handles by holding `Running` clones and awaiting them in `start`.
/// - Spawn long-lived work with [`SubsysCtx::spawn`] and make it wind down when the ctx token is
///   cancelled. `spawn` intentionally follows [`TaskTracker`] semantics: closing the tracker lets
///   the supervisor wait for the current task set, but does not reject later inserts by itself, so
///   handles that keep a cloned ctx for dynamic work must gate new work on cancellation.
/// - Components with their own lifecycle contract integrate via [`SubsysCtx::on_shutdown`]
///   instead.
/// - The supervisor owns lifecycle, not health: spawned tasks are cancelled and drained at
///   shutdown but never watched for panics or early exits. Handling worker errors — and deciding
///   whether one is fatal — is the subsystem's own job; a subsystem whose failure should stop the
///   node can hold the root trigger token as ordinary configuration and cancel it.
pub trait Subsystem: Send + Sized + 'static {
    /// Cloneable capability handle given to dependents. Use `()` if the subsystem exposes no
    /// capabilities.
    type Handle: Clone + Send + Sync + 'static;

    /// Stable name used in logs and the dependency graph.
    fn name(&self) -> &'static str;

    /// Bring the subsystem up; return once it is ready to serve.
    fn start(self, ctx: SubsysCtx) -> impl Future<Output = anyhow::Result<Self::Handle>> + Send;
}

/// Per-subsystem lifecycle context handed to [`Subsystem::start`].
///
/// Cheap to clone; subsystems that spawn tasks dynamically at runtime may keep a clone around,
/// but should reject new dynamic work once [`Self::shutdown_token`] is cancelled.
#[derive(Clone)]
pub struct SubsysCtx {
    name: &'static str,
    shutdown: CancellationToken,
    tasks: TaskTracker,
}

lyquor_primitives::debug_struct_name!(SubsysCtx);

impl SubsysCtx {
    /// Token cancelled when this subsystem must stop. Hand out `child_token()`s to inner
    /// components that take a `CancellationToken`.
    pub fn shutdown_token(&self) -> &CancellationToken {
        &self.shutdown
    }

    /// Spawn a task owned by this subsystem. The supervisor cancels the shutdown token and then
    /// waits for the tracked task set at this subsystem's turn in the shutdown sequence.
    ///
    /// Like [`TaskTracker::spawn`], this does not reject inserts after the tracker has been
    /// closed. Runtime request paths that keep a cloned [`SubsysCtx`] should check
    /// [`CancellationToken::is_cancelled`] before spawning new work.
    pub fn spawn<F>(&self, task: F) -> tokio::task::JoinHandle<F::Output>
    where
        F: Future + Send + 'static,
        F::Output: Send + 'static,
    {
        self.tasks.spawn(task)
    }

    /// Run `teardown` at this subsystem's turn in the shutdown sequence, once its token is
    /// cancelled; the subsystem counts as stopped only after it finishes. This is how components
    /// with their own lifecycle contract integrate:
    /// `ctx.on_shutdown(async move { owner.shutdown().await })`.
    ///
    /// Multiple teardowns run concurrently — sequence dependent steps inside one future.
    pub fn on_shutdown<F>(&self, teardown: F)
    where
        F: Future<Output = ()> + Send + 'static,
    {
        let stop = self.shutdown.clone();
        self.tasks.spawn(async move {
            stop.cancelled().await;
            teardown.await;
        });
    }
}

/// Identifier of a subsystem or resource in a [`Supervisor`]'s dependency graph.
///
/// Ids are only meaningful for the supervisor that issued them. They are not globally tagged with
/// a supervisor identity; mixing ids from different supervisors is unsupported wiring.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct SubsysId(usize);

/// A spawned subsystem (or registered resource): a promise of its capability handle plus its
/// graph id.
///
/// Cheap to clone regardless of the handle type; subsystem configurations hold `Running` clones
/// of their dependencies and await [`Self::handle`] inside `start`.
pub struct Running<H> {
    id: SubsysId,
    name: &'static str,
    cell: HandleCell<H>,
}

impl<H> Clone for Running<H> {
    fn clone(&self) -> Self {
        Self {
            id: self.id,
            name: self.name,
            cell: self.cell.clone(),
        }
    }
}

impl<H> std::fmt::Debug for Running<H> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Running")
            .field("name", &self.name)
            .finish_non_exhaustive()
    }
}

impl<H: Clone> Running<H> {
    /// Graph id of this subsystem, used to declare dependencies on it.
    pub fn id(&self) -> SubsysId {
        self.id
    }

    /// Wait until the subsystem is ready and clone its capability handle out.
    ///
    /// If the subsystem fails to start this never resolves: the failure surfaces through
    /// [`Supervisor::wait_all_ready`], and [`Supervisor::shutdown`] aborts starts parked here.
    /// Await handles inside `start`; in the composition root, only after `wait_all_ready`.
    pub async fn handle(&self) -> H {
        self.cell.wait().await.clone3()
    }
}

/// One subsystem or resource registered with the supervisor.
struct Entry {
    name: &'static str,
    /// Stops-before edges (dependencies plus [`Supervisor::shutdown_order`] overrides): this
    /// entry stops before each of them.
    deps: Vec<usize>,
    /// Cancelled at this entry's turn in the shutdown sequence.
    stop: CancellationToken,
    /// The entry's lifecycle task plus everything spawned via its [`SubsysCtx`]; drained by the
    /// supervisor at this entry's turn in the shutdown sequence.
    tasks: TaskTracker,
    ready: ReadyCell,
    /// Passive value with no tasks: never stopped, present only for ordering.
    resource: bool,
}

/// Owner of the subsystem dependency graph; derives the shutdown sequence from it.
///
/// The root [`CancellationToken`] passed to [`Supervisor::new`] is only the *trigger*: cancelling
/// it aborts starts that are still pending and is the signal to run [`Supervisor::shutdown`].
/// Each subsystem has its own token, cancelled at its turn in the shutdown sequence.
///
/// Registration takes `&mut self` and is sequential; the spawned subsystems themselves start
/// concurrently, gated only by the [`Running::handle`]s they await.
pub struct Supervisor {
    root: CancellationToken,
    entries: Vec<Entry>,
}

lyquor_primitives::debug_struct_name!(Supervisor);

impl Supervisor {
    /// Create a supervisor triggered by `root`.
    pub fn new(root: CancellationToken) -> Self {
        Self {
            root,
            entries: Vec::new(),
        }
    }

    /// Clone of the root trigger token, for wiring external triggers such as signal handlers.
    pub fn shutdown_token(&self) -> CancellationToken {
        self.root.clone()
    }

    /// Completes once shutdown has been requested.
    pub async fn wait_shutdown_triggered(&self) {
        self.root.cancelled().await;
    }

    /// Register a subsystem with dependency edges to `deps` and immediately launch its `start` in
    /// a lifecycle task. Returns a [`Running`] promise of its handle.
    ///
    /// If shutdown is triggered while the start is pending, the start is abandoned and whatever
    /// work it already put on its ctx is cleaned up at its turn during [`Self::shutdown`].
    ///
    /// Must be called from within a tokio runtime.
    ///
    /// # Panics
    /// Panics if a dependency id is out of range for this supervisor. Ids from another supervisor
    /// are unsupported and may alias an in-range entry because ids are plain graph indices.
    pub fn spawn<S: Subsystem>(&mut self, subsys: S, deps: impl IntoIterator<Item = SubsysId>) -> Running<S::Handle> {
        let name = subsys.name();
        let deps = self.collect_deps(name, deps);
        let cell: HandleCell<S::Handle> = Arc::new(SetOnce::new());
        let ready: ReadyCell = Arc::new(SetOnce::new());
        let ctx = SubsysCtx {
            name,
            shutdown: CancellationToken::new(),
            tasks: TaskTracker::new(),
        };
        let stop = ctx.shutdown.clone();
        let tasks = ctx.tasks.clone();
        tasks.spawn(run_subsystem(
            subsys,
            ctx,
            self.root.clone(),
            cell.clone(),
            ready.clone(),
        ));
        let id = SubsysId(self.entries.len());
        self.entries.push(Entry {
            name,
            deps,
            stop,
            tasks,
            ready,
            resource: false,
        });
        Running { id, name, cell }
    }

    /// Record a passive value (store, registry, resolver, ...) as a graph node.
    ///
    /// Resources own no tasks and are never stopped, but ordering constraints propagate through
    /// them: if task `A` depends on resource `R` and `R` depends on task `B`, then `A` stops
    /// before `B`. The returned [`Running::handle`] resolves immediately.
    ///
    /// # Panics
    /// Panics if a dependency id is out of range for this supervisor. Ids from another supervisor
    /// are unsupported and may alias an in-range entry because ids are plain graph indices.
    pub fn resource<T>(&mut self, name: &'static str, value: T, deps: impl IntoIterator<Item = SubsysId>) -> Running<T>
    where
        T: Clone + Send + Sync + 'static,
    {
        let deps = self.collect_deps(name, deps);
        let cell = Arc::new(SetOnce::new());
        cell.set(value).ok().expect("fresh cell is settable");
        let ready: ReadyCell = Arc::new(SetOnce::new());
        ready.set(Ok(())).expect("fresh cell is settable");
        let id = SubsysId(self.entries.len());
        self.entries.push(Entry {
            name,
            deps,
            stop: CancellationToken::new(),
            tasks: TaskTracker::new(),
            ready,
            resource: true,
        });
        Running { id, name, cell }
    }

    /// Add a shutdown-ordering edge with no capability meaning: `first` is cancelled and drained
    /// strictly before `later`. Use sparingly, for flow constraints the capability graph cannot
    /// express (e.g. "the DB flusher stops only after the state writers stopped").
    ///
    /// # Panics
    /// Panics if the edge contradicts the existing stop order (would form a cycle) or if an id is
    /// out of range for this supervisor. Ids from another supervisor are unsupported and may alias
    /// an in-range entry because ids are plain graph indices.
    pub fn shutdown_order(&mut self, first: SubsysId, later: SubsysId) {
        assert!(
            first.0 < self.entries.len() && later.0 < self.entries.len(),
            "shutdown_order: SubsysId is out of range for this supervisor graph"
        );
        assert_ne!(
            first.0, later.0,
            "shutdown_order: a subsystem cannot stop before itself"
        );
        assert!(
            !stops_before(&self.entries, later.0, first.0),
            "shutdown_order({} before {}) contradicts the existing stop order",
            self.entries[first.0].name,
            self.entries[later.0].name,
        );
        self.entries[first.0].deps.push(later.0);
    }

    /// Wait until every spawned subsystem is ready, failing fast on the first one whose start
    /// fails. On error the caller should run [`Self::shutdown`] to unwind the subsystems that did
    /// start.
    pub async fn wait_all_ready(&self) -> anyhow::Result<()> {
        let waits15 = self.entries.iter()15.map15(|entry| async {
            match entry.ready.wait().await {
                Ok(()) => Ok(()),
                Err(error) => Err(anyhow::anyhow!("subsystem {} failed to start: {error:#}", entry.name)),
            }
        });
        futures::future::try_join_all(waits).await.map(|_| ())
    }

    /// Stop all subsystems by walking the graph in reverse: each one is cancelled and drained
    /// (its [`SubsysCtx::on_shutdown`] teardowns and spawned tasks) as soon as every subsystem
    /// depending on it has drained, so independent chains stop in parallel. Also triggers the
    /// root token, so starts that are still pending abort.
    ///
    /// Dropping the supervisor without calling this (e.g. `?` on a failed
    /// [`Self::wait_all_ready`]) cancels every subsystem's token as a safety net, but waits for
    /// nothing; `shutdown` is the ordered, drained teardown.
    ///
    /// This future consumes the supervisor and is intended to be awaited to completion. It is not
    /// a cancel-safe operation to race in `tokio::select!` and then drop; enforce external
    /// shutdown deadlines at a higher level if the process needs one.
    pub async fn shutdown(mut self) {
        self.root.cancel();
        let started = Instant::now();
        // Taken out so the `Drop` safety net sees an already-emptied supervisor.
        let entries = std::mem::take(&mut self.entries);
        // dependents[e] = entries that must stop before entry e
        let mut dependents = vec![Vec::new(); entries.len()];
        for (index, entry) in entries.iter()14.enumerate14() {
            for &dep14 in &entry.deps {
                dependents[dep].push(index);
            }
        }
        // Mirrors the ReadyCell on the way down: set once an entry has fully stopped.
        let stopped14: Vec<SetOnce<()>>14 = entries.iter()14.map14(|_| SetOnce::new()).collect14();
        let waits14 = entries.iter()14.enumerate14().map14(|(index, entry)| {
            let blockers = &dependents[index];
            let stopped = &stopped;
            async move {
                for &blocker14 in blockers {
                    stopped[blocker].wait().await;
                }
                if !entry.resource {
                    entry.stop.cancel();
                    stop_entry(entry.name, entry.tasks.clone()).await;
                }
                let _ = stopped[index].set(());
            }
        });
        futures::future::join_all(waits).await;
        tracing::info!(
            elapsed_ms = started.elapsed().as_millis() as u64,
            "all subsystems stopped"
        );
    }

    fn collect_deps(&self, name: &'static str, deps: impl IntoIterator<Item = SubsysId>) -> Vec<usize> {
        let deps: Vec<usize> = deps.into_iter().map(|id| id.0).collect();
        for &dep14 in &deps {
            assert!(
                dep < self.entries.len(),
                "subsystem {name}: dependency SubsysId is out of range for this supervisor graph"
            );
        }
        deps
    }
}

/// Safety net for supervisors dropped without [`Supervisor::shutdown`] (e.g. `?` on a failed
/// [`Supervisor::wait_all_ready`]): cancel every subsystem so no work outlives its owner. This
/// only signals — nothing is drained or waited for.
impl Drop for Supervisor {
    fn drop(&mut self) {
        if self.entries.is_empty() {
            return;
        }
        tracing::warn!("supervisor dropped without shutdown; cancelling subsystems without draining them");
        self.root.cancel();
        for entry in &self.entries2 {
            entry.stop.cancel();
        }
    }
}

/// Runs one subsystem's `start` (abandoning it if shutdown is triggered first) and publishes the
/// handle or the failure. Lives on the entry's own tracker like any ctx-spawned task.
async fn run_subsystem<S: Subsystem>(
    subsys: S, ctx: SubsysCtx, root: CancellationToken, cell: HandleCell<S::Handle>, ready: ReadyCell,
) {
    let name = ctx.name;
    // If this task unwinds (subsystem start panicked) before publishing readiness, the guard
    // publishes failure so `wait_all_ready` does not hang.
    let mut guard = PublishOnDrop {
        ready: &ready,
        armed: true,
    };

    let publish31 = if root.is_cancelled() {
        Err(anyhow::anyhow!("shutdown triggered before start"))
    } else {
        tracing::debug!(subsystem = name, "starting subsystem");
        let started = Instant::now();
        let start = async {
            tokio::select! {
                biased;
                result27 = subsys.start(ctx.clone()) => result27,
                () = root.cancelled() => Err(anyhow::anyhow!("shutdown triggered while subsystem was starting"))3,
            }
        };
        let result30 = warn_if_slow(name, "start", start).await;
        if result.is_ok() {
            tracing::info!(
                subsystem = name,
                elapsed_ms = started.elapsed().as_millis() as u64,
                "subsystem ready"
            );
        }
        result
    };
    guard.armed = false;
    match publish {
        Ok(handle) => {
            let _ = cell.set(handle);
            let _ = ready.set(Ok(()));
        }
        Err(error) => {
            tracing::error!(
                subsystem = name,
                error = format!("{error:#}"),
                "subsystem failed to start"
            );
            let _ = ready.set(Err(error));
        }
    }
}

/// Publishes failure if the lifecycle task unwinds before publishing readiness.
struct PublishOnDrop<'a> {
    ready: &'a SetOnce<anyhow::Result<()>>,
    armed: bool,
}

impl Drop for PublishOnDrop<'_> {
    fn drop(&mut self) {
        if self.armed {
            let _ = self.ready.set(Err(anyhow::anyhow!("panicked while starting")));
        }
    }
}

/// Drain an entry's task tracker: its ctx-spawned tasks and [`SubsysCtx::on_shutdown`]
/// teardowns (which wake on the just-cancelled stop token).
async fn stop_entry(name: &'static str, tasks: TaskTracker) {
    let started = Instant::now();
    tracing::debug!(subsystem = name, "stopping subsystem");
    tasks.close();
    warn_if_slow(name, "stop", tasks.wait()).await;
    tracing::debug!(
        subsystem = name,
        elapsed_ms = started.elapsed().as_millis() as u64,
        "subsystem stopped"
    );
}

/// Await `operation`, logging a warning every [`SLOW_WARN_INTERVAL`] until it completes.
async fn warn_if_slow<T>(name: &'static str, phase: &'static str, operation: impl Future<Output = T>) -> T {
    let started = Instant::now();
    tokio::pin!(operation);
    loop {
        match tokio::time::timeout(SLOW_WARN_INTERVAL, &mut operation).await {
            Ok(output) => break output,
            Err(_) => tracing::warn!(
                subsystem = name,
                phase,
                waited_secs = started.elapsed().as_secs(),
                "subsystem is taking a long time"
            ),
        }
    }
}

/// True if `from` transitively stops before `to` under the current graph.
fn stops_before(entries: &[Entry], from: usize, to: usize) -> bool {
    let mut seen = vec![false; entries.len()];
    let mut stack = vec![from];
    while let Some(current3) = stack.pop() {
        if current == to {
            return true;
        }
        if std::mem::replace(&mut seen[current], true) {
            continue;
        }
        stack.extend(entries[current].deps.iter().copied());
    }
    false
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::{AtomicUsize, Ordering};

    use lyquor_test::test;
    use tokio::sync::{mpsc, oneshot};

    use super::*;

    type Events = mpsc::UnboundedSender<String>;

    /// Reports lifecycle events, optionally waits for a dependency handle and/or a start gate
    /// inside `start`, and can hold up its stop until a stop gate fires.
    struct Tracked {
        name: &'static str,
        events: Events,
        needs: Option<Running<()>>,
        start_gate: Option<oneshot::Receiver<()>>,
        stop_gate: Option<oneshot::Receiver<()>>,
    }

    impl Tracked {
        fn new(name: &'static str, events: &Events) -> Self {
            Self {
                name,
                events: events.clone(),
                needs: None,
                start_gate: None,
                stop_gate: None,
            }
        }

        fn needs(mut self, dep: &Running<()>) -> Self {
            self.needs = Some(dep.clone());
            self
        }

        fn start_gated(mut self) -> (Self, oneshot::Sender<()>) {
            let (release, gate) = oneshot::channel();
            self.start_gate = Some(gate);
            (self, release)
        }

        fn stop_gated(mut self) -> (Self, oneshot::Sender<()>) {
            let (release, gate) = oneshot::channel();
            self.stop_gate = Some(gate);
            (self, release)
        }
    }

    impl Subsystem for Tracked {
        type Handle = ();

        fn name(&self) -> &'static str {
            self.name
        }

        async fn start(self, ctx: SubsysCtx) -> anyhow::Result<()> {
            let Self {
                name,
                events,
                needs,
                start_gate,
                stop_gate,
            } = self;
            events.send(format!("{name}:starting")).unwrap();
            if let Some(dep3) = needs {
                dep.handle().await;
            }
            if let Some(gate3) = start_gate {
                gate.await.ok();
            }
            events.send(format!("{name}:started")).unwrap();
            let token = ctx.shutdown_token().clone();
            ctx.spawn(async move {
                token.cancelled().await;
                events.send(format!("{name}:cancelled")).unwrap();
                if let Some(gate9) = stop_gate {
                    let _ = gate.await;
                }
                events.send(format!("{name}:stopped")).unwrap();
            });
            Ok(())
        }
    }

    async fn expect_event(events: &mut mpsc::UnboundedReceiver<String>, expected: &str) {
        assert_eq!(events.recv().await.as_deref(), Some(expected));
    }

    async fn expect_events_any_order(events: &mut mpsc::UnboundedReceiver<String>, expected: [&str; 2]) {
        let mut received = [events.recv().await.unwrap(), events.recv().await.unwrap()];
        received.sort();
        let mut expected = expected.map(str::to_string);
        expected.sort();
        assert_eq!(received, expected);
    }

    fn drain(events: &mut mpsc::UnboundedReceiver<String>) {
        while events.try_recv().is_ok() {}37
    }

    #[test(tokio::test)]
    async fn independent_subsystems_start_concurrently() {
        let (events_tx, mut events) = mpsc::unbounded_channel();
        let mut sup = Supervisor::new(CancellationToken::new());
        let (a_subsys, release_a) = Tracked::new("a", &events_tx).start_gated();
        let (b_subsys, release_b) = Tracked::new("b", &events_tx).start_gated();
        sup.spawn(a_subsys, []);
        sup.spawn(b_subsys, []);

        // Both starts are in flight at once: each reports `starting` while neither has been
        // allowed to finish. A supervisor that awaited each start in turn would hang here.
        expect_events_any_order(&mut events, ["a:starting", "b:starting"]).await;
        release_a.send(()).unwrap();
        release_b.send(()).unwrap();
        sup.wait_all_ready().await.unwrap();
        expect_events_any_order(&mut events, ["a:started", "b:started"]).await;
        sup.shutdown().await;
    }

    #[test(tokio::test)]
    async fn dependent_start_waits_for_dependency_handle() {
        let (events_tx, mut events) = mpsc::unbounded_channel();
        let mut sup = Supervisor::new(CancellationToken::new());
        let (a_subsys, release_a) = Tracked::new("a", &events_tx).start_gated();
        let a = sup.spawn(a_subsys, []);
        sup.spawn(Tracked::new("b", &events_tx).needs(&a), [a.id()]);

        expect_events_any_order(&mut events, ["a:starting", "b:starting"]).await;
        // b is parked on a's handle; nothing further happens until a is released.
        assert!(events.try_recv().is_err(), "b must not start before a is ready");
        release_a.send(()).unwrap();
        sup.wait_all_ready().await.unwrap();
        expect_event(&mut events, "a:started").await;
        expect_event(&mut events, "b:started").await;
        sup.shutdown().await;
    }

    struct Producer;

    impl Subsystem for Producer {
        type Handle = Arc<AtomicUsize>;

        fn name(&self) -> &'static str {
            "producer"
        }

        async fn start(self, _ctx: SubsysCtx) -> anyhow::Result<Self::Handle> {
            Ok(Arc::new(AtomicUsize::new(1)))
        }
    }

    struct Consumer {
        producer: Running<Arc<AtomicUsize>>,
    }

    impl Subsystem for Consumer {
        type Handle = ();

        fn name(&self) -> &'static str {
            "consumer"
        }

        async fn start(self, _ctx: SubsysCtx) -> anyhow::Result<()> {
            let counter = self.producer.handle().await;
            counter.fetch_add(1, Ordering::SeqCst);
            Ok(())
        }
    }

    #[test(tokio::test)]
    async fn handles_flow_between_subsystems() {
        let mut sup = Supervisor::new(CancellationToken::new());
        let producer = sup.spawn(Producer, []);
        sup.spawn(
            Consumer {
                producer: producer.clone(),
            },
            [producer.id()],
        );

        sup.wait_all_ready().await.unwrap();
        assert_eq!(producer.handle().await.load(Ordering::SeqCst), 2);
        sup.shutdown().await;
    }

    #[test(tokio::test)]
    async fn shutdown_stops_dependents_before_dependencies() {
        let (events_tx, mut events) = mpsc::unbounded_channel();
        let mut sup = Supervisor::new(CancellationToken::new());
        let a = sup.spawn(Tracked::new("a", &events_tx), []);
        let (b_subsys, release_b) = Tracked::new("b", &events_tx).stop_gated();
        let b = sup.spawn(b_subsys, [a.id()]);
        let (c_subsys, release_c) = Tracked::new("c", &events_tx).stop_gated();
        sup.spawn(c_subsys, [b.id()]);
        sup.wait_all_ready().await.unwrap();
        drain(&mut events);

        let shutdown = tokio::spawn(sup.shutdown());
        expect_event(&mut events, "c:cancelled").await;
        assert!(events.try_recv().is_err(), "b must not be cancelled while c drains");
        release_c.send(()).unwrap();
        expect_event(&mut events, "c:stopped").await;
        expect_event(&mut events, "b:cancelled").await;
        assert!(events.try_recv().is_err(), "a must not be cancelled while b drains");
        release_b.send(()).unwrap();
        expect_event(&mut events, "b:stopped").await;
        expect_event(&mut events, "a:cancelled").await;
        expect_event(&mut events, "a:stopped").await;
        shutdown.await.unwrap();
    }

    #[test(tokio::test)]
    async fn layer_members_are_cancelled_together() {
        let (events_tx, mut events) = mpsc::unbounded_channel();
        let mut sup = Supervisor::new(CancellationToken::new());
        let (x_subsys, release_x) = Tracked::new("x", &events_tx).stop_gated();
        let (y_subsys, release_y) = Tracked::new("y", &events_tx).stop_gated();
        sup.spawn(x_subsys, []);
        sup.spawn(y_subsys, []);
        sup.wait_all_ready().await.unwrap();
        drain(&mut events);

        let shutdown = tokio::spawn(sup.shutdown());
        // Both leaves are cancelled even though neither has finished stopping yet.
        expect_events_any_order(&mut events, ["x:cancelled", "y:cancelled"]).await;
        release_x.send(()).unwrap();
        release_y.send(()).unwrap();
        shutdown.await.unwrap();
    }

    #[test(tokio::test)]
    async fn independent_chains_stop_without_waiting_for_each_other() {
        let (events_tx, mut events) = mpsc::unbounded_channel();
        let mut sup = Supervisor::new(CancellationToken::new());
        let a = sup.spawn(Tracked::new("a", &events_tx), []);
        sup.spawn(Tracked::new("b", &events_tx), [a.id()]);
        let x = sup.spawn(Tracked::new("x", &events_tx), []);
        let (y_subsys, release_y) = Tracked::new("y", &events_tx).stop_gated();
        sup.spawn(y_subsys, [x.id()]);
        sup.wait_all_ready().await.unwrap();
        drain(&mut events);

        // The a<-b chain drains to completion while y (same "layer" as b) is still held open:
        // layered shutdown would park a behind the b/y layer barrier.
        let shutdown = tokio::spawn(sup.shutdown());
        let mut seen = Vec::new();
        while seen.last().map(String::as_str) != Some("a:stopped") {
            seen.push(events.recv().await.unwrap());
        }
        assert!(
            !seen.iter().any(|event| event.starts_with("x:")),
            "x stopped early: {seen:?}"
        );
        assert!(
            !seen.contains(&"y:stopped".to_string()),
            "y drained unexpectedly: {seen:?}"
        );
        release_y.send(()).unwrap();
        expect_event(&mut events, "y:stopped").await;
        expect_event(&mut events, "x:cancelled").await;
        expect_event(&mut events, "x:stopped").await;
        shutdown.await.unwrap();
    }

    #[test(tokio::test)]
    async fn resources_link_shutdown_ordering_transitively() {
        let (events_tx, mut events) = mpsc::unbounded_channel();
        let mut sup = Supervisor::new(CancellationToken::new());
        let b = sup.spawn(Tracked::new("b", &events_tx), []);
        let store = sup.resource("store", (), [b.id()]);
        let (a_subsys, release_a) = Tracked::new("a", &events_tx).stop_gated();
        sup.spawn(a_subsys, [store.id()]);
        sup.wait_all_ready().await.unwrap();
        drain(&mut events);

        let shutdown = tokio::spawn(sup.shutdown());
        expect_event(&mut events, "a:cancelled").await;
        assert!(events.try_recv().is_err(), "b must wait for a to stop");
        release_a.send(()).unwrap();
        expect_event(&mut events, "a:stopped").await;
        expect_event(&mut events, "b:cancelled").await;
        expect_event(&mut events, "b:stopped").await;
        shutdown.await.unwrap();
    }

    #[test(tokio::test)]
    async fn shutdown_order_stops_first_before_later() {
        let (events_tx, mut events) = mpsc::unbounded_channel();
        let mut sup = Supervisor::new(CancellationToken::new());
        let (x_subsys, release_x) = Tracked::new("x", &events_tx).stop_gated();
        let x = sup.spawn(x_subsys, []);
        let y = sup.spawn(Tracked::new("y", &events_tx), []);
        sup.shutdown_order(x.id(), y.id());
        sup.wait_all_ready().await.unwrap();
        drain(&mut events);

        let shutdown = tokio::spawn(sup.shutdown());
        expect_event(&mut events, "x:cancelled").await;
        assert!(events.try_recv().is_err(), "y must wait for x to stop");
        release_x.send(()).unwrap();
        expect_event(&mut events, "x:stopped").await;
        expect_event(&mut events, "y:cancelled").await;
        expect_event(&mut events, "y:stopped").await;
        shutdown.await.unwrap();
    }

    #[test]
    fn shutdown_order_rejects_contradictions() {
        let mut sup = Supervisor::new(CancellationToken::new());
        let a = sup.resource("a", (), []);
        let b = sup.resource("b", (), [a.id()]); // b stops before a
        let panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            sup.shutdown_order(a.id(), b.id()); // a stops before b: contradiction
        }))
        .unwrap_err();
        let message = panic.downcast_ref::<String>().cloned().unwrap_or_default();
        assert!(message.contains("contradicts"), "unexpected panic message: {message}");
    }

    #[test(tokio::test)]
    async fn entry_waits_for_all_of_its_dependents() {
        let (events_tx, mut events) = mpsc::unbounded_channel();
        let mut sup = Supervisor::new(CancellationToken::new());
        let x = sup.spawn(Tracked::new("x", &events_tx), []);
        let (y_subsys, release_y) = Tracked::new("y", &events_tx).stop_gated();
        sup.spawn(y_subsys, [x.id()]);
        let (z_subsys, release_z) = Tracked::new("z", &events_tx).stop_gated();
        sup.spawn(z_subsys, [x.id()]);
        sup.wait_all_ready().await.unwrap();
        drain(&mut events);

        let shutdown = tokio::spawn(sup.shutdown());
        expect_events_any_order(&mut events, ["y:cancelled", "z:cancelled"]).await;
        release_y.send(()).unwrap();
        expect_event(&mut events, "y:stopped").await;
        assert!(events.try_recv().is_err(), "x must wait for z as well");
        release_z.send(()).unwrap();
        expect_event(&mut events, "z:stopped").await;
        expect_event(&mut events, "x:cancelled").await;
        expect_event(&mut events, "x:stopped").await;
        shutdown.await.unwrap();
    }

    struct FailAfterSpawning {
        events: Events,
    }

    impl Subsystem for FailAfterSpawning {
        type Handle = ();

        fn name(&self) -> &'static str {
            "failing"
        }

        async fn start(self, ctx: SubsysCtx) -> anyhow::Result<()> {
            let events = self.events.clone();
            let token = ctx.shutdown_token().clone();
            ctx.spawn(async move {
                token.cancelled().await;
                events.send("failing:task-cancelled".into()).unwrap();
            });
            let events = self.events;
            ctx.on_shutdown(async move {
                events.send("failing:teardown".into()).unwrap();
            });
            anyhow::bail!("boom")
        }
    }

    #[test(tokio::test)]
    async fn failed_start_fails_the_boot_and_unwinds_partial_work() {
        let (events_tx, mut events) = mpsc::unbounded_channel();
        let mut sup = Supervisor::new(CancellationToken::new());
        let a = sup.spawn(Tracked::new("a", &events_tx), []);
        let failing = sup.spawn(
            FailAfterSpawning {
                events: events_tx.clone(),
            },
            [a.id()],
        );
        // A dependent of the failed subsystem stays parked on its handle; the boot barrier
        // reports the root cause and shutdown aborts the parked start.
        sup.spawn(Tracked::new("dependent", &events_tx).needs(&failing), [failing.id()]);

        let error = sup.wait_all_ready().await.unwrap_err();
        assert!(
            error.to_string().contains("subsystem failing failed to start: boom"),
            "unexpected error: {error}"
        );
        drain(&mut events);

        sup.shutdown().await;
        // The failed subsystem's partial work is cleaned up in its layer (its teardown and task
        // run concurrently), then `a` stops.
        expect_events_any_order(&mut events, ["failing:teardown", "failing:task-cancelled"]).await;
        expect_event(&mut events, "a:cancelled").await;
        expect_event(&mut events, "a:stopped").await;
    }

    #[test(tokio::test)]
    async fn dropping_the_supervisor_cancels_subsystems() {
        let (events_tx, mut events) = mpsc::unbounded_channel();
        let mut sup = Supervisor::new(CancellationToken::new());
        sup.spawn(Tracked::new("a", &events_tx), []);
        sup.spawn(
            FailAfterSpawning {
                events: events_tx.clone(),
            },
            [],
        );
        sup.wait_all_ready().await.unwrap_err();
        drain(&mut events);

        // The natural `?` path: the supervisor is dropped without shutdown(). Every subsystem's
        // partial work must still be cancelled rather than outliving its owner.
        drop(sup);
        let mut seen: Vec<String> = Vec::new();
        for _ in 0..4 {
            seen.push(events.recv().await.unwrap());
        }
        seen.sort();
        assert_eq!(
            seen,
            ["a:cancelled", "a:stopped", "failing:task-cancelled", "failing:teardown"]
        );
    }

    struct PanicsWhileStarting;

    impl Subsystem for PanicsWhileStarting {
        type Handle = ();

        fn name(&self) -> &'static str {
            "panicking"
        }

        async fn start(self, _ctx: SubsysCtx) -> anyhow::Result<()> {
            panic!("boom");
        }
    }

    #[test(tokio::test)]
    async fn panicking_start_fails_the_boot_instead_of_hanging() {
        let (events_tx, _events) = mpsc::unbounded_channel();
        let mut sup = Supervisor::new(CancellationToken::new());
        let panicking = sup.spawn(PanicsWhileStarting, []);
        sup.spawn(
            Tracked::new("dependent", &events_tx).needs(&panicking),
            [panicking.id()],
        );

        let error = sup.wait_all_ready().await.unwrap_err();
        assert!(
            error.to_string().contains("panicking failed to start: panicked"),
            "unexpected error: {error}"
        );
        sup.shutdown().await;
    }

    struct HangsWhileStarting {
        events: Events,
    }

    impl Subsystem for HangsWhileStarting {
        type Handle = ();

        fn name(&self) -> &'static str {
            "hanging"
        }

        async fn start(self, ctx: SubsysCtx) -> anyhow::Result<()> {
            let Self { events } = self;
            events.send("hanging:entered".into()).unwrap();
            let token = ctx.shutdown_token().clone();
            ctx.spawn(async move {
                token.cancelled().await;
                events.send("hanging:task-cancelled".into()).unwrap();
            });
            std::future::pending::<()>().await;
            Ok(())
        }
    }

    #[test(tokio::test)]
    async fn triggered_shutdown_aborts_pending_start() {
        let (events_tx, mut events) = mpsc::unbounded_channel();
        let mut sup = Supervisor::new(CancellationToken::new());
        sup.spawn(
            HangsWhileStarting {
                events: events_tx.clone(),
            },
            [],
        );
        expect_event(&mut events, "hanging:entered").await;

        sup.shutdown_token().cancel();
        let error = sup.wait_all_ready().await.unwrap_err();
        assert!(
            error.to_string().contains("shutdown triggered"),
            "unexpected error: {error}"
        );

        // The abandoned subsystem was still registered; shutdown cleans up its spawned work.
        sup.shutdown().await;
        expect_event(&mut events, "hanging:task-cancelled").await;
    }

    struct OwnerTeardown {
        events: Events,
    }

    impl Subsystem for OwnerTeardown {
        type Handle = ();

        fn name(&self) -> &'static str {
            "owner"
        }

        async fn start(self, ctx: SubsysCtx) -> anyhow::Result<()> {
            let (torn_down_tx, torn_down) = oneshot::channel();
            let events = self.events.clone();
            ctx.on_shutdown(async move {
                events.send("owner:teardown".into()).unwrap();
                let _ = torn_down_tx.send(());
            });
            let events = self.events;
            ctx.spawn(async move {
                // Exits only after the teardown ran: shutdown would hang here otherwise, proving
                // the layer waits for on_shutdown teardowns like any other tracked task.
                let _ = torn_down.await;
                events.send("owner:task-exited".into()).unwrap();
            });
            Ok(())
        }
    }

    #[test(tokio::test)]
    async fn on_shutdown_runs_at_the_stop_turn() {
        let (events_tx, mut events) = mpsc::unbounded_channel();
        let mut sup = Supervisor::new(CancellationToken::new());
        sup.spawn(
            OwnerTeardown {
                events: events_tx.clone(),
            },
            [],
        );
        sup.wait_all_ready().await.unwrap();
        assert!(events.try_recv().is_err(), "teardown must not run before shutdown");

        sup.shutdown().await;
        expect_event(&mut events, "owner:teardown").await;
        expect_event(&mut events, "owner:task-exited").await;
    }

    #[test(tokio::test)]
    async fn spawn_after_trigger_fails_fast() {
        let (events_tx, _events) = mpsc::unbounded_channel();
        let mut sup = Supervisor::new(CancellationToken::new());
        sup.shutdown_token().cancel();
        sup.spawn(Tracked::new("late", &events_tx), []);
        let error = sup.wait_all_ready().await.unwrap_err();
        assert!(
            error.to_string().contains("subsystem late failed to start"),
            "unexpected error: {error}"
        );
        sup.shutdown().await;
    }
}

Coverage Report

Created: 2026-07-07 06:35