saluki_core/runtime/
supervisor.rs

1use std::{
2    future::Future,
3    pin::Pin,
4    sync::{
5        atomic::{AtomicU64, AtomicUsize, Ordering},
6        Arc, Mutex,
7    },
8    time::Duration,
9};
10
11use async_trait::async_trait;
12use saluki_common::collections::FastHashMap;
13use saluki_common::sync::shutdown::ShutdownHandle;
14use saluki_error::GenericError;
15use snafu::{OptionExt as _, Snafu};
16use tokio::{
17    pin,
18    runtime::Handle,
19    select,
20    sync::{mpsc, oneshot},
21};
22use tracing::{debug, error, warn};
23
24use super::{
25    dedicated::{spawn_dedicated_runtime, RuntimeConfiguration, RuntimeMode},
26    restart::{RestartAction, RestartMode, RestartState, RestartStrategy, RestartType},
27    worker_state::WorkerState,
28};
29use crate::runtime::{
30    process::{Process, ProcessExt as _},
31    state::DataspaceRegistry,
32};
33
34/// A `Future` that represents the execution of a supervised process.
35pub type SupervisorFuture = Pin<Box<dyn Future<Output = Result<(), GenericError>> + Send>>;
36
37/// A `Future` that represents the full lifecycle of a worker, including initialization.
38///
39/// Unlike [`SupervisorFuture`], which only represents the runtime phase, this future first performs async
40/// initialization and then runs the worker. This allows initialization to happen concurrently when multiple workers are
41/// spawned, and keeps the supervisor loop responsive to shutdown signals during initialization.
42pub(super) type WorkerFuture = Pin<Box<dyn Future<Output = Result<(), WorkerError>> + Send>>;
43
44/// Worker lifecycle errors.
45///
46/// Distinguishes between initialization failures (which shouldn't trigger restart logic) and runtime failures (which
47/// are eligible for restart).
48#[derive(Debug)]
49pub(super) enum WorkerError {
50    /// The worker failed during async initialization.
51    ///
52    /// The optional `child_name` carries the name of the original failing child when the error originates from a
53    /// nested supervisor. This allows the parent to include it in its own `FailedToInitialize` error for better
54    /// diagnostics across supervision tree levels.
55    Initialization {
56        child_name: Option<String>,
57        source: InitializationError,
58    },
59
60    /// The worker failed during runtime execution.
61    Runtime(GenericError),
62
63    /// The worker was a nested supervisor that completed a requested shutdown after forcefully aborting one or more of
64    /// its own workers.
65    ///
66    /// Carried as a distinct variant (rather than collapsed into [`Runtime`][WorkerError::Runtime]) so the parent's
67    /// shutdown drain can recover the structured count and merge it into its own tally, aggregating forced aborts up
68    /// the supervision tree.
69    ShutdownTimedOut {
70        /// The number of workers the nested supervisor forcefully aborted, summed across its own supervision tree.
71        aborted: usize,
72    },
73}
74
75impl From<SupervisorError> for WorkerError {
76    fn from(err: SupervisorError) -> Self {
77        match err {
78            // Propagate initialization failures so the parent supervisor does NOT attempt to restart.
79            // Preserve the original child name so the parent can include it in diagnostics.
80            SupervisorError::FailedToInitialize { child_name, source } => WorkerError::Initialization {
81                child_name: Some(child_name),
82                source,
83            },
84            // Preserve the structured abort count so the parent can merge it into its own shutdown tally.
85            SupervisorError::ShutdownTimedOut { aborted } => WorkerError::ShutdownTimedOut { aborted },
86            // All other supervisor errors (shutdown, no children, invalid name) are runtime-level.
87            other => WorkerError::Runtime(other.into()),
88        }
89    }
90}
91
92/// Process errors.
93#[derive(Debug, Snafu)]
94pub enum ProcessError {
95    /// The child process was aborted by the supervisor.
96    #[snafu(display("Child process was aborted by the supervisor."))]
97    Aborted,
98
99    /// The child process panicked.
100    #[snafu(display("Child process panicked."))]
101    Panicked,
102
103    /// The child process terminated with an error.
104    #[snafu(display("Child process terminated with an error: {}", source))]
105    Terminated {
106        /// The error that caused the termination.
107        source: GenericError,
108    },
109}
110
111/// Initialization errors.
112///
113/// Initialization errors are distinct from runtime errors: they indicate that a process couldn't be started at all
114/// (for example, failed to bind a port, missing configuration). These errors don't trigger restart logic; instead, they
115/// immediately propagate up and fail the supervisor.
116#[derive(Debug, Snafu)]
117#[snafu(context(suffix(false)))]
118pub enum InitializationError {
119    /// The process couldn't be initialized due to an error.
120    #[snafu(display("Process failed to initialize: {}", source))]
121    Failed {
122        /// The underlying error that caused initialization to fail.
123        source: GenericError,
124    },
125}
126
127impl From<GenericError> for InitializationError {
128    fn from(source: GenericError) -> Self {
129        Self::Failed { source }
130    }
131}
132
133/// Strategy for shutting down a process.
134#[derive(Clone, Copy, Debug)]
135pub enum ShutdownStrategy {
136    /// Waits for the configured duration for the process to exit, and then forcefully aborts it otherwise.
137    Graceful(Duration),
138
139    /// Forcefully aborts the process without waiting.
140    Brutal,
141}
142
143/// Policy for automatically shutting a supervisor down based on the termination of its _significant_ children.
144///
145/// A significant child (see [`ChildSpecification::with_significant`]) is one whose termination -- when it isn't restarted -- can
146/// drive the supervisor to shut down. This mirrors Erlang/OTP's `auto_shutdown` supervisor flag, and is how an
147/// unexpected (or intentional) child exit cascades into the supervisor stopping, and thus propagating up the tree,
148/// without that child being restarted.
149#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
150pub enum AutoShutdown {
151    /// Never shut down automatically; significant children have no special effect. This is the default.
152    #[default]
153    Never,
154
155    /// Shut down as soon as _any_ significant child terminates without being restarted.
156    AnySignificant,
157
158    /// Shut down once _all_ significant children have terminated without being restarted.
159    AllSignificant,
160}
161
162/// How a supervisor shuts its children down.
163#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
164pub enum ShutdownMode {
165    /// Shut children down one at a time, in reverse order of starting (last-started first).
166    ///
167    /// This is the default, and is appropriate when later children may depend on earlier ones: each child is fully
168    /// stopped before the next is signalled.
169    #[default]
170    Ordered,
171
172    /// Shut all children down at once and wait for them concurrently.
173    ///
174    /// Total shutdown time is bounded by the slowest child rather than the sum of all children, which suits large,
175    /// independent child sets -- for example, one task per network connection.
176    Concurrent,
177}
178
179/// A supervisable process.
180#[async_trait]
181pub trait Supervisable: Send + Sync {
182    /// Returns the name of the process.
183    fn name(&self) -> &str;
184
185    /// Returns the shutdown strategy for the process.
186    fn shutdown_strategy(&self) -> ShutdownStrategy {
187        ShutdownStrategy::Graceful(Duration::from_secs(5))
188    }
189
190    /// Initializes the process asynchronously.
191    ///
192    /// During initialization, any resources or configuration for the process can be created asynchronously, and the
193    /// same runtime that's used for running the process is used for initialization. The resulting future is expected to
194    /// complete as soon as reasonably possible after `shutdown` resolves.
195    ///
196    /// **Important:** The `process_shutdown` signal must be moved into the returned [`SupervisorFuture`] so the worker
197    /// can respond to supervisor-initiated shutdown. If `process_shutdown` is dropped during initialization, the worker
198    /// will be unable to shut down gracefully and will be forcefully aborted after the shutdown timeout.
199    ///
200    /// # Errors
201    ///
202    /// If the process can't be initialized, an error is returned.
203    async fn initialize(&self, process_shutdown: ShutdownHandle) -> Result<SupervisorFuture, InitializationError>;
204}
205
206/// Supervisor errors.
207#[derive(Debug, Snafu)]
208#[snafu(context(suffix(false)))]
209pub enum SupervisorError {
210    /// Supervisor or worker name is invalid.
211    #[snafu(display("Invalid name for supervisor or worker: '{}'", name))]
212    InvalidName {
213        /// The name of the supervisor is invalid.
214        name: String,
215    },
216
217    /// A child process failed to initialize.
218    ///
219    /// This error indicates that a child couldn't complete its async initialization. This is distinct from runtime
220    /// failures and doesn't trigger restart logic.
221    #[snafu(display("Child process '{}' failed to initialize: {}", child_name, source))]
222    FailedToInitialize {
223        /// The name of the child that failed to initialize.
224        child_name: String,
225
226        /// The underlying initialization error.
227        source: InitializationError,
228    },
229
230    /// The supervisor exceeded its restart limits and was forced to shutdown.
231    #[snafu(display("Supervisor has exceeded restart limits and was forced to shutdown."))]
232    Shutdown,
233
234    /// The supervisor shut down because a significant child terminated.
235    ///
236    /// See [`AutoShutdown`] and [`ChildSpecification::with_significant`]. The supervisor stopped, and drained its remaining
237    /// children, because a child marked significant terminated without being restarted.
238    #[snafu(display("Supervisor shut down after a significant child terminated."))]
239    SignificantChildExited,
240
241    /// The supervisor completed a requested shutdown, but one or more workers ignored graceful shutdown and had to be
242    /// forcefully aborted after exceeding their shutdown timeout.
243    ///
244    /// The shutdown itself was requested and otherwise orderly; this variant exists so that having to forcefully stop a
245    /// worker is surfaced as a failure rather than reported as a clean shutdown. The count aggregates forced aborts
246    /// across the entire supervision tree: a parent merges in the counts reported by any child supervisors that also
247    /// timed out, so the value observed at the root supervisor is the total number of workers that had to be
248    /// force-stopped tree-wide.
249    #[snafu(display(
250        "Shutdown completed uncleanly: {} worker(s) were forcefully aborted after exceeding their shutdown timeout.",
251        aborted
252    ))]
253    ShutdownTimedOut {
254        /// The number of workers that had to be forcefully aborted.
255        aborted: usize,
256    },
257}
258
259/// A specification for a process to be added to a [`Supervisor`].
260///
261/// A child specification describes how the supervisor should create and manage a child: the underlying future that
262/// represents the process, along with metadata such as its name and shutdown strategy. All processes in a supervisor,
263/// whether a worker or a (nested) supervisor, are represented by a [`ChildSpecification`].
264///
265/// Generally, callers should prefer to use [`add_worker`][Supervisor::add_worker] directly, which can accept either
266/// [`Supervisor`] or any value that implements [`Supervisable`], without needing to explicitly create a
267/// [`ChildSpecification`]. This is preferred as it is more concise but also will ensure that relevant settings are
268/// configured properly for the given worker type, such as using the proper shutdown strategy for supervisors to allow
269/// for complete, graceful shutdown.
270///
271/// If more control is needed, [`ChildSpecification::worker`] can be used to create a specification directly, allowing
272/// access to configuring those more advanced settings. This is currently only valid for worker processes, as
273/// supervisors have no additional user-configurable settings.
274pub struct ChildSpecification<S = WorkerSpec> {
275    spec_inner: S,
276}
277
278/// Child specification state for a worker.
279pub struct WorkerSpec {
280    worker: Arc<dyn Supervisable>,
281    config: ChildConfig,
282}
283
284/// Child specification state for a supervisor.
285pub struct SupervisorSpec {
286    supervisor: Supervisor,
287}
288
289impl ChildSpecification<WorkerSpec> {
290    /// Creates a specification for the given worker.
291    pub fn worker<T: Supervisable + 'static>(worker: T) -> Self {
292        Self {
293            spec_inner: WorkerSpec {
294                worker: Arc::new(worker),
295                config: ChildConfig::default(),
296            },
297        }
298    }
299
300    /// Creates a specification for a worker that can only run once.
301    ///
302    /// This function is shorthand for calling [`worker`][Self::worker] followed by
303    /// [`with_restart_type`][Self::with_restart_type] set to [`RestartType::Temporary`][RestartType::Temporary].
304    pub fn one_shot_worker<T: Supervisable + 'static>(worker: T) -> Self {
305        Self::worker(worker).with_restart_type(RestartType::Temporary)
306    }
307
308    /// Sets the restart policy for this worker.
309    ///
310    /// Defaults to [`RestartType::Permanent`].
311    #[must_use]
312    pub fn with_restart_type(mut self, restart_type: RestartType) -> Self {
313        self.spec_inner.config.restart = restart_type;
314        self
315    }
316
317    /// Sets whether this worker is _significant_.
318    ///
319    /// A significant worker's termination (when it isn't restarted) can drive the supervisor to shut down, per the
320    /// supervisor's [`AutoShutdown`] policy. Only meaningful for non-permanent workers, since a permanent worker is
321    /// always restarted and so never terminates without being restarted.
322    #[must_use]
323    pub fn with_significant(mut self, significant: bool) -> Self {
324        self.spec_inner.config.significant = significant;
325        self
326    }
327
328    /// Runs this worker on the given Tokio runtime rather than the supervisor's own runtime.
329    ///
330    /// By default, a worker runs on whatever runtime its supervisor runs on. Use this for compute-heavy workers that
331    /// shouldn't contend with the supervisor's runtime -- for example, a topology component offloading encoding work
332    /// onto a shared worker pool.
333    ///
334    /// Note that this only affects where the worker's task is spawned. Supervision itself -- shutdown signalling,
335    /// restart handling, and abort-on-timeout -- is unchanged, and is still driven from the supervisor's runtime.
336    #[must_use]
337    pub fn with_runtime(mut self, handle: Handle) -> Self {
338        self.spec_inner.config.runtime = Some(handle);
339        self
340    }
341
342    /// Overrides the shutdown strategy for this worker.
343    ///
344    /// By default, a worker's strategy comes from [`Supervisable::shutdown_strategy`], which itself defaults to
345    /// `Graceful(5s)`. Use this when the grace period depends on where the worker is used rather than on the worker
346    /// type: a worker that a component drains during its own shutdown needs at least as long as the component itself,
347    /// otherwise it is forcefully aborted while the component is still waiting on it.
348    #[must_use]
349    pub fn with_shutdown_strategy(mut self, strategy: ShutdownStrategy) -> Self {
350        self.spec_inner.config.shutdown_strategy = Some(strategy);
351        self
352    }
353
354    /// Lowers this worker specification into its type-erased child and configuration.
355    fn into_worker_parts(self) -> (SupervisedChild, ChildConfig) {
356        (SupervisedChild::Worker(self.spec_inner.worker), self.spec_inner.config)
357    }
358}
359
360impl<T> From<T> for ChildSpecification<WorkerSpec>
361where
362    T: Supervisable + 'static,
363{
364    fn from(worker: T) -> Self {
365        Self::worker(worker)
366    }
367}
368
369impl From<Supervisor> for ChildSpecification<SupervisorSpec> {
370    fn from(supervisor: Supervisor) -> Self {
371        Self {
372            spec_inner: SupervisorSpec { supervisor },
373        }
374    }
375}
376
377mod sealed {
378    pub trait Sealed {}
379}
380
381impl sealed::Sealed for WorkerSpec {}
382impl sealed::Sealed for SupervisorSpec {}
383
384/// Child specification state.
385///
386/// This trait is sealed -- it cannot be implemented outside of this crate -- and is implemented only for
387/// [`WorkerSpec`] and [`SupervisorSpec`]. It exists so that [`Supervisor::add_worker`] can accept a
388/// [`ChildSpecification`] in either state (as well as bare workers and supervisors) while lowering each into the
389/// supervisor's internal representation.
390pub trait ChildState: sealed::Sealed + Sized {
391    #[doc(hidden)]
392    fn register(spec: ChildSpecification<Self>, supervisor: &mut Supervisor);
393}
394
395impl ChildState for WorkerSpec {
396    fn register(spec: ChildSpecification<Self>, supervisor: &mut Supervisor) {
397        let (child, config) = spec.into_worker_parts();
398        supervisor.push_child(ChildEntry {
399            spec: child,
400            config,
401            dynamic: false,
402        });
403    }
404}
405
406impl ChildState for SupervisorSpec {
407    fn register(spec: ChildSpecification<Self>, supervisor: &mut Supervisor) {
408        supervisor.push_child(ChildEntry {
409            spec: SupervisedChild::Supervisor(spec.spec_inner.supervisor),
410            config: ChildConfig::default(),
411            dynamic: false,
412        });
413    }
414}
415
416/// The type-erased, runnable form of a child: either a worker or a nested supervisor.
417pub(super) enum SupervisedChild {
418    Worker(Arc<dyn Supervisable>),
419    Supervisor(Supervisor),
420}
421
422impl SupervisedChild {
423    /// Returns whether this child is a nested supervisor rather than a leaf worker.
424    pub(super) fn is_supervisor(&self) -> bool {
425        matches!(self, Self::Supervisor(_))
426    }
427
428    fn process_type(&self) -> &'static str {
429        match self {
430            Self::Worker(_) => "worker",
431            Self::Supervisor(_) => "supervisor",
432        }
433    }
434
435    fn name(&self) -> &str {
436        match self {
437            Self::Worker(worker) => worker.name(),
438            Self::Supervisor(supervisor) => &supervisor.supervisor_id,
439        }
440    }
441
442    pub(super) fn shutdown_strategy(&self) -> ShutdownStrategy {
443        match self {
444            Self::Worker(worker) => worker.shutdown_strategy(),
445
446            // Supervisors should always be given as much time as necessary shutdown down gracefully to ensure that the
447            // entire supervision subtree can be shutdown cleanly.
448            Self::Supervisor(_) => ShutdownStrategy::Graceful(Duration::MAX),
449        }
450    }
451
452    pub(super) fn create_process(&self, parent_process: &Process) -> Result<Process, SupervisorError> {
453        match self {
454            Self::Worker(worker) => Process::worker(worker.name(), parent_process).context(InvalidName {
455                name: worker.name().to_string(),
456            }),
457            Self::Supervisor(sup) => {
458                Process::supervisor(&sup.supervisor_id, Some(parent_process)).context(InvalidName {
459                    name: sup.supervisor_id.to_string(),
460                })
461            }
462        }
463    }
464
465    pub(super) fn create_worker_future(
466        &self, process: Process, process_shutdown: ShutdownHandle,
467    ) -> Result<WorkerFuture, SupervisorError> {
468        match self {
469            Self::Worker(worker) => {
470                let worker = Arc::clone(worker);
471                Ok(Box::pin(async move {
472                    let run_future =
473                        worker
474                            .initialize(process_shutdown)
475                            .await
476                            .map_err(|source| WorkerError::Initialization {
477                                child_name: None,
478                                source,
479                            })?;
480                    run_future.await.map_err(WorkerError::Runtime)
481                }))
482            }
483            Self::Supervisor(sup) => {
484                match sup.runtime_mode() {
485                    RuntimeMode::Ambient => {
486                        // Run on the parent's ambient runtime.
487                        Ok(sup.as_nested_process(process, process_shutdown))
488                    }
489                    RuntimeMode::Dedicated(config) => {
490                        // Spawn in a dedicated runtime on a new OS thread, passing the parent's
491                        // dataspace so the nested supervisor inherits it across the thread boundary.
492                        let child_name = sup.supervisor_id.to_string();
493                        let dataspace = process.dataspace().clone();
494                        let handle =
495                            spawn_dedicated_runtime(sup.inner_clone(), config.clone(), process_shutdown, dataspace)
496                                .map_err(|e| SupervisorError::FailedToInitialize {
497                                    child_name,
498                                    source: e.into(),
499                                })?;
500
501                        Ok(Box::pin(async move { handle.await.map_err(WorkerError::from) }))
502                    }
503                }
504            }
505        }
506    }
507}
508
509impl Clone for SupervisedChild {
510    fn clone(&self) -> Self {
511        match self {
512            Self::Worker(worker) => Self::Worker(Arc::clone(worker)),
513            Self::Supervisor(supervisor) => Self::Supervisor(supervisor.inner_clone()),
514        }
515    }
516}
517
518/// Per-child configuration: its [`RestartType`], whether it is _significant_ (see [`AutoShutdown`]), and any
519/// per-child overrides for runtime placement and shutdown strategy.
520///
521/// Defaults to a permanent, non-significant child that runs on the supervisor's own runtime and takes its shutdown
522/// strategy from [`Supervisable::shutdown_strategy`]. On a worker, this is set through
523/// [`ChildSpecification::with_restart_type`], [`ChildSpecification::with_significant`],
524/// [`ChildSpecification::with_runtime`], and [`ChildSpecification::with_shutdown_strategy`].
525#[derive(Clone, Debug)]
526pub(super) struct ChildConfig {
527    restart: RestartType,
528    significant: bool,
529
530    /// Runtime to spawn the child on. `None` means the supervisor's own runtime.
531    runtime: Option<Handle>,
532
533    /// Shutdown strategy override. `None` means defer to [`Supervisable::shutdown_strategy`].
534    shutdown_strategy: Option<ShutdownStrategy>,
535}
536
537impl ChildConfig {
538    /// Returns the runtime the child should be spawned on, if it isn't the supervisor's own.
539    pub(super) fn runtime(&self) -> Option<&Handle> {
540        self.runtime.as_ref()
541    }
542
543    /// Returns the child's shutdown strategy override, if one was configured.
544    pub(super) fn shutdown_strategy(&self) -> Option<ShutdownStrategy> {
545        self.shutdown_strategy
546    }
547}
548
549impl Default for ChildConfig {
550    fn default() -> Self {
551        Self {
552            restart: RestartType::Permanent,
553            significant: false,
554            runtime: None,
555            shutdown_strategy: None,
556        }
557    }
558}
559
560/// A registered child: its specification together with the configuration chosen at registration time.
561#[derive(Clone)]
562struct ChildEntry {
563    spec: SupervisedChild,
564    config: ChildConfig,
565    /// Whether this child was added dynamically (via [`SupervisorHandle`]) rather than statically before the run. Used
566    /// to maintain the dynamic-children gauge.
567    dynamic: bool,
568}
569
570/// Identifier for a child managed by a [`Supervisor`].
571///
572/// Returned by [`SupervisorHandle::spawn`] for dynamically spawned children. Unique within a single process for the
573/// lifetime of a supervisor run.
574#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
575pub struct ChildId(u64);
576
577impl ChildId {
578    /// Returns the raw numeric value of this identifier.
579    pub const fn as_u64(self) -> u64 {
580        self.0
581    }
582}
583
584/// Error returned when spawning a dynamic child on a [`Supervisor`] fails.
585#[derive(Debug, Snafu)]
586pub enum SpawnError {
587    /// The supervisor isn't currently running, so it can't accept the spawn.
588    ///
589    /// Returned when the supervisor hasn't started yet, is between restarts, or has shut down -- and also if the run
590    /// ends after the request is accepted but before the child is started. To add children before the supervisor
591    /// starts, configure them statically with [`Supervisor::add_worker`] instead.
592    #[snafu(display("supervisor is gone"))]
593    SupervisorGone,
594
595    /// The supervisor was running but rejected the spawn (for example, an invalid child name).
596    ///
597    /// Unlike [`SupervisorGone`](Self::SupervisorGone), the supervisor accepted the request and then couldn't start the
598    /// child; the underlying error is preserved as the source.
599    #[snafu(display("supervisor rejected the spawn: {}", source))]
600    Rejected {
601        /// The underlying error that caused the spawn to be rejected.
602        source: GenericError,
603    },
604}
605
606/// A dynamic spawn request sent from a [`SupervisorHandle`] to the running supervisor.
607struct PendingSpawn {
608    id: u64,
609    spec: SupervisedChild,
610    config: ChildConfig,
611    ack: oneshot::Sender<Result<(), SpawnError>>,
612}
613
614/// Capacity of the per-run channel that carries dynamic spawn requests from handles to the running supervisor.
615///
616/// Each request is short-lived -- the supervisor processes it and signals the waiting caller promptly -- so this only
617/// bounds how many spawns can be in flight before a caller's send applies backpressure.
618const DYNAMIC_SPAWN_CHANNEL_CAPACITY: usize = 1024;
619
620/// A handle for spawning dynamic children on a running [`Supervisor`].
621///
622/// Obtained from [`Supervisor::handle`]. Handles are cheap to clone and can be shared across tasks. Spawning is async:
623/// the request is handed to the running supervisor and the call returns once the child has been started. If the
624/// supervisor isn't currently running, spawning returns [`SpawnError::SupervisorGone`].
625#[derive(Clone)]
626pub struct SupervisorHandle {
627    name: Arc<str>,
628    // The currently running supervisor publishes its command sender here so handles can reach the live run; it's
629    // cleared when no run is active, at which point spawns observe `SupervisorGone`.
630    current_tx: Arc<Mutex<Option<mpsc::Sender<PendingSpawn>>>>,
631    id_counter: Arc<AtomicU64>,
632    active: Arc<AtomicUsize>,
633}
634
635impl SupervisorHandle {
636    /// Returns the name of the supervisor this handle refers to.
637    pub fn name(&self) -> &str {
638        &self.name
639    }
640
641    /// Spawns a new dynamic worker.
642    ///
643    /// Dynamic workers are temporary children that are not restarted by the supervisor when they die or when the
644    /// supervisor itself is restarted. They are useful for short-lived, non-critical background tasks that require
645    /// structured concurrency: the process should be cancelled when the supervisor itself is restarted or terminated,
646    /// and so on.
647    ///
648    /// Use [`spawn_with`](Self::spawn_with) to configure the child's restart policy or significance.
649    ///
650    /// # Errors
651    ///
652    /// If the supervisor isn't current running, or if the child specification is invalid, an error is returned.
653    pub async fn spawn<T: Supervisable + 'static>(&self, worker: T) -> Result<ChildId, SpawnError> {
654        self.spawn_with(ChildSpecification::worker(worker).with_restart_type(RestartType::Temporary))
655            .await
656    }
657
658    /// Spawns a new dynamic child from a fully configured [`ChildSpecification`].
659    ///
660    /// Dynamic workers are temporary children that are not restarted by the supervisor when they die or when the
661    /// supervisor itself is restarted. They are useful for short-lived, non-critical background tasks that require
662    /// structured concurrency: the process should be cancelled when the supervisor itself is restarted or terminated,
663    /// and so on.
664    ///
665    /// This method allows for configuring more advanced aspects of the child process, such as its restart type and
666    /// significance.
667    ///
668    /// # Errors
669    ///
670    /// If the supervisor isn't current running, or if the child specification is invalid, an error is returned.
671    pub async fn spawn_with(&self, spec: ChildSpecification<WorkerSpec>) -> Result<ChildId, SpawnError> {
672        let id = self.id_counter.fetch_add(1, Ordering::Relaxed);
673        let (spec, config) = spec.into_worker_parts();
674        let (ack_tx, ack_rx) = oneshot::channel();
675        self.send(PendingSpawn {
676            id,
677            spec,
678            config,
679            ack: ack_tx,
680        })
681        .await?;
682
683        // Wait for the supervisor to start (or reject) the child. A dropped ack channel means the run ended before it
684        // got to us, which is indistinguishable from `SupervisorGone` to the caller.
685        ack_rx
686            .await
687            .map_err(|_| SpawnError::SupervisorGone)?
688            .map(|()| ChildId(id))
689    }
690
691    /// Returns whether the supervisor is currently running.
692    pub fn is_running(&self) -> bool {
693        self.current_tx.lock().unwrap().is_some()
694    }
695
696    /// Returns the number of dynamic children currently running under the supervisor.
697    pub fn active_children(&self) -> usize {
698        self.active.load(Ordering::Relaxed)
699    }
700
701    /// Hands a spawn request to the currently running supervisor, applying backpressure if its channel is full.
702    ///
703    /// Returns [`SpawnError::SupervisorGone`] if no run is active, or if the run ends before the request is accepted.
704    async fn send(&self, spawn: PendingSpawn) -> Result<(), SpawnError> {
705        // Clone the sender out from under the lock so we don't hold the (synchronous) mutex guard across the await.
706        let tx = self.current_tx.lock().unwrap().clone();
707        match tx {
708            Some(tx) => tx.send(spawn).await.map_err(|_| SpawnError::SupervisorGone),
709            None => Err(SpawnError::SupervisorGone),
710        }
711    }
712}
713
714/// Supervises a set of workers.
715///
716/// # Workers
717///
718/// All workers are defined through implementation of the [`Supervisable`] trait, which provides the logic for both
719/// creating the underlying worker future that's spawned, as well as other metadata, such as the worker's name, how the
720/// worker should be shutdown, and so on.
721///
722/// Supervisors also (indirectly) implement the [`Supervisable`] trait, allowing them to be supervised by other
723/// supervisors in order to construct _supervision trees_.
724///
725/// # Instrumentation
726///
727/// Supervisors automatically create their own allocation group
728/// ([`TrackingAllocator`][saluki_common::resource_tracking::TrackingAllocator]), which is used to track both the memory
729/// usage of the supervisor itself and its children. Additionally, individual worker processes are wrapped in a
730/// dedicated [`tracing::Span`] to allow tracing the causal relationship between arbitrary code and the worker executing
731/// it, and statistics about task polls (poll count, poll duration) are collected.
732///
733/// # Restart Strategies
734///
735/// As the main purpose of a supervisor, restart behavior is fully configurable. A number of restart strategies are
736/// available, which generally relate to the purpose of the supervisor: whether the workers being managed are
737/// independent or interdependent.
738///
739/// All restart strategies are configured through [`RestartStrategy`], which has more information on the available
740/// strategies and configuration settings.
741pub struct Supervisor {
742    supervisor_id: Arc<str>,
743    child_specs: Vec<ChildEntry>,
744    restart_strategy: RestartStrategy,
745    auto_shutdown: AutoShutdown,
746    shutdown_mode: ShutdownMode,
747    shutdown_budget: Option<Duration>,
748    runtime_mode: RuntimeMode,
749    // Shared across clones (a nested supervisor is cloned each time it runs) and across all handles. While a run is
750    // active it holds that run's spawn-command sender so handles can reach the live supervisor; it's `None` whenever no
751    // run is active, at which point spawns observe `SupervisorGone`. Doubles as the `is_running` signal.
752    current_tx: Arc<Mutex<Option<mpsc::Sender<PendingSpawn>>>>,
753    id_counter: Arc<AtomicU64>,
754    // Number of dynamic children currently running, shared with handles so it can be surfaced as a gauge.
755    active: Arc<AtomicUsize>,
756}
757
758impl Supervisor {
759    /// Creates an empty `Supervisor` with the default restart strategy.
760    pub fn new<S: AsRef<str>>(supervisor_id: S) -> Result<Self, SupervisorError> {
761        // We try to throw an error about invalid names as early as possible. This is a manual check, so we might still
762        // encounter an error later when actually running the supervisor, but this is a good first step to catch the
763        // bulk of invalid names.
764        if supervisor_id.as_ref().is_empty() {
765            return Err(SupervisorError::InvalidName {
766                name: supervisor_id.as_ref().to_string(),
767            });
768        }
769
770        Ok(Self {
771            supervisor_id: supervisor_id.as_ref().into(),
772            child_specs: Vec::new(),
773            restart_strategy: RestartStrategy::default(),
774            auto_shutdown: AutoShutdown::default(),
775            shutdown_mode: ShutdownMode::default(),
776            shutdown_budget: None,
777            runtime_mode: RuntimeMode::default(),
778            current_tx: Arc::new(Mutex::new(None)),
779            id_counter: Arc::new(AtomicU64::new(0)),
780            active: Arc::new(AtomicUsize::new(0)),
781        })
782    }
783
784    /// Returns the supervisor's ID.
785    pub fn id(&self) -> &str {
786        &self.supervisor_id
787    }
788
789    /// Sets the restart strategy for the supervisor.
790    pub fn with_restart_strategy(mut self, strategy: RestartStrategy) -> Self {
791        self.restart_strategy = strategy;
792        self
793    }
794
795    /// Sets the supervisor's automatic-shutdown policy.
796    ///
797    /// Controls whether the termination of _significant_ children (see [`ChildSpecification::with_significant`]) drives the
798    /// supervisor to shut down. Defaults to [`AutoShutdown::Never`].
799    pub fn with_auto_shutdown(mut self, auto_shutdown: AutoShutdown) -> Self {
800        self.auto_shutdown = auto_shutdown;
801        self
802    }
803
804    /// Sets the supervisor's shutdown mode. See [`ShutdownMode`]. Defaults to [`ShutdownMode::Ordered`].
805    pub fn with_shutdown_mode(mut self, mode: ShutdownMode) -> Self {
806        self.shutdown_mode = mode;
807        self
808    }
809
810    /// Bounds how long this supervisor waits for its worker children during shutdown.
811    ///
812    /// Without a budget, a supervisor waits as long as each child's own [`ShutdownStrategy`] allows, and waits
813    /// indefinitely for any child that has no finite deadline of its own. A budget makes the supervisor responsible for
814    /// the deadline instead: children need no individual timeouts, and whatever is still running when the budget
815    /// elapses is forcefully aborted -- each one named in the logs, and counted in the resulting
816    /// [`SupervisorError::ShutdownTimedOut`].
817    ///
818    /// The budget is a ceiling, not a replacement: a child that also carries its own finite deadline is still held to
819    /// whichever elapses first.
820    ///
821    /// Two kinds of child are outside it. A nested supervisor is never cut off by its parent's budget -- it bounds its
822    /// own subtree, and aborting it would both truncate that drain and, for a supervisor running on a dedicated
823    /// runtime, fail to stop it at all. A [`ShutdownStrategy::Brutal`] child is aborted up front and never waited on.
824    /// Neither can a budget bound work that ignores cancellation, since an abort only takes effect at an await point.
825    ///
826    /// Use this where one deadline for a whole subtree is more meaningful than a guess per worker -- a topology
827    /// component and its background tasks, for instance, where what matters is that the component as a whole stops in
828    /// time.
829    #[must_use]
830    pub fn with_shutdown_budget(mut self, budget: Duration) -> Self {
831        self.shutdown_budget = Some(budget);
832        self
833    }
834
835    /// Returns a handle for spawning dynamic children on this supervisor while it runs.
836    ///
837    /// The handle can be created before the supervisor starts and cloned freely. Spawns only succeed while the
838    /// supervisor is actually running; if it hasn't started yet, is between restarts, or has shut down, they return
839    /// [`SpawnError::SupervisorGone`].
840    pub fn handle(&self) -> SupervisorHandle {
841        SupervisorHandle {
842            name: Arc::clone(&self.supervisor_id),
843            current_tx: Arc::clone(&self.current_tx),
844            id_counter: Arc::clone(&self.id_counter),
845            active: Arc::clone(&self.active),
846        }
847    }
848
849    /// Configures this supervisor to run in a dedicated runtime.
850    ///
851    /// When this supervisor is added as a child to another supervisor, it will spawn its own OS threads and Tokio
852    /// runtime instead of running on the parent's ambient runtime.
853    ///
854    /// This provides runtime isolation, which can be useful for:
855    /// - CPU-bound work that shouldn't block the parent's runtime
856    /// - Isolating failures in one part of the system
857    /// - Using different runtime configurations (for example, single-threaded vs multi-threaded)
858    pub fn with_dedicated_runtime(mut self, config: RuntimeConfiguration) -> Self {
859        self.runtime_mode = RuntimeMode::Dedicated(config);
860        self
861    }
862
863    /// Returns the runtime mode for this supervisor.
864    pub(crate) fn runtime_mode(&self) -> &RuntimeMode {
865        &self.runtime_mode
866    }
867
868    /// Adds a worker (or nested supervisor) to the supervisor.
869    ///
870    /// A worker can be anything that implements the [`Supervisable`] trait. A [`Supervisor`] can also be added as a
871    /// worker and managed in a nested fashion, known as a supervision tree.
872    ///
873    /// See [`ChildSpecification`] for more details on how workers are represented internally and what options are
874    /// available to configure.
875    pub fn add_worker<S, T>(&mut self, child: T)
876    where
877        S: ChildState,
878        T: Into<ChildSpecification<S>>,
879    {
880        S::register(child.into(), self);
881    }
882
883    fn push_child(&mut self, entry: ChildEntry) {
884        debug!(
885            supervisor_id = %self.supervisor_id,
886            "Adding new static child process #{}. ({}, {}, {:?})",
887            self.child_specs.len(),
888            entry.spec.process_type(),
889            entry.spec.name(),
890            entry.config,
891        );
892        self.child_specs.push(entry);
893    }
894
895    fn spawn_static_children(
896        &self, children: &mut FastHashMap<u64, ChildEntry>, worker_state: &mut WorkerState,
897    ) -> Result<(), SupervisorError> {
898        debug!(supervisor_id = %self.supervisor_id, "Spawning all static child processes.");
899        for entry in &self.child_specs {
900            let id = self.id_counter.fetch_add(1, Ordering::Relaxed);
901            worker_state.add_worker(id, &entry.spec, &entry.config)?;
902            children.insert(id, entry.clone());
903        }
904
905        Ok(())
906    }
907
908    /// Respawns children after a one-for-all restart, honoring each child's [`RestartType`].
909    ///
910    /// Every child except [`RestartType::Temporary`] is restarted, matching Erlang/OTP: a group restart restarts all
911    /// permanent and transient children -- regardless of how they last exited, including a transient child that had
912    /// already exited cleanly -- but never temporary children, which are shut down with the group and not brought back.
913    /// A transient child's "restart only on abnormal exit" rule governs its _own_ termination, not a group restart
914    /// driven by a sibling. Dynamic children are not restored (they are lost on a supervisor-level restart).
915    fn respawn_children_one_for_all(
916        &self, children: &mut FastHashMap<u64, ChildEntry>, worker_state: &mut WorkerState,
917    ) -> Result<(), SupervisorError> {
918        debug!(supervisor_id = %self.supervisor_id, "Restarting all eligible static child processes.");
919        for entry in &self.child_specs {
920            // Temporary children are never restarted by a group restart (matching OTP): they are shut down with the
921            // group but not brought back.
922            if entry.config.restart == RestartType::Temporary {
923                continue;
924            }
925            let id = self.id_counter.fetch_add(1, Ordering::Relaxed);
926            worker_state.add_worker(id, &entry.spec, &entry.config)?;
927            children.insert(id, entry.clone());
928        }
929
930        Ok(())
931    }
932
933    /// Spawns one dynamic child into the running supervisor's worker set and roster, signaling the requesting handle.
934    fn spawn_dynamic_child(
935        &self, spawn: PendingSpawn, worker_state: &mut WorkerState, children: &mut FastHashMap<u64, ChildEntry>,
936        significant_remaining: &mut usize,
937    ) {
938        let PendingSpawn { id, spec, config, ack } = spawn;
939        let entry = ChildEntry {
940            spec,
941            config,
942            dynamic: true,
943        };
944        match worker_state.add_worker(id, &entry.spec, &entry.config) {
945            Ok(()) => {
946                if entry.config.significant {
947                    *significant_remaining += 1;
948                }
949                self.active.fetch_add(1, Ordering::Relaxed);
950                children.insert(id, entry);
951                let _ = ack.send(Ok(()));
952            }
953            Err(e) => {
954                // Registration failed (e.g. an invalid child name). Report it to the waiting caller as `Rejected` --
955                // distinct from `SupervisorGone` -- so the underlying cause isn't lost.
956                error!(supervisor_id = %self.supervisor_id, error = %e, "Failed to spawn dynamic child.");
957                let _ = ack.send(Err(SpawnError::Rejected { source: e.into() }));
958            }
959        }
960    }
961
962    async fn run_inner(&self, process: Process, process_shutdown: ShutdownHandle) -> Result<(), SupervisorError> {
963        // Publish a fresh command channel for this run so handles can spawn dynamic children into it; while it's set,
964        // handles observe us as running.
965        let (cmd_tx, cmd_rx) = mpsc::channel(DYNAMIC_SPAWN_CHANNEL_CAPACITY);
966        *self.current_tx.lock().unwrap() = Some(cmd_tx);
967
968        let result = self.supervise(process, process_shutdown, cmd_rx).await;
969
970        // The run is over. Clear the sender so later spawns observe `SupervisorGone`, and reset the dynamic-children
971        // gauge. Dropping the receiver (owned by `supervise`) already rejected anything still in flight.
972        *self.current_tx.lock().unwrap() = None;
973        self.active.store(0, Ordering::Relaxed);
974        result
975    }
976
977    async fn supervise(
978        &self, process: Process, process_shutdown: ShutdownHandle, mut cmd_rx: mpsc::Receiver<PendingSpawn>,
979    ) -> Result<(), SupervisorError> {
980        let mut restart_state = RestartState::new(self.restart_strategy);
981        let mut worker_state = WorkerState::new(process, self.shutdown_mode, self.shutdown_budget);
982
983        // The live roster of children -- both static (seeded below) and dynamic (added via the handle) -- keyed by a
984        // stable id. A restart re-runs a child by id; a child that isn't restarted is removed from the roster.
985        let mut children: FastHashMap<u64, ChildEntry> = FastHashMap::default();
986
987        // Spawn the static children. Initialization is folded into each worker's task, so this returns immediately --
988        // children initialize concurrently in the background.
989        self.spawn_static_children(&mut children, &mut worker_state)?;
990
991        // Track how many significant children are still running, for `AutoShutdown` evaluation.
992        let mut significant_remaining = children.values().filter(|entry| entry.config.significant).count();
993
994        // Now we supervise.
995        pin!(process_shutdown);
996
997        let outcome = loop {
998            select! {
999                // Shutdown takes priority so a flood of dynamic spawns can't starve it.
1000                biased;
1001
1002                // Shutdown has been triggered; break out of the loop with a clean outcome and tear down below. (We
1003                // can't touch `cmd_rx` in any arm's handler -- the `recv` arm below borrows it for the whole
1004                // `select!` -- so all teardown happens after the loop.)
1005                _ = &mut process_shutdown => break Ok(()),
1006
1007                // A handle asked us to spawn a dynamic child. The published sender keeps the channel open for the whole
1008                // run, so `recv` only yields `None` once we close it during teardown.
1009                spawn = cmd_rx.recv() => {
1010                    if let Some(spawn) = spawn {
1011                        self.spawn_dynamic_child(spawn, &mut worker_state, &mut children, &mut significant_remaining);
1012                    }
1013                }
1014
1015                (child_id, worker_result) = worker_state.wait_for_next_worker() => {
1016                    // Pull out what we need from the roster before we mutate it.
1017                    let (child_name, config, dynamic) = {
1018                        let entry = children.get(&child_id).expect("completed worker must be present in the roster");
1019                        (entry.spec.name().to_string(), entry.config.clone(), entry.dynamic)
1020                    };
1021
1022                    // Initialization failures are not eligible for restart -- they propagate immediately.
1023                    if let Err(WorkerError::Initialization { child_name: inner, source }) = worker_result {
1024                        // If the error came from a nested supervisor, include the original child name to make the error
1025                        // chain more informative (e.g., "ctrl-pln/privileged-api").
1026                        let full_name = match inner {
1027                            Some(inner) => format!("{}/{}", child_name, inner),
1028                            None => child_name.clone(),
1029                        };
1030
1031                        error!(supervisor_id = %self.supervisor_id, worker_name = full_name, "Child process failed to initialize: {}", source);
1032                        break Err(SupervisorError::FailedToInitialize { child_name: full_name, source });
1033                    }
1034
1035                    // A worker exited abnormally if it returned an error, panicked, or was aborted; a clean exit is
1036                    // `Ok(())`. Together with the worker's restart policy, this determines whether we restart it.
1037                    let abnormal = worker_result.is_err();
1038                    let worker_result = worker_result.map_err(|e| match e {
1039                        WorkerError::Runtime(e) => ProcessError::Terminated { source: e },
1040                        WorkerError::Initialization { .. } => unreachable!("handled above"),
1041                        // A nested supervisor only reports `ShutdownTimedOut` while draining, which is driven by its own
1042                        // `process_shutdown` -- and that fires only when *this* supervisor is itself draining it, i.e.
1043                        // from `shutdown_workers` below, never from this main-loop arm. Treat it as a runtime
1044                        // termination defensively rather than asserting unreachable.
1045                        WorkerError::ShutdownTimedOut { aborted } => ProcessError::Terminated {
1046                            source: SupervisorError::ShutdownTimedOut { aborted }.into(),
1047                        },
1048                    });
1049
1050                    if !config.restart.should_restart(abnormal) {
1051                        // Not eligible for restart given how it exited. Drop it from the roster, and free its slot/gauge
1052                        // if it was dynamic. Crucially, we do NOT consult `evaluate_restart` here: non-restarts must not
1053                        // consume the restart-intensity budget, otherwise a steady stream of terminating temporary
1054                        // children would eventually trip the limit and tear the supervisor (and its siblings) down.
1055                        //
1056                        // An abnormal exit is reported at `warn` rather than `debug`: a child that isn't restarted --
1057                        // every dynamically-spawned child, in practice -- has no other path back to its owner, so this
1058                        // is the only place its failure is surfaced.
1059                        if abnormal {
1060                            warn!(supervisor_id = %self.supervisor_id, worker_name = %child_name, restart = ?config.restart, ?worker_result, "Child process exited with an error and is not eligible for restart.");
1061                        } else {
1062                            debug!(supervisor_id = %self.supervisor_id, worker_name = %child_name, restart = ?config.restart, "Child process exited and is not eligible for restart.");
1063                        }
1064                        children.remove(&child_id);
1065                        if dynamic {
1066                            self.active.fetch_sub(1, Ordering::Relaxed);
1067                        }
1068
1069                        // A significant child terminating without restart can drive the supervisor to shut down, per its
1070                        // `AutoShutdown` policy -- cascading an unexpected (or intentional) child exit into the
1071                        // supervisor stopping and propagating up the tree.
1072                        if config.significant {
1073                            significant_remaining = significant_remaining.saturating_sub(1);
1074                            let auto_shutdown = match self.auto_shutdown {
1075                                AutoShutdown::Never => false,
1076                                AutoShutdown::AnySignificant => true,
1077                                AutoShutdown::AllSignificant => significant_remaining == 0,
1078                            };
1079                            if auto_shutdown {
1080                                warn!(supervisor_id = %self.supervisor_id, worker_name = %child_name, ?worker_result, "Significant child terminated; shutting down supervisor.");
1081                                break Err(SupervisorError::SignificantChildExited);
1082                            }
1083                        }
1084                    } else {
1085                        match restart_state.evaluate_restart() {
1086                            RestartAction::Restart(mode) => match mode {
1087                                RestartMode::OneForOne => {
1088                                    warn!(supervisor_id = %self.supervisor_id, worker_name = %child_name, ?worker_result, "Child process terminated, restarting.");
1089                                    let spec = children.get(&child_id).expect("present for restart").spec.clone();
1090                                    if let Err(e) = worker_state.add_worker(child_id, &spec, &config) {
1091                                        break Err(e);
1092                                    }
1093                                }
1094                                RestartMode::OneForAll => {
1095                                    warn!(supervisor_id = %self.supervisor_id, worker_name = %child_name, ?worker_result, "Child process terminated, restarting all processes.");
1096                                    // This drain is part of a restart, not a shutdown: any forced aborts here are
1097                                    // already logged per-worker, and the supervisor keeps running, so the count does
1098                                    // not feed the unclean-shutdown signal.
1099                                    let _ = worker_state.shutdown_workers().await;
1100                                    // A one-for-all restart resets to the static roster; dynamic children are not
1101                                    // restored (they're lost on a supervisor-level restart, matching Erlang/OTP), and
1102                                    // temporary children are not restarted.
1103                                    children.clear();
1104                                    self.active.store(0, Ordering::Relaxed);
1105                                    let respawn = self.respawn_children_one_for_all(&mut children, &mut worker_state);
1106                                    if let Err(e) = respawn {
1107                                        break Err(e);
1108                                    }
1109                                    significant_remaining =
1110                                        children.values().filter(|entry| entry.config.significant).count();
1111                                }
1112                            },
1113                            RestartAction::Shutdown => {
1114                                error!(supervisor_id = %self.supervisor_id, worker_name = %child_name, ?worker_result, "Supervisor shutting down due to restart limits.");
1115                                break Err(SupervisorError::Shutdown);
1116                            }
1117                        }
1118                    }
1119                }
1120            }
1121        };
1122
1123        // The run is ending -- either cleanly (shutdown was signalled) or with an error (a child failed to initialize
1124        // or restart, the restart limit was exceeded, or a significant child exited). On every path: stop accepting
1125        // spawns and reject anything still queued -- rather
1126        // than starting children only to tear them down immediately -- then shut down all children. Closing the channel
1127        // before the (possibly slow) shutdown also unblocks any handle parked on a full channel, so a spawn racing the
1128        // teardown observes `SupervisorGone` promptly instead of hanging until shutdown finishes.
1129        cmd_rx.close();
1130        while let Ok(spawn) = cmd_rx.try_recv() {
1131            let _ = spawn.ack.send(Err(SpawnError::SupervisorGone));
1132        }
1133        let aborted = worker_state.shutdown_workers().await;
1134
1135        // A requested shutdown that nonetheless had to forcefully abort one or more workers (here or anywhere in the
1136        // subtree below us) is surfaced as an unclean shutdown so it propagates up the tree rather than being reported
1137        // as success. An outcome that already carries an error (initialization, restart limit, significant child)
1138        // takes precedence -- that's the root cause -- and the forced aborts are left to the per-worker warnings.
1139        match outcome {
1140            Ok(()) if aborted > 0 => {
1141                warn!(supervisor_id = %self.supervisor_id, aborted, "Shutdown completed uncleanly; workers were forcefully aborted.");
1142                Err(SupervisorError::ShutdownTimedOut { aborted })
1143            }
1144            outcome => outcome,
1145        }
1146    }
1147
1148    fn as_nested_process(&self, process: Process, process_shutdown: ShutdownHandle) -> WorkerFuture {
1149        // Simple wrapper around `run_inner` to satisfy the return type signature needed when running the supervisor as
1150        // a nested child process in another supervisor.
1151        debug!(supervisor_id = %self.supervisor_id, "Nested supervisor starting.");
1152
1153        // Create a standalone clone of ourselves so we can fulfill the future signature.
1154        let sup = self.inner_clone();
1155
1156        Box::pin(async move {
1157            sup.run_inner(process, process_shutdown)
1158                .await
1159                .map_err(WorkerError::from)
1160        })
1161    }
1162
1163    /// Runs the supervisor forever.
1164    ///
1165    /// # Errors
1166    ///
1167    /// If the supervisor exceeds its restart limits, or fails to initialize a child process, an error is returned.
1168    pub async fn run(&mut self) -> Result<(), SupervisorError> {
1169        // Create a no-op `ShutdownHandle` to satisfy the `run_inner` function. This is never used since we want to run
1170        // forever, but we need to satisfy the signature.
1171        let process_shutdown = ShutdownHandle::noop();
1172        let process = Process::supervisor(&self.supervisor_id, None).context(InvalidName {
1173            name: self.supervisor_id.to_string(),
1174        })?;
1175
1176        debug!(supervisor_id = %self.supervisor_id, "Supervisor starting.");
1177        self.run_inner(process.clone(), process_shutdown)
1178            .into_process_future(process)
1179            .await
1180    }
1181
1182    /// Runs the supervisor until shutdown is triggered.
1183    ///
1184    /// When `shutdown` resolves, the supervisor will shutdown all child processes according to their shutdown strategy,
1185    /// and then return.
1186    ///
1187    /// # Errors
1188    ///
1189    /// If the supervisor exceeds its restart limits, or fails to initialize a child process, an error is returned.
1190    pub async fn run_with_shutdown<F: Future + Send + 'static>(&mut self, shutdown: F) -> Result<(), SupervisorError> {
1191        // Drive the caller-provided shutdown future into a trigger so the supervisor can begin shutting down its
1192        // children once `shutdown` resolves. The trigger fires at most once (guarded), and otherwise fires on drop if
1193        // the supervisor returns on its own first.
1194        let (shutdown_coordinator, shutdown_handle) = ShutdownHandle::paired();
1195        let run = self.run_with_shutdown_inner(shutdown_handle, None);
1196        pin!(run, shutdown);
1197
1198        let mut shutdown_coordinator = Some(shutdown_coordinator);
1199        loop {
1200            select! {
1201                result = &mut run => return result,
1202                _ = &mut shutdown, if shutdown_coordinator.is_some() => {
1203                    shutdown_coordinator.take().expect("coordinator present per select guard").shutdown();
1204                }
1205            }
1206        }
1207    }
1208
1209    /// Runs the supervisor until the given `ShutdownHandle` signal is received.
1210    ///
1211    /// This is an internal variant of `run_with_shutdown` that takes a `ShutdownHandle` directly, used when spawning
1212    /// supervisors in dedicated runtimes where the shutdown signal is already wrapped in a `ShutdownHandle`.
1213    ///
1214    /// If `dataspace` is provided, the supervisor will use it instead of creating a new one. This is used to propagate
1215    /// the parent's dataspace across OS thread boundaries for dedicated runtimes.
1216    ///
1217    /// # Errors
1218    ///
1219    /// If the supervisor exceeds its restart limits, or fails to initialize a child process, an error is returned.
1220    pub(crate) async fn run_with_shutdown_inner(
1221        &mut self, process_shutdown: ShutdownHandle, dataspace: Option<DataspaceRegistry>,
1222    ) -> Result<(), SupervisorError> {
1223        let process =
1224            Process::supervisor_with_dataspace(&self.supervisor_id, None, dataspace).context(InvalidName {
1225                name: self.supervisor_id.to_string(),
1226            })?;
1227
1228        debug!(supervisor_id = %self.supervisor_id, "Supervisor starting.");
1229        self.run_inner(process.clone(), process_shutdown)
1230            .into_process_future(process)
1231            .await
1232    }
1233
1234    fn inner_clone(&self) -> Self {
1235        // This is no different than if we just implemented `Clone` directly, but it allows us to avoid exposing a
1236        // _public_ implementation of `Clone`, which we don't want normal users to be able to do. We only need this
1237        // internally to support nested supervisors.
1238        Self {
1239            supervisor_id: Arc::clone(&self.supervisor_id),
1240            child_specs: self.child_specs.clone(),
1241            restart_strategy: self.restart_strategy,
1242            auto_shutdown: self.auto_shutdown,
1243            shutdown_mode: self.shutdown_mode,
1244            shutdown_budget: self.shutdown_budget,
1245            runtime_mode: self.runtime_mode.clone(),
1246            current_tx: Arc::clone(&self.current_tx),
1247            id_counter: Arc::clone(&self.id_counter),
1248            active: Arc::clone(&self.active),
1249        }
1250    }
1251}
1252
1253#[cfg(test)]
1254mod tests {
1255    use std::{
1256        future::pending,
1257        sync::atomic::{AtomicBool, AtomicUsize, Ordering},
1258    };
1259
1260    use async_trait::async_trait;
1261    use saluki_common::sync::shutdown::ShutdownCoordinator;
1262    use saluki_metrics::test::TestRecorder;
1263    use tokio::{
1264        sync::oneshot,
1265        task::JoinHandle,
1266        time::{sleep, timeout},
1267    };
1268
1269    use super::*;
1270    use crate::runtime::noninterruptible_worker;
1271    use crate::test_support::wait_until;
1272
1273    /// Behavior for a mock worker during initialization.
1274    #[derive(Clone)]
1275    enum InitBehavior {
1276        /// Initialization succeeds immediately.
1277        Instant,
1278
1279        /// Initialization takes the given duration before succeeding.
1280        Slow(Duration),
1281
1282        /// Initialization fails with the given message.
1283        Fail(&'static str),
1284    }
1285
1286    /// Behavior for a mock worker during runtime (after initialization).
1287    #[derive(Clone)]
1288    enum RunBehavior {
1289        /// Runs until shutdown is received.
1290        UntilShutdown,
1291
1292        /// Fails with the given error message after the given delay.
1293        FailAfter(Duration, &'static str),
1294
1295        /// Completes successfully after the given delay.
1296        CompleteAfter(Duration),
1297
1298        /// On shutdown, sleeps for the given duration before exiting (to exercise concurrent draining).
1299        SlowShutdown(Duration),
1300
1301        /// Ignores shutdown entirely and runs forever (to exercise abort-at-deadline).
1302        IgnoreShutdown,
1303
1304        /// Panics after the given delay, unless shutdown arrives first.
1305        PanicAfter(Duration),
1306    }
1307
1308    /// A configurable mock worker for testing supervisor behavior.
1309    struct MockWorker {
1310        name: &'static str,
1311        init_behavior: InitBehavior,
1312        run_behavior: RunBehavior,
1313        start_count: Arc<AtomicUsize>,
1314        finish_count: Arc<AtomicUsize>,
1315        brutal_shutdown: bool,
1316        graceful_timeout: Duration,
1317    }
1318
1319    impl MockWorker {
1320        /// Creates a worker that runs until shutdown.
1321        fn long_running(name: &'static str) -> Self {
1322            Self {
1323                name,
1324                init_behavior: InitBehavior::Instant,
1325                run_behavior: RunBehavior::UntilShutdown,
1326                start_count: Arc::new(AtomicUsize::new(0)),
1327                finish_count: Arc::new(AtomicUsize::new(0)),
1328                brutal_shutdown: false,
1329                graceful_timeout: Duration::from_millis(500),
1330            }
1331        }
1332
1333        /// Creates a worker that fails after the given delay.
1334        fn failing(name: &'static str, delay: Duration) -> Self {
1335            Self {
1336                name,
1337                init_behavior: InitBehavior::Instant,
1338                run_behavior: RunBehavior::FailAfter(delay, "worker failed"),
1339                start_count: Arc::new(AtomicUsize::new(0)),
1340                finish_count: Arc::new(AtomicUsize::new(0)),
1341                brutal_shutdown: false,
1342                graceful_timeout: Duration::from_millis(500),
1343            }
1344        }
1345
1346        /// Creates a worker that completes successfully after the given delay.
1347        fn completing(name: &'static str, delay: Duration) -> Self {
1348            Self {
1349                name,
1350                init_behavior: InitBehavior::Instant,
1351                run_behavior: RunBehavior::CompleteAfter(delay),
1352                start_count: Arc::new(AtomicUsize::new(0)),
1353                finish_count: Arc::new(AtomicUsize::new(0)),
1354                brutal_shutdown: false,
1355                graceful_timeout: Duration::from_millis(500),
1356            }
1357        }
1358
1359        /// Creates a worker that sleeps for `delay` after observing shutdown before exiting.
1360        fn slow_shutdown(name: &'static str, delay: Duration) -> Self {
1361            Self {
1362                name,
1363                init_behavior: InitBehavior::Instant,
1364                run_behavior: RunBehavior::SlowShutdown(delay),
1365                start_count: Arc::new(AtomicUsize::new(0)),
1366                finish_count: Arc::new(AtomicUsize::new(0)),
1367                brutal_shutdown: false,
1368                graceful_timeout: Duration::from_millis(500),
1369            }
1370        }
1371
1372        /// Creates a worker that never reacts to shutdown.
1373        fn ignore_shutdown(name: &'static str) -> Self {
1374            Self {
1375                name,
1376                init_behavior: InitBehavior::Instant,
1377                run_behavior: RunBehavior::IgnoreShutdown,
1378                start_count: Arc::new(AtomicUsize::new(0)),
1379                finish_count: Arc::new(AtomicUsize::new(0)),
1380                brutal_shutdown: false,
1381                graceful_timeout: Duration::from_millis(500),
1382            }
1383        }
1384
1385        /// Creates a worker that panics after the given delay.
1386        fn panicking(name: &'static str, delay: Duration) -> Self {
1387            Self {
1388                name,
1389                init_behavior: InitBehavior::Instant,
1390                run_behavior: RunBehavior::PanicAfter(delay),
1391                start_count: Arc::new(AtomicUsize::new(0)),
1392                finish_count: Arc::new(AtomicUsize::new(0)),
1393                brutal_shutdown: false,
1394                graceful_timeout: Duration::from_millis(500),
1395            }
1396        }
1397
1398        /// Creates a worker that fails during initialization.
1399        fn init_failure(name: &'static str) -> Self {
1400            Self {
1401                name,
1402                init_behavior: InitBehavior::Fail("init failed"),
1403                run_behavior: RunBehavior::UntilShutdown,
1404                start_count: Arc::new(AtomicUsize::new(0)),
1405                finish_count: Arc::new(AtomicUsize::new(0)),
1406                brutal_shutdown: false,
1407                graceful_timeout: Duration::from_millis(500),
1408            }
1409        }
1410
1411        /// Creates a worker with slow initialization.
1412        fn slow_init(name: &'static str, init_delay: Duration) -> Self {
1413            Self {
1414                name,
1415                init_behavior: InitBehavior::Slow(init_delay),
1416                run_behavior: RunBehavior::UntilShutdown,
1417                start_count: Arc::new(AtomicUsize::new(0)),
1418                finish_count: Arc::new(AtomicUsize::new(0)),
1419                brutal_shutdown: false,
1420                graceful_timeout: Duration::from_millis(500),
1421            }
1422        }
1423
1424        /// Returns a shared handle to the start count for this worker.
1425        ///
1426        /// The start count ticks up the instant the worker's run future begins executing, which is *before* any
1427        /// programmed delay elapses. It records that the worker started (or was restarted), not that it ran to any
1428        /// particular outcome.
1429        fn start_count(&self) -> Arc<AtomicUsize> {
1430            Arc::clone(&self.start_count)
1431        }
1432
1433        /// Returns a shared handle to the finish count for this worker.
1434        ///
1435        /// The finish count ticks up only when the worker runs to its *own* programmed terminal state -- a
1436        /// [`RunBehavior::FailAfter`] failure or a [`RunBehavior::CompleteAfter`] completion -- and not when it is cut
1437        /// short by shutdown. Tests use it to wait for a worker to actually fail or complete (rather than merely
1438        /// start) before asserting on restart behavior, so the failure/completion path is genuinely exercised.
1439        fn finish_count(&self) -> Arc<AtomicUsize> {
1440            Arc::clone(&self.finish_count)
1441        }
1442
1443        /// Configures this worker to use a `Brutal` shutdown strategy (immediate abort, no graceful wait).
1444        fn with_brutal_shutdown(mut self) -> Self {
1445            self.brutal_shutdown = true;
1446            self
1447        }
1448
1449        /// Overrides the worker's graceful shutdown timeout (defaults to 500 milliseconds).
1450        fn with_graceful_timeout(mut self, timeout: Duration) -> Self {
1451            self.graceful_timeout = timeout;
1452            self
1453        }
1454    }
1455
1456    #[async_trait]
1457    impl Supervisable for MockWorker {
1458        fn name(&self) -> &str {
1459            self.name
1460        }
1461
1462        fn shutdown_strategy(&self) -> ShutdownStrategy {
1463            if self.brutal_shutdown {
1464                ShutdownStrategy::Brutal
1465            } else {
1466                ShutdownStrategy::Graceful(self.graceful_timeout)
1467            }
1468        }
1469
1470        async fn initialize(&self, process_shutdown: ShutdownHandle) -> Result<SupervisorFuture, InitializationError> {
1471            match &self.init_behavior {
1472                InitBehavior::Instant => {}
1473                InitBehavior::Slow(delay) => {
1474                    sleep(*delay).await;
1475                }
1476                InitBehavior::Fail(msg) => {
1477                    return Err(InitializationError::Failed {
1478                        source: GenericError::msg(*msg),
1479                    });
1480                }
1481            }
1482
1483            let start_count = Arc::clone(&self.start_count);
1484            let finish_count = Arc::clone(&self.finish_count);
1485            let run_behavior = self.run_behavior.clone();
1486
1487            Ok(Box::pin(async move {
1488                start_count.fetch_add(1, Ordering::SeqCst);
1489
1490                match run_behavior {
1491                    RunBehavior::UntilShutdown => {
1492                        process_shutdown.await;
1493                        Ok(())
1494                    }
1495                    RunBehavior::FailAfter(delay, msg) => {
1496                        select! {
1497                            _ = sleep(delay) => {
1498                                // Ran to our own programmed failure rather than being cut short by shutdown; record
1499                                // it so tests can wait for the failure to actually happen before asserting.
1500                                finish_count.fetch_add(1, Ordering::SeqCst);
1501                                Err(GenericError::msg(msg))
1502                            }
1503                            _ = process_shutdown => {
1504                                Ok(())
1505                            }
1506                        }
1507                    }
1508                    RunBehavior::CompleteAfter(delay) => {
1509                        select! {
1510                            _ = sleep(delay) => {
1511                                // Ran to our own programmed completion rather than being cut short by shutdown.
1512                                finish_count.fetch_add(1, Ordering::SeqCst);
1513                                Ok(())
1514                            }
1515                            _ = process_shutdown => Ok(()),
1516                        }
1517                    }
1518                    RunBehavior::SlowShutdown(delay) => {
1519                        process_shutdown.await;
1520                        sleep(delay).await;
1521                        Ok(())
1522                    }
1523                    RunBehavior::IgnoreShutdown => {
1524                        // Hold the handle (so the supervisor counts us as outstanding) but never react to it.
1525                        let _hold = process_shutdown;
1526                        pending().await
1527                    }
1528                    RunBehavior::PanicAfter(delay) => {
1529                        select! {
1530                            _ = sleep(delay) => panic!("worker panicked"),
1531                            _ = process_shutdown => Ok(()),
1532                        }
1533                    }
1534                }
1535            }))
1536        }
1537    }
1538
1539    /// Helper: run a supervisor with a oneshot-based shutdown trigger.
1540    ///
1541    /// Returns the shutdown sender and a join handle for the run. The supervisor is polled to a running state (its
1542    /// static children spawned) via readiness polling rather than a blind startup sleep, so callers can rely on it
1543    /// being live on return.
1544    async fn run_supervisor_with_trigger(
1545        supervisor: Supervisor,
1546    ) -> (oneshot::Sender<()>, JoinHandle<Result<(), SupervisorError>>) {
1547        // Grab a handle before moving the supervisor into the run task so we can observe when it actually starts.
1548        let sup_handle = supervisor.handle();
1549        let mut supervisor = supervisor;
1550
1551        let (tx, rx) = oneshot::channel();
1552        let handle = tokio::spawn(async move { supervisor.run_with_shutdown(rx).await });
1553
1554        wait_until("supervisor is running", || sup_handle.is_running()).await;
1555        (tx, handle)
1556    }
1557
1558    /// Helper: awaits a spawned supervisor run to completion under a bounded timeout, unwrapping the join.
1559    ///
1560    /// Collapses the `timeout(..).await.unwrap().unwrap()` suffix repeated across the restart/shutdown tests into one
1561    /// call with useful panic messages.
1562    async fn join_supervisor(handle: JoinHandle<Result<(), SupervisorError>>) -> Result<(), SupervisorError> {
1563        timeout(Duration::from_secs(2), handle)
1564            .await
1565            .expect("supervisor should exit promptly")
1566            .expect("supervisor task should not panic")
1567    }
1568
1569    // -- Supervisor run mode tests ---------------------------------------------------------
1570
1571    #[tokio::test]
1572    async fn standalone_supervisor_shuts_down_cleanly() {
1573        let mut sup = Supervisor::new("test-sup").unwrap();
1574        sup.add_worker(MockWorker::long_running("worker1"));
1575        sup.add_worker(MockWorker::long_running("worker2"));
1576
1577        let (tx, handle) = run_supervisor_with_trigger(sup).await;
1578        tx.send(()).unwrap();
1579
1580        let result = join_supervisor(handle).await;
1581        assert!(result.is_ok());
1582    }
1583
1584    #[tokio::test]
1585    async fn nested_supervisor_shuts_down_cleanly() {
1586        let mut child_sup = Supervisor::new("child-sup").unwrap();
1587        child_sup.add_worker(MockWorker::long_running("inner-worker"));
1588
1589        let mut parent_sup = Supervisor::new("parent-sup").unwrap();
1590        parent_sup.add_worker(MockWorker::long_running("outer-worker"));
1591        parent_sup.add_worker(child_sup);
1592
1593        let (tx, handle) = run_supervisor_with_trigger(parent_sup).await;
1594        tx.send(()).unwrap();
1595
1596        let result = join_supervisor(handle).await;
1597        assert!(result.is_ok());
1598    }
1599
1600    #[tokio::test]
1601    async fn empty_supervisor_idles_until_shutdown() {
1602        // A supervisor with no static children is valid: it idles, waiting for dynamic children, and shuts down
1603        // cleanly when signalled. (Before dynamic children were folded in, this returned a `NoChildren` error.)
1604        let sup = Supervisor::new("empty-sup").unwrap();
1605
1606        let (tx, handle) = run_supervisor_with_trigger(sup).await;
1607        assert!(!handle.is_finished(), "an empty supervisor must idle rather than exit");
1608
1609        tx.send(()).unwrap();
1610        let result = join_supervisor(handle).await;
1611        assert!(result.is_ok());
1612    }
1613
1614    // -- Child restart behavior tests ------------------------------------------------------
1615
1616    #[tokio::test]
1617    async fn one_for_one_restarts_only_failed_child() {
1618        let failing = MockWorker::failing("failing-worker", Duration::from_millis(50));
1619        let failing_count = failing.start_count();
1620
1621        let stable = MockWorker::long_running("stable-worker");
1622        let stable_count = stable.start_count();
1623
1624        let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
1625            RestartStrategy::one_to_one().with_intensity_and_period(20, Duration::from_secs(10)),
1626        );
1627        sup.add_worker(stable);
1628        sup.add_worker(failing);
1629
1630        let (tx, handle) = run_supervisor_with_trigger(sup).await;
1631
1632        // Wait until the failing worker has actually been restarted (its second start), then shut down.
1633        wait_until("the failing worker has been restarted", || {
1634            failing_count.load(Ordering::SeqCst) >= 2
1635        })
1636        .await;
1637        let _ = tx.send(());
1638
1639        let result = join_supervisor(handle).await;
1640        assert!(result.is_ok());
1641
1642        // The failing worker should have been started multiple times.
1643        assert!(
1644            failing_count.load(Ordering::SeqCst) >= 2,
1645            "failing worker should have been restarted"
1646        );
1647        // The stable worker should only have been started once (never restarted).
1648        assert_eq!(
1649            stable_count.load(Ordering::SeqCst),
1650            1,
1651            "stable worker should not have been restarted"
1652        );
1653    }
1654
1655    #[tokio::test]
1656    async fn one_for_all_restarts_all_children() {
1657        let failing = MockWorker::failing("failing-worker", Duration::from_millis(50));
1658        let failing_count = failing.start_count();
1659
1660        let stable = MockWorker::long_running("stable-worker");
1661        let stable_count = stable.start_count();
1662
1663        let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
1664            RestartStrategy::one_for_all().with_intensity_and_period(20, Duration::from_secs(10)),
1665        );
1666        sup.add_worker(stable);
1667        sup.add_worker(failing);
1668
1669        let (tx, handle) = run_supervisor_with_trigger(sup).await;
1670
1671        // Wait until a one-for-all cycle has restarted both workers (each on its second start), then shut down.
1672        wait_until("both workers have been restarted", || {
1673            failing_count.load(Ordering::SeqCst) >= 2 && stable_count.load(Ordering::SeqCst) >= 2
1674        })
1675        .await;
1676        let _ = tx.send(());
1677
1678        let result = join_supervisor(handle).await;
1679        assert!(result.is_ok());
1680
1681        // Both workers should have been started multiple times.
1682        assert!(
1683            failing_count.load(Ordering::SeqCst) >= 2,
1684            "failing worker should have been restarted"
1685        );
1686        assert!(
1687            stable_count.load(Ordering::SeqCst) >= 2,
1688            "stable worker should also have been restarted"
1689        );
1690    }
1691
1692    #[tokio::test]
1693    async fn one_for_all_does_not_restart_temporary_children() {
1694        // A permanent worker that fails repeatedly drives one-for-all restarts; a temporary sibling is shut down with
1695        // the group on each cycle but, per OTP semantics, must never be brought back.
1696        let failing = MockWorker::failing("failing-worker", Duration::from_millis(50));
1697        let failing_count = failing.start_count();
1698
1699        let temp = MockWorker::long_running("temp-worker");
1700        let temp_count = temp.start_count();
1701
1702        let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
1703            RestartStrategy::one_for_all().with_intensity_and_period(20, Duration::from_secs(10)),
1704        );
1705        sup.add_worker(ChildSpecification::worker(temp).with_restart_type(RestartType::Temporary));
1706        sup.add_worker(failing);
1707
1708        let (tx, handle) = run_supervisor_with_trigger(sup).await;
1709
1710        // Wait until the permanent worker has driven at least one one-for-all restart, then shut down.
1711        wait_until("the permanent worker has been restarted", || {
1712            failing_count.load(Ordering::SeqCst) >= 2
1713        })
1714        .await;
1715        let _ = tx.send(());
1716
1717        let result = join_supervisor(handle).await;
1718        assert!(result.is_ok());
1719        assert!(
1720            failing_count.load(Ordering::SeqCst) >= 2,
1721            "permanent worker should have been restarted by one-for-all"
1722        );
1723        assert_eq!(
1724            temp_count.load(Ordering::SeqCst),
1725            1,
1726            "temporary child must not be restarted by a one-for-all group restart"
1727        );
1728    }
1729
1730    #[tokio::test]
1731    async fn one_for_all_restarts_transient_children() {
1732        // A transient child that exits cleanly is not restarted on its own, but a one-for-all restart triggered by a
1733        // sibling restarts it anyway -- matching OTP, where only temporary children are exempt from group restarts.
1734        let transient = MockWorker::completing("transient-worker", Duration::from_millis(30));
1735        let transient_count = transient.start_count();
1736
1737        // Fails after the transient has already exited cleanly, so the group restart is what brings the transient back.
1738        let failing = MockWorker::failing("failing-worker", Duration::from_millis(80));
1739
1740        let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
1741            RestartStrategy::one_for_all().with_intensity_and_period(20, Duration::from_secs(10)),
1742        );
1743        sup.add_worker(ChildSpecification::worker(transient).with_restart_type(RestartType::Transient));
1744        sup.add_worker(failing);
1745
1746        let (tx, handle) = run_supervisor_with_trigger(sup).await;
1747
1748        wait_until("the transient worker has been restarted by the group", || {
1749            transient_count.load(Ordering::SeqCst) >= 2
1750        })
1751        .await;
1752        let _ = tx.send(());
1753
1754        let result = join_supervisor(handle).await;
1755        assert!(result.is_ok());
1756        assert!(
1757            transient_count.load(Ordering::SeqCst) >= 2,
1758            "transient child must be restarted by a one-for-all group restart, even after a clean exit"
1759        );
1760    }
1761
1762    #[tokio::test]
1763    async fn transient_abnormal_exit_triggers_one_for_all() {
1764        // A transient child's *own* abnormal exit is restartable, so under one-for-all it triggers a whole-group
1765        // restart -- the sibling is restarted too, not just the transient.
1766        let transient = MockWorker::failing("transient-worker", Duration::from_millis(50));
1767        let transient_count = transient.start_count();
1768
1769        let stable = MockWorker::long_running("stable-worker");
1770        let stable_count = stable.start_count();
1771
1772        let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
1773            RestartStrategy::one_for_all().with_intensity_and_period(20, Duration::from_secs(10)),
1774        );
1775        sup.add_worker(ChildSpecification::worker(transient).with_restart_type(RestartType::Transient));
1776        sup.add_worker(stable);
1777
1778        let (tx, handle) = run_supervisor_with_trigger(sup).await;
1779
1780        wait_until("the abnormal exit has restarted both workers", || {
1781            transient_count.load(Ordering::SeqCst) >= 2 && stable_count.load(Ordering::SeqCst) >= 2
1782        })
1783        .await;
1784        let _ = tx.send(());
1785
1786        let result = join_supervisor(handle).await;
1787        assert!(result.is_ok());
1788        assert!(
1789            transient_count.load(Ordering::SeqCst) >= 2,
1790            "transient worker must be restarted after its own abnormal exit"
1791        );
1792        assert!(
1793            stable_count.load(Ordering::SeqCst) >= 2,
1794            "the transient's abnormal exit must trigger a one-for-all that also restarts the sibling"
1795        );
1796    }
1797
1798    #[tokio::test]
1799    async fn restart_limit_exceeded_shuts_down_supervisor() {
1800        let mut sup = Supervisor::new("test-sup")
1801            .unwrap()
1802            .with_restart_strategy(RestartStrategy::one_to_one().with_intensity_and_period(1, Duration::from_secs(10)));
1803        // This worker fails immediately, which will exhaust the restart budget quickly.
1804        sup.add_worker(MockWorker::failing("fast-fail", Duration::ZERO));
1805
1806        let (tx, rx) = oneshot::channel::<()>();
1807        let handle = tokio::spawn(async move { sup.run_with_shutdown(rx).await });
1808
1809        let result = join_supervisor(handle).await;
1810        drop(tx);
1811
1812        assert!(matches!(result, Err(SupervisorError::Shutdown)));
1813    }
1814
1815    // -- Restart type tests ----------------------------------------------------------------
1816
1817    #[tokio::test]
1818    async fn temporary_child_is_not_restarted() {
1819        // A temporary worker that fails quickly, alongside a long-running worker that keeps the supervisor alive.
1820        let temp = MockWorker::failing("temp-worker", Duration::from_millis(50));
1821        let temp_started = temp.start_count();
1822        let temp_failed = temp.finish_count();
1823
1824        let stable = MockWorker::long_running("stable-worker");
1825
1826        let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
1827            RestartStrategy::one_to_one().with_intensity_and_period(20, Duration::from_secs(10)),
1828        );
1829        sup.add_worker(stable);
1830        sup.add_worker(ChildSpecification::worker(temp).with_restart_type(RestartType::Temporary));
1831
1832        let (tx, handle) = run_supervisor_with_trigger(sup).await;
1833
1834        // Wait for the worker to *actually fail*, not merely start. `start_count` ticks up the instant the worker
1835        // begins running -- well before its 50ms failure -- so shutting down as soon as it reached 1 would tear the
1836        // supervisor down before the failure -> no-restart path ever ran, hiding a regression that restarted a
1837        // temporary child (or charged the failure against restart intensity). `finish_count` ticks only once the
1838        // worker runs to its own failure, so waiting on it genuinely exercises that path before we shut down.
1839        wait_until("the temporary worker has failed once", || {
1840            temp_failed.load(Ordering::SeqCst) == 1
1841        })
1842        .await;
1843        let _ = tx.send(());
1844
1845        let result = join_supervisor(handle).await;
1846        assert!(result.is_ok());
1847        assert_eq!(
1848            temp_started.load(Ordering::SeqCst),
1849            1,
1850            "temporary worker must not be restarted after it fails"
1851        );
1852    }
1853
1854    #[tokio::test]
1855    async fn transient_child_is_not_restarted_on_clean_exit() {
1856        let transient = MockWorker::completing("transient-worker", Duration::from_millis(50));
1857        let transient_started = transient.start_count();
1858        let transient_finished = transient.finish_count();
1859
1860        let stable = MockWorker::long_running("stable-worker");
1861
1862        let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
1863            RestartStrategy::one_to_one().with_intensity_and_period(20, Duration::from_secs(10)),
1864        );
1865        sup.add_worker(stable);
1866        sup.add_worker(ChildSpecification::worker(transient).with_restart_type(RestartType::Transient));
1867
1868        let (tx, handle) = run_supervisor_with_trigger(sup).await;
1869
1870        // Wait for the worker to *actually complete*, not merely start: `start_count` ticks the instant it begins
1871        // running, so shutting down as soon as it reached 1 would drive the supervisor's teardown before the clean
1872        // exit -> no-restart path ran, hiding a regression that restarted a transient child after a clean exit.
1873        // `finish_count` ticks only once the worker runs to its own completion.
1874        wait_until("the transient worker has completed once", || {
1875            transient_finished.load(Ordering::SeqCst) == 1
1876        })
1877        .await;
1878        let _ = tx.send(());
1879
1880        let result = join_supervisor(handle).await;
1881        assert!(result.is_ok());
1882        assert_eq!(
1883            transient_started.load(Ordering::SeqCst),
1884            1,
1885            "transient worker must not be restarted after a clean exit"
1886        );
1887    }
1888
1889    #[tokio::test]
1890    async fn transient_child_is_restarted_on_failure() {
1891        let transient = MockWorker::failing("transient-worker", Duration::from_millis(50));
1892        let transient_count = transient.start_count();
1893
1894        let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
1895            RestartStrategy::one_to_one().with_intensity_and_period(20, Duration::from_secs(10)),
1896        );
1897        sup.add_worker(ChildSpecification::worker(transient).with_restart_type(RestartType::Transient));
1898
1899        let (tx, handle) = run_supervisor_with_trigger(sup).await;
1900
1901        wait_until("the transient worker has been restarted", || {
1902            transient_count.load(Ordering::SeqCst) >= 2
1903        })
1904        .await;
1905        let _ = tx.send(());
1906
1907        let result = join_supervisor(handle).await;
1908        assert!(result.is_ok());
1909        assert!(
1910            transient_count.load(Ordering::SeqCst) >= 2,
1911            "transient worker must be restarted after an abnormal exit"
1912        );
1913    }
1914
1915    #[tokio::test]
1916    async fn permanent_child_is_restarted_on_clean_exit() {
1917        // A permanent worker that completes cleanly must still be restarted -- this is what distinguishes
1918        // `Permanent` from `Transient`, which is left stopped after a clean exit.
1919        let permanent = MockWorker::completing("permanent-worker", Duration::from_millis(50));
1920        let permanent_count = permanent.start_count();
1921
1922        let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
1923            RestartStrategy::one_to_one().with_intensity_and_period(20, Duration::from_secs(10)),
1924        );
1925        // Added with the default restart policy, which is `Permanent`.
1926        sup.add_worker(permanent);
1927
1928        let (tx, handle) = run_supervisor_with_trigger(sup).await;
1929
1930        wait_until("the permanent worker has been restarted", || {
1931            permanent_count.load(Ordering::SeqCst) >= 2
1932        })
1933        .await;
1934        let _ = tx.send(());
1935
1936        let result = join_supervisor(handle).await;
1937        assert!(result.is_ok());
1938        assert!(
1939            permanent_count.load(Ordering::SeqCst) >= 2,
1940            "permanent worker must be restarted even after a clean exit"
1941        );
1942    }
1943
1944    #[tokio::test]
1945    async fn temporary_failures_do_not_consume_restart_intensity() {
1946        // With intensity=1, two *restartable* failures within the period would shut the supervisor down. Here several
1947        // temporary workers all fail quickly. Because temporary exits aren't eligible for restart, they must not consume
1948        // the restart-intensity budget, and the supervisor must stay up.
1949        let mut sup = Supervisor::new("test-sup")
1950            .unwrap()
1951            .with_restart_strategy(RestartStrategy::one_to_one().with_intensity_and_period(1, Duration::from_secs(10)));
1952
1953        let workers = [
1954            MockWorker::failing("temp-0", Duration::from_millis(20)),
1955            MockWorker::failing("temp-1", Duration::from_millis(20)),
1956            MockWorker::failing("temp-2", Duration::from_millis(20)),
1957            MockWorker::failing("temp-3", Duration::from_millis(20)),
1958            MockWorker::failing("temp-4", Duration::from_millis(20)),
1959        ];
1960        let started: Vec<_> = workers.iter().map(|w| w.start_count()).collect();
1961        let failed: Vec<_> = workers.iter().map(|w| w.finish_count()).collect();
1962        for worker in workers {
1963            sup.add_worker(ChildSpecification::worker(worker).with_restart_type(RestartType::Temporary));
1964        }
1965        // A long-running worker so the supervisor doesn't simply idle once the temporaries are gone.
1966        sup.add_worker(MockWorker::long_running("stable-worker"));
1967
1968        let (tx, handle) = run_supervisor_with_trigger(sup).await;
1969        // Wait for every temporary worker to *actually fail* on its own. Keying off `start_count` would let shutdown
1970        // cut them short before their failures ran, so the supervisor would never get the chance to (mis)charge those
1971        // failures against its intensity=1 budget -- hiding the very regression this guards against.
1972        wait_until("every temporary worker has failed once", || {
1973            failed.iter().all(|c| c.load(Ordering::SeqCst) == 1)
1974        })
1975        .await;
1976        let _ = tx.send(());
1977
1978        let result = join_supervisor(handle).await;
1979        assert!(
1980            result.is_ok(),
1981            "supervisor must not trip its restart limit on temporary exits"
1982        );
1983        for count in started {
1984            assert_eq!(
1985                count.load(Ordering::SeqCst),
1986                1,
1987                "each temporary worker runs exactly once"
1988            );
1989        }
1990    }
1991
1992    #[tokio::test]
1993    async fn transient_clean_exits_do_not_consume_restart_intensity() {
1994        // With intensity=1, two *restartable* exits within the period would shut the supervisor down. Here several
1995        // transient workers all complete cleanly. A transient child's clean exit isn't eligible for restart, so it
1996        // must not consume the restart-intensity budget, and the supervisor must stay up.
1997        let mut sup = Supervisor::new("test-sup")
1998            .unwrap()
1999            .with_restart_strategy(RestartStrategy::one_to_one().with_intensity_and_period(1, Duration::from_secs(10)));
2000
2001        let workers = [
2002            MockWorker::completing("transient-0", Duration::from_millis(20)),
2003            MockWorker::completing("transient-1", Duration::from_millis(20)),
2004            MockWorker::completing("transient-2", Duration::from_millis(20)),
2005            MockWorker::completing("transient-3", Duration::from_millis(20)),
2006            MockWorker::completing("transient-4", Duration::from_millis(20)),
2007        ];
2008        let started: Vec<_> = workers.iter().map(|w| w.start_count()).collect();
2009        let finished: Vec<_> = workers.iter().map(|w| w.finish_count()).collect();
2010        for worker in workers {
2011            sup.add_worker(ChildSpecification::worker(worker).with_restart_type(RestartType::Transient));
2012        }
2013        // A long-running worker so the supervisor doesn't simply idle once the transients have completed.
2014        sup.add_worker(MockWorker::long_running("stable-worker"));
2015
2016        let (tx, handle) = run_supervisor_with_trigger(sup).await;
2017        // Wait for every transient to *actually complete* on its own. Keying off `start_count` would let shutdown cut
2018        // the workers short before their clean exits ran, so the supervisor would never get the chance to (mis)charge
2019        // those exits against its intensity=1 budget -- hiding the very regression this guards against.
2020        wait_until("every transient worker has completed once", || {
2021            finished.iter().all(|c| c.load(Ordering::SeqCst) == 1)
2022        })
2023        .await;
2024        let _ = tx.send(());
2025
2026        let result = join_supervisor(handle).await;
2027        assert!(
2028            result.is_ok(),
2029            "supervisor must not trip its restart limit on clean transient exits"
2030        );
2031        for count in started {
2032            assert_eq!(
2033                count.load(Ordering::SeqCst),
2034                1,
2035                "each transient worker runs exactly once"
2036            );
2037        }
2038    }
2039
2040    #[tokio::test]
2041    async fn supervisor_idles_when_all_temporary_children_exit() {
2042        // When every static child is temporary and they all exit, the worker set drains. The supervisor must not panic
2043        // or exit on its own; it must keep running and remain able to accept new (dynamic) work until shutdown is
2044        // triggered.
2045        let temp_a = MockWorker::completing("temp-a", Duration::from_millis(10));
2046        let a_finished = temp_a.finish_count();
2047        let temp_b = MockWorker::completing("temp-b", Duration::from_millis(10));
2048        let b_finished = temp_b.finish_count();
2049
2050        let mut sup = Supervisor::new("test-sup").unwrap();
2051        let handle = sup.handle();
2052        sup.add_worker(ChildSpecification::worker(temp_a).with_restart_type(RestartType::Temporary));
2053        sup.add_worker(ChildSpecification::worker(temp_b).with_restart_type(RestartType::Temporary));
2054
2055        let (tx, run) = run_supervisor_with_trigger(sup).await;
2056
2057        // Wait for both temporary children to actually complete -- draining the worker set to empty -- before probing.
2058        // Keying off `start_count` could spawn the probe child before the set ever emptied, letting a supervisor that
2059        // (wrongly) exited once its last child left slip through.
2060        wait_until("both temporary children have completed", || {
2061            a_finished.load(Ordering::SeqCst) == 1 && b_finished.load(Ordering::SeqCst) == 1
2062        })
2063        .await;
2064
2065        // The supervisor must still be alive after its worker set empties: spawning a new dynamic child succeeds and
2066        // runs, which is only possible if the supervise loop kept running rather than exiting when the last child left.
2067        let dynamic = MockWorker::long_running("late-comer");
2068        let dynamic_count = dynamic.start_count();
2069        handle
2070            .spawn(dynamic)
2071            .await
2072            .expect("supervisor must still accept work after its children drain");
2073        wait_until("the late dynamic child has started", || {
2074            dynamic_count.load(Ordering::SeqCst) == 1
2075        })
2076        .await;
2077        assert!(
2078            handle.is_running(),
2079            "supervisor must keep running after all temporary children exit"
2080        );
2081
2082        tx.send(()).unwrap();
2083        let result = join_supervisor(run).await;
2084        assert!(result.is_ok());
2085    }
2086
2087    // -- Significant child / auto-shutdown tests -------------------------------------------
2088
2089    #[tokio::test]
2090    async fn significant_child_drives_auto_shutdown() {
2091        // With `AnySignificant`, a significant child terminating (even cleanly, and without being restarted) must
2092        // shut the supervisor down, surfacing the significant-exit error.
2093        let mut sup = Supervisor::new("test-sup")
2094            .unwrap()
2095            .with_auto_shutdown(AutoShutdown::AnySignificant);
2096        sup.add_worker(MockWorker::long_running("stable"));
2097        sup.add_worker(
2098            ChildSpecification::worker(MockWorker::completing("significant", Duration::from_millis(50)))
2099                .with_restart_type(RestartType::Temporary)
2100                .with_significant(true),
2101        );
2102
2103        // Hold the shutdown sender so the only thing that can stop the supervisor is the significant child.
2104        let (_tx, rx) = oneshot::channel::<()>();
2105        let result = timeout(Duration::from_secs(2), sup.run_with_shutdown(rx))
2106            .await
2107            .unwrap();
2108        assert!(matches!(result, Err(SupervisorError::SignificantChildExited)));
2109    }
2110
2111    #[tokio::test]
2112    async fn non_significant_exit_does_not_auto_shutdown() {
2113        // Even with `AnySignificant` set, a non-significant child exiting must not shut the supervisor down.
2114        let plain = MockWorker::completing("plain", Duration::from_millis(10));
2115        let plain_finished = plain.finish_count();
2116
2117        let mut sup = Supervisor::new("test-sup")
2118            .unwrap()
2119            .with_auto_shutdown(AutoShutdown::AnySignificant);
2120        let handle = sup.handle();
2121        sup.add_worker(MockWorker::long_running("stable"));
2122        sup.add_worker(ChildSpecification::worker(plain).with_restart_type(RestartType::Temporary));
2123
2124        let (tx, run) = run_supervisor_with_trigger(sup).await;
2125
2126        // Let the non-significant child actually run to completion -- not merely start. Its completion is what could
2127        // (wrongly) trip `AnySignificant`, so we must observe the real exit before probing liveness; keying off
2128        // `start_count` could assert before the completion was ever processed.
2129        wait_until("the non-significant child has completed", || {
2130            plain_finished.load(Ordering::SeqCst) == 1
2131        })
2132        .await;
2133
2134        // The supervisor must still be alive after the non-significant child exits (had it been treated as
2135        // significant, `AnySignificant` would have torn the supervisor down). Spawning a dynamic child and observing
2136        // it start proves the supervise loop is still running.
2137        let dynamic = MockWorker::long_running("late-comer");
2138        let dynamic_count = dynamic.start_count();
2139        handle
2140            .spawn(dynamic)
2141            .await
2142            .expect("supervisor must still accept work after a non-significant child exits");
2143        wait_until("the late dynamic child has started", || {
2144            dynamic_count.load(Ordering::SeqCst) == 1
2145        })
2146        .await;
2147        assert!(
2148            handle.is_running(),
2149            "a non-significant child exiting must not trigger auto-shutdown"
2150        );
2151
2152        tx.send(()).unwrap();
2153        let result = join_supervisor(run).await;
2154        assert!(result.is_ok());
2155    }
2156
2157    #[tokio::test]
2158    async fn all_significant_waits_for_last() {
2159        // With `AllSignificant`, the supervisor shuts down only once *all* significant children have terminated.
2160        let mut sup = Supervisor::new("test-sup")
2161            .unwrap()
2162            .with_auto_shutdown(AutoShutdown::AllSignificant);
2163        sup.add_worker(
2164            ChildSpecification::worker(MockWorker::completing("sig-a", Duration::from_millis(50)))
2165                .with_restart_type(RestartType::Temporary)
2166                .with_significant(true),
2167        );
2168        sup.add_worker(
2169            ChildSpecification::worker(MockWorker::completing("sig-b", Duration::from_millis(250)))
2170                .with_restart_type(RestartType::Temporary)
2171                .with_significant(true),
2172        );
2173
2174        let (_tx, rx) = oneshot::channel::<()>();
2175        let start = std::time::Instant::now();
2176        let result = timeout(Duration::from_secs(2), sup.run_with_shutdown(rx))
2177            .await
2178            .unwrap();
2179        let elapsed = start.elapsed();
2180
2181        assert!(matches!(result, Err(SupervisorError::SignificantChildExited)));
2182        // The first significant child exits at ~50ms but must NOT trigger shutdown; only the second (~250ms) does.
2183        assert!(
2184            elapsed >= Duration::from_millis(200),
2185            "auto-shutdown must wait for all significant children (took {elapsed:?})"
2186        );
2187    }
2188
2189    // -- Initialization failure tests ------------------------------------------------------
2190
2191    #[tokio::test]
2192    async fn init_failure_propagates_with_child_name() {
2193        let mut sup = Supervisor::new("test-sup").unwrap();
2194        sup.add_worker(MockWorker::long_running("good-worker"));
2195        sup.add_worker(MockWorker::init_failure("bad-worker"));
2196
2197        let (_tx, rx) = oneshot::channel::<()>();
2198        let result = timeout(Duration::from_secs(2), sup.run_with_shutdown(rx))
2199            .await
2200            .unwrap();
2201
2202        match result {
2203            Err(SupervisorError::FailedToInitialize { child_name, .. }) => {
2204                assert_eq!(child_name, "bad-worker");
2205            }
2206            other => panic!("expected FailedToInitialize, got: {:?}", other),
2207        }
2208    }
2209
2210    #[tokio::test]
2211    async fn init_failure_does_not_trigger_restart() {
2212        let init_fail = MockWorker::init_failure("bad-worker");
2213        let start_count = init_fail.start_count();
2214
2215        let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
2216            RestartStrategy::one_to_one().with_intensity_and_period(10, Duration::from_secs(10)),
2217        );
2218        sup.add_worker(init_fail);
2219
2220        let (_tx, rx) = oneshot::channel::<()>();
2221        let result = timeout(Duration::from_secs(2), sup.run_with_shutdown(rx))
2222            .await
2223            .unwrap();
2224
2225        assert!(matches!(result, Err(SupervisorError::FailedToInitialize { .. })));
2226        // The worker never got past init, so start_count should be 0.
2227        assert_eq!(start_count.load(Ordering::SeqCst), 0);
2228    }
2229
2230    // -- Shutdown responsiveness tests -----------------------------------------------------
2231
2232    #[tokio::test]
2233    async fn shutdown_completes_promptly_in_steady_state() {
2234        let mut sup = Supervisor::new("test-sup").unwrap();
2235        sup.add_worker(MockWorker::long_running("worker1"));
2236        sup.add_worker(MockWorker::long_running("worker2"));
2237
2238        let (tx, handle) = run_supervisor_with_trigger(sup).await;
2239        tx.send(()).unwrap();
2240
2241        // Shutdown should complete well within 1 second (workers respond to shutdown signal immediately).
2242        let result = timeout(Duration::from_secs(1), handle).await;
2243        assert!(result.is_ok(), "shutdown should complete promptly");
2244    }
2245
2246    #[tokio::test]
2247    async fn shutdown_during_slow_init_completes_promptly() {
2248        let mut sup = Supervisor::new("test-sup").unwrap();
2249        // This worker takes 30 seconds to initialize — but we'll trigger shutdown immediately.
2250        sup.add_worker(MockWorker::slow_init("slow-worker", Duration::from_secs(30)));
2251
2252        let (tx, rx) = oneshot::channel();
2253        let handle = tokio::spawn(async move { sup.run_with_shutdown(rx).await });
2254
2255        // Give the supervisor just enough time to spawn the task, then trigger shutdown.
2256        sleep(Duration::from_millis(20)).await;
2257        tx.send(()).unwrap();
2258
2259        // Shutdown should complete quickly even though the worker hasn't finished initializing.
2260        // The supervisor loop sees the shutdown signal and aborts the still-initializing task.
2261        let result = timeout(Duration::from_secs(2), handle).await;
2262        assert!(result.is_ok(), "shutdown during slow init should complete promptly");
2263    }
2264
2265    // -- Dynamic children tests ------------------------------------------------------------
2266
2267    #[tokio::test]
2268    async fn dynamic_children_spawn_after_start() {
2269        let sup = Supervisor::new("dyn-sup").unwrap();
2270        let handle = sup.handle();
2271        let (tx, run) = run_supervisor_with_trigger(sup).await;
2272        wait_until("supervisor is running", || handle.is_running()).await;
2273
2274        let c1 = MockWorker::long_running("c1");
2275        let c2 = MockWorker::long_running("c2");
2276        let c1_count = c1.start_count();
2277        let c2_count = c2.start_count();
2278        handle.spawn(c1).await.unwrap();
2279        handle.spawn(c2).await.unwrap();
2280
2281        wait_until("both dynamic children have started", || {
2282            c1_count.load(Ordering::SeqCst) == 1 && c2_count.load(Ordering::SeqCst) == 1
2283        })
2284        .await;
2285        assert_eq!(handle.active_children(), 2);
2286
2287        tx.send(()).unwrap();
2288        let result = join_supervisor(run).await;
2289        assert!(result.is_ok());
2290        assert_eq!(
2291            handle.active_children(),
2292            0,
2293            "all dynamic children must be drained on shutdown"
2294        );
2295    }
2296
2297    #[tokio::test]
2298    async fn temporary_dynamic_child_failure_is_isolated() {
2299        // A dynamic child added with the default config (temporary, not significant) is fault-isolated: its failure is
2300        // reaped and removed without restarting it or disturbing the supervisor or its siblings.
2301        let sup = Supervisor::new("dyn-sup").unwrap();
2302        let handle = sup.handle();
2303        let (tx, run) = run_supervisor_with_trigger(sup).await;
2304        wait_until("supervisor is running", || handle.is_running()).await;
2305
2306        let failing = MockWorker::failing("boom", Duration::from_millis(20));
2307        let failing_count = failing.start_count();
2308        handle.spawn(failing).await.unwrap();
2309        wait_until("the failing dynamic child has run once", || {
2310            failing_count.load(Ordering::SeqCst) == 1
2311        })
2312        .await;
2313        wait_until("all dynamic children have drained", || handle.active_children() == 0).await;
2314
2315        sleep(Duration::from_millis(50)).await;
2316        assert!(
2317            handle.is_running(),
2318            "supervisor stays up after an isolated child failure"
2319        );
2320        assert_eq!(
2321            failing_count.load(Ordering::SeqCst),
2322            1,
2323            "a temporary child is never restarted"
2324        );
2325
2326        // It still accepts new children.
2327        handle.spawn(MockWorker::long_running("c2")).await.unwrap();
2328        wait_until("one dynamic child is running", || handle.active_children() == 1).await;
2329
2330        tx.send(()).unwrap();
2331        let result = join_supervisor(run).await;
2332        assert!(result.is_ok());
2333    }
2334
2335    #[tokio::test]
2336    async fn temporary_dynamic_child_panic_is_isolated() {
2337        // A panicking temporary, non-significant child is isolated exactly like an error exit.
2338        let sup = Supervisor::new("dyn-sup").unwrap();
2339        let handle = sup.handle();
2340        let (tx, run) = run_supervisor_with_trigger(sup).await;
2341        wait_until("supervisor is running", || handle.is_running()).await;
2342
2343        handle
2344            .spawn(MockWorker::panicking("boom", Duration::from_millis(20)))
2345            .await
2346            .unwrap();
2347        wait_until("all dynamic children have drained", || handle.active_children() == 0).await;
2348
2349        sleep(Duration::from_millis(50)).await;
2350        assert!(handle.is_running(), "supervisor stays up after an isolated child panic");
2351
2352        tx.send(()).unwrap();
2353        let result = join_supervisor(run).await;
2354        assert!(result.is_ok());
2355    }
2356
2357    #[tokio::test]
2358    async fn significant_dynamic_child_failure_shuts_down_supervisor() {
2359        // A dynamic child added as significant, under `AutoShutdown::AnySignificant`, drives the supervisor to shut
2360        // down when it terminates -- the opt-in mechanism that replaces the old escalate-on-error behavior.
2361        let sup = Supervisor::new("dyn-sup")
2362            .unwrap()
2363            .with_auto_shutdown(AutoShutdown::AnySignificant);
2364        let handle = sup.handle();
2365        let (_tx, run) = run_supervisor_with_trigger(sup).await;
2366        wait_until("supervisor is running", || handle.is_running()).await;
2367
2368        handle
2369            .spawn_with(
2370                ChildSpecification::worker(MockWorker::failing("boom", Duration::from_millis(20)))
2371                    .with_restart_type(RestartType::Temporary)
2372                    .with_significant(true),
2373            )
2374            .await
2375            .unwrap();
2376
2377        let result = join_supervisor(run).await;
2378        assert!(matches!(result, Err(SupervisorError::SignificantChildExited)));
2379    }
2380
2381    #[tokio::test]
2382    async fn dynamic_spawn_fails_before_start_and_after_shutdown() {
2383        let sup = Supervisor::new("dyn-sup").unwrap();
2384        let handle = sup.handle();
2385
2386        // Before the supervisor is running there's nothing to accept the spawn, so it's rejected outright (static
2387        // children should be configured up front via `add_worker` instead).
2388        assert!(!handle.is_running());
2389        let err = handle
2390            .spawn(MockWorker::long_running("before-start"))
2391            .await
2392            .unwrap_err();
2393        assert!(matches!(err, SpawnError::SupervisorGone));
2394
2395        // Once it's running, spawns succeed.
2396        let (tx, run) = run_supervisor_with_trigger(sup).await;
2397        wait_until("supervisor is running", || handle.is_running()).await;
2398        let worker = MockWorker::long_running("after-start");
2399        let started = worker.start_count();
2400        handle.spawn(worker).await.unwrap();
2401        wait_until("the dynamic child has started", || started.load(Ordering::SeqCst) == 1).await;
2402
2403        tx.send(()).unwrap();
2404        let result = join_supervisor(run).await;
2405        assert!(result.is_ok());
2406
2407        // Once the supervisor has shut down, the run is gone and spawns are rejected again.
2408        wait_until("the supervisor has stopped", || !handle.is_running()).await;
2409        let err = handle
2410            .spawn(MockWorker::long_running("after-shutdown"))
2411            .await
2412            .unwrap_err();
2413        assert!(matches!(err, SpawnError::SupervisorGone));
2414    }
2415
2416    #[tokio::test]
2417    async fn dynamic_spawn_returns_after_registration() {
2418        let sup = Supervisor::new("dyn-sup").unwrap();
2419        let handle = sup.handle();
2420        let (tx, run) = run_supervisor_with_trigger(sup).await;
2421        wait_until("supervisor is running", || handle.is_running()).await;
2422
2423        let worker = MockWorker::long_running("c");
2424        let started = worker.start_count();
2425        let id = handle.spawn(worker).await.unwrap();
2426        // No static children, so the first dynamic child takes id 0.
2427        assert_eq!(id.as_u64(), 0);
2428        wait_until("the dynamic child has started", || started.load(Ordering::SeqCst) == 1).await;
2429
2430        tx.send(()).unwrap();
2431        let result = join_supervisor(run).await;
2432        assert!(result.is_ok());
2433    }
2434
2435    #[tokio::test]
2436    async fn dynamic_spawn_rejects_invalid_child_name() {
2437        // While running, a spawn that fails registration (here, an empty/invalid child name) is reported as
2438        // `Rejected` with the underlying cause -- not `SupervisorGone`, which means the supervisor isn't running.
2439        let sup = Supervisor::new("dyn-sup").unwrap();
2440        let handle = sup.handle();
2441        let (tx, run) = run_supervisor_with_trigger(sup).await;
2442        wait_until("supervisor is running", || handle.is_running()).await;
2443
2444        let err = handle.spawn(MockWorker::long_running("")).await.unwrap_err();
2445        assert!(matches!(err, SpawnError::Rejected { .. }), "got {err:?}");
2446
2447        // The supervisor stays up and still accepts valid children.
2448        assert!(handle.is_running());
2449        handle.spawn(MockWorker::long_running("ok")).await.unwrap();
2450        wait_until("one dynamic child is running", || handle.active_children() == 1).await;
2451
2452        tx.send(()).unwrap();
2453        let result = join_supervisor(run).await;
2454        assert!(result.is_ok());
2455    }
2456
2457    #[tokio::test]
2458    async fn concurrent_shutdown_drains_many_children_quickly() {
2459        const CHILDREN: usize = 500;
2460        const SHUTDOWN_DELAY: Duration = Duration::from_millis(50);
2461
2462        let sup = Supervisor::new("dyn-sup")
2463            .unwrap()
2464            .with_shutdown_mode(ShutdownMode::Concurrent);
2465        let handle = sup.handle();
2466        let (tx, run) = run_supervisor_with_trigger(sup).await;
2467        wait_until("supervisor is running", || handle.is_running()).await;
2468
2469        for _ in 0..CHILDREN {
2470            handle
2471                .spawn(MockWorker::slow_shutdown("conn", SHUTDOWN_DELAY))
2472                .await
2473                .unwrap();
2474        }
2475        wait_until("all dynamic children are running", || {
2476            handle.active_children() == CHILDREN
2477        })
2478        .await;
2479
2480        // Each child sleeps after observing shutdown. Concurrent shutdown drains them all in roughly one delay; an
2481        // ordered shutdown would take CHILDREN * delay (25s here). Assert it finishes well under that.
2482        let start = std::time::Instant::now();
2483        tx.send(()).unwrap();
2484        let result = timeout(Duration::from_secs(5), run).await.unwrap().unwrap();
2485        let elapsed = start.elapsed();
2486
2487        assert!(result.is_ok());
2488        assert_eq!(handle.active_children(), 0, "active count must return to zero");
2489        assert!(
2490            elapsed < Duration::from_secs(2),
2491            "shutdown must be concurrent (took {elapsed:?})"
2492        );
2493    }
2494
2495    #[tokio::test]
2496    async fn concurrent_shutdown_aborts_unresponsive_children() {
2497        let sup = Supervisor::new("dyn-sup")
2498            .unwrap()
2499            .with_shutdown_mode(ShutdownMode::Concurrent);
2500        let handle = sup.handle();
2501        let (tx, run) = run_supervisor_with_trigger(sup).await;
2502        wait_until("supervisor is running", || handle.is_running()).await;
2503
2504        handle.spawn(MockWorker::ignore_shutdown("stuck")).await.unwrap();
2505        wait_until("one dynamic child is running", || handle.active_children() == 1).await;
2506
2507        // The child never reacts to shutdown, so it must be aborted once its graceful deadline (500ms) elapses rather
2508        // than hanging the supervisor.
2509        let start = std::time::Instant::now();
2510        tx.send(()).unwrap();
2511        let result = join_supervisor(run).await;
2512        let elapsed = start.elapsed();
2513
2514        // Forcefully aborting an unresponsive child is surfaced as an unclean shutdown rather than reported as success.
2515        assert!(
2516            matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 1 })),
2517            "aborting a stuck child must surface as an unclean shutdown, got {result:?}"
2518        );
2519        assert_eq!(handle.active_children(), 0);
2520        assert!(
2521            elapsed < Duration::from_secs(1),
2522            "stuck child must be aborted at the deadline (took {elapsed:?})"
2523        );
2524    }
2525
2526    #[tokio::test]
2527    async fn concurrent_shutdown_honors_per_child_deadline() {
2528        // Each child must be aborted at its OWN graceful deadline, not a single shared one. A responsive child with an
2529        // effectively-infinite timeout (modeling a nested supervisor, which uses `Graceful(Duration::MAX)`) coexists
2530        // with an unresponsive child with a short timeout. Under a shared `max` deadline the short-timeout child would
2531        // never be aborted (the shared deadline would be `MAX`) and shutdown would hang.
2532        let sup = Supervisor::new("dyn-sup")
2533            .unwrap()
2534            .with_shutdown_mode(ShutdownMode::Concurrent);
2535        let handle = sup.handle();
2536        let (tx, run) = run_supervisor_with_trigger(sup).await;
2537        wait_until("supervisor is running", || handle.is_running()).await;
2538
2539        // Responds to shutdown promptly, but its deadline is effectively infinite.
2540        handle
2541            .spawn(MockWorker::long_running("responsive").with_graceful_timeout(Duration::MAX))
2542            .await
2543            .unwrap();
2544        // Never responds; must be aborted at its own short deadline.
2545        handle
2546            .spawn(MockWorker::ignore_shutdown("stuck").with_graceful_timeout(Duration::from_millis(200)))
2547            .await
2548            .unwrap();
2549        wait_until("both dynamic children are running", || handle.active_children() == 2).await;
2550
2551        let start = std::time::Instant::now();
2552        tx.send(()).unwrap();
2553        let result = join_supervisor(run).await;
2554        let elapsed = start.elapsed();
2555
2556        // Only the stuck child is aborted (the responsive one exits cleanly), so the unclean-shutdown tally is exactly 1.
2557        assert!(
2558            matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 1 })),
2559            "aborting the stuck child must surface as an unclean shutdown with a count of 1, got {result:?}"
2560        );
2561        assert_eq!(handle.active_children(), 0);
2562        assert!(
2563            elapsed < Duration::from_secs(1),
2564            "stuck child must be aborted at its own deadline despite an infinite-timeout sibling (took {elapsed:?})"
2565        );
2566    }
2567
2568    #[tokio::test]
2569    async fn ordered_shutdown_aborts_unresponsive_child() {
2570        // Under the default `ShutdownMode::Ordered`, a child that never reacts to shutdown must be aborted once its
2571        // graceful deadline (500ms) elapses, rather than hanging the supervisor.
2572        let mut sup = Supervisor::new("test-sup").unwrap();
2573        sup.add_worker(MockWorker::ignore_shutdown("stuck"));
2574
2575        let (tx, handle) = run_supervisor_with_trigger(sup).await;
2576
2577        let start = std::time::Instant::now();
2578        tx.send(()).unwrap();
2579        let result = join_supervisor(handle).await;
2580        let elapsed = start.elapsed();
2581
2582        assert!(
2583            matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 1 })),
2584            "aborting a stuck child under ordered shutdown must surface as an unclean shutdown, got {result:?}"
2585        );
2586        assert!(
2587            elapsed < Duration::from_secs(1),
2588            "unresponsive child must be aborted at its deadline under ordered shutdown (took {elapsed:?})"
2589        );
2590    }
2591
2592    #[tokio::test]
2593    async fn brutal_shutdown_aborts_child_immediately() {
2594        // A child with a `Brutal` shutdown strategy is aborted at once on shutdown, with no graceful wait -- so even a
2595        // child that ignores shutdown is torn down promptly rather than after the graceful deadline.
2596        let mut sup = Supervisor::new("test-sup").unwrap();
2597        sup.add_worker(MockWorker::ignore_shutdown("brutal-stuck").with_brutal_shutdown());
2598
2599        let (tx, handle) = run_supervisor_with_trigger(sup).await;
2600
2601        let start = std::time::Instant::now();
2602        tx.send(()).unwrap();
2603        let result = join_supervisor(handle).await;
2604        let elapsed = start.elapsed();
2605
2606        // A brutal abort is the configured, expected way to stop this child -- not a graceful-timeout overrun -- so it
2607        // is NOT counted toward the unclean-shutdown tally, and the shutdown reports success.
2608        assert!(result.is_ok());
2609        assert!(
2610            elapsed < Duration::from_millis(200),
2611            "brutal-shutdown child must be aborted immediately, not after a graceful wait (took {elapsed:?})"
2612        );
2613    }
2614
2615    #[tokio::test]
2616    async fn shutdown_timeout_aborts_aggregate_to_root() {
2617        // Forced aborts must surface as an unclean shutdown and aggregate up the tree: a supervisor adds the workers it
2618        // aborts directly to the counts reported by any child supervisors that also timed out. Here the parent aborts
2619        // one direct child and a nested supervisor aborts one of its own, so the root observes a total of 2.
2620        let mut child_sup = Supervisor::new("child-sup")
2621            .unwrap()
2622            .with_shutdown_mode(ShutdownMode::Concurrent);
2623        child_sup
2624            .add_worker(MockWorker::ignore_shutdown("child-stuck").with_graceful_timeout(Duration::from_millis(200)));
2625
2626        let mut parent_sup = Supervisor::new("parent-sup")
2627            .unwrap()
2628            .with_shutdown_mode(ShutdownMode::Concurrent);
2629        parent_sup
2630            .add_worker(MockWorker::ignore_shutdown("parent-stuck").with_graceful_timeout(Duration::from_millis(200)));
2631        parent_sup.add_worker(MockWorker::long_running("parent-clean"));
2632        parent_sup.add_worker(child_sup);
2633
2634        let (tx, handle) = run_supervisor_with_trigger(parent_sup).await;
2635        tx.send(()).unwrap();
2636
2637        let result = join_supervisor(handle).await;
2638        assert!(
2639            matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 2 })),
2640            "forced aborts must aggregate across the tree (1 direct + 1 nested), got {result:?}"
2641        );
2642    }
2643
2644    // -- Restart-policy edge cases ---------------------------------------------------------
2645
2646    #[tokio::test]
2647    async fn restart_intensity_zero_shuts_down_on_first_failure() {
2648        // A restart intensity of zero means the supervisor gives up the moment any restartable child fails: it shuts
2649        // down on the very first failure without ever restarting the worker. (See `RestartState::evaluate_restart`,
2650        // which short-circuits to `Shutdown` when intensity is zero.)
2651        let worker = MockWorker::failing("boom", Duration::from_millis(20));
2652        let start_count = worker.start_count();
2653
2654        let mut sup = Supervisor::new("test-sup")
2655            .unwrap()
2656            .with_restart_strategy(RestartStrategy::new(RestartMode::OneForOne, 0, Duration::from_secs(5)));
2657        sup.add_worker(worker);
2658
2659        let (_tx, rx) = oneshot::channel::<()>();
2660        let result = timeout(Duration::from_secs(2), sup.run_with_shutdown(rx))
2661            .await
2662            .unwrap();
2663
2664        assert!(matches!(result, Err(SupervisorError::Shutdown)));
2665        assert_eq!(
2666            start_count.load(Ordering::SeqCst),
2667            1,
2668            "with intensity zero the worker must run exactly once and never be restarted"
2669        );
2670    }
2671
2672    #[tokio::test]
2673    async fn one_for_all_restart_loses_dynamic_children() {
2674        // Documented one-for-all semantics: a group restart resets to the static roster only -- dynamic children are
2675        // NOT restored (they're lost on a supervisor-level restart, matching Erlang/OTP). A permanent static worker
2676        // that keeps failing drives repeated one-for-all restarts; a dynamic child spawned before the first restart
2677        // must be torn down and never brought back.
2678        let failing = MockWorker::failing("failing-static", Duration::from_millis(50));
2679        let failing_count = failing.start_count();
2680
2681        let sup = Supervisor::new("dyn-sup").unwrap().with_restart_strategy(
2682            RestartStrategy::one_for_all().with_intensity_and_period(20, Duration::from_secs(10)),
2683        );
2684        let handle = sup.handle();
2685        let mut sup = sup;
2686        sup.add_worker(failing);
2687
2688        let (tx, run) = run_supervisor_with_trigger(sup).await;
2689
2690        // Spawn a long-running dynamic child and wait for it to be running.
2691        let dynamic = MockWorker::long_running("dynamic");
2692        let dynamic_count = dynamic.start_count();
2693        handle.spawn(dynamic).await.expect("should spawn dynamic child");
2694        wait_until("the dynamic child is running", || handle.active_children() == 1).await;
2695
2696        // Let the static worker drive at least one one-for-all restart (its second start).
2697        wait_until("the static worker has been restarted", || {
2698            failing_count.load(Ordering::SeqCst) >= 2
2699        })
2700        .await;
2701
2702        // The one-for-all restart must have discarded the dynamic child: the active count returns to zero, and the
2703        // dynamic child ran exactly once (it was never restored).
2704        wait_until("the dynamic child has been discarded", || handle.active_children() == 0).await;
2705        assert_eq!(
2706            dynamic_count.load(Ordering::SeqCst),
2707            1,
2708            "a dynamic child must be lost -- not restored -- across a one-for-all restart"
2709        );
2710
2711        tx.send(()).unwrap();
2712        let result = join_supervisor(run).await;
2713        assert!(result.is_ok());
2714    }
2715
2716    // -- Dedicated-runtime tests -----------------------------------------------------------
2717
2718    #[tokio::test]
2719    async fn dedicated_single_threaded_runtime_runs_nested_worker_and_shuts_down_cleanly() {
2720        // A nested supervisor configured with a dedicated single-threaded runtime spawns its own OS thread and Tokio
2721        // runtime (via `spawn_dedicated_runtime`). Its worker must run there, and a shutdown signalled by the parent
2722        // must propagate across the thread boundary and tear it down cleanly.
2723        let worker = MockWorker::long_running("dedicated-worker");
2724        let worker_count = worker.start_count();
2725
2726        let mut child_sup = Supervisor::new("child-sup")
2727            .unwrap()
2728            .with_dedicated_runtime(RuntimeConfiguration::single_threaded());
2729        child_sup.add_worker(worker);
2730
2731        let mut parent_sup = Supervisor::new("parent-sup").unwrap();
2732        parent_sup.add_worker(child_sup);
2733
2734        let (tx, handle) = run_supervisor_with_trigger(parent_sup).await;
2735
2736        // The worker starts on the dedicated runtime's own thread.
2737        wait_until("the dedicated worker has started", || {
2738            worker_count.load(Ordering::SeqCst) == 1
2739        })
2740        .await;
2741
2742        tx.send(()).unwrap();
2743        let result = join_supervisor(handle).await;
2744        assert!(
2745            result.is_ok(),
2746            "dedicated-runtime supervisor should shut down cleanly, got {result:?}"
2747        );
2748    }
2749
2750    #[tokio::test]
2751    async fn dedicated_multi_threaded_runtime_runs_nested_worker() {
2752        // The same nested-dedicated flow, but exercising the multi-threaded dedicated runtime builder path.
2753        let worker = MockWorker::long_running("dedicated-worker");
2754        let worker_count = worker.start_count();
2755
2756        let mut child_sup = Supervisor::new("child-sup")
2757            .unwrap()
2758            .with_dedicated_runtime(RuntimeConfiguration::multi_threaded(2));
2759        child_sup.add_worker(worker);
2760
2761        let mut parent_sup = Supervisor::new("parent-sup").unwrap();
2762        parent_sup.add_worker(child_sup);
2763
2764        let (tx, handle) = run_supervisor_with_trigger(parent_sup).await;
2765        wait_until("the dedicated worker has started", || {
2766            worker_count.load(Ordering::SeqCst) == 1
2767        })
2768        .await;
2769
2770        tx.send(()).unwrap();
2771        let result = join_supervisor(handle).await;
2772        assert!(
2773            result.is_ok(),
2774            "multi-threaded dedicated-runtime supervisor should shut down cleanly, got {result:?}"
2775        );
2776    }
2777
2778    #[tokio::test]
2779    async fn dedicated_runtime_forced_abort_aggregates_to_root() {
2780        // A worker inside a dedicated-runtime nested supervisor that ignores shutdown must be forcefully aborted at its
2781        // deadline, and that abort tally must survive the OS-thread boundary (`DedicatedRuntimeHandle` -> `WorkerError`)
2782        // and be observed by the root supervisor as `ShutdownTimedOut`.
2783        let stuck = MockWorker::ignore_shutdown("stuck").with_graceful_timeout(Duration::from_millis(200));
2784        let stuck_count = stuck.start_count();
2785
2786        let mut child_sup = Supervisor::new("child-sup")
2787            .unwrap()
2788            .with_dedicated_runtime(RuntimeConfiguration::single_threaded())
2789            .with_shutdown_mode(ShutdownMode::Concurrent);
2790        child_sup.add_worker(stuck);
2791
2792        let mut parent_sup = Supervisor::new("parent-sup").unwrap();
2793        parent_sup.add_worker(child_sup);
2794
2795        let (tx, handle) = run_supervisor_with_trigger(parent_sup).await;
2796
2797        // Make sure the stuck worker is actually running on the dedicated runtime before signalling shutdown, so the
2798        // forced-abort path (rather than an early exit) is what we exercise.
2799        wait_until("the stuck worker has started", || {
2800            stuck_count.load(Ordering::SeqCst) == 1
2801        })
2802        .await;
2803
2804        tx.send(()).unwrap();
2805        let result = join_supervisor(handle).await;
2806        assert!(
2807            matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 1 })),
2808            "a stuck worker in a dedicated runtime must surface as an unclean shutdown aggregated to the root, got {result:?}"
2809        );
2810    }
2811
2812    // -- Per-child override tests ----------------------------------------------------------
2813
2814    #[tokio::test]
2815    async fn child_with_runtime_override_runs_on_that_runtime() {
2816        // `with_runtime` places an individual child's task on a caller-provided runtime instead of the supervisor's
2817        // own. The child reports the name of the thread it's actually running on, which must belong to that runtime.
2818        let child_runtime = tokio::runtime::Builder::new_multi_thread()
2819            .worker_threads(1)
2820            .thread_name("child-rt-test")
2821            .enable_all()
2822            .build()
2823            .expect("should build child runtime");
2824
2825        let (thread_tx, thread_rx) = oneshot::channel();
2826        let worker = noninterruptible_worker("placed", move |shutdown| async move {
2827            let thread_name = std::thread::current().name().unwrap_or_default().to_string();
2828            let _ = thread_tx.send(thread_name);
2829            shutdown.await;
2830        });
2831
2832        let mut sup = Supervisor::new("test-sup").unwrap();
2833        sup.add_worker(
2834            ChildSpecification::worker(worker)
2835                .with_restart_type(RestartType::Temporary)
2836                .with_runtime(child_runtime.handle().clone()),
2837        );
2838
2839        let (tx, handle) = run_supervisor_with_trigger(sup).await;
2840
2841        let thread_name = timeout(Duration::from_secs(2), thread_rx)
2842            .await
2843            .expect("child should report its thread promptly")
2844            .expect("child should not be dropped before reporting");
2845        assert!(
2846            thread_name.starts_with("child-rt-test"),
2847            "child must run on the runtime given to `with_runtime`, but ran on thread {thread_name:?}"
2848        );
2849
2850        tx.send(()).unwrap();
2851        assert!(join_supervisor(handle).await.is_ok());
2852
2853        // Dropping a `Runtime` from within an async context panics, so tear it down without blocking.
2854        child_runtime.shutdown_background();
2855    }
2856
2857    #[tokio::test]
2858    async fn child_shutdown_strategy_override_takes_precedence_over_worker() {
2859        // `with_shutdown_strategy` overrides what the worker reports for itself. The worker below asks for a 30-second
2860        // grace period and then ignores shutdown entirely; the override cuts that to 50ms, so the supervisor must
2861        // abort it and report an unclean shutdown well inside `join_supervisor`'s two-second bound. Without the
2862        // override taking precedence, this test times out.
2863        let worker = noninterruptible_worker("stuck", |_shutdown| std::future::pending::<()>())
2864            .with_shutdown_timeout(Duration::from_secs(30));
2865
2866        let mut sup = Supervisor::new("test-sup").unwrap();
2867        sup.add_worker(
2868            ChildSpecification::worker(worker)
2869                .with_restart_type(RestartType::Temporary)
2870                .with_shutdown_strategy(ShutdownStrategy::Graceful(Duration::from_millis(50))),
2871        );
2872
2873        let (tx, handle) = run_supervisor_with_trigger(sup).await;
2874        tx.send(()).unwrap();
2875
2876        let result = join_supervisor(handle).await;
2877        assert!(
2878            matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 1 })),
2879            "the overridden 50ms deadline should have aborted the stuck child, got {result:?}"
2880        );
2881    }
2882
2883    /// Builds two children modelling an owner that drains a background task during shutdown.
2884    ///
2885    /// `stuck` holds a shutdown handle and never releases it voluntarily, so the only way it goes away is a forced
2886    /// abort. `waiter` blocks on that handle being dropped, standing in for a component's `shutdown_and_wait`. The
2887    /// returned flag records whether `waiter` ran to completion rather than being aborted itself.
2888    fn build_drain_pair(
2889        stuck_strategy: ShutdownStrategy, waiter_strategy: ShutdownStrategy,
2890    ) -> (Supervisor, Arc<AtomicBool>) {
2891        let mut coordinator = ShutdownCoordinator::default();
2892        let held_handle = coordinator.register();
2893
2894        let stuck = noninterruptible_worker("stuck", move |_shutdown| async move {
2895            // Hold the handle for as long as this future lives, and ignore shutdown entirely.
2896            let _held = held_handle;
2897            pending::<()>().await;
2898        });
2899
2900        let waiter_finished = Arc::new(AtomicBool::new(false));
2901        let finished = Arc::clone(&waiter_finished);
2902        let waiter = noninterruptible_worker("waiter", move |shutdown| async move {
2903            shutdown.await;
2904            coordinator.shutdown_and_wait().await;
2905            finished.store(true, Ordering::SeqCst);
2906        });
2907
2908        let mut sup = Supervisor::new("test-sup")
2909            .unwrap()
2910            .with_shutdown_mode(ShutdownMode::Concurrent);
2911        sup.add_worker(
2912            ChildSpecification::worker(stuck)
2913                .with_restart_type(RestartType::Temporary)
2914                .with_shutdown_strategy(stuck_strategy),
2915        );
2916        sup.add_worker(
2917            ChildSpecification::worker(waiter)
2918                .with_restart_type(RestartType::Temporary)
2919                .with_shutdown_strategy(waiter_strategy),
2920        );
2921
2922        (sup, waiter_finished)
2923    }
2924
2925    #[tokio::test]
2926    async fn shorter_child_deadline_releases_a_waiting_sibling() {
2927        // Aborting a stuck child drops the shutdown handle it was holding, which is what releases anything waiting on
2928        // it. A child bounded more tightly than its waiter therefore stays recoverable: the child is aborted, the
2929        // waiter unblocks and finishes cleanly, and only one abort is reported.
2930        let (sup, waiter_finished) = build_drain_pair(
2931            ShutdownStrategy::Graceful(Duration::from_millis(100)),
2932            ShutdownStrategy::Graceful(Duration::from_secs(1)),
2933        );
2934
2935        let (tx, handle) = run_supervisor_with_trigger(sup).await;
2936        tx.send(()).unwrap();
2937
2938        let result = join_supervisor(handle).await;
2939        assert!(
2940            matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 1 })),
2941            "only the stuck child should have been aborted, got {result:?}"
2942        );
2943        assert!(
2944            waiter_finished.load(Ordering::SeqCst),
2945            "the waiter should have been released by the stuck child's abort and run to completion"
2946        );
2947    }
2948
2949    #[tokio::test]
2950    async fn equal_child_deadlines_abort_the_waiter_too() {
2951        // The counterpart: concurrent shutdown computes every deadline from one shared instant, so identical timeouts
2952        // elapse in the same pass and the waiter is aborted alongside the child it was waiting on. This is also what a
2953        // shutdown budget does to a whole subtree, which is why the budget is set at a level where losing the entire
2954        // group at once is the intended outcome.
2955        let (sup, waiter_finished) = build_drain_pair(
2956            ShutdownStrategy::Graceful(Duration::from_millis(100)),
2957            ShutdownStrategy::Graceful(Duration::from_millis(100)),
2958        );
2959
2960        let (tx, handle) = run_supervisor_with_trigger(sup).await;
2961        tx.send(()).unwrap();
2962
2963        let result = join_supervisor(handle).await;
2964        assert!(
2965            matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 2 })),
2966            "both children should have been aborted together, got {result:?}"
2967        );
2968        assert!(
2969            !waiter_finished.load(Ordering::SeqCst),
2970            "the waiter should have been aborted mid-wait, not completed"
2971        );
2972    }
2973
2974    // -- Shutdown budget tests -------------------------------------------------------------
2975
2976    #[tokio::test]
2977    async fn budget_bounds_children_that_have_no_deadline_of_their_own() {
2978        // Without a budget, `Graceful(Duration::MAX)` children are waited on indefinitely and a stuck one hangs
2979        // shutdown forever. A budget makes the supervisor responsible for the deadline instead, and each child it has
2980        // to abort is still counted individually -- so an overrun says how many tasks were responsible, not merely
2981        // that the group as a whole overran.
2982        let mut sup = Supervisor::new("test-sup")
2983            .unwrap()
2984            .with_shutdown_mode(ShutdownMode::Concurrent)
2985            .with_shutdown_budget(Duration::from_millis(100));
2986
2987        for name in ["stuck_one", "stuck_two"] {
2988            sup.add_worker(
2989                ChildSpecification::worker(noninterruptible_worker(name, |_shutdown| pending::<()>()))
2990                    .with_restart_type(RestartType::Temporary)
2991                    .with_shutdown_strategy(ShutdownStrategy::Graceful(Duration::MAX)),
2992            );
2993        }
2994
2995        let (tx, handle) = run_supervisor_with_trigger(sup).await;
2996        tx.send(()).unwrap();
2997
2998        let result = join_supervisor(handle).await;
2999        assert!(
3000            matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 2 })),
3001            "the budget should have aborted both deadline-less children, got {result:?}"
3002        );
3003    }
3004
3005    #[tokio::test]
3006    async fn budget_does_not_delay_children_that_stop_on_their_own() {
3007        // The budget is a ceiling, not a wait. Asserting on elapsed time rather than merely on success is what makes
3008        // this meaningful: a supervisor that waited out its budget regardless would still report `Ok`.
3009        let mut sup = Supervisor::new("test-sup")
3010            .unwrap()
3011            .with_shutdown_mode(ShutdownMode::Concurrent)
3012            .with_shutdown_budget(Duration::from_secs(30));
3013        sup.add_worker(
3014            ChildSpecification::one_shot_worker(noninterruptible_worker("prompt", |shutdown| shutdown))
3015                .with_shutdown_strategy(ShutdownStrategy::Graceful(Duration::MAX)),
3016        );
3017
3018        let (tx, handle) = run_supervisor_with_trigger(sup).await;
3019        let started = tokio::time::Instant::now();
3020        tx.send(()).unwrap();
3021
3022        assert!(join_supervisor(handle).await.is_ok());
3023        let elapsed = started.elapsed();
3024        assert!(
3025            elapsed < Duration::from_millis(500),
3026            "shutdown should finish as soon as the child does, not burn the budget; took {elapsed:?}"
3027        );
3028    }
3029
3030    #[tokio::test]
3031    async fn child_deadline_shorter_than_budget_still_wins() {
3032        // A child that carries its own finite deadline is held to whichever elapses first, so a component can still
3033        // bound one particular task more tightly than the budget covering the rest.
3034        let mut sup = Supervisor::new("test-sup")
3035            .unwrap()
3036            .with_shutdown_mode(ShutdownMode::Concurrent)
3037            .with_shutdown_budget(Duration::from_secs(30));
3038        sup.add_worker(
3039            ChildSpecification::worker(noninterruptible_worker("stuck", |_shutdown| pending::<()>()))
3040                .with_restart_type(RestartType::Temporary)
3041                .with_shutdown_strategy(ShutdownStrategy::Graceful(Duration::from_millis(100))),
3042        );
3043
3044        let (tx, handle) = run_supervisor_with_trigger(sup).await;
3045        tx.send(()).unwrap();
3046
3047        // Again bounded at two seconds: if the 30-second budget had won, this would time out instead.
3048        let result = join_supervisor(handle).await;
3049        assert!(
3050            matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 1 })),
3051            "the child's own 100ms deadline should have won over the budget, got {result:?}"
3052        );
3053    }
3054
3055    #[tokio::test]
3056    async fn every_worker_records_poll_metrics() {
3057        // Poll timing is a property of being supervised, not something a child opts into, so a plain statically
3058        // registered worker gets it too -- tagged with its fully qualified process name.
3059        let recorder = TestRecorder::default();
3060        let _guard = metrics::set_default_local_recorder(&recorder);
3061
3062        // The recorder has to be installed before the worker spawns: its metric handles are resolved once, at spawn.
3063        let mut sup = Supervisor::new("metrics_sup").unwrap();
3064        sup.add_worker(ChildSpecification::one_shot_worker(noninterruptible_worker(
3065            "timed",
3066            |shutdown| shutdown,
3067        )));
3068
3069        let (tx, handle) = run_supervisor_with_trigger(sup).await;
3070        tx.send(()).unwrap();
3071        assert!(join_supervisor(handle).await.is_ok());
3072
3073        let polls = recorder.counter(("runtime_task_poll_count", &[("task_name", "metrics_sup.timed")]));
3074        assert!(
3075            polls.is_some_and(|polls| polls > 0),
3076            "a supervised worker should have recorded poll metrics, got {polls:?}"
3077        );
3078    }
3079
3080    #[tokio::test]
3081    async fn budget_bounds_the_whole_drain_in_ordered_mode() {
3082        // Ordered shutdown stops children one at a time, so a budget has to cover the sequence as a whole rather than
3083        // resetting per child. Three children that each ignore a 10-second deadline must all be aborted at the shared
3084        // 150ms budget, not 30 seconds later.
3085        let mut sup = Supervisor::new("test-sup")
3086            .unwrap()
3087            .with_shutdown_mode(ShutdownMode::Ordered)
3088            .with_shutdown_budget(Duration::from_millis(150));
3089
3090        for name in ["stuck_one", "stuck_two", "stuck_three"] {
3091            sup.add_worker(
3092                ChildSpecification::one_shot_worker(noninterruptible_worker(name, |_shutdown| pending::<()>()))
3093                    .with_shutdown_strategy(ShutdownStrategy::Graceful(Duration::from_secs(10))),
3094            );
3095        }
3096
3097        let (tx, handle) = run_supervisor_with_trigger(sup).await;
3098        tx.send(()).unwrap();
3099
3100        let result = join_supervisor(handle).await;
3101        assert!(
3102            matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 3 })),
3103            "the budget should have bounded the whole ordered drain, got {result:?}"
3104        );
3105    }
3106
3107    #[tokio::test]
3108    async fn budget_of_duration_max_is_treated_as_no_budget() {
3109        // `Duration::MAX` is the natural spelling of "no ceiling" and used to panic the supervisor task on an instant
3110        // overflow.
3111        let mut sup = Supervisor::new("test-sup")
3112            .unwrap()
3113            .with_shutdown_mode(ShutdownMode::Concurrent)
3114            .with_shutdown_budget(Duration::MAX);
3115        sup.add_worker(ChildSpecification::one_shot_worker(noninterruptible_worker(
3116            "prompt",
3117            |shutdown| shutdown,
3118        )));
3119
3120        let (tx, handle) = run_supervisor_with_trigger(sup).await;
3121        tx.send(()).unwrap();
3122        assert!(join_supervisor(handle).await.is_ok());
3123    }
3124
3125    #[tokio::test]
3126    async fn near_max_child_timeout_does_not_panic_in_ordered_mode() {
3127        // The ordered path used to pass the timeout straight to `sleep`, which clamps. Resolving it to an instant
3128        // instead made anything just under `Duration::MAX` overflow.
3129        let mut sup = Supervisor::new("test-sup").unwrap();
3130        sup.add_worker(
3131            ChildSpecification::one_shot_worker(noninterruptible_worker("prompt", |shutdown| shutdown))
3132                .with_shutdown_strategy(ShutdownStrategy::Graceful(Duration::MAX - Duration::from_nanos(1))),
3133        );
3134
3135        let (tx, handle) = run_supervisor_with_trigger(sup).await;
3136        tx.send(()).unwrap();
3137        assert!(join_supervisor(handle).await.is_ok());
3138    }
3139
3140    #[tokio::test]
3141    async fn budget_does_not_cut_off_a_nested_supervisor_mid_drain() {
3142        // A nested supervisor bounds its own subtree, so a parent's budget must not abort it: doing so truncates the
3143        // subtree's drain, discards its abort tally, and -- for a supervisor on a dedicated runtime, whose work is on
3144        // another OS thread -- reports it as stopped without actually stopping it.
3145        let drained = Arc::new(AtomicBool::new(false));
3146        let child_drained = Arc::clone(&drained);
3147
3148        let mut nested = Supervisor::new("nested").unwrap();
3149        nested.add_worker(
3150            ChildSpecification::one_shot_worker(noninterruptible_worker("slow", move |shutdown| async move {
3151                shutdown.await;
3152                sleep(Duration::from_millis(300)).await;
3153                child_drained.store(true, Ordering::SeqCst);
3154            }))
3155            .with_shutdown_strategy(ShutdownStrategy::Graceful(Duration::from_secs(10))),
3156        );
3157
3158        let mut parent = Supervisor::new("parent")
3159            .unwrap()
3160            .with_shutdown_mode(ShutdownMode::Concurrent)
3161            .with_shutdown_budget(Duration::from_millis(50));
3162        parent.add_worker(nested);
3163
3164        let (tx, handle) = run_supervisor_with_trigger(parent).await;
3165        tx.send(()).unwrap();
3166
3167        let result = join_supervisor(handle).await;
3168        assert!(
3169            drained.load(Ordering::SeqCst),
3170            "the nested subtree should have drained rather than being cut off by the parent's budget: {result:?}"
3171        );
3172        assert!(
3173            result.is_ok(),
3174            "the nested drain finished in time, so shutdown was clean: {result:?}"
3175        );
3176    }
3177}