1use std::{
2 future::Future,
3 pin::Pin,
4 sync::{
5 atomic::{AtomicU64, AtomicUsize, Ordering},
6 Arc, Mutex,
7 },
8 time::Duration,
9};
10
11use async_trait::async_trait;
12use saluki_common::collections::FastHashMap;
13use saluki_common::sync::shutdown::ShutdownHandle;
14use saluki_error::GenericError;
15use snafu::{OptionExt as _, Snafu};
16use tokio::{
17 pin,
18 runtime::Handle,
19 select,
20 sync::{mpsc, oneshot},
21};
22use tracing::{debug, error, warn};
23
24use super::{
25 dedicated::{spawn_dedicated_runtime, RuntimeConfiguration, RuntimeMode},
26 restart::{RestartAction, RestartMode, RestartState, RestartStrategy, RestartType},
27 worker_state::WorkerState,
28};
29use crate::runtime::{
30 process::{Process, ProcessExt as _},
31 state::DataspaceRegistry,
32};
33
34pub type SupervisorFuture = Pin<Box<dyn Future<Output = Result<(), GenericError>> + Send>>;
36
37pub(super) type WorkerFuture = Pin<Box<dyn Future<Output = Result<(), WorkerError>> + Send>>;
43
44#[derive(Debug)]
49pub(super) enum WorkerError {
50 Initialization {
56 child_name: Option<String>,
57 source: InitializationError,
58 },
59
60 Runtime(GenericError),
62
63 ShutdownTimedOut {
70 aborted: usize,
72 },
73}
74
75impl From<SupervisorError> for WorkerError {
76 fn from(err: SupervisorError) -> Self {
77 match err {
78 SupervisorError::FailedToInitialize { child_name, source } => WorkerError::Initialization {
81 child_name: Some(child_name),
82 source,
83 },
84 SupervisorError::ShutdownTimedOut { aborted } => WorkerError::ShutdownTimedOut { aborted },
86 other => WorkerError::Runtime(other.into()),
88 }
89 }
90}
91
92#[derive(Debug, Snafu)]
94pub enum ProcessError {
95 #[snafu(display("Child process was aborted by the supervisor."))]
97 Aborted,
98
99 #[snafu(display("Child process panicked."))]
101 Panicked,
102
103 #[snafu(display("Child process terminated with an error: {}", source))]
105 Terminated {
106 source: GenericError,
108 },
109}
110
111#[derive(Debug, Snafu)]
117#[snafu(context(suffix(false)))]
118pub enum InitializationError {
119 #[snafu(display("Process failed to initialize: {}", source))]
121 Failed {
122 source: GenericError,
124 },
125}
126
127impl From<GenericError> for InitializationError {
128 fn from(source: GenericError) -> Self {
129 Self::Failed { source }
130 }
131}
132
133#[derive(Clone, Copy, Debug)]
135pub enum ShutdownStrategy {
136 Graceful(Duration),
138
139 Brutal,
141}
142
143#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
150pub enum AutoShutdown {
151 #[default]
153 Never,
154
155 AnySignificant,
157
158 AllSignificant,
160}
161
162#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
164pub enum ShutdownMode {
165 #[default]
170 Ordered,
171
172 Concurrent,
177}
178
179#[async_trait]
181pub trait Supervisable: Send + Sync {
182 fn name(&self) -> &str;
184
185 fn shutdown_strategy(&self) -> ShutdownStrategy {
187 ShutdownStrategy::Graceful(Duration::from_secs(5))
188 }
189
190 async fn initialize(&self, process_shutdown: ShutdownHandle) -> Result<SupervisorFuture, InitializationError>;
204}
205
206#[derive(Debug, Snafu)]
208#[snafu(context(suffix(false)))]
209pub enum SupervisorError {
210 #[snafu(display("Invalid name for supervisor or worker: '{}'", name))]
212 InvalidName {
213 name: String,
215 },
216
217 #[snafu(display("Child process '{}' failed to initialize: {}", child_name, source))]
222 FailedToInitialize {
223 child_name: String,
225
226 source: InitializationError,
228 },
229
230 #[snafu(display("Supervisor has exceeded restart limits and was forced to shutdown."))]
232 Shutdown,
233
234 #[snafu(display("Supervisor shut down after a significant child terminated."))]
239 SignificantChildExited,
240
241 #[snafu(display(
250 "Shutdown completed uncleanly: {} worker(s) were forcefully aborted after exceeding their shutdown timeout.",
251 aborted
252 ))]
253 ShutdownTimedOut {
254 aborted: usize,
256 },
257}
258
259pub struct ChildSpecification<S = WorkerSpec> {
275 spec_inner: S,
276}
277
278pub struct WorkerSpec {
280 worker: Arc<dyn Supervisable>,
281 config: ChildConfig,
282}
283
284pub struct SupervisorSpec {
286 supervisor: Supervisor,
287}
288
289impl ChildSpecification<WorkerSpec> {
290 pub fn worker<T: Supervisable + 'static>(worker: T) -> Self {
292 Self {
293 spec_inner: WorkerSpec {
294 worker: Arc::new(worker),
295 config: ChildConfig::default(),
296 },
297 }
298 }
299
300 pub fn one_shot_worker<T: Supervisable + 'static>(worker: T) -> Self {
305 Self::worker(worker).with_restart_type(RestartType::Temporary)
306 }
307
308 #[must_use]
312 pub fn with_restart_type(mut self, restart_type: RestartType) -> Self {
313 self.spec_inner.config.restart = restart_type;
314 self
315 }
316
317 #[must_use]
323 pub fn with_significant(mut self, significant: bool) -> Self {
324 self.spec_inner.config.significant = significant;
325 self
326 }
327
328 #[must_use]
337 pub fn with_runtime(mut self, handle: Handle) -> Self {
338 self.spec_inner.config.runtime = Some(handle);
339 self
340 }
341
342 #[must_use]
349 pub fn with_shutdown_strategy(mut self, strategy: ShutdownStrategy) -> Self {
350 self.spec_inner.config.shutdown_strategy = Some(strategy);
351 self
352 }
353
354 fn into_worker_parts(self) -> (SupervisedChild, ChildConfig) {
356 (SupervisedChild::Worker(self.spec_inner.worker), self.spec_inner.config)
357 }
358}
359
360impl<T> From<T> for ChildSpecification<WorkerSpec>
361where
362 T: Supervisable + 'static,
363{
364 fn from(worker: T) -> Self {
365 Self::worker(worker)
366 }
367}
368
369impl From<Supervisor> for ChildSpecification<SupervisorSpec> {
370 fn from(supervisor: Supervisor) -> Self {
371 Self {
372 spec_inner: SupervisorSpec { supervisor },
373 }
374 }
375}
376
377mod sealed {
378 pub trait Sealed {}
379}
380
381impl sealed::Sealed for WorkerSpec {}
382impl sealed::Sealed for SupervisorSpec {}
383
384pub trait ChildState: sealed::Sealed + Sized {
391 #[doc(hidden)]
392 fn register(spec: ChildSpecification<Self>, supervisor: &mut Supervisor);
393}
394
395impl ChildState for WorkerSpec {
396 fn register(spec: ChildSpecification<Self>, supervisor: &mut Supervisor) {
397 let (child, config) = spec.into_worker_parts();
398 supervisor.push_child(ChildEntry {
399 spec: child,
400 config,
401 dynamic: false,
402 });
403 }
404}
405
406impl ChildState for SupervisorSpec {
407 fn register(spec: ChildSpecification<Self>, supervisor: &mut Supervisor) {
408 supervisor.push_child(ChildEntry {
409 spec: SupervisedChild::Supervisor(spec.spec_inner.supervisor),
410 config: ChildConfig::default(),
411 dynamic: false,
412 });
413 }
414}
415
416pub(super) enum SupervisedChild {
418 Worker(Arc<dyn Supervisable>),
419 Supervisor(Supervisor),
420}
421
422impl SupervisedChild {
423 pub(super) fn is_supervisor(&self) -> bool {
425 matches!(self, Self::Supervisor(_))
426 }
427
428 fn process_type(&self) -> &'static str {
429 match self {
430 Self::Worker(_) => "worker",
431 Self::Supervisor(_) => "supervisor",
432 }
433 }
434
435 fn name(&self) -> &str {
436 match self {
437 Self::Worker(worker) => worker.name(),
438 Self::Supervisor(supervisor) => &supervisor.supervisor_id,
439 }
440 }
441
442 pub(super) fn shutdown_strategy(&self) -> ShutdownStrategy {
443 match self {
444 Self::Worker(worker) => worker.shutdown_strategy(),
445
446 Self::Supervisor(_) => ShutdownStrategy::Graceful(Duration::MAX),
449 }
450 }
451
452 pub(super) fn create_process(&self, parent_process: &Process) -> Result<Process, SupervisorError> {
453 match self {
454 Self::Worker(worker) => Process::worker(worker.name(), parent_process).context(InvalidName {
455 name: worker.name().to_string(),
456 }),
457 Self::Supervisor(sup) => {
458 Process::supervisor(&sup.supervisor_id, Some(parent_process)).context(InvalidName {
459 name: sup.supervisor_id.to_string(),
460 })
461 }
462 }
463 }
464
465 pub(super) fn create_worker_future(
466 &self, process: Process, process_shutdown: ShutdownHandle,
467 ) -> Result<WorkerFuture, SupervisorError> {
468 match self {
469 Self::Worker(worker) => {
470 let worker = Arc::clone(worker);
471 Ok(Box::pin(async move {
472 let run_future =
473 worker
474 .initialize(process_shutdown)
475 .await
476 .map_err(|source| WorkerError::Initialization {
477 child_name: None,
478 source,
479 })?;
480 run_future.await.map_err(WorkerError::Runtime)
481 }))
482 }
483 Self::Supervisor(sup) => {
484 match sup.runtime_mode() {
485 RuntimeMode::Ambient => {
486 Ok(sup.as_nested_process(process, process_shutdown))
488 }
489 RuntimeMode::Dedicated(config) => {
490 let child_name = sup.supervisor_id.to_string();
493 let dataspace = process.dataspace().clone();
494 let handle =
495 spawn_dedicated_runtime(sup.inner_clone(), config.clone(), process_shutdown, dataspace)
496 .map_err(|e| SupervisorError::FailedToInitialize {
497 child_name,
498 source: e.into(),
499 })?;
500
501 Ok(Box::pin(async move { handle.await.map_err(WorkerError::from) }))
502 }
503 }
504 }
505 }
506 }
507}
508
509impl Clone for SupervisedChild {
510 fn clone(&self) -> Self {
511 match self {
512 Self::Worker(worker) => Self::Worker(Arc::clone(worker)),
513 Self::Supervisor(supervisor) => Self::Supervisor(supervisor.inner_clone()),
514 }
515 }
516}
517
518#[derive(Clone, Debug)]
526pub(super) struct ChildConfig {
527 restart: RestartType,
528 significant: bool,
529
530 runtime: Option<Handle>,
532
533 shutdown_strategy: Option<ShutdownStrategy>,
535}
536
537impl ChildConfig {
538 pub(super) fn runtime(&self) -> Option<&Handle> {
540 self.runtime.as_ref()
541 }
542
543 pub(super) fn shutdown_strategy(&self) -> Option<ShutdownStrategy> {
545 self.shutdown_strategy
546 }
547}
548
549impl Default for ChildConfig {
550 fn default() -> Self {
551 Self {
552 restart: RestartType::Permanent,
553 significant: false,
554 runtime: None,
555 shutdown_strategy: None,
556 }
557 }
558}
559
560#[derive(Clone)]
562struct ChildEntry {
563 spec: SupervisedChild,
564 config: ChildConfig,
565 dynamic: bool,
568}
569
570#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
575pub struct ChildId(u64);
576
577impl ChildId {
578 pub const fn as_u64(self) -> u64 {
580 self.0
581 }
582}
583
584#[derive(Debug, Snafu)]
586pub enum SpawnError {
587 #[snafu(display("supervisor is gone"))]
593 SupervisorGone,
594
595 #[snafu(display("supervisor rejected the spawn: {}", source))]
600 Rejected {
601 source: GenericError,
603 },
604}
605
606struct PendingSpawn {
608 id: u64,
609 spec: SupervisedChild,
610 config: ChildConfig,
611 ack: oneshot::Sender<Result<(), SpawnError>>,
612}
613
614const DYNAMIC_SPAWN_CHANNEL_CAPACITY: usize = 1024;
619
620#[derive(Clone)]
626pub struct SupervisorHandle {
627 name: Arc<str>,
628 current_tx: Arc<Mutex<Option<mpsc::Sender<PendingSpawn>>>>,
631 id_counter: Arc<AtomicU64>,
632 active: Arc<AtomicUsize>,
633}
634
635impl SupervisorHandle {
636 pub fn name(&self) -> &str {
638 &self.name
639 }
640
641 pub async fn spawn<T: Supervisable + 'static>(&self, worker: T) -> Result<ChildId, SpawnError> {
654 self.spawn_with(ChildSpecification::worker(worker).with_restart_type(RestartType::Temporary))
655 .await
656 }
657
658 pub async fn spawn_with(&self, spec: ChildSpecification<WorkerSpec>) -> Result<ChildId, SpawnError> {
672 let id = self.id_counter.fetch_add(1, Ordering::Relaxed);
673 let (spec, config) = spec.into_worker_parts();
674 let (ack_tx, ack_rx) = oneshot::channel();
675 self.send(PendingSpawn {
676 id,
677 spec,
678 config,
679 ack: ack_tx,
680 })
681 .await?;
682
683 ack_rx
686 .await
687 .map_err(|_| SpawnError::SupervisorGone)?
688 .map(|()| ChildId(id))
689 }
690
691 pub fn is_running(&self) -> bool {
693 self.current_tx.lock().unwrap().is_some()
694 }
695
696 pub fn active_children(&self) -> usize {
698 self.active.load(Ordering::Relaxed)
699 }
700
701 async fn send(&self, spawn: PendingSpawn) -> Result<(), SpawnError> {
705 let tx = self.current_tx.lock().unwrap().clone();
707 match tx {
708 Some(tx) => tx.send(spawn).await.map_err(|_| SpawnError::SupervisorGone),
709 None => Err(SpawnError::SupervisorGone),
710 }
711 }
712}
713
714pub struct Supervisor {
742 supervisor_id: Arc<str>,
743 child_specs: Vec<ChildEntry>,
744 restart_strategy: RestartStrategy,
745 auto_shutdown: AutoShutdown,
746 shutdown_mode: ShutdownMode,
747 shutdown_budget: Option<Duration>,
748 runtime_mode: RuntimeMode,
749 current_tx: Arc<Mutex<Option<mpsc::Sender<PendingSpawn>>>>,
753 id_counter: Arc<AtomicU64>,
754 active: Arc<AtomicUsize>,
756}
757
758impl Supervisor {
759 pub fn new<S: AsRef<str>>(supervisor_id: S) -> Result<Self, SupervisorError> {
761 if supervisor_id.as_ref().is_empty() {
765 return Err(SupervisorError::InvalidName {
766 name: supervisor_id.as_ref().to_string(),
767 });
768 }
769
770 Ok(Self {
771 supervisor_id: supervisor_id.as_ref().into(),
772 child_specs: Vec::new(),
773 restart_strategy: RestartStrategy::default(),
774 auto_shutdown: AutoShutdown::default(),
775 shutdown_mode: ShutdownMode::default(),
776 shutdown_budget: None,
777 runtime_mode: RuntimeMode::default(),
778 current_tx: Arc::new(Mutex::new(None)),
779 id_counter: Arc::new(AtomicU64::new(0)),
780 active: Arc::new(AtomicUsize::new(0)),
781 })
782 }
783
784 pub fn id(&self) -> &str {
786 &self.supervisor_id
787 }
788
789 pub fn with_restart_strategy(mut self, strategy: RestartStrategy) -> Self {
791 self.restart_strategy = strategy;
792 self
793 }
794
795 pub fn with_auto_shutdown(mut self, auto_shutdown: AutoShutdown) -> Self {
800 self.auto_shutdown = auto_shutdown;
801 self
802 }
803
804 pub fn with_shutdown_mode(mut self, mode: ShutdownMode) -> Self {
806 self.shutdown_mode = mode;
807 self
808 }
809
810 #[must_use]
830 pub fn with_shutdown_budget(mut self, budget: Duration) -> Self {
831 self.shutdown_budget = Some(budget);
832 self
833 }
834
835 pub fn handle(&self) -> SupervisorHandle {
841 SupervisorHandle {
842 name: Arc::clone(&self.supervisor_id),
843 current_tx: Arc::clone(&self.current_tx),
844 id_counter: Arc::clone(&self.id_counter),
845 active: Arc::clone(&self.active),
846 }
847 }
848
849 pub fn with_dedicated_runtime(mut self, config: RuntimeConfiguration) -> Self {
859 self.runtime_mode = RuntimeMode::Dedicated(config);
860 self
861 }
862
863 pub(crate) fn runtime_mode(&self) -> &RuntimeMode {
865 &self.runtime_mode
866 }
867
868 pub fn add_worker<S, T>(&mut self, child: T)
876 where
877 S: ChildState,
878 T: Into<ChildSpecification<S>>,
879 {
880 S::register(child.into(), self);
881 }
882
883 fn push_child(&mut self, entry: ChildEntry) {
884 debug!(
885 supervisor_id = %self.supervisor_id,
886 "Adding new static child process #{}. ({}, {}, {:?})",
887 self.child_specs.len(),
888 entry.spec.process_type(),
889 entry.spec.name(),
890 entry.config,
891 );
892 self.child_specs.push(entry);
893 }
894
895 fn spawn_static_children(
896 &self, children: &mut FastHashMap<u64, ChildEntry>, worker_state: &mut WorkerState,
897 ) -> Result<(), SupervisorError> {
898 debug!(supervisor_id = %self.supervisor_id, "Spawning all static child processes.");
899 for entry in &self.child_specs {
900 let id = self.id_counter.fetch_add(1, Ordering::Relaxed);
901 worker_state.add_worker(id, &entry.spec, &entry.config)?;
902 children.insert(id, entry.clone());
903 }
904
905 Ok(())
906 }
907
908 fn respawn_children_one_for_all(
916 &self, children: &mut FastHashMap<u64, ChildEntry>, worker_state: &mut WorkerState,
917 ) -> Result<(), SupervisorError> {
918 debug!(supervisor_id = %self.supervisor_id, "Restarting all eligible static child processes.");
919 for entry in &self.child_specs {
920 if entry.config.restart == RestartType::Temporary {
923 continue;
924 }
925 let id = self.id_counter.fetch_add(1, Ordering::Relaxed);
926 worker_state.add_worker(id, &entry.spec, &entry.config)?;
927 children.insert(id, entry.clone());
928 }
929
930 Ok(())
931 }
932
933 fn spawn_dynamic_child(
935 &self, spawn: PendingSpawn, worker_state: &mut WorkerState, children: &mut FastHashMap<u64, ChildEntry>,
936 significant_remaining: &mut usize,
937 ) {
938 let PendingSpawn { id, spec, config, ack } = spawn;
939 let entry = ChildEntry {
940 spec,
941 config,
942 dynamic: true,
943 };
944 match worker_state.add_worker(id, &entry.spec, &entry.config) {
945 Ok(()) => {
946 if entry.config.significant {
947 *significant_remaining += 1;
948 }
949 self.active.fetch_add(1, Ordering::Relaxed);
950 children.insert(id, entry);
951 let _ = ack.send(Ok(()));
952 }
953 Err(e) => {
954 error!(supervisor_id = %self.supervisor_id, error = %e, "Failed to spawn dynamic child.");
957 let _ = ack.send(Err(SpawnError::Rejected { source: e.into() }));
958 }
959 }
960 }
961
962 async fn run_inner(&self, process: Process, process_shutdown: ShutdownHandle) -> Result<(), SupervisorError> {
963 let (cmd_tx, cmd_rx) = mpsc::channel(DYNAMIC_SPAWN_CHANNEL_CAPACITY);
966 *self.current_tx.lock().unwrap() = Some(cmd_tx);
967
968 let result = self.supervise(process, process_shutdown, cmd_rx).await;
969
970 *self.current_tx.lock().unwrap() = None;
973 self.active.store(0, Ordering::Relaxed);
974 result
975 }
976
977 async fn supervise(
978 &self, process: Process, process_shutdown: ShutdownHandle, mut cmd_rx: mpsc::Receiver<PendingSpawn>,
979 ) -> Result<(), SupervisorError> {
980 let mut restart_state = RestartState::new(self.restart_strategy);
981 let mut worker_state = WorkerState::new(process, self.shutdown_mode, self.shutdown_budget);
982
983 let mut children: FastHashMap<u64, ChildEntry> = FastHashMap::default();
986
987 self.spawn_static_children(&mut children, &mut worker_state)?;
990
991 let mut significant_remaining = children.values().filter(|entry| entry.config.significant).count();
993
994 pin!(process_shutdown);
996
997 let outcome = loop {
998 select! {
999 biased;
1001
1002 _ = &mut process_shutdown => break Ok(()),
1006
1007 spawn = cmd_rx.recv() => {
1010 if let Some(spawn) = spawn {
1011 self.spawn_dynamic_child(spawn, &mut worker_state, &mut children, &mut significant_remaining);
1012 }
1013 }
1014
1015 (child_id, worker_result) = worker_state.wait_for_next_worker() => {
1016 let (child_name, config, dynamic) = {
1018 let entry = children.get(&child_id).expect("completed worker must be present in the roster");
1019 (entry.spec.name().to_string(), entry.config.clone(), entry.dynamic)
1020 };
1021
1022 if let Err(WorkerError::Initialization { child_name: inner, source }) = worker_result {
1024 let full_name = match inner {
1027 Some(inner) => format!("{}/{}", child_name, inner),
1028 None => child_name.clone(),
1029 };
1030
1031 error!(supervisor_id = %self.supervisor_id, worker_name = full_name, "Child process failed to initialize: {}", source);
1032 break Err(SupervisorError::FailedToInitialize { child_name: full_name, source });
1033 }
1034
1035 let abnormal = worker_result.is_err();
1038 let worker_result = worker_result.map_err(|e| match e {
1039 WorkerError::Runtime(e) => ProcessError::Terminated { source: e },
1040 WorkerError::Initialization { .. } => unreachable!("handled above"),
1041 WorkerError::ShutdownTimedOut { aborted } => ProcessError::Terminated {
1046 source: SupervisorError::ShutdownTimedOut { aborted }.into(),
1047 },
1048 });
1049
1050 if !config.restart.should_restart(abnormal) {
1051 if abnormal {
1060 warn!(supervisor_id = %self.supervisor_id, worker_name = %child_name, restart = ?config.restart, ?worker_result, "Child process exited with an error and is not eligible for restart.");
1061 } else {
1062 debug!(supervisor_id = %self.supervisor_id, worker_name = %child_name, restart = ?config.restart, "Child process exited and is not eligible for restart.");
1063 }
1064 children.remove(&child_id);
1065 if dynamic {
1066 self.active.fetch_sub(1, Ordering::Relaxed);
1067 }
1068
1069 if config.significant {
1073 significant_remaining = significant_remaining.saturating_sub(1);
1074 let auto_shutdown = match self.auto_shutdown {
1075 AutoShutdown::Never => false,
1076 AutoShutdown::AnySignificant => true,
1077 AutoShutdown::AllSignificant => significant_remaining == 0,
1078 };
1079 if auto_shutdown {
1080 warn!(supervisor_id = %self.supervisor_id, worker_name = %child_name, ?worker_result, "Significant child terminated; shutting down supervisor.");
1081 break Err(SupervisorError::SignificantChildExited);
1082 }
1083 }
1084 } else {
1085 match restart_state.evaluate_restart() {
1086 RestartAction::Restart(mode) => match mode {
1087 RestartMode::OneForOne => {
1088 warn!(supervisor_id = %self.supervisor_id, worker_name = %child_name, ?worker_result, "Child process terminated, restarting.");
1089 let spec = children.get(&child_id).expect("present for restart").spec.clone();
1090 if let Err(e) = worker_state.add_worker(child_id, &spec, &config) {
1091 break Err(e);
1092 }
1093 }
1094 RestartMode::OneForAll => {
1095 warn!(supervisor_id = %self.supervisor_id, worker_name = %child_name, ?worker_result, "Child process terminated, restarting all processes.");
1096 let _ = worker_state.shutdown_workers().await;
1100 children.clear();
1104 self.active.store(0, Ordering::Relaxed);
1105 let respawn = self.respawn_children_one_for_all(&mut children, &mut worker_state);
1106 if let Err(e) = respawn {
1107 break Err(e);
1108 }
1109 significant_remaining =
1110 children.values().filter(|entry| entry.config.significant).count();
1111 }
1112 },
1113 RestartAction::Shutdown => {
1114 error!(supervisor_id = %self.supervisor_id, worker_name = %child_name, ?worker_result, "Supervisor shutting down due to restart limits.");
1115 break Err(SupervisorError::Shutdown);
1116 }
1117 }
1118 }
1119 }
1120 }
1121 };
1122
1123 cmd_rx.close();
1130 while let Ok(spawn) = cmd_rx.try_recv() {
1131 let _ = spawn.ack.send(Err(SpawnError::SupervisorGone));
1132 }
1133 let aborted = worker_state.shutdown_workers().await;
1134
1135 match outcome {
1140 Ok(()) if aborted > 0 => {
1141 warn!(supervisor_id = %self.supervisor_id, aborted, "Shutdown completed uncleanly; workers were forcefully aborted.");
1142 Err(SupervisorError::ShutdownTimedOut { aborted })
1143 }
1144 outcome => outcome,
1145 }
1146 }
1147
1148 fn as_nested_process(&self, process: Process, process_shutdown: ShutdownHandle) -> WorkerFuture {
1149 debug!(supervisor_id = %self.supervisor_id, "Nested supervisor starting.");
1152
1153 let sup = self.inner_clone();
1155
1156 Box::pin(async move {
1157 sup.run_inner(process, process_shutdown)
1158 .await
1159 .map_err(WorkerError::from)
1160 })
1161 }
1162
1163 pub async fn run(&mut self) -> Result<(), SupervisorError> {
1169 let process_shutdown = ShutdownHandle::noop();
1172 let process = Process::supervisor(&self.supervisor_id, None).context(InvalidName {
1173 name: self.supervisor_id.to_string(),
1174 })?;
1175
1176 debug!(supervisor_id = %self.supervisor_id, "Supervisor starting.");
1177 self.run_inner(process.clone(), process_shutdown)
1178 .into_process_future(process)
1179 .await
1180 }
1181
1182 pub async fn run_with_shutdown<F: Future + Send + 'static>(&mut self, shutdown: F) -> Result<(), SupervisorError> {
1191 let (shutdown_coordinator, shutdown_handle) = ShutdownHandle::paired();
1195 let run = self.run_with_shutdown_inner(shutdown_handle, None);
1196 pin!(run, shutdown);
1197
1198 let mut shutdown_coordinator = Some(shutdown_coordinator);
1199 loop {
1200 select! {
1201 result = &mut run => return result,
1202 _ = &mut shutdown, if shutdown_coordinator.is_some() => {
1203 shutdown_coordinator.take().expect("coordinator present per select guard").shutdown();
1204 }
1205 }
1206 }
1207 }
1208
1209 pub(crate) async fn run_with_shutdown_inner(
1221 &mut self, process_shutdown: ShutdownHandle, dataspace: Option<DataspaceRegistry>,
1222 ) -> Result<(), SupervisorError> {
1223 let process =
1224 Process::supervisor_with_dataspace(&self.supervisor_id, None, dataspace).context(InvalidName {
1225 name: self.supervisor_id.to_string(),
1226 })?;
1227
1228 debug!(supervisor_id = %self.supervisor_id, "Supervisor starting.");
1229 self.run_inner(process.clone(), process_shutdown)
1230 .into_process_future(process)
1231 .await
1232 }
1233
1234 fn inner_clone(&self) -> Self {
1235 Self {
1239 supervisor_id: Arc::clone(&self.supervisor_id),
1240 child_specs: self.child_specs.clone(),
1241 restart_strategy: self.restart_strategy,
1242 auto_shutdown: self.auto_shutdown,
1243 shutdown_mode: self.shutdown_mode,
1244 shutdown_budget: self.shutdown_budget,
1245 runtime_mode: self.runtime_mode.clone(),
1246 current_tx: Arc::clone(&self.current_tx),
1247 id_counter: Arc::clone(&self.id_counter),
1248 active: Arc::clone(&self.active),
1249 }
1250 }
1251}
1252
1253#[cfg(test)]
1254mod tests {
1255 use std::{
1256 future::pending,
1257 sync::atomic::{AtomicBool, AtomicUsize, Ordering},
1258 };
1259
1260 use async_trait::async_trait;
1261 use saluki_common::sync::shutdown::ShutdownCoordinator;
1262 use saluki_metrics::test::TestRecorder;
1263 use tokio::{
1264 sync::oneshot,
1265 task::JoinHandle,
1266 time::{sleep, timeout},
1267 };
1268
1269 use super::*;
1270 use crate::runtime::noninterruptible_worker;
1271 use crate::test_support::wait_until;
1272
1273 #[derive(Clone)]
1275 enum InitBehavior {
1276 Instant,
1278
1279 Slow(Duration),
1281
1282 Fail(&'static str),
1284 }
1285
1286 #[derive(Clone)]
1288 enum RunBehavior {
1289 UntilShutdown,
1291
1292 FailAfter(Duration, &'static str),
1294
1295 CompleteAfter(Duration),
1297
1298 SlowShutdown(Duration),
1300
1301 IgnoreShutdown,
1303
1304 PanicAfter(Duration),
1306 }
1307
1308 struct MockWorker {
1310 name: &'static str,
1311 init_behavior: InitBehavior,
1312 run_behavior: RunBehavior,
1313 start_count: Arc<AtomicUsize>,
1314 finish_count: Arc<AtomicUsize>,
1315 brutal_shutdown: bool,
1316 graceful_timeout: Duration,
1317 }
1318
1319 impl MockWorker {
1320 fn long_running(name: &'static str) -> Self {
1322 Self {
1323 name,
1324 init_behavior: InitBehavior::Instant,
1325 run_behavior: RunBehavior::UntilShutdown,
1326 start_count: Arc::new(AtomicUsize::new(0)),
1327 finish_count: Arc::new(AtomicUsize::new(0)),
1328 brutal_shutdown: false,
1329 graceful_timeout: Duration::from_millis(500),
1330 }
1331 }
1332
1333 fn failing(name: &'static str, delay: Duration) -> Self {
1335 Self {
1336 name,
1337 init_behavior: InitBehavior::Instant,
1338 run_behavior: RunBehavior::FailAfter(delay, "worker failed"),
1339 start_count: Arc::new(AtomicUsize::new(0)),
1340 finish_count: Arc::new(AtomicUsize::new(0)),
1341 brutal_shutdown: false,
1342 graceful_timeout: Duration::from_millis(500),
1343 }
1344 }
1345
1346 fn completing(name: &'static str, delay: Duration) -> Self {
1348 Self {
1349 name,
1350 init_behavior: InitBehavior::Instant,
1351 run_behavior: RunBehavior::CompleteAfter(delay),
1352 start_count: Arc::new(AtomicUsize::new(0)),
1353 finish_count: Arc::new(AtomicUsize::new(0)),
1354 brutal_shutdown: false,
1355 graceful_timeout: Duration::from_millis(500),
1356 }
1357 }
1358
1359 fn slow_shutdown(name: &'static str, delay: Duration) -> Self {
1361 Self {
1362 name,
1363 init_behavior: InitBehavior::Instant,
1364 run_behavior: RunBehavior::SlowShutdown(delay),
1365 start_count: Arc::new(AtomicUsize::new(0)),
1366 finish_count: Arc::new(AtomicUsize::new(0)),
1367 brutal_shutdown: false,
1368 graceful_timeout: Duration::from_millis(500),
1369 }
1370 }
1371
1372 fn ignore_shutdown(name: &'static str) -> Self {
1374 Self {
1375 name,
1376 init_behavior: InitBehavior::Instant,
1377 run_behavior: RunBehavior::IgnoreShutdown,
1378 start_count: Arc::new(AtomicUsize::new(0)),
1379 finish_count: Arc::new(AtomicUsize::new(0)),
1380 brutal_shutdown: false,
1381 graceful_timeout: Duration::from_millis(500),
1382 }
1383 }
1384
1385 fn panicking(name: &'static str, delay: Duration) -> Self {
1387 Self {
1388 name,
1389 init_behavior: InitBehavior::Instant,
1390 run_behavior: RunBehavior::PanicAfter(delay),
1391 start_count: Arc::new(AtomicUsize::new(0)),
1392 finish_count: Arc::new(AtomicUsize::new(0)),
1393 brutal_shutdown: false,
1394 graceful_timeout: Duration::from_millis(500),
1395 }
1396 }
1397
1398 fn init_failure(name: &'static str) -> Self {
1400 Self {
1401 name,
1402 init_behavior: InitBehavior::Fail("init failed"),
1403 run_behavior: RunBehavior::UntilShutdown,
1404 start_count: Arc::new(AtomicUsize::new(0)),
1405 finish_count: Arc::new(AtomicUsize::new(0)),
1406 brutal_shutdown: false,
1407 graceful_timeout: Duration::from_millis(500),
1408 }
1409 }
1410
1411 fn slow_init(name: &'static str, init_delay: Duration) -> Self {
1413 Self {
1414 name,
1415 init_behavior: InitBehavior::Slow(init_delay),
1416 run_behavior: RunBehavior::UntilShutdown,
1417 start_count: Arc::new(AtomicUsize::new(0)),
1418 finish_count: Arc::new(AtomicUsize::new(0)),
1419 brutal_shutdown: false,
1420 graceful_timeout: Duration::from_millis(500),
1421 }
1422 }
1423
1424 fn start_count(&self) -> Arc<AtomicUsize> {
1430 Arc::clone(&self.start_count)
1431 }
1432
1433 fn finish_count(&self) -> Arc<AtomicUsize> {
1440 Arc::clone(&self.finish_count)
1441 }
1442
1443 fn with_brutal_shutdown(mut self) -> Self {
1445 self.brutal_shutdown = true;
1446 self
1447 }
1448
1449 fn with_graceful_timeout(mut self, timeout: Duration) -> Self {
1451 self.graceful_timeout = timeout;
1452 self
1453 }
1454 }
1455
1456 #[async_trait]
1457 impl Supervisable for MockWorker {
1458 fn name(&self) -> &str {
1459 self.name
1460 }
1461
1462 fn shutdown_strategy(&self) -> ShutdownStrategy {
1463 if self.brutal_shutdown {
1464 ShutdownStrategy::Brutal
1465 } else {
1466 ShutdownStrategy::Graceful(self.graceful_timeout)
1467 }
1468 }
1469
1470 async fn initialize(&self, process_shutdown: ShutdownHandle) -> Result<SupervisorFuture, InitializationError> {
1471 match &self.init_behavior {
1472 InitBehavior::Instant => {}
1473 InitBehavior::Slow(delay) => {
1474 sleep(*delay).await;
1475 }
1476 InitBehavior::Fail(msg) => {
1477 return Err(InitializationError::Failed {
1478 source: GenericError::msg(*msg),
1479 });
1480 }
1481 }
1482
1483 let start_count = Arc::clone(&self.start_count);
1484 let finish_count = Arc::clone(&self.finish_count);
1485 let run_behavior = self.run_behavior.clone();
1486
1487 Ok(Box::pin(async move {
1488 start_count.fetch_add(1, Ordering::SeqCst);
1489
1490 match run_behavior {
1491 RunBehavior::UntilShutdown => {
1492 process_shutdown.await;
1493 Ok(())
1494 }
1495 RunBehavior::FailAfter(delay, msg) => {
1496 select! {
1497 _ = sleep(delay) => {
1498 finish_count.fetch_add(1, Ordering::SeqCst);
1501 Err(GenericError::msg(msg))
1502 }
1503 _ = process_shutdown => {
1504 Ok(())
1505 }
1506 }
1507 }
1508 RunBehavior::CompleteAfter(delay) => {
1509 select! {
1510 _ = sleep(delay) => {
1511 finish_count.fetch_add(1, Ordering::SeqCst);
1513 Ok(())
1514 }
1515 _ = process_shutdown => Ok(()),
1516 }
1517 }
1518 RunBehavior::SlowShutdown(delay) => {
1519 process_shutdown.await;
1520 sleep(delay).await;
1521 Ok(())
1522 }
1523 RunBehavior::IgnoreShutdown => {
1524 let _hold = process_shutdown;
1526 pending().await
1527 }
1528 RunBehavior::PanicAfter(delay) => {
1529 select! {
1530 _ = sleep(delay) => panic!("worker panicked"),
1531 _ = process_shutdown => Ok(()),
1532 }
1533 }
1534 }
1535 }))
1536 }
1537 }
1538
1539 async fn run_supervisor_with_trigger(
1545 supervisor: Supervisor,
1546 ) -> (oneshot::Sender<()>, JoinHandle<Result<(), SupervisorError>>) {
1547 let sup_handle = supervisor.handle();
1549 let mut supervisor = supervisor;
1550
1551 let (tx, rx) = oneshot::channel();
1552 let handle = tokio::spawn(async move { supervisor.run_with_shutdown(rx).await });
1553
1554 wait_until("supervisor is running", || sup_handle.is_running()).await;
1555 (tx, handle)
1556 }
1557
1558 async fn join_supervisor(handle: JoinHandle<Result<(), SupervisorError>>) -> Result<(), SupervisorError> {
1563 timeout(Duration::from_secs(2), handle)
1564 .await
1565 .expect("supervisor should exit promptly")
1566 .expect("supervisor task should not panic")
1567 }
1568
1569 #[tokio::test]
1572 async fn standalone_supervisor_shuts_down_cleanly() {
1573 let mut sup = Supervisor::new("test-sup").unwrap();
1574 sup.add_worker(MockWorker::long_running("worker1"));
1575 sup.add_worker(MockWorker::long_running("worker2"));
1576
1577 let (tx, handle) = run_supervisor_with_trigger(sup).await;
1578 tx.send(()).unwrap();
1579
1580 let result = join_supervisor(handle).await;
1581 assert!(result.is_ok());
1582 }
1583
1584 #[tokio::test]
1585 async fn nested_supervisor_shuts_down_cleanly() {
1586 let mut child_sup = Supervisor::new("child-sup").unwrap();
1587 child_sup.add_worker(MockWorker::long_running("inner-worker"));
1588
1589 let mut parent_sup = Supervisor::new("parent-sup").unwrap();
1590 parent_sup.add_worker(MockWorker::long_running("outer-worker"));
1591 parent_sup.add_worker(child_sup);
1592
1593 let (tx, handle) = run_supervisor_with_trigger(parent_sup).await;
1594 tx.send(()).unwrap();
1595
1596 let result = join_supervisor(handle).await;
1597 assert!(result.is_ok());
1598 }
1599
1600 #[tokio::test]
1601 async fn empty_supervisor_idles_until_shutdown() {
1602 let sup = Supervisor::new("empty-sup").unwrap();
1605
1606 let (tx, handle) = run_supervisor_with_trigger(sup).await;
1607 assert!(!handle.is_finished(), "an empty supervisor must idle rather than exit");
1608
1609 tx.send(()).unwrap();
1610 let result = join_supervisor(handle).await;
1611 assert!(result.is_ok());
1612 }
1613
1614 #[tokio::test]
1617 async fn one_for_one_restarts_only_failed_child() {
1618 let failing = MockWorker::failing("failing-worker", Duration::from_millis(50));
1619 let failing_count = failing.start_count();
1620
1621 let stable = MockWorker::long_running("stable-worker");
1622 let stable_count = stable.start_count();
1623
1624 let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
1625 RestartStrategy::one_to_one().with_intensity_and_period(20, Duration::from_secs(10)),
1626 );
1627 sup.add_worker(stable);
1628 sup.add_worker(failing);
1629
1630 let (tx, handle) = run_supervisor_with_trigger(sup).await;
1631
1632 wait_until("the failing worker has been restarted", || {
1634 failing_count.load(Ordering::SeqCst) >= 2
1635 })
1636 .await;
1637 let _ = tx.send(());
1638
1639 let result = join_supervisor(handle).await;
1640 assert!(result.is_ok());
1641
1642 assert!(
1644 failing_count.load(Ordering::SeqCst) >= 2,
1645 "failing worker should have been restarted"
1646 );
1647 assert_eq!(
1649 stable_count.load(Ordering::SeqCst),
1650 1,
1651 "stable worker should not have been restarted"
1652 );
1653 }
1654
1655 #[tokio::test]
1656 async fn one_for_all_restarts_all_children() {
1657 let failing = MockWorker::failing("failing-worker", Duration::from_millis(50));
1658 let failing_count = failing.start_count();
1659
1660 let stable = MockWorker::long_running("stable-worker");
1661 let stable_count = stable.start_count();
1662
1663 let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
1664 RestartStrategy::one_for_all().with_intensity_and_period(20, Duration::from_secs(10)),
1665 );
1666 sup.add_worker(stable);
1667 sup.add_worker(failing);
1668
1669 let (tx, handle) = run_supervisor_with_trigger(sup).await;
1670
1671 wait_until("both workers have been restarted", || {
1673 failing_count.load(Ordering::SeqCst) >= 2 && stable_count.load(Ordering::SeqCst) >= 2
1674 })
1675 .await;
1676 let _ = tx.send(());
1677
1678 let result = join_supervisor(handle).await;
1679 assert!(result.is_ok());
1680
1681 assert!(
1683 failing_count.load(Ordering::SeqCst) >= 2,
1684 "failing worker should have been restarted"
1685 );
1686 assert!(
1687 stable_count.load(Ordering::SeqCst) >= 2,
1688 "stable worker should also have been restarted"
1689 );
1690 }
1691
1692 #[tokio::test]
1693 async fn one_for_all_does_not_restart_temporary_children() {
1694 let failing = MockWorker::failing("failing-worker", Duration::from_millis(50));
1697 let failing_count = failing.start_count();
1698
1699 let temp = MockWorker::long_running("temp-worker");
1700 let temp_count = temp.start_count();
1701
1702 let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
1703 RestartStrategy::one_for_all().with_intensity_and_period(20, Duration::from_secs(10)),
1704 );
1705 sup.add_worker(ChildSpecification::worker(temp).with_restart_type(RestartType::Temporary));
1706 sup.add_worker(failing);
1707
1708 let (tx, handle) = run_supervisor_with_trigger(sup).await;
1709
1710 wait_until("the permanent worker has been restarted", || {
1712 failing_count.load(Ordering::SeqCst) >= 2
1713 })
1714 .await;
1715 let _ = tx.send(());
1716
1717 let result = join_supervisor(handle).await;
1718 assert!(result.is_ok());
1719 assert!(
1720 failing_count.load(Ordering::SeqCst) >= 2,
1721 "permanent worker should have been restarted by one-for-all"
1722 );
1723 assert_eq!(
1724 temp_count.load(Ordering::SeqCst),
1725 1,
1726 "temporary child must not be restarted by a one-for-all group restart"
1727 );
1728 }
1729
1730 #[tokio::test]
1731 async fn one_for_all_restarts_transient_children() {
1732 let transient = MockWorker::completing("transient-worker", Duration::from_millis(30));
1735 let transient_count = transient.start_count();
1736
1737 let failing = MockWorker::failing("failing-worker", Duration::from_millis(80));
1739
1740 let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
1741 RestartStrategy::one_for_all().with_intensity_and_period(20, Duration::from_secs(10)),
1742 );
1743 sup.add_worker(ChildSpecification::worker(transient).with_restart_type(RestartType::Transient));
1744 sup.add_worker(failing);
1745
1746 let (tx, handle) = run_supervisor_with_trigger(sup).await;
1747
1748 wait_until("the transient worker has been restarted by the group", || {
1749 transient_count.load(Ordering::SeqCst) >= 2
1750 })
1751 .await;
1752 let _ = tx.send(());
1753
1754 let result = join_supervisor(handle).await;
1755 assert!(result.is_ok());
1756 assert!(
1757 transient_count.load(Ordering::SeqCst) >= 2,
1758 "transient child must be restarted by a one-for-all group restart, even after a clean exit"
1759 );
1760 }
1761
1762 #[tokio::test]
1763 async fn transient_abnormal_exit_triggers_one_for_all() {
1764 let transient = MockWorker::failing("transient-worker", Duration::from_millis(50));
1767 let transient_count = transient.start_count();
1768
1769 let stable = MockWorker::long_running("stable-worker");
1770 let stable_count = stable.start_count();
1771
1772 let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
1773 RestartStrategy::one_for_all().with_intensity_and_period(20, Duration::from_secs(10)),
1774 );
1775 sup.add_worker(ChildSpecification::worker(transient).with_restart_type(RestartType::Transient));
1776 sup.add_worker(stable);
1777
1778 let (tx, handle) = run_supervisor_with_trigger(sup).await;
1779
1780 wait_until("the abnormal exit has restarted both workers", || {
1781 transient_count.load(Ordering::SeqCst) >= 2 && stable_count.load(Ordering::SeqCst) >= 2
1782 })
1783 .await;
1784 let _ = tx.send(());
1785
1786 let result = join_supervisor(handle).await;
1787 assert!(result.is_ok());
1788 assert!(
1789 transient_count.load(Ordering::SeqCst) >= 2,
1790 "transient worker must be restarted after its own abnormal exit"
1791 );
1792 assert!(
1793 stable_count.load(Ordering::SeqCst) >= 2,
1794 "the transient's abnormal exit must trigger a one-for-all that also restarts the sibling"
1795 );
1796 }
1797
1798 #[tokio::test]
1799 async fn restart_limit_exceeded_shuts_down_supervisor() {
1800 let mut sup = Supervisor::new("test-sup")
1801 .unwrap()
1802 .with_restart_strategy(RestartStrategy::one_to_one().with_intensity_and_period(1, Duration::from_secs(10)));
1803 sup.add_worker(MockWorker::failing("fast-fail", Duration::ZERO));
1805
1806 let (tx, rx) = oneshot::channel::<()>();
1807 let handle = tokio::spawn(async move { sup.run_with_shutdown(rx).await });
1808
1809 let result = join_supervisor(handle).await;
1810 drop(tx);
1811
1812 assert!(matches!(result, Err(SupervisorError::Shutdown)));
1813 }
1814
1815 #[tokio::test]
1818 async fn temporary_child_is_not_restarted() {
1819 let temp = MockWorker::failing("temp-worker", Duration::from_millis(50));
1821 let temp_started = temp.start_count();
1822 let temp_failed = temp.finish_count();
1823
1824 let stable = MockWorker::long_running("stable-worker");
1825
1826 let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
1827 RestartStrategy::one_to_one().with_intensity_and_period(20, Duration::from_secs(10)),
1828 );
1829 sup.add_worker(stable);
1830 sup.add_worker(ChildSpecification::worker(temp).with_restart_type(RestartType::Temporary));
1831
1832 let (tx, handle) = run_supervisor_with_trigger(sup).await;
1833
1834 wait_until("the temporary worker has failed once", || {
1840 temp_failed.load(Ordering::SeqCst) == 1
1841 })
1842 .await;
1843 let _ = tx.send(());
1844
1845 let result = join_supervisor(handle).await;
1846 assert!(result.is_ok());
1847 assert_eq!(
1848 temp_started.load(Ordering::SeqCst),
1849 1,
1850 "temporary worker must not be restarted after it fails"
1851 );
1852 }
1853
1854 #[tokio::test]
1855 async fn transient_child_is_not_restarted_on_clean_exit() {
1856 let transient = MockWorker::completing("transient-worker", Duration::from_millis(50));
1857 let transient_started = transient.start_count();
1858 let transient_finished = transient.finish_count();
1859
1860 let stable = MockWorker::long_running("stable-worker");
1861
1862 let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
1863 RestartStrategy::one_to_one().with_intensity_and_period(20, Duration::from_secs(10)),
1864 );
1865 sup.add_worker(stable);
1866 sup.add_worker(ChildSpecification::worker(transient).with_restart_type(RestartType::Transient));
1867
1868 let (tx, handle) = run_supervisor_with_trigger(sup).await;
1869
1870 wait_until("the transient worker has completed once", || {
1875 transient_finished.load(Ordering::SeqCst) == 1
1876 })
1877 .await;
1878 let _ = tx.send(());
1879
1880 let result = join_supervisor(handle).await;
1881 assert!(result.is_ok());
1882 assert_eq!(
1883 transient_started.load(Ordering::SeqCst),
1884 1,
1885 "transient worker must not be restarted after a clean exit"
1886 );
1887 }
1888
1889 #[tokio::test]
1890 async fn transient_child_is_restarted_on_failure() {
1891 let transient = MockWorker::failing("transient-worker", Duration::from_millis(50));
1892 let transient_count = transient.start_count();
1893
1894 let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
1895 RestartStrategy::one_to_one().with_intensity_and_period(20, Duration::from_secs(10)),
1896 );
1897 sup.add_worker(ChildSpecification::worker(transient).with_restart_type(RestartType::Transient));
1898
1899 let (tx, handle) = run_supervisor_with_trigger(sup).await;
1900
1901 wait_until("the transient worker has been restarted", || {
1902 transient_count.load(Ordering::SeqCst) >= 2
1903 })
1904 .await;
1905 let _ = tx.send(());
1906
1907 let result = join_supervisor(handle).await;
1908 assert!(result.is_ok());
1909 assert!(
1910 transient_count.load(Ordering::SeqCst) >= 2,
1911 "transient worker must be restarted after an abnormal exit"
1912 );
1913 }
1914
1915 #[tokio::test]
1916 async fn permanent_child_is_restarted_on_clean_exit() {
1917 let permanent = MockWorker::completing("permanent-worker", Duration::from_millis(50));
1920 let permanent_count = permanent.start_count();
1921
1922 let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
1923 RestartStrategy::one_to_one().with_intensity_and_period(20, Duration::from_secs(10)),
1924 );
1925 sup.add_worker(permanent);
1927
1928 let (tx, handle) = run_supervisor_with_trigger(sup).await;
1929
1930 wait_until("the permanent worker has been restarted", || {
1931 permanent_count.load(Ordering::SeqCst) >= 2
1932 })
1933 .await;
1934 let _ = tx.send(());
1935
1936 let result = join_supervisor(handle).await;
1937 assert!(result.is_ok());
1938 assert!(
1939 permanent_count.load(Ordering::SeqCst) >= 2,
1940 "permanent worker must be restarted even after a clean exit"
1941 );
1942 }
1943
1944 #[tokio::test]
1945 async fn temporary_failures_do_not_consume_restart_intensity() {
1946 let mut sup = Supervisor::new("test-sup")
1950 .unwrap()
1951 .with_restart_strategy(RestartStrategy::one_to_one().with_intensity_and_period(1, Duration::from_secs(10)));
1952
1953 let workers = [
1954 MockWorker::failing("temp-0", Duration::from_millis(20)),
1955 MockWorker::failing("temp-1", Duration::from_millis(20)),
1956 MockWorker::failing("temp-2", Duration::from_millis(20)),
1957 MockWorker::failing("temp-3", Duration::from_millis(20)),
1958 MockWorker::failing("temp-4", Duration::from_millis(20)),
1959 ];
1960 let started: Vec<_> = workers.iter().map(|w| w.start_count()).collect();
1961 let failed: Vec<_> = workers.iter().map(|w| w.finish_count()).collect();
1962 for worker in workers {
1963 sup.add_worker(ChildSpecification::worker(worker).with_restart_type(RestartType::Temporary));
1964 }
1965 sup.add_worker(MockWorker::long_running("stable-worker"));
1967
1968 let (tx, handle) = run_supervisor_with_trigger(sup).await;
1969 wait_until("every temporary worker has failed once", || {
1973 failed.iter().all(|c| c.load(Ordering::SeqCst) == 1)
1974 })
1975 .await;
1976 let _ = tx.send(());
1977
1978 let result = join_supervisor(handle).await;
1979 assert!(
1980 result.is_ok(),
1981 "supervisor must not trip its restart limit on temporary exits"
1982 );
1983 for count in started {
1984 assert_eq!(
1985 count.load(Ordering::SeqCst),
1986 1,
1987 "each temporary worker runs exactly once"
1988 );
1989 }
1990 }
1991
1992 #[tokio::test]
1993 async fn transient_clean_exits_do_not_consume_restart_intensity() {
1994 let mut sup = Supervisor::new("test-sup")
1998 .unwrap()
1999 .with_restart_strategy(RestartStrategy::one_to_one().with_intensity_and_period(1, Duration::from_secs(10)));
2000
2001 let workers = [
2002 MockWorker::completing("transient-0", Duration::from_millis(20)),
2003 MockWorker::completing("transient-1", Duration::from_millis(20)),
2004 MockWorker::completing("transient-2", Duration::from_millis(20)),
2005 MockWorker::completing("transient-3", Duration::from_millis(20)),
2006 MockWorker::completing("transient-4", Duration::from_millis(20)),
2007 ];
2008 let started: Vec<_> = workers.iter().map(|w| w.start_count()).collect();
2009 let finished: Vec<_> = workers.iter().map(|w| w.finish_count()).collect();
2010 for worker in workers {
2011 sup.add_worker(ChildSpecification::worker(worker).with_restart_type(RestartType::Transient));
2012 }
2013 sup.add_worker(MockWorker::long_running("stable-worker"));
2015
2016 let (tx, handle) = run_supervisor_with_trigger(sup).await;
2017 wait_until("every transient worker has completed once", || {
2021 finished.iter().all(|c| c.load(Ordering::SeqCst) == 1)
2022 })
2023 .await;
2024 let _ = tx.send(());
2025
2026 let result = join_supervisor(handle).await;
2027 assert!(
2028 result.is_ok(),
2029 "supervisor must not trip its restart limit on clean transient exits"
2030 );
2031 for count in started {
2032 assert_eq!(
2033 count.load(Ordering::SeqCst),
2034 1,
2035 "each transient worker runs exactly once"
2036 );
2037 }
2038 }
2039
2040 #[tokio::test]
2041 async fn supervisor_idles_when_all_temporary_children_exit() {
2042 let temp_a = MockWorker::completing("temp-a", Duration::from_millis(10));
2046 let a_finished = temp_a.finish_count();
2047 let temp_b = MockWorker::completing("temp-b", Duration::from_millis(10));
2048 let b_finished = temp_b.finish_count();
2049
2050 let mut sup = Supervisor::new("test-sup").unwrap();
2051 let handle = sup.handle();
2052 sup.add_worker(ChildSpecification::worker(temp_a).with_restart_type(RestartType::Temporary));
2053 sup.add_worker(ChildSpecification::worker(temp_b).with_restart_type(RestartType::Temporary));
2054
2055 let (tx, run) = run_supervisor_with_trigger(sup).await;
2056
2057 wait_until("both temporary children have completed", || {
2061 a_finished.load(Ordering::SeqCst) == 1 && b_finished.load(Ordering::SeqCst) == 1
2062 })
2063 .await;
2064
2065 let dynamic = MockWorker::long_running("late-comer");
2068 let dynamic_count = dynamic.start_count();
2069 handle
2070 .spawn(dynamic)
2071 .await
2072 .expect("supervisor must still accept work after its children drain");
2073 wait_until("the late dynamic child has started", || {
2074 dynamic_count.load(Ordering::SeqCst) == 1
2075 })
2076 .await;
2077 assert!(
2078 handle.is_running(),
2079 "supervisor must keep running after all temporary children exit"
2080 );
2081
2082 tx.send(()).unwrap();
2083 let result = join_supervisor(run).await;
2084 assert!(result.is_ok());
2085 }
2086
2087 #[tokio::test]
2090 async fn significant_child_drives_auto_shutdown() {
2091 let mut sup = Supervisor::new("test-sup")
2094 .unwrap()
2095 .with_auto_shutdown(AutoShutdown::AnySignificant);
2096 sup.add_worker(MockWorker::long_running("stable"));
2097 sup.add_worker(
2098 ChildSpecification::worker(MockWorker::completing("significant", Duration::from_millis(50)))
2099 .with_restart_type(RestartType::Temporary)
2100 .with_significant(true),
2101 );
2102
2103 let (_tx, rx) = oneshot::channel::<()>();
2105 let result = timeout(Duration::from_secs(2), sup.run_with_shutdown(rx))
2106 .await
2107 .unwrap();
2108 assert!(matches!(result, Err(SupervisorError::SignificantChildExited)));
2109 }
2110
2111 #[tokio::test]
2112 async fn non_significant_exit_does_not_auto_shutdown() {
2113 let plain = MockWorker::completing("plain", Duration::from_millis(10));
2115 let plain_finished = plain.finish_count();
2116
2117 let mut sup = Supervisor::new("test-sup")
2118 .unwrap()
2119 .with_auto_shutdown(AutoShutdown::AnySignificant);
2120 let handle = sup.handle();
2121 sup.add_worker(MockWorker::long_running("stable"));
2122 sup.add_worker(ChildSpecification::worker(plain).with_restart_type(RestartType::Temporary));
2123
2124 let (tx, run) = run_supervisor_with_trigger(sup).await;
2125
2126 wait_until("the non-significant child has completed", || {
2130 plain_finished.load(Ordering::SeqCst) == 1
2131 })
2132 .await;
2133
2134 let dynamic = MockWorker::long_running("late-comer");
2138 let dynamic_count = dynamic.start_count();
2139 handle
2140 .spawn(dynamic)
2141 .await
2142 .expect("supervisor must still accept work after a non-significant child exits");
2143 wait_until("the late dynamic child has started", || {
2144 dynamic_count.load(Ordering::SeqCst) == 1
2145 })
2146 .await;
2147 assert!(
2148 handle.is_running(),
2149 "a non-significant child exiting must not trigger auto-shutdown"
2150 );
2151
2152 tx.send(()).unwrap();
2153 let result = join_supervisor(run).await;
2154 assert!(result.is_ok());
2155 }
2156
2157 #[tokio::test]
2158 async fn all_significant_waits_for_last() {
2159 let mut sup = Supervisor::new("test-sup")
2161 .unwrap()
2162 .with_auto_shutdown(AutoShutdown::AllSignificant);
2163 sup.add_worker(
2164 ChildSpecification::worker(MockWorker::completing("sig-a", Duration::from_millis(50)))
2165 .with_restart_type(RestartType::Temporary)
2166 .with_significant(true),
2167 );
2168 sup.add_worker(
2169 ChildSpecification::worker(MockWorker::completing("sig-b", Duration::from_millis(250)))
2170 .with_restart_type(RestartType::Temporary)
2171 .with_significant(true),
2172 );
2173
2174 let (_tx, rx) = oneshot::channel::<()>();
2175 let start = std::time::Instant::now();
2176 let result = timeout(Duration::from_secs(2), sup.run_with_shutdown(rx))
2177 .await
2178 .unwrap();
2179 let elapsed = start.elapsed();
2180
2181 assert!(matches!(result, Err(SupervisorError::SignificantChildExited)));
2182 assert!(
2184 elapsed >= Duration::from_millis(200),
2185 "auto-shutdown must wait for all significant children (took {elapsed:?})"
2186 );
2187 }
2188
2189 #[tokio::test]
2192 async fn init_failure_propagates_with_child_name() {
2193 let mut sup = Supervisor::new("test-sup").unwrap();
2194 sup.add_worker(MockWorker::long_running("good-worker"));
2195 sup.add_worker(MockWorker::init_failure("bad-worker"));
2196
2197 let (_tx, rx) = oneshot::channel::<()>();
2198 let result = timeout(Duration::from_secs(2), sup.run_with_shutdown(rx))
2199 .await
2200 .unwrap();
2201
2202 match result {
2203 Err(SupervisorError::FailedToInitialize { child_name, .. }) => {
2204 assert_eq!(child_name, "bad-worker");
2205 }
2206 other => panic!("expected FailedToInitialize, got: {:?}", other),
2207 }
2208 }
2209
2210 #[tokio::test]
2211 async fn init_failure_does_not_trigger_restart() {
2212 let init_fail = MockWorker::init_failure("bad-worker");
2213 let start_count = init_fail.start_count();
2214
2215 let mut sup = Supervisor::new("test-sup").unwrap().with_restart_strategy(
2216 RestartStrategy::one_to_one().with_intensity_and_period(10, Duration::from_secs(10)),
2217 );
2218 sup.add_worker(init_fail);
2219
2220 let (_tx, rx) = oneshot::channel::<()>();
2221 let result = timeout(Duration::from_secs(2), sup.run_with_shutdown(rx))
2222 .await
2223 .unwrap();
2224
2225 assert!(matches!(result, Err(SupervisorError::FailedToInitialize { .. })));
2226 assert_eq!(start_count.load(Ordering::SeqCst), 0);
2228 }
2229
2230 #[tokio::test]
2233 async fn shutdown_completes_promptly_in_steady_state() {
2234 let mut sup = Supervisor::new("test-sup").unwrap();
2235 sup.add_worker(MockWorker::long_running("worker1"));
2236 sup.add_worker(MockWorker::long_running("worker2"));
2237
2238 let (tx, handle) = run_supervisor_with_trigger(sup).await;
2239 tx.send(()).unwrap();
2240
2241 let result = timeout(Duration::from_secs(1), handle).await;
2243 assert!(result.is_ok(), "shutdown should complete promptly");
2244 }
2245
2246 #[tokio::test]
2247 async fn shutdown_during_slow_init_completes_promptly() {
2248 let mut sup = Supervisor::new("test-sup").unwrap();
2249 sup.add_worker(MockWorker::slow_init("slow-worker", Duration::from_secs(30)));
2251
2252 let (tx, rx) = oneshot::channel();
2253 let handle = tokio::spawn(async move { sup.run_with_shutdown(rx).await });
2254
2255 sleep(Duration::from_millis(20)).await;
2257 tx.send(()).unwrap();
2258
2259 let result = timeout(Duration::from_secs(2), handle).await;
2262 assert!(result.is_ok(), "shutdown during slow init should complete promptly");
2263 }
2264
2265 #[tokio::test]
2268 async fn dynamic_children_spawn_after_start() {
2269 let sup = Supervisor::new("dyn-sup").unwrap();
2270 let handle = sup.handle();
2271 let (tx, run) = run_supervisor_with_trigger(sup).await;
2272 wait_until("supervisor is running", || handle.is_running()).await;
2273
2274 let c1 = MockWorker::long_running("c1");
2275 let c2 = MockWorker::long_running("c2");
2276 let c1_count = c1.start_count();
2277 let c2_count = c2.start_count();
2278 handle.spawn(c1).await.unwrap();
2279 handle.spawn(c2).await.unwrap();
2280
2281 wait_until("both dynamic children have started", || {
2282 c1_count.load(Ordering::SeqCst) == 1 && c2_count.load(Ordering::SeqCst) == 1
2283 })
2284 .await;
2285 assert_eq!(handle.active_children(), 2);
2286
2287 tx.send(()).unwrap();
2288 let result = join_supervisor(run).await;
2289 assert!(result.is_ok());
2290 assert_eq!(
2291 handle.active_children(),
2292 0,
2293 "all dynamic children must be drained on shutdown"
2294 );
2295 }
2296
2297 #[tokio::test]
2298 async fn temporary_dynamic_child_failure_is_isolated() {
2299 let sup = Supervisor::new("dyn-sup").unwrap();
2302 let handle = sup.handle();
2303 let (tx, run) = run_supervisor_with_trigger(sup).await;
2304 wait_until("supervisor is running", || handle.is_running()).await;
2305
2306 let failing = MockWorker::failing("boom", Duration::from_millis(20));
2307 let failing_count = failing.start_count();
2308 handle.spawn(failing).await.unwrap();
2309 wait_until("the failing dynamic child has run once", || {
2310 failing_count.load(Ordering::SeqCst) == 1
2311 })
2312 .await;
2313 wait_until("all dynamic children have drained", || handle.active_children() == 0).await;
2314
2315 sleep(Duration::from_millis(50)).await;
2316 assert!(
2317 handle.is_running(),
2318 "supervisor stays up after an isolated child failure"
2319 );
2320 assert_eq!(
2321 failing_count.load(Ordering::SeqCst),
2322 1,
2323 "a temporary child is never restarted"
2324 );
2325
2326 handle.spawn(MockWorker::long_running("c2")).await.unwrap();
2328 wait_until("one dynamic child is running", || handle.active_children() == 1).await;
2329
2330 tx.send(()).unwrap();
2331 let result = join_supervisor(run).await;
2332 assert!(result.is_ok());
2333 }
2334
2335 #[tokio::test]
2336 async fn temporary_dynamic_child_panic_is_isolated() {
2337 let sup = Supervisor::new("dyn-sup").unwrap();
2339 let handle = sup.handle();
2340 let (tx, run) = run_supervisor_with_trigger(sup).await;
2341 wait_until("supervisor is running", || handle.is_running()).await;
2342
2343 handle
2344 .spawn(MockWorker::panicking("boom", Duration::from_millis(20)))
2345 .await
2346 .unwrap();
2347 wait_until("all dynamic children have drained", || handle.active_children() == 0).await;
2348
2349 sleep(Duration::from_millis(50)).await;
2350 assert!(handle.is_running(), "supervisor stays up after an isolated child panic");
2351
2352 tx.send(()).unwrap();
2353 let result = join_supervisor(run).await;
2354 assert!(result.is_ok());
2355 }
2356
2357 #[tokio::test]
2358 async fn significant_dynamic_child_failure_shuts_down_supervisor() {
2359 let sup = Supervisor::new("dyn-sup")
2362 .unwrap()
2363 .with_auto_shutdown(AutoShutdown::AnySignificant);
2364 let handle = sup.handle();
2365 let (_tx, run) = run_supervisor_with_trigger(sup).await;
2366 wait_until("supervisor is running", || handle.is_running()).await;
2367
2368 handle
2369 .spawn_with(
2370 ChildSpecification::worker(MockWorker::failing("boom", Duration::from_millis(20)))
2371 .with_restart_type(RestartType::Temporary)
2372 .with_significant(true),
2373 )
2374 .await
2375 .unwrap();
2376
2377 let result = join_supervisor(run).await;
2378 assert!(matches!(result, Err(SupervisorError::SignificantChildExited)));
2379 }
2380
2381 #[tokio::test]
2382 async fn dynamic_spawn_fails_before_start_and_after_shutdown() {
2383 let sup = Supervisor::new("dyn-sup").unwrap();
2384 let handle = sup.handle();
2385
2386 assert!(!handle.is_running());
2389 let err = handle
2390 .spawn(MockWorker::long_running("before-start"))
2391 .await
2392 .unwrap_err();
2393 assert!(matches!(err, SpawnError::SupervisorGone));
2394
2395 let (tx, run) = run_supervisor_with_trigger(sup).await;
2397 wait_until("supervisor is running", || handle.is_running()).await;
2398 let worker = MockWorker::long_running("after-start");
2399 let started = worker.start_count();
2400 handle.spawn(worker).await.unwrap();
2401 wait_until("the dynamic child has started", || started.load(Ordering::SeqCst) == 1).await;
2402
2403 tx.send(()).unwrap();
2404 let result = join_supervisor(run).await;
2405 assert!(result.is_ok());
2406
2407 wait_until("the supervisor has stopped", || !handle.is_running()).await;
2409 let err = handle
2410 .spawn(MockWorker::long_running("after-shutdown"))
2411 .await
2412 .unwrap_err();
2413 assert!(matches!(err, SpawnError::SupervisorGone));
2414 }
2415
2416 #[tokio::test]
2417 async fn dynamic_spawn_returns_after_registration() {
2418 let sup = Supervisor::new("dyn-sup").unwrap();
2419 let handle = sup.handle();
2420 let (tx, run) = run_supervisor_with_trigger(sup).await;
2421 wait_until("supervisor is running", || handle.is_running()).await;
2422
2423 let worker = MockWorker::long_running("c");
2424 let started = worker.start_count();
2425 let id = handle.spawn(worker).await.unwrap();
2426 assert_eq!(id.as_u64(), 0);
2428 wait_until("the dynamic child has started", || started.load(Ordering::SeqCst) == 1).await;
2429
2430 tx.send(()).unwrap();
2431 let result = join_supervisor(run).await;
2432 assert!(result.is_ok());
2433 }
2434
2435 #[tokio::test]
2436 async fn dynamic_spawn_rejects_invalid_child_name() {
2437 let sup = Supervisor::new("dyn-sup").unwrap();
2440 let handle = sup.handle();
2441 let (tx, run) = run_supervisor_with_trigger(sup).await;
2442 wait_until("supervisor is running", || handle.is_running()).await;
2443
2444 let err = handle.spawn(MockWorker::long_running("")).await.unwrap_err();
2445 assert!(matches!(err, SpawnError::Rejected { .. }), "got {err:?}");
2446
2447 assert!(handle.is_running());
2449 handle.spawn(MockWorker::long_running("ok")).await.unwrap();
2450 wait_until("one dynamic child is running", || handle.active_children() == 1).await;
2451
2452 tx.send(()).unwrap();
2453 let result = join_supervisor(run).await;
2454 assert!(result.is_ok());
2455 }
2456
2457 #[tokio::test]
2458 async fn concurrent_shutdown_drains_many_children_quickly() {
2459 const CHILDREN: usize = 500;
2460 const SHUTDOWN_DELAY: Duration = Duration::from_millis(50);
2461
2462 let sup = Supervisor::new("dyn-sup")
2463 .unwrap()
2464 .with_shutdown_mode(ShutdownMode::Concurrent);
2465 let handle = sup.handle();
2466 let (tx, run) = run_supervisor_with_trigger(sup).await;
2467 wait_until("supervisor is running", || handle.is_running()).await;
2468
2469 for _ in 0..CHILDREN {
2470 handle
2471 .spawn(MockWorker::slow_shutdown("conn", SHUTDOWN_DELAY))
2472 .await
2473 .unwrap();
2474 }
2475 wait_until("all dynamic children are running", || {
2476 handle.active_children() == CHILDREN
2477 })
2478 .await;
2479
2480 let start = std::time::Instant::now();
2483 tx.send(()).unwrap();
2484 let result = timeout(Duration::from_secs(5), run).await.unwrap().unwrap();
2485 let elapsed = start.elapsed();
2486
2487 assert!(result.is_ok());
2488 assert_eq!(handle.active_children(), 0, "active count must return to zero");
2489 assert!(
2490 elapsed < Duration::from_secs(2),
2491 "shutdown must be concurrent (took {elapsed:?})"
2492 );
2493 }
2494
2495 #[tokio::test]
2496 async fn concurrent_shutdown_aborts_unresponsive_children() {
2497 let sup = Supervisor::new("dyn-sup")
2498 .unwrap()
2499 .with_shutdown_mode(ShutdownMode::Concurrent);
2500 let handle = sup.handle();
2501 let (tx, run) = run_supervisor_with_trigger(sup).await;
2502 wait_until("supervisor is running", || handle.is_running()).await;
2503
2504 handle.spawn(MockWorker::ignore_shutdown("stuck")).await.unwrap();
2505 wait_until("one dynamic child is running", || handle.active_children() == 1).await;
2506
2507 let start = std::time::Instant::now();
2510 tx.send(()).unwrap();
2511 let result = join_supervisor(run).await;
2512 let elapsed = start.elapsed();
2513
2514 assert!(
2516 matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 1 })),
2517 "aborting a stuck child must surface as an unclean shutdown, got {result:?}"
2518 );
2519 assert_eq!(handle.active_children(), 0);
2520 assert!(
2521 elapsed < Duration::from_secs(1),
2522 "stuck child must be aborted at the deadline (took {elapsed:?})"
2523 );
2524 }
2525
2526 #[tokio::test]
2527 async fn concurrent_shutdown_honors_per_child_deadline() {
2528 let sup = Supervisor::new("dyn-sup")
2533 .unwrap()
2534 .with_shutdown_mode(ShutdownMode::Concurrent);
2535 let handle = sup.handle();
2536 let (tx, run) = run_supervisor_with_trigger(sup).await;
2537 wait_until("supervisor is running", || handle.is_running()).await;
2538
2539 handle
2541 .spawn(MockWorker::long_running("responsive").with_graceful_timeout(Duration::MAX))
2542 .await
2543 .unwrap();
2544 handle
2546 .spawn(MockWorker::ignore_shutdown("stuck").with_graceful_timeout(Duration::from_millis(200)))
2547 .await
2548 .unwrap();
2549 wait_until("both dynamic children are running", || handle.active_children() == 2).await;
2550
2551 let start = std::time::Instant::now();
2552 tx.send(()).unwrap();
2553 let result = join_supervisor(run).await;
2554 let elapsed = start.elapsed();
2555
2556 assert!(
2558 matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 1 })),
2559 "aborting the stuck child must surface as an unclean shutdown with a count of 1, got {result:?}"
2560 );
2561 assert_eq!(handle.active_children(), 0);
2562 assert!(
2563 elapsed < Duration::from_secs(1),
2564 "stuck child must be aborted at its own deadline despite an infinite-timeout sibling (took {elapsed:?})"
2565 );
2566 }
2567
2568 #[tokio::test]
2569 async fn ordered_shutdown_aborts_unresponsive_child() {
2570 let mut sup = Supervisor::new("test-sup").unwrap();
2573 sup.add_worker(MockWorker::ignore_shutdown("stuck"));
2574
2575 let (tx, handle) = run_supervisor_with_trigger(sup).await;
2576
2577 let start = std::time::Instant::now();
2578 tx.send(()).unwrap();
2579 let result = join_supervisor(handle).await;
2580 let elapsed = start.elapsed();
2581
2582 assert!(
2583 matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 1 })),
2584 "aborting a stuck child under ordered shutdown must surface as an unclean shutdown, got {result:?}"
2585 );
2586 assert!(
2587 elapsed < Duration::from_secs(1),
2588 "unresponsive child must be aborted at its deadline under ordered shutdown (took {elapsed:?})"
2589 );
2590 }
2591
2592 #[tokio::test]
2593 async fn brutal_shutdown_aborts_child_immediately() {
2594 let mut sup = Supervisor::new("test-sup").unwrap();
2597 sup.add_worker(MockWorker::ignore_shutdown("brutal-stuck").with_brutal_shutdown());
2598
2599 let (tx, handle) = run_supervisor_with_trigger(sup).await;
2600
2601 let start = std::time::Instant::now();
2602 tx.send(()).unwrap();
2603 let result = join_supervisor(handle).await;
2604 let elapsed = start.elapsed();
2605
2606 assert!(result.is_ok());
2609 assert!(
2610 elapsed < Duration::from_millis(200),
2611 "brutal-shutdown child must be aborted immediately, not after a graceful wait (took {elapsed:?})"
2612 );
2613 }
2614
2615 #[tokio::test]
2616 async fn shutdown_timeout_aborts_aggregate_to_root() {
2617 let mut child_sup = Supervisor::new("child-sup")
2621 .unwrap()
2622 .with_shutdown_mode(ShutdownMode::Concurrent);
2623 child_sup
2624 .add_worker(MockWorker::ignore_shutdown("child-stuck").with_graceful_timeout(Duration::from_millis(200)));
2625
2626 let mut parent_sup = Supervisor::new("parent-sup")
2627 .unwrap()
2628 .with_shutdown_mode(ShutdownMode::Concurrent);
2629 parent_sup
2630 .add_worker(MockWorker::ignore_shutdown("parent-stuck").with_graceful_timeout(Duration::from_millis(200)));
2631 parent_sup.add_worker(MockWorker::long_running("parent-clean"));
2632 parent_sup.add_worker(child_sup);
2633
2634 let (tx, handle) = run_supervisor_with_trigger(parent_sup).await;
2635 tx.send(()).unwrap();
2636
2637 let result = join_supervisor(handle).await;
2638 assert!(
2639 matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 2 })),
2640 "forced aborts must aggregate across the tree (1 direct + 1 nested), got {result:?}"
2641 );
2642 }
2643
2644 #[tokio::test]
2647 async fn restart_intensity_zero_shuts_down_on_first_failure() {
2648 let worker = MockWorker::failing("boom", Duration::from_millis(20));
2652 let start_count = worker.start_count();
2653
2654 let mut sup = Supervisor::new("test-sup")
2655 .unwrap()
2656 .with_restart_strategy(RestartStrategy::new(RestartMode::OneForOne, 0, Duration::from_secs(5)));
2657 sup.add_worker(worker);
2658
2659 let (_tx, rx) = oneshot::channel::<()>();
2660 let result = timeout(Duration::from_secs(2), sup.run_with_shutdown(rx))
2661 .await
2662 .unwrap();
2663
2664 assert!(matches!(result, Err(SupervisorError::Shutdown)));
2665 assert_eq!(
2666 start_count.load(Ordering::SeqCst),
2667 1,
2668 "with intensity zero the worker must run exactly once and never be restarted"
2669 );
2670 }
2671
2672 #[tokio::test]
2673 async fn one_for_all_restart_loses_dynamic_children() {
2674 let failing = MockWorker::failing("failing-static", Duration::from_millis(50));
2679 let failing_count = failing.start_count();
2680
2681 let sup = Supervisor::new("dyn-sup").unwrap().with_restart_strategy(
2682 RestartStrategy::one_for_all().with_intensity_and_period(20, Duration::from_secs(10)),
2683 );
2684 let handle = sup.handle();
2685 let mut sup = sup;
2686 sup.add_worker(failing);
2687
2688 let (tx, run) = run_supervisor_with_trigger(sup).await;
2689
2690 let dynamic = MockWorker::long_running("dynamic");
2692 let dynamic_count = dynamic.start_count();
2693 handle.spawn(dynamic).await.expect("should spawn dynamic child");
2694 wait_until("the dynamic child is running", || handle.active_children() == 1).await;
2695
2696 wait_until("the static worker has been restarted", || {
2698 failing_count.load(Ordering::SeqCst) >= 2
2699 })
2700 .await;
2701
2702 wait_until("the dynamic child has been discarded", || handle.active_children() == 0).await;
2705 assert_eq!(
2706 dynamic_count.load(Ordering::SeqCst),
2707 1,
2708 "a dynamic child must be lost -- not restored -- across a one-for-all restart"
2709 );
2710
2711 tx.send(()).unwrap();
2712 let result = join_supervisor(run).await;
2713 assert!(result.is_ok());
2714 }
2715
2716 #[tokio::test]
2719 async fn dedicated_single_threaded_runtime_runs_nested_worker_and_shuts_down_cleanly() {
2720 let worker = MockWorker::long_running("dedicated-worker");
2724 let worker_count = worker.start_count();
2725
2726 let mut child_sup = Supervisor::new("child-sup")
2727 .unwrap()
2728 .with_dedicated_runtime(RuntimeConfiguration::single_threaded());
2729 child_sup.add_worker(worker);
2730
2731 let mut parent_sup = Supervisor::new("parent-sup").unwrap();
2732 parent_sup.add_worker(child_sup);
2733
2734 let (tx, handle) = run_supervisor_with_trigger(parent_sup).await;
2735
2736 wait_until("the dedicated worker has started", || {
2738 worker_count.load(Ordering::SeqCst) == 1
2739 })
2740 .await;
2741
2742 tx.send(()).unwrap();
2743 let result = join_supervisor(handle).await;
2744 assert!(
2745 result.is_ok(),
2746 "dedicated-runtime supervisor should shut down cleanly, got {result:?}"
2747 );
2748 }
2749
2750 #[tokio::test]
2751 async fn dedicated_multi_threaded_runtime_runs_nested_worker() {
2752 let worker = MockWorker::long_running("dedicated-worker");
2754 let worker_count = worker.start_count();
2755
2756 let mut child_sup = Supervisor::new("child-sup")
2757 .unwrap()
2758 .with_dedicated_runtime(RuntimeConfiguration::multi_threaded(2));
2759 child_sup.add_worker(worker);
2760
2761 let mut parent_sup = Supervisor::new("parent-sup").unwrap();
2762 parent_sup.add_worker(child_sup);
2763
2764 let (tx, handle) = run_supervisor_with_trigger(parent_sup).await;
2765 wait_until("the dedicated worker has started", || {
2766 worker_count.load(Ordering::SeqCst) == 1
2767 })
2768 .await;
2769
2770 tx.send(()).unwrap();
2771 let result = join_supervisor(handle).await;
2772 assert!(
2773 result.is_ok(),
2774 "multi-threaded dedicated-runtime supervisor should shut down cleanly, got {result:?}"
2775 );
2776 }
2777
2778 #[tokio::test]
2779 async fn dedicated_runtime_forced_abort_aggregates_to_root() {
2780 let stuck = MockWorker::ignore_shutdown("stuck").with_graceful_timeout(Duration::from_millis(200));
2784 let stuck_count = stuck.start_count();
2785
2786 let mut child_sup = Supervisor::new("child-sup")
2787 .unwrap()
2788 .with_dedicated_runtime(RuntimeConfiguration::single_threaded())
2789 .with_shutdown_mode(ShutdownMode::Concurrent);
2790 child_sup.add_worker(stuck);
2791
2792 let mut parent_sup = Supervisor::new("parent-sup").unwrap();
2793 parent_sup.add_worker(child_sup);
2794
2795 let (tx, handle) = run_supervisor_with_trigger(parent_sup).await;
2796
2797 wait_until("the stuck worker has started", || {
2800 stuck_count.load(Ordering::SeqCst) == 1
2801 })
2802 .await;
2803
2804 tx.send(()).unwrap();
2805 let result = join_supervisor(handle).await;
2806 assert!(
2807 matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 1 })),
2808 "a stuck worker in a dedicated runtime must surface as an unclean shutdown aggregated to the root, got {result:?}"
2809 );
2810 }
2811
2812 #[tokio::test]
2815 async fn child_with_runtime_override_runs_on_that_runtime() {
2816 let child_runtime = tokio::runtime::Builder::new_multi_thread()
2819 .worker_threads(1)
2820 .thread_name("child-rt-test")
2821 .enable_all()
2822 .build()
2823 .expect("should build child runtime");
2824
2825 let (thread_tx, thread_rx) = oneshot::channel();
2826 let worker = noninterruptible_worker("placed", move |shutdown| async move {
2827 let thread_name = std::thread::current().name().unwrap_or_default().to_string();
2828 let _ = thread_tx.send(thread_name);
2829 shutdown.await;
2830 });
2831
2832 let mut sup = Supervisor::new("test-sup").unwrap();
2833 sup.add_worker(
2834 ChildSpecification::worker(worker)
2835 .with_restart_type(RestartType::Temporary)
2836 .with_runtime(child_runtime.handle().clone()),
2837 );
2838
2839 let (tx, handle) = run_supervisor_with_trigger(sup).await;
2840
2841 let thread_name = timeout(Duration::from_secs(2), thread_rx)
2842 .await
2843 .expect("child should report its thread promptly")
2844 .expect("child should not be dropped before reporting");
2845 assert!(
2846 thread_name.starts_with("child-rt-test"),
2847 "child must run on the runtime given to `with_runtime`, but ran on thread {thread_name:?}"
2848 );
2849
2850 tx.send(()).unwrap();
2851 assert!(join_supervisor(handle).await.is_ok());
2852
2853 child_runtime.shutdown_background();
2855 }
2856
2857 #[tokio::test]
2858 async fn child_shutdown_strategy_override_takes_precedence_over_worker() {
2859 let worker = noninterruptible_worker("stuck", |_shutdown| std::future::pending::<()>())
2864 .with_shutdown_timeout(Duration::from_secs(30));
2865
2866 let mut sup = Supervisor::new("test-sup").unwrap();
2867 sup.add_worker(
2868 ChildSpecification::worker(worker)
2869 .with_restart_type(RestartType::Temporary)
2870 .with_shutdown_strategy(ShutdownStrategy::Graceful(Duration::from_millis(50))),
2871 );
2872
2873 let (tx, handle) = run_supervisor_with_trigger(sup).await;
2874 tx.send(()).unwrap();
2875
2876 let result = join_supervisor(handle).await;
2877 assert!(
2878 matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 1 })),
2879 "the overridden 50ms deadline should have aborted the stuck child, got {result:?}"
2880 );
2881 }
2882
2883 fn build_drain_pair(
2889 stuck_strategy: ShutdownStrategy, waiter_strategy: ShutdownStrategy,
2890 ) -> (Supervisor, Arc<AtomicBool>) {
2891 let mut coordinator = ShutdownCoordinator::default();
2892 let held_handle = coordinator.register();
2893
2894 let stuck = noninterruptible_worker("stuck", move |_shutdown| async move {
2895 let _held = held_handle;
2897 pending::<()>().await;
2898 });
2899
2900 let waiter_finished = Arc::new(AtomicBool::new(false));
2901 let finished = Arc::clone(&waiter_finished);
2902 let waiter = noninterruptible_worker("waiter", move |shutdown| async move {
2903 shutdown.await;
2904 coordinator.shutdown_and_wait().await;
2905 finished.store(true, Ordering::SeqCst);
2906 });
2907
2908 let mut sup = Supervisor::new("test-sup")
2909 .unwrap()
2910 .with_shutdown_mode(ShutdownMode::Concurrent);
2911 sup.add_worker(
2912 ChildSpecification::worker(stuck)
2913 .with_restart_type(RestartType::Temporary)
2914 .with_shutdown_strategy(stuck_strategy),
2915 );
2916 sup.add_worker(
2917 ChildSpecification::worker(waiter)
2918 .with_restart_type(RestartType::Temporary)
2919 .with_shutdown_strategy(waiter_strategy),
2920 );
2921
2922 (sup, waiter_finished)
2923 }
2924
2925 #[tokio::test]
2926 async fn shorter_child_deadline_releases_a_waiting_sibling() {
2927 let (sup, waiter_finished) = build_drain_pair(
2931 ShutdownStrategy::Graceful(Duration::from_millis(100)),
2932 ShutdownStrategy::Graceful(Duration::from_secs(1)),
2933 );
2934
2935 let (tx, handle) = run_supervisor_with_trigger(sup).await;
2936 tx.send(()).unwrap();
2937
2938 let result = join_supervisor(handle).await;
2939 assert!(
2940 matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 1 })),
2941 "only the stuck child should have been aborted, got {result:?}"
2942 );
2943 assert!(
2944 waiter_finished.load(Ordering::SeqCst),
2945 "the waiter should have been released by the stuck child's abort and run to completion"
2946 );
2947 }
2948
2949 #[tokio::test]
2950 async fn equal_child_deadlines_abort_the_waiter_too() {
2951 let (sup, waiter_finished) = build_drain_pair(
2956 ShutdownStrategy::Graceful(Duration::from_millis(100)),
2957 ShutdownStrategy::Graceful(Duration::from_millis(100)),
2958 );
2959
2960 let (tx, handle) = run_supervisor_with_trigger(sup).await;
2961 tx.send(()).unwrap();
2962
2963 let result = join_supervisor(handle).await;
2964 assert!(
2965 matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 2 })),
2966 "both children should have been aborted together, got {result:?}"
2967 );
2968 assert!(
2969 !waiter_finished.load(Ordering::SeqCst),
2970 "the waiter should have been aborted mid-wait, not completed"
2971 );
2972 }
2973
2974 #[tokio::test]
2977 async fn budget_bounds_children_that_have_no_deadline_of_their_own() {
2978 let mut sup = Supervisor::new("test-sup")
2983 .unwrap()
2984 .with_shutdown_mode(ShutdownMode::Concurrent)
2985 .with_shutdown_budget(Duration::from_millis(100));
2986
2987 for name in ["stuck_one", "stuck_two"] {
2988 sup.add_worker(
2989 ChildSpecification::worker(noninterruptible_worker(name, |_shutdown| pending::<()>()))
2990 .with_restart_type(RestartType::Temporary)
2991 .with_shutdown_strategy(ShutdownStrategy::Graceful(Duration::MAX)),
2992 );
2993 }
2994
2995 let (tx, handle) = run_supervisor_with_trigger(sup).await;
2996 tx.send(()).unwrap();
2997
2998 let result = join_supervisor(handle).await;
2999 assert!(
3000 matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 2 })),
3001 "the budget should have aborted both deadline-less children, got {result:?}"
3002 );
3003 }
3004
3005 #[tokio::test]
3006 async fn budget_does_not_delay_children_that_stop_on_their_own() {
3007 let mut sup = Supervisor::new("test-sup")
3010 .unwrap()
3011 .with_shutdown_mode(ShutdownMode::Concurrent)
3012 .with_shutdown_budget(Duration::from_secs(30));
3013 sup.add_worker(
3014 ChildSpecification::one_shot_worker(noninterruptible_worker("prompt", |shutdown| shutdown))
3015 .with_shutdown_strategy(ShutdownStrategy::Graceful(Duration::MAX)),
3016 );
3017
3018 let (tx, handle) = run_supervisor_with_trigger(sup).await;
3019 let started = tokio::time::Instant::now();
3020 tx.send(()).unwrap();
3021
3022 assert!(join_supervisor(handle).await.is_ok());
3023 let elapsed = started.elapsed();
3024 assert!(
3025 elapsed < Duration::from_millis(500),
3026 "shutdown should finish as soon as the child does, not burn the budget; took {elapsed:?}"
3027 );
3028 }
3029
3030 #[tokio::test]
3031 async fn child_deadline_shorter_than_budget_still_wins() {
3032 let mut sup = Supervisor::new("test-sup")
3035 .unwrap()
3036 .with_shutdown_mode(ShutdownMode::Concurrent)
3037 .with_shutdown_budget(Duration::from_secs(30));
3038 sup.add_worker(
3039 ChildSpecification::worker(noninterruptible_worker("stuck", |_shutdown| pending::<()>()))
3040 .with_restart_type(RestartType::Temporary)
3041 .with_shutdown_strategy(ShutdownStrategy::Graceful(Duration::from_millis(100))),
3042 );
3043
3044 let (tx, handle) = run_supervisor_with_trigger(sup).await;
3045 tx.send(()).unwrap();
3046
3047 let result = join_supervisor(handle).await;
3049 assert!(
3050 matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 1 })),
3051 "the child's own 100ms deadline should have won over the budget, got {result:?}"
3052 );
3053 }
3054
3055 #[tokio::test]
3056 async fn every_worker_records_poll_metrics() {
3057 let recorder = TestRecorder::default();
3060 let _guard = metrics::set_default_local_recorder(&recorder);
3061
3062 let mut sup = Supervisor::new("metrics_sup").unwrap();
3064 sup.add_worker(ChildSpecification::one_shot_worker(noninterruptible_worker(
3065 "timed",
3066 |shutdown| shutdown,
3067 )));
3068
3069 let (tx, handle) = run_supervisor_with_trigger(sup).await;
3070 tx.send(()).unwrap();
3071 assert!(join_supervisor(handle).await.is_ok());
3072
3073 let polls = recorder.counter(("runtime_task_poll_count", &[("task_name", "metrics_sup.timed")]));
3074 assert!(
3075 polls.is_some_and(|polls| polls > 0),
3076 "a supervised worker should have recorded poll metrics, got {polls:?}"
3077 );
3078 }
3079
3080 #[tokio::test]
3081 async fn budget_bounds_the_whole_drain_in_ordered_mode() {
3082 let mut sup = Supervisor::new("test-sup")
3086 .unwrap()
3087 .with_shutdown_mode(ShutdownMode::Ordered)
3088 .with_shutdown_budget(Duration::from_millis(150));
3089
3090 for name in ["stuck_one", "stuck_two", "stuck_three"] {
3091 sup.add_worker(
3092 ChildSpecification::one_shot_worker(noninterruptible_worker(name, |_shutdown| pending::<()>()))
3093 .with_shutdown_strategy(ShutdownStrategy::Graceful(Duration::from_secs(10))),
3094 );
3095 }
3096
3097 let (tx, handle) = run_supervisor_with_trigger(sup).await;
3098 tx.send(()).unwrap();
3099
3100 let result = join_supervisor(handle).await;
3101 assert!(
3102 matches!(result, Err(SupervisorError::ShutdownTimedOut { aborted: 3 })),
3103 "the budget should have bounded the whole ordered drain, got {result:?}"
3104 );
3105 }
3106
3107 #[tokio::test]
3108 async fn budget_of_duration_max_is_treated_as_no_budget() {
3109 let mut sup = Supervisor::new("test-sup")
3112 .unwrap()
3113 .with_shutdown_mode(ShutdownMode::Concurrent)
3114 .with_shutdown_budget(Duration::MAX);
3115 sup.add_worker(ChildSpecification::one_shot_worker(noninterruptible_worker(
3116 "prompt",
3117 |shutdown| shutdown,
3118 )));
3119
3120 let (tx, handle) = run_supervisor_with_trigger(sup).await;
3121 tx.send(()).unwrap();
3122 assert!(join_supervisor(handle).await.is_ok());
3123 }
3124
3125 #[tokio::test]
3126 async fn near_max_child_timeout_does_not_panic_in_ordered_mode() {
3127 let mut sup = Supervisor::new("test-sup").unwrap();
3130 sup.add_worker(
3131 ChildSpecification::one_shot_worker(noninterruptible_worker("prompt", |shutdown| shutdown))
3132 .with_shutdown_strategy(ShutdownStrategy::Graceful(Duration::MAX - Duration::from_nanos(1))),
3133 );
3134
3135 let (tx, handle) = run_supervisor_with_trigger(sup).await;
3136 tx.send(()).unwrap();
3137 assert!(join_supervisor(handle).await.is_ok());
3138 }
3139
3140 #[tokio::test]
3141 async fn budget_does_not_cut_off_a_nested_supervisor_mid_drain() {
3142 let drained = Arc::new(AtomicBool::new(false));
3146 let child_drained = Arc::clone(&drained);
3147
3148 let mut nested = Supervisor::new("nested").unwrap();
3149 nested.add_worker(
3150 ChildSpecification::one_shot_worker(noninterruptible_worker("slow", move |shutdown| async move {
3151 shutdown.await;
3152 sleep(Duration::from_millis(300)).await;
3153 child_drained.store(true, Ordering::SeqCst);
3154 }))
3155 .with_shutdown_strategy(ShutdownStrategy::Graceful(Duration::from_secs(10))),
3156 );
3157
3158 let mut parent = Supervisor::new("parent")
3159 .unwrap()
3160 .with_shutdown_mode(ShutdownMode::Concurrent)
3161 .with_shutdown_budget(Duration::from_millis(50));
3162 parent.add_worker(nested);
3163
3164 let (tx, handle) = run_supervisor_with_trigger(parent).await;
3165 tx.send(()).unwrap();
3166
3167 let result = join_supervisor(handle).await;
3168 assert!(
3169 drained.load(Ordering::SeqCst),
3170 "the nested subtree should have drained rather than being cut off by the parent's budget: {result:?}"
3171 );
3172 assert!(
3173 result.is_ok(),
3174 "the nested drain finished in time, so shutdown was clean: {result:?}"
3175 );
3176 }
3177}