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