Skip to main content

actor_kit/
supervisor.rs

1//! Erlang-style supervisor trees for hierarchical actor supervision.
2//!
3//! This module implements the supervisor pattern from Erlang/OTP, providing
4//! fault-tolerant actor management with automatic restart policies and
5//! hierarchical supervision trees.
6//!
7//! # Overview
8//!
9//! Supervisors manage child actors with configurable restart strategies:
10//!
11//! - **One-for-one**: Restart only the failed child
12//! - **One-for-all**: Restart all children when one fails
13//! - **Rest-for-one**: Restart failed child and those started after it
14//! - **Simple one-for-one**: No restart, just acknowledge failure
15//!
16//! # Example
17//!
18//! ```ignore
19//! use actor_kit::supervisor::{
20//!     SupervisorTree, SupervisionStrategy, ChildSpec, RestartPolicy,
21//! };
22//! use std::time::Duration;
23//!
24//! // Create a supervisor tree with one-for-one strategy
25//! let mut tree = SupervisorTree::new(
26//!     SupervisionStrategy::one_for_one(5, Duration::from_secs(60))
27//! );
28//!
29//! // Start a child actor under the root supervisor
30//! let child_id = tree.start_child_under(
31//!     tree.root(),
32//!     ChildSpec {
33//!         name: "worker-1".into(),
34//!         actor_config: Default::default(),
35//!         restart_policy: RestartPolicy::Permanent,
36//!         shutdown_timeout: Duration::from_secs(30),
37//!         significant: false,
38//!     },
39//! )?;
40//! ```
41
42use parking_lot::RwLock;
43use std::collections::HashMap;
44use std::sync::Arc;
45use std::time::{Duration, Instant};
46
47use crate::{ActorId, Signal};
48use crate::{Error, Result};
49
50/// Maximum restart count threshold before considering it excessive.
51const MAX_RESTART_THRESHOLD: u32 = 100;
52
53/// How a supervisor should handle child failures.
54#[derive(Debug, Clone, PartialEq, Eq)]
55pub enum SupervisionStrategy {
56    /// Restart only the failed child.
57    ///
58    /// When a child crashes, only that child is restarted.
59    /// Other children are unaffected.
60    OneForOne {
61        /// Maximum number of restarts allowed within the time window.
62        max_restarts: u32,
63        /// Time window for counting restarts.
64        within: Duration,
65    },
66    /// Restart all children when one fails.
67    ///
68    /// When a child crashes, all children are terminated and restarted.
69    /// Use when children are tightly coupled.
70    OneForAll {
71        /// Maximum number of restarts allowed within the time window.
72        max_restarts: u32,
73        /// Time window for counting restarts.
74        within: Duration,
75    },
76    /// Restart failed child and those started after it.
77    ///
78    /// When a child crashes, that child and all children started after it
79    /// are terminated and restarted. Children started before are unaffected.
80    /// Use for dependent chains of actors.
81    RestForOne {
82        /// Maximum number of restarts allowed within the time window.
83        max_restarts: u32,
84        /// Time window for counting restarts.
85        within: Duration,
86    },
87    /// Don't restart, just acknowledge failure.
88    ///
89    /// No automatic restart is performed. Useful for dynamic, temporary
90    /// children where failures are expected.
91    SimpleOneForOne,
92}
93
94impl SupervisionStrategy {
95    /// Create a one-for-one strategy.
96    pub fn one_for_one(max_restarts: u32, within: Duration) -> Self {
97        Self::OneForOne {
98            max_restarts,
99            within,
100        }
101    }
102
103    /// Create a one-for-all strategy.
104    pub fn one_for_all(max_restarts: u32, within: Duration) -> Self {
105        Self::OneForAll {
106            max_restarts,
107            within,
108        }
109    }
110
111    /// Create a rest-for-one strategy.
112    pub fn rest_for_one(max_restarts: u32, within: Duration) -> Self {
113        Self::RestForOne {
114            max_restarts,
115            within,
116        }
117    }
118
119    /// Create a simple one-for-one strategy.
120    pub fn simple_one_for_one() -> Self {
121        Self::SimpleOneForOne
122    }
123
124    /// Get the max restarts limit for this strategy
125    #[cfg(test)]
126    pub(crate) fn max_restarts(&self) -> Option<u32> {
127        match self {
128            Self::OneForOne { max_restarts, .. }
129            | Self::OneForAll { max_restarts, .. }
130            | Self::RestForOne { max_restarts, .. } => Some(*max_restarts),
131            Self::SimpleOneForOne => None,
132        }
133    }
134
135    fn within(&self) -> Option<Duration> {
136        match self {
137            Self::OneForOne { within, .. }
138            | Self::OneForAll { within, .. }
139            | Self::RestForOne { within, .. } => Some(*within),
140            Self::SimpleOneForOne => None,
141        }
142    }
143}
144
145impl Default for SupervisionStrategy {
146    fn default() -> Self {
147        Self::one_for_one(5, Duration::from_secs(60))
148    }
149}
150
151/// Restart policy for a supervised child.
152#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
153pub enum RestartPolicy {
154    /// Always restart the child, regardless of exit reason.
155    #[default]
156    Permanent,
157    /// Restart only on abnormal termination (error, killed, signaled).
158    Transient,
159    /// Never restart the child.
160    Temporary,
161}
162
163/// Why an actor terminated.
164#[derive(Debug, Clone, PartialEq, Eq)]
165pub enum ExitReason {
166    /// Graceful shutdown completed successfully.
167    Normal,
168    /// Supervisor requested shutdown.
169    Shutdown,
170    /// Received a termination signal.
171    Signaled(Signal),
172    /// Actor failed with an error.
173    Error(String),
174    /// Actor was forcefully killed.
175    Killed,
176}
177
178impl ExitReason {
179    /// Check if this is a normal termination.
180    pub fn is_normal(&self) -> bool {
181        matches!(self, Self::Normal)
182    }
183
184    /// Check if this is an abnormal termination (should trigger restart for Transient).
185    pub fn is_abnormal(&self) -> bool {
186        !self.is_normal()
187    }
188}
189
190impl std::fmt::Display for ExitReason {
191    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
192        match self {
193            Self::Normal => write!(f, "normal"),
194            Self::Shutdown => write!(f, "shutdown"),
195            Self::Signaled(signal) => write!(f, "signaled({:?})", signal),
196            Self::Error(msg) => write!(f, "error({})", msg),
197            Self::Killed => write!(f, "killed"),
198        }
199    }
200}
201
202/// Specification for a supervised child actor.
203#[derive(Debug, Clone)]
204pub struct ChildSpec {
205    /// Unique name for the child within this supervisor.
206    pub name: String,
207    /// Startup and behavioral configuration passed to the child actor on spawn,
208    /// such as mailbox capacity, dispatch priority, and initialization parameters.
209    pub actor_config: ActorConfig,
210    /// When to restart the child.
211    pub restart_policy: RestartPolicy,
212    /// How long to wait for graceful shutdown before force-killing.
213    pub shutdown_timeout: Duration,
214    /// If true, abnormal termination escalates to parent supervisor.
215    pub significant: bool,
216}
217
218impl ChildSpec {
219    /// Create a new child spec with the given name.
220    pub fn new(name: impl Into<String>) -> Self {
221        Self {
222            name: name.into(),
223            actor_config: ActorConfig::default(),
224            restart_policy: RestartPolicy::default(),
225            shutdown_timeout: Duration::from_secs(30),
226            significant: false,
227        }
228    }
229
230    /// Set the restart policy.
231    pub fn restart_policy(mut self, policy: RestartPolicy) -> Self {
232        self.restart_policy = policy;
233        self
234    }
235
236    /// Set the shutdown timeout.
237    pub fn shutdown_timeout(mut self, timeout: Duration) -> Self {
238        self.shutdown_timeout = timeout;
239        self
240    }
241
242    /// Mark this child as significant (termination escalates).
243    pub fn significant(mut self, significant: bool) -> Self {
244        self.significant = significant;
245        self
246    }
247}
248
249/// Startup and behavioral configuration for a child actor, controlling parameters
250/// such as mailbox capacity, dispatch priority, and initialization arguments.
251#[derive(Debug, Clone, Default)]
252pub struct ActorConfig {
253    /// Reserved for additional configuration options (e.g., mailbox bounds,
254    /// dispatch priority, custom metadata) in future releases.
255    _future: (),
256}
257
258/// State of a supervised child.
259#[derive(Debug, Clone)]
260pub enum ChildState {
261    /// Child is running normally.
262    Running,
263    /// Child is being stopped.
264    Stopping,
265    /// Child has stopped.
266    Stopped {
267        /// Why the child stopped.
268        reason: ExitReason,
269    },
270    /// Child is being restarted.
271    Restarting {
272        /// Current restart attempt number.
273        attempt: u32,
274    },
275}
276
277impl ChildState {
278    /// Check if the child is in an active state.
279    pub fn is_active(&self) -> bool {
280        matches!(self, Self::Running | Self::Restarting { .. })
281    }
282}
283
284/// A child actor being supervised.
285#[derive(Debug, Clone)]
286pub struct SupervisedChild {
287    /// Actor ID.
288    pub id: ActorId,
289    /// Child specification.
290    pub spec: ChildSpec,
291    /// Current state.
292    pub state: ChildState,
293    /// Order in which this child was started (for rest-for-one).
294    pub start_order: u32,
295}
296
297/// Statistics about a supervisor's children.
298#[derive(Debug, Clone, Default)]
299pub struct SupervisorStats {
300    /// Total number of children.
301    pub total: usize,
302    /// Number of running children.
303    pub running: usize,
304    /// Number of stopped children.
305    pub stopped: usize,
306    /// Number of restarting children.
307    pub restarting: usize,
308    /// Total restarts across all children.
309    pub total_restarts: u64,
310}
311
312/// Action to take when max restarts are exceeded.
313#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
314pub enum EscalationAction {
315    /// Escalate to parent supervisor.
316    #[default]
317    Escalate,
318    /// Shutdown the entire node.
319    ShutdownNode,
320    /// Ignore and stop restarting.
321    GiveUp,
322}
323
324/// Error type for supervisor operations.
325#[derive(Debug, Clone)]
326pub enum SupervisorError {
327    /// Child with this name already exists.
328    ChildExists(String),
329    /// Child with this name not found.
330    ChildNotFound(String),
331    /// Maximum restarts exceeded.
332    MaxRestartsExceeded {
333        /// Child name.
334        name: String,
335        /// Restart count in window.
336        count: u32,
337        /// Max allowed.
338        max: u32,
339    },
340    /// Supervisor not found.
341    SupervisorNotFound(ActorId),
342    /// Invalid operation for current state.
343    InvalidState {
344        /// Child name.
345        name: String,
346        /// Current state.
347        state: String,
348        /// Attempted operation.
349        operation: String,
350    },
351    /// Parent supervisor ID is invalid.
352    InvalidParent(ActorId),
353}
354
355impl std::fmt::Display for SupervisorError {
356    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
357        match self {
358            Self::ChildExists(name) => write!(f, "child '{}' already exists", name),
359            Self::ChildNotFound(name) => write!(f, "child '{}' not found", name),
360            Self::MaxRestartsExceeded { name, count, max } => {
361                write!(
362                    f,
363                    "child '{}' exceeded max restarts ({}/{})",
364                    name, count, max
365                )
366            }
367            Self::SupervisorNotFound(id) => write!(f, "supervisor {:?} not found", id),
368            Self::InvalidState {
369                name,
370                state,
371                operation,
372            } => {
373                write!(
374                    f,
375                    "invalid operation '{}' on child '{}' in state '{}'",
376                    operation, name, state
377                )
378            }
379            Self::InvalidParent(id) => write!(f, "invalid parent supervisor {:?}", id),
380        }
381    }
382}
383
384impl From<SupervisorError> for Error {
385    fn from(e: SupervisorError) -> Error {
386        Error::actor(e.to_string())
387    }
388}
389
390/// A supervisor that manages child actors.
391#[derive(Debug)]
392pub struct Supervisor {
393    /// Supervisor's own ID.
394    id: ActorId,
395    /// Supervision strategy.
396    strategy: SupervisionStrategy,
397    /// Child actors by name.
398    children: HashMap<String, SupervisedChild>,
399    /// Restart timestamps per child for rate limiting.
400    restart_counts: HashMap<String, Vec<Instant>>,
401    /// Parent supervisor ID (None for root).
402    parent: Option<ActorId>,
403    /// Next start order number.
404    next_start_order: u32,
405    /// Total restarts across all children.
406    total_restarts: u64,
407    /// Escalation action when max restarts exceeded.
408    escalation_action: EscalationAction,
409}
410
411impl Supervisor {
412    /// Create a new supervisor with the given strategy.
413    pub fn new(strategy: SupervisionStrategy) -> Self {
414        Self {
415            id: ActorId::new(),
416            strategy,
417            children: HashMap::new(),
418            restart_counts: HashMap::new(),
419            parent: None,
420            next_start_order: 0,
421            total_restarts: 0,
422            escalation_action: EscalationAction::default(),
423        }
424    }
425
426    /// Create a new supervisor with a parent.
427    pub fn with_parent(strategy: SupervisionStrategy, parent: ActorId) -> Self {
428        Self {
429            id: ActorId::new(),
430            strategy,
431            children: HashMap::new(),
432            restart_counts: HashMap::new(),
433            parent: Some(parent),
434            next_start_order: 0,
435            total_restarts: 0,
436            escalation_action: EscalationAction::default(),
437        }
438    }
439
440    /// Get the supervisor's ID.
441    pub fn id(&self) -> ActorId {
442        self.id
443    }
444
445    /// Get the parent supervisor ID.
446    pub fn parent(&self) -> Option<ActorId> {
447        self.parent
448    }
449
450    /// Set the escalation action.
451    pub fn set_escalation_action(&mut self, action: EscalationAction) {
452        self.escalation_action = action;
453    }
454
455    /// Start a new child actor.
456    pub fn start_child(&mut self, spec: ChildSpec) -> Result<ActorId> {
457        if self.children.contains_key(&spec.name) {
458            return Err(SupervisorError::ChildExists(spec.name.clone()).into());
459        }
460
461        let id = ActorId::new();
462        let start_order = self.next_start_order;
463        self.next_start_order = self.next_start_order.saturating_add(1);
464
465        let child = SupervisedChild {
466            id,
467            spec,
468            state: ChildState::Running,
469            start_order,
470        };
471
472        let name = child.spec.name.clone();
473        self.children.insert(name, child);
474
475        Ok(id)
476    }
477
478    /// Stop a child actor gracefully.
479    pub fn stop_child(&mut self, name: &str) -> Result<()> {
480        let child = self
481            .children
482            .get_mut(name)
483            .ok_or_else(|| SupervisorError::ChildNotFound(name.to_string()))?;
484
485        match &child.state {
486            ChildState::Running | ChildState::Restarting { .. } => {
487                child.state = ChildState::Stopping;
488                Ok(())
489            }
490            ChildState::Stopping => Ok(()),
491            ChildState::Stopped { .. } => Ok(()),
492        }
493    }
494
495    /// Terminate a child actor with a specific reason.
496    pub fn terminate_child(&mut self, name: &str, reason: ExitReason) -> Result<()> {
497        let child = self
498            .children
499            .get_mut(name)
500            .ok_or_else(|| SupervisorError::ChildNotFound(name.to_string()))?;
501
502        child.state = ChildState::Stopped { reason };
503        Ok(())
504    }
505
506    /// Get a list of all children.
507    pub fn which_children(&self) -> Vec<&SupervisedChild> {
508        self.children.values().collect()
509    }
510
511    /// Get statistics about the supervisor.
512    pub fn count_children(&self) -> SupervisorStats {
513        let mut running = 0;
514        let mut stopped = 0;
515        let mut restarting = 0;
516
517        for child in self.children.values() {
518            match &child.state {
519                ChildState::Running => running += 1,
520                ChildState::Stopped { .. } => stopped += 1,
521                ChildState::Restarting { .. } => restarting += 1,
522                ChildState::Stopping => stopped += 1,
523            }
524        }
525
526        SupervisorStats {
527            total: self.children.len(),
528            running,
529            stopped,
530            restarting,
531            total_restarts: self.total_restarts,
532        }
533    }
534
535    /// Handle a child exit event.
536    pub async fn handle_child_exit(&mut self, name: &str, reason: ExitReason) -> Result<()> {
537        let child = self
538            .children
539            .get(name)
540            .ok_or_else(|| SupervisorError::ChildNotFound(name.to_string()))?;
541
542        let should_restart = match child.spec.restart_policy {
543            RestartPolicy::Permanent => true,
544            RestartPolicy::Transient => reason.is_abnormal(),
545            RestartPolicy::Temporary => false,
546        };
547
548        let is_significant = child.spec.significant && reason.is_abnormal();
549
550        if is_significant {
551            if let Some(parent_id) = self.parent {
552                tracing::warn!(
553                    "Significant child '{}' exited, escalating to parent {:?}",
554                    name,
555                    parent_id
556                );
557            }
558        }
559
560        if !should_restart {
561            if let Some(child) = self.children.get_mut(name) {
562                child.state = ChildState::Stopped { reason };
563            }
564            return Ok(());
565        }
566
567        match &self.strategy {
568            SupervisionStrategy::SimpleOneForOne => {
569                if let Some(child) = self.children.get_mut(name) {
570                    child.state = ChildState::Stopped { reason };
571                }
572                Ok(())
573            }
574            SupervisionStrategy::OneForOne {
575                max_restarts,
576                within,
577            } => {
578                self.restart_one_for_one(name, reason, *max_restarts, *within)
579                    .await
580            }
581            SupervisionStrategy::OneForAll {
582                max_restarts,
583                within,
584            } => {
585                self.restart_one_for_all(name, reason, *max_restarts, *within)
586                    .await
587            }
588            SupervisionStrategy::RestForOne {
589                max_restarts,
590                within,
591            } => {
592                self.restart_rest_for_one(name, reason, *max_restarts, *within)
593                    .await
594            }
595        }
596    }
597
598    async fn restart_one_for_one(
599        &mut self,
600        name: &str,
601        reason: ExitReason,
602        max_restarts: u32,
603        within: Duration,
604    ) -> Result<()> {
605        if !self.check_restart_allowed(name, max_restarts, within) {
606            return self.handle_max_restarts_exceeded(name, max_restarts);
607        }
608
609        self.record_restart(name);
610        self.total_restarts = self.total_restarts.saturating_add(1);
611
612        if let Some(child) = self.children.get_mut(name) {
613            let attempt = self
614                .restart_counts
615                .get(name)
616                .map(|v| v.len() as u32)
617                .unwrap_or(1);
618            tracing::info!(
619                "Restarting child '{}' (attempt {}) after {:?}",
620                name,
621                attempt,
622                reason
623            );
624            child.state = ChildState::Restarting { attempt };
625            child.id = ActorId::new();
626        }
627
628        Ok(())
629    }
630
631    async fn restart_one_for_all(
632        &mut self,
633        name: &str,
634        reason: ExitReason,
635        max_restarts: u32,
636        within: Duration,
637    ) -> Result<()> {
638        if !self.check_restart_allowed(name, max_restarts, within) {
639            return self.handle_max_restarts_exceeded(name, max_restarts);
640        }
641
642        self.record_restart(name);
643        self.total_restarts = self.total_restarts.saturating_add(1);
644
645        tracing::info!(
646            "Restarting all children after '{}' exited with {:?}",
647            name,
648            reason
649        );
650
651        for child in self.children.values_mut() {
652            let attempt = self
653                .restart_counts
654                .get(&child.spec.name)
655                .map(|v| v.len() as u32)
656                .unwrap_or(1);
657            child.state = ChildState::Restarting { attempt };
658            child.id = ActorId::new();
659        }
660
661        Ok(())
662    }
663
664    async fn restart_rest_for_one(
665        &mut self,
666        name: &str,
667        _reason: ExitReason,
668        max_restarts: u32,
669        within: Duration,
670    ) -> Result<()> {
671        if !self.check_restart_allowed(name, max_restarts, within) {
672            return self.handle_max_restarts_exceeded(name, max_restarts);
673        }
674
675        let failed_order = self
676            .children
677            .get(name)
678            .map(|c| c.start_order)
679            .ok_or_else(|| SupervisorError::ChildNotFound(name.to_string()))?;
680
681        self.record_restart(name);
682        self.total_restarts = self.total_restarts.saturating_add(1);
683
684        tracing::info!(
685            "Restarting child '{}' and siblings started after it (rest-for-one)",
686            name
687        );
688
689        for child in self.children.values_mut() {
690            if child.start_order >= failed_order {
691                let attempt = self
692                    .restart_counts
693                    .get(&child.spec.name)
694                    .map(|v| v.len() as u32)
695                    .unwrap_or(1);
696                child.state = ChildState::Restarting { attempt };
697                child.id = ActorId::new();
698            }
699        }
700
701        Ok(())
702    }
703
704    fn check_restart_allowed(&self, name: &str, max_restarts: u32, within: Duration) -> bool {
705        if max_restarts == 0 {
706            return true;
707        }
708
709        let now = Instant::now();
710        let restarts = self.restart_counts.get(name);
711
712        if let Some(timestamps) = restarts {
713            let recent = timestamps
714                .iter()
715                .filter(|&&t| now.duration_since(t) < within)
716                .count();
717
718            let recent_u32 = u32::try_from(recent).unwrap_or(MAX_RESTART_THRESHOLD);
719            recent_u32 < max_restarts
720        } else {
721            true
722        }
723    }
724
725    fn record_restart(&mut self, name: &str) {
726        let now = Instant::now();
727        let timestamps = self.restart_counts.entry(name.to_string()).or_default();
728        timestamps.push(now);
729
730        if let Some(within) = self.strategy.within() {
731            timestamps.retain(|&t| now.duration_since(t) < within);
732        }
733    }
734
735    fn handle_max_restarts_exceeded(&mut self, name: &str, max_restarts: u32) -> Result<()> {
736        let count = self
737            .restart_counts
738            .get(name)
739            .map(|v| u32::try_from(v.len()).unwrap_or(MAX_RESTART_THRESHOLD))
740            .unwrap_or(0);
741
742        tracing::error!(
743            "Child '{}' exceeded max restarts ({}/{})",
744            name,
745            count,
746            max_restarts
747        );
748
749        if let Some(child) = self.children.get_mut(name) {
750            child.state = ChildState::Stopped {
751                reason: ExitReason::Error(format!(
752                    "max restarts exceeded ({}/{})",
753                    count, max_restarts
754                )),
755            };
756        }
757
758        match self.escalation_action {
759            EscalationAction::Escalate => {
760                if let Some(parent_id) = self.parent {
761                    tracing::warn!("Escalating max restarts failure to parent {:?}", parent_id);
762                }
763                Err(SupervisorError::MaxRestartsExceeded {
764                    name: name.to_string(),
765                    count,
766                    max: max_restarts,
767                }
768                .into())
769            }
770            EscalationAction::ShutdownNode => {
771                tracing::error!("Initiating node shutdown due to max restarts exceeded");
772                Err(SupervisorError::MaxRestartsExceeded {
773                    name: name.to_string(),
774                    count,
775                    max: max_restarts,
776                }
777                .into())
778            }
779            EscalationAction::GiveUp => {
780                tracing::warn!("Giving up on child '{}' after max restarts", name);
781                Ok(())
782            }
783        }
784    }
785
786    /// Remove a child from supervision.
787    pub fn remove_child(&mut self, name: &str) -> Result<SupervisedChild> {
788        self.children
789            .remove(name)
790            .ok_or_else(|| SupervisorError::ChildNotFound(name.to_string()).into())
791    }
792
793    /// Get a child by name.
794    pub fn get_child(&self, name: &str) -> Option<&SupervisedChild> {
795        self.children.get(name)
796    }
797
798    /// Get a child by name mutably.
799    pub fn get_child_mut(&mut self, name: &str) -> Option<&mut SupervisedChild> {
800        self.children.get_mut(name)
801    }
802
803    /// Check if a child exists.
804    pub fn has_child(&self, name: &str) -> bool {
805        self.children.contains_key(name)
806    }
807
808    /// Get the number of children.
809    pub fn child_count(&self) -> usize {
810        self.children.len()
811    }
812
813    /// Mark a restarting child as running.
814    pub fn mark_child_running(&mut self, name: &str) -> Result<()> {
815        let child = self
816            .children
817            .get_mut(name)
818            .ok_or_else(|| SupervisorError::ChildNotFound(name.to_string()))?;
819
820        match &child.state {
821            ChildState::Restarting { .. } | ChildState::Running => {
822                child.state = ChildState::Running;
823                Ok(())
824            }
825            ChildState::Stopping | ChildState::Stopped { .. } => {
826                Err(SupervisorError::InvalidState {
827                    name: name.to_string(),
828                    state: format!("{:?}", child.state),
829                    operation: "mark_running".to_string(),
830                }
831                .into())
832            }
833        }
834    }
835}
836
837/// Hierarchical supervisor tree structure.
838pub struct SupervisorTree {
839    /// Root supervisor ID.
840    root: ActorId,
841    /// All supervisors by ID.
842    supervisors: HashMap<ActorId, Supervisor>,
843    /// Parent-child relationships for supervisors.
844    supervisor_children: HashMap<ActorId, Vec<ActorId>>,
845}
846
847impl SupervisorTree {
848    /// Create a new supervisor tree with a root supervisor.
849    pub fn new(root_strategy: SupervisionStrategy) -> Self {
850        let root_supervisor = Supervisor::new(root_strategy);
851        let root_id = root_supervisor.id();
852
853        let mut supervisors = HashMap::new();
854        supervisors.insert(root_id, root_supervisor);
855
856        let mut supervisor_children = HashMap::new();
857        supervisor_children.insert(root_id, Vec::new());
858
859        Self {
860            root: root_id,
861            supervisors,
862            supervisor_children,
863        }
864    }
865
866    /// Get the root supervisor ID.
867    pub fn root(&self) -> ActorId {
868        self.root
869    }
870
871    /// Add a new supervisor under a parent.
872    pub fn add_supervisor(
873        &mut self,
874        parent: ActorId,
875        strategy: SupervisionStrategy,
876    ) -> Result<ActorId> {
877        if !self.supervisors.contains_key(&parent) {
878            return Err(SupervisorError::SupervisorNotFound(parent).into());
879        }
880
881        let supervisor = Supervisor::with_parent(strategy, parent);
882        let id = supervisor.id();
883
884        self.supervisors.insert(id, supervisor);
885        self.supervisor_children.insert(id, Vec::new());
886
887        if let Some(children) = self.supervisor_children.get_mut(&parent) {
888            children.push(id);
889        }
890
891        Ok(id)
892    }
893
894    /// Start a child actor under a specific supervisor.
895    pub fn start_child_under(&mut self, supervisor: ActorId, spec: ChildSpec) -> Result<ActorId> {
896        let sup = self
897            .supervisors
898            .get_mut(&supervisor)
899            .ok_or(SupervisorError::SupervisorNotFound(supervisor))?;
900
901        sup.start_child(spec)
902    }
903
904    /// Get a supervisor by ID.
905    pub fn get_supervisor(&self, id: &ActorId) -> Option<&Supervisor> {
906        self.supervisors.get(id)
907    }
908
909    /// Get a supervisor by ID mutably.
910    pub fn get_supervisor_mut(&mut self, id: &ActorId) -> Option<&mut Supervisor> {
911        self.supervisors.get_mut(id)
912    }
913
914    /// Terminate a subtree starting from a given supervisor.
915    pub async fn terminate_tree(&mut self, root_id: &ActorId) -> Result<()> {
916        if !self.supervisors.contains_key(root_id) {
917            return Err(SupervisorError::SupervisorNotFound(*root_id).into());
918        }
919
920        let to_terminate = self.collect_subtree(root_id);
921
922        for id in to_terminate.iter().rev() {
923            if let Some(supervisor) = self.supervisors.get(id) {
924                let child_names: Vec<String> = supervisor
925                    .which_children()
926                    .iter()
927                    .map(|c| c.spec.name.clone())
928                    .collect();
929
930                for name in child_names {
931                    if let Some(sup) = self.supervisors.get_mut(id) {
932                        let _ = sup.terminate_child(&name, ExitReason::Shutdown);
933                    }
934                }
935            }
936        }
937
938        for id in &to_terminate {
939            if *id != self.root {
940                self.supervisors.remove(id);
941                self.supervisor_children.remove(id);
942
943                for children in self.supervisor_children.values_mut() {
944                    children.retain(|child_id| child_id != id);
945                }
946            }
947        }
948
949        Ok(())
950    }
951
952    fn collect_subtree(&self, root_id: &ActorId) -> Vec<ActorId> {
953        let mut result = Vec::new();
954        self.collect_subtree_recursive(root_id, &mut result);
955        result
956    }
957
958    fn collect_subtree_recursive(&self, id: &ActorId, result: &mut Vec<ActorId>) {
959        result.push(*id);
960
961        if let Some(children) = self.supervisor_children.get(id) {
962            for child in children {
963                self.collect_subtree_recursive(child, result);
964            }
965        }
966    }
967
968    /// Handle a child exit in the appropriate supervisor.
969    pub async fn handle_child_exit(
970        &mut self,
971        supervisor_id: ActorId,
972        child_name: &str,
973        reason: ExitReason,
974    ) -> Result<()> {
975        let supervisor = self
976            .supervisors
977            .get_mut(&supervisor_id)
978            .ok_or(SupervisorError::SupervisorNotFound(supervisor_id))?;
979
980        supervisor.handle_child_exit(child_name, reason).await
981    }
982
983    /// Get statistics for the entire tree.
984    pub fn tree_stats(&self) -> SupervisorTreeStats {
985        let mut total_children = 0;
986        let mut running_children = 0;
987        let mut stopped_children = 0;
988        let mut restarting_children = 0;
989        let mut total_restarts = 0;
990
991        for supervisor in self.supervisors.values() {
992            let sup_stats = supervisor.count_children();
993            total_children += sup_stats.total;
994            running_children += sup_stats.running;
995            stopped_children += sup_stats.stopped;
996            restarting_children += sup_stats.restarting;
997            total_restarts += sup_stats.total_restarts;
998        }
999
1000        SupervisorTreeStats {
1001            supervisor_count: self.supervisors.len(),
1002            total_children,
1003            running_children,
1004            stopped_children,
1005            restarting_children,
1006            total_restarts,
1007        }
1008    }
1009
1010    /// Get the depth of the tree.
1011    pub fn tree_depth(&self) -> usize {
1012        self.calculate_depth(&self.root, 0)
1013    }
1014
1015    fn calculate_depth(&self, id: &ActorId, current_depth: usize) -> usize {
1016        let children = self.supervisor_children.get(id);
1017        let child_depths: Vec<usize> = children
1018            .map(|c| {
1019                c.iter()
1020                    .map(|child_id| self.calculate_depth(child_id, current_depth + 1))
1021                    .collect()
1022            })
1023            .unwrap_or_default();
1024
1025        child_depths.into_iter().max().unwrap_or(current_depth)
1026    }
1027}
1028
1029/// Statistics for the entire supervisor tree.
1030#[derive(Debug, Clone, Default)]
1031pub struct SupervisorTreeStats {
1032    /// Number of supervisors.
1033    pub supervisor_count: usize,
1034    /// Total children across all supervisors.
1035    pub total_children: usize,
1036    /// Running children.
1037    pub running_children: usize,
1038    /// Stopped children.
1039    pub stopped_children: usize,
1040    /// Restarting children.
1041    pub restarting_children: usize,
1042    /// Total restarts.
1043    pub total_restarts: u64,
1044}
1045
1046/// Thread-safe wrapper for supervisor operations.
1047#[derive(Clone)]
1048pub struct SupervisorHandle {
1049    inner: Arc<RwLock<Supervisor>>,
1050}
1051
1052impl SupervisorHandle {
1053    /// Create a new handle wrapping a supervisor.
1054    pub fn new(supervisor: Supervisor) -> Self {
1055        Self {
1056            inner: Arc::new(RwLock::new(supervisor)),
1057        }
1058    }
1059
1060    /// Get the supervisor ID.
1061    pub fn id(&self) -> ActorId {
1062        self.inner.read().id()
1063    }
1064
1065    /// Start a child actor.
1066    pub fn start_child(&self, spec: ChildSpec) -> Result<ActorId> {
1067        self.inner.write().start_child(spec)
1068    }
1069
1070    /// Stop a child actor.
1071    pub fn stop_child(&self, name: &str) -> Result<()> {
1072        self.inner.write().stop_child(name)
1073    }
1074
1075    /// Get child statistics.
1076    pub fn count_children(&self) -> SupervisorStats {
1077        self.inner.read().count_children()
1078    }
1079
1080    /// Get all children.
1081    pub fn which_children(&self) -> Vec<SupervisedChild> {
1082        self.inner
1083            .read()
1084            .which_children()
1085            .into_iter()
1086            .cloned()
1087            .collect()
1088    }
1089}
1090
1091// ---------------------------------------------------------------------------
1092// Graceful Degradation
1093// ---------------------------------------------------------------------------
1094
1095/// Resource pressure level reported by a [`ResourceMonitor`].
1096#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
1097pub enum PressureLevel {
1098    /// System is operating within normal resource bounds.
1099    #[default]
1100    Normal,
1101    /// Resource usage is elevated; non-essential work should be throttled.
1102    Elevated,
1103    /// Resources are critically constrained; spawn new work is rejected.
1104    Critical,
1105}
1106
1107impl std::fmt::Display for PressureLevel {
1108    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1109        match self {
1110            Self::Normal => write!(f, "normal"),
1111            Self::Elevated => write!(f, "elevated"),
1112            Self::Critical => write!(f, "critical"),
1113        }
1114    }
1115}
1116
1117/// Trait for monitoring system resource pressure.
1118///
1119/// Implementations report the current [`PressureLevel`] so that supervisors
1120/// can make graceful degradation decisions.
1121pub trait ResourceMonitor: Send + Sync {
1122    /// Return the current resource pressure level.
1123    fn current_pressure(&self) -> PressureLevel;
1124}
1125
1126/// A static resource monitor that always reports the same pressure level.
1127///
1128/// Useful for testing and for nodes where resource monitoring is delegated
1129/// to an external system.
1130pub struct StaticResourceMonitor {
1131    level: PressureLevel,
1132}
1133
1134impl StaticResourceMonitor {
1135    /// Create a monitor that always reports `level`.
1136    pub fn new(level: PressureLevel) -> Self {
1137        Self { level }
1138    }
1139}
1140
1141impl ResourceMonitor for StaticResourceMonitor {
1142    fn current_pressure(&self) -> PressureLevel {
1143        self.level
1144    }
1145}
1146
1147/// A resource monitor that reports `Elevated` when a closure returns `true`
1148/// and `Critical` when a second closure returns `true`.
1149pub struct FnResourceMonitor {
1150    elevated_fn: Box<dyn Fn() -> bool + Send + Sync>,
1151    critical_fn: Box<dyn Fn() -> bool + Send + Sync>,
1152}
1153
1154impl FnResourceMonitor {
1155    /// Create a function-based resource monitor.
1156    pub fn new(
1157        elevated_fn: Box<dyn Fn() -> bool + Send + Sync>,
1158        critical_fn: Box<dyn Fn() -> bool + Send + Sync>,
1159    ) -> Self {
1160        Self {
1161            elevated_fn,
1162            critical_fn,
1163        }
1164    }
1165}
1166
1167impl ResourceMonitor for FnResourceMonitor {
1168    fn current_pressure(&self) -> PressureLevel {
1169        if (self.critical_fn)() {
1170            PressureLevel::Critical
1171        } else if (self.elevated_fn)() {
1172            PressureLevel::Elevated
1173        } else {
1174            PressureLevel::Normal
1175        }
1176    }
1177}
1178
1179/// Result of a graceful degradation decision.
1180#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1181pub enum DegradationDecision {
1182    /// Allow the operation to proceed normally.
1183    Allow,
1184    /// Throttle: reduce batch sizes or concurrency.
1185    Throttle,
1186    /// Reject: the operation cannot proceed under current pressure.
1187    Reject,
1188}
1189
1190impl DegradationDecision {
1191    /// Returns `true` if the operation is allowed (possibly throttled).
1192    pub fn is_allowed(self) -> bool {
1193        matches!(self, Self::Allow | Self::Throttle)
1194    }
1195}
1196
1197/// Graceful degradation controller that wraps a [`ResourceMonitor`] and
1198/// makes admission-control decisions.
1199pub struct DegradationController<M: ResourceMonitor> {
1200    monitor: M,
1201    elevated_batch_factor: f64,
1202}
1203
1204impl<M: ResourceMonitor> DegradationController<M> {
1205    /// Create a new controller with the given resource monitor.
1206    ///
1207    /// `elevated_batch_factor` is multiplied by the normal batch size when
1208    /// pressure is elevated (e.g. `0.5` halves the batch).
1209    pub fn new(monitor: M, elevated_batch_factor: f64) -> Self {
1210        Self {
1211            monitor,
1212            elevated_batch_factor: elevated_batch_factor.clamp(0.0, 1.0),
1213        }
1214    }
1215
1216    /// Return the current pressure level.
1217    pub fn pressure(&self) -> PressureLevel {
1218        self.monitor.current_pressure()
1219    }
1220
1221    /// Return a reference to the underlying monitor.
1222    pub fn monitor(&self) -> &M {
1223        &self.monitor
1224    }
1225
1226    /// Return the elevated batch factor.
1227    pub fn elevated_batch_factor(&self) -> f64 {
1228        self.elevated_batch_factor
1229    }
1230
1231    /// Decide whether a new child actor spawn should be admitted.
1232    pub fn admit_spawn(&self) -> DegradationDecision {
1233        match self.monitor.current_pressure() {
1234            PressureLevel::Normal => DegradationDecision::Allow,
1235            PressureLevel::Elevated => DegradationDecision::Throttle,
1236            PressureLevel::Critical => DegradationDecision::Reject,
1237        }
1238    }
1239
1240    /// Compute the effective batch size given normal and current pressure.
1241    pub fn effective_batch_size(&self, normal: usize) -> usize {
1242        match self.monitor.current_pressure() {
1243            PressureLevel::Normal => normal,
1244            PressureLevel::Elevated => {
1245                let scaled = (normal as f64 * self.elevated_batch_factor) as usize;
1246                scaled.max(1)
1247            }
1248            PressureLevel::Critical => 1,
1249        }
1250    }
1251
1252    /// Decide whether a non-essential actor should be paused.
1253    pub fn should_pause_non_essential(&self) -> bool {
1254        match self.monitor.current_pressure() {
1255            PressureLevel::Normal => false,
1256            PressureLevel::Elevated => true,
1257            PressureLevel::Critical => true,
1258        }
1259    }
1260}
1261
1262#[cfg(test)]
1263mod tests {
1264    use super::*;
1265
1266    #[test]
1267    fn test_supervision_strategy_constructors() {
1268        let s1 = SupervisionStrategy::one_for_one(5, Duration::from_secs(60));
1269        assert_eq!(s1.max_restarts(), Some(5));
1270        assert_eq!(s1.within(), Some(Duration::from_secs(60)));
1271
1272        let s2 = SupervisionStrategy::one_for_all(3, Duration::from_secs(300));
1273        assert_eq!(s2.max_restarts(), Some(3));
1274
1275        let s3 = SupervisionStrategy::rest_for_one(10, Duration::from_secs(120));
1276        assert_eq!(s3.max_restarts(), Some(10));
1277
1278        let s4 = SupervisionStrategy::simple_one_for_one();
1279        assert_eq!(s4.max_restarts(), None);
1280    }
1281
1282    #[test]
1283    fn test_exit_reason_checks() {
1284        assert!(ExitReason::Normal.is_normal());
1285        assert!(!ExitReason::Normal.is_abnormal());
1286
1287        assert!(!ExitReason::Killed.is_normal());
1288        assert!(ExitReason::Killed.is_abnormal());
1289
1290        assert!(ExitReason::Error("test".to_string()).is_abnormal());
1291        assert!(ExitReason::Shutdown.is_abnormal());
1292        assert!(ExitReason::Signaled(Signal::Pause).is_abnormal());
1293    }
1294
1295    #[test]
1296    fn test_child_spec_builder() {
1297        let spec = ChildSpec::new("test-child")
1298            .restart_policy(RestartPolicy::Transient)
1299            .shutdown_timeout(Duration::from_secs(60))
1300            .significant(true);
1301
1302        assert_eq!(spec.name, "test-child");
1303        assert_eq!(spec.restart_policy, RestartPolicy::Transient);
1304        assert_eq!(spec.shutdown_timeout, Duration::from_secs(60));
1305        assert!(spec.significant);
1306    }
1307
1308    #[test]
1309    fn test_supervisor_start_child() {
1310        let mut supervisor = Supervisor::new(SupervisionStrategy::default());
1311
1312        let spec = ChildSpec::new("child-1");
1313        let id = supervisor.start_child(spec).unwrap();
1314
1315        assert!(supervisor.has_child("child-1"));
1316        assert_eq!(supervisor.get_child("child-1").map(|c| c.id), Some(id));
1317        assert_eq!(supervisor.child_count(), 1);
1318    }
1319
1320    #[test]
1321    fn test_supervisor_duplicate_child() {
1322        let mut supervisor = Supervisor::new(SupervisionStrategy::default());
1323
1324        let spec1 = ChildSpec::new("child-1");
1325        let _ = supervisor.start_child(spec1).unwrap();
1326
1327        let spec2 = ChildSpec::new("child-1");
1328        let result = supervisor.start_child(spec2);
1329
1330        assert!(result.is_err());
1331    }
1332
1333    #[test]
1334    fn test_supervisor_stats() {
1335        let mut supervisor = Supervisor::new(SupervisionStrategy::default());
1336
1337        supervisor.start_child(ChildSpec::new("child-1")).unwrap();
1338        supervisor.start_child(ChildSpec::new("child-2")).unwrap();
1339        supervisor.start_child(ChildSpec::new("child-3")).unwrap();
1340
1341        let stats = supervisor.count_children();
1342        assert_eq!(stats.total, 3);
1343        assert_eq!(stats.running, 3);
1344        assert_eq!(stats.stopped, 0);
1345    }
1346
1347    #[tokio::test]
1348    async fn test_supervisor_stop_child() {
1349        let mut supervisor = Supervisor::new(SupervisionStrategy::default());
1350
1351        supervisor.start_child(ChildSpec::new("child-1")).unwrap();
1352        supervisor.stop_child("child-1").unwrap();
1353
1354        let child = supervisor.get_child("child-1").unwrap();
1355        matches!(
1356            child.state,
1357            ChildState::Stopping | ChildState::Stopped { .. }
1358        );
1359    }
1360
1361    #[tokio::test]
1362    async fn test_supervisor_terminate_child() {
1363        let mut supervisor = Supervisor::new(SupervisionStrategy::default());
1364
1365        supervisor.start_child(ChildSpec::new("child-1")).unwrap();
1366        supervisor
1367            .terminate_child("child-1", ExitReason::Error("test error".to_string()))
1368            .unwrap();
1369
1370        let child = supervisor.get_child("child-1").unwrap();
1371        match &child.state {
1372            ChildState::Stopped { reason } => {
1373                assert!(matches!(reason, ExitReason::Error(_)));
1374            }
1375            _ => panic!("Expected stopped state"),
1376        }
1377    }
1378
1379    #[tokio::test]
1380    async fn test_supervisor_handle_exit_permanent_restart() {
1381        let mut supervisor =
1382            Supervisor::new(SupervisionStrategy::one_for_one(5, Duration::from_secs(60)));
1383
1384        supervisor
1385            .start_child(ChildSpec::new("child-1").restart_policy(RestartPolicy::Permanent))
1386            .unwrap();
1387
1388        supervisor
1389            .handle_child_exit("child-1", ExitReason::Error("crashed".to_string()))
1390            .await
1391            .unwrap();
1392
1393        let child = supervisor.get_child("child-1").unwrap();
1394        match &child.state {
1395            ChildState::Restarting { attempt } => {
1396                assert_eq!(*attempt, 1);
1397            }
1398            _ => panic!("Expected restarting state"),
1399        }
1400
1401        let stats = supervisor.count_children();
1402        assert_eq!(stats.total_restarts, 1);
1403    }
1404
1405    #[tokio::test]
1406    async fn test_supervisor_handle_exit_temporary_no_restart() {
1407        let mut supervisor =
1408            Supervisor::new(SupervisionStrategy::one_for_one(5, Duration::from_secs(60)));
1409
1410        supervisor
1411            .start_child(ChildSpec::new("child-1").restart_policy(RestartPolicy::Temporary))
1412            .unwrap();
1413
1414        supervisor
1415            .handle_child_exit("child-1", ExitReason::Error("crashed".to_string()))
1416            .await
1417            .unwrap();
1418
1419        let child = supervisor.get_child("child-1").unwrap();
1420        match &child.state {
1421            ChildState::Stopped { .. } => {}
1422            _ => panic!("Expected stopped state"),
1423        }
1424    }
1425
1426    #[tokio::test]
1427    async fn test_supervisor_handle_exit_transient_normal_no_restart() {
1428        let mut supervisor =
1429            Supervisor::new(SupervisionStrategy::one_for_one(5, Duration::from_secs(60)));
1430
1431        supervisor
1432            .start_child(ChildSpec::new("child-1").restart_policy(RestartPolicy::Transient))
1433            .unwrap();
1434
1435        supervisor
1436            .handle_child_exit("child-1", ExitReason::Normal)
1437            .await
1438            .unwrap();
1439
1440        let child = supervisor.get_child("child-1").unwrap();
1441        match &child.state {
1442            ChildState::Stopped { reason } => {
1443                assert!(reason.is_normal());
1444            }
1445            _ => panic!("Expected stopped state"),
1446        }
1447    }
1448
1449    #[tokio::test]
1450    async fn test_supervisor_handle_exit_transient_abnormal_restart() {
1451        let mut supervisor =
1452            Supervisor::new(SupervisionStrategy::one_for_one(5, Duration::from_secs(60)));
1453
1454        supervisor
1455            .start_child(ChildSpec::new("child-1").restart_policy(RestartPolicy::Transient))
1456            .unwrap();
1457
1458        supervisor
1459            .handle_child_exit("child-1", ExitReason::Killed)
1460            .await
1461            .unwrap();
1462
1463        let child = supervisor.get_child("child-1").unwrap();
1464        matches!(child.state, ChildState::Restarting { .. });
1465    }
1466
1467    #[tokio::test]
1468    async fn test_supervisor_max_restarts_exceeded() {
1469        let mut supervisor =
1470            Supervisor::new(SupervisionStrategy::one_for_one(2, Duration::from_secs(60)));
1471
1472        supervisor
1473            .start_child(ChildSpec::new("child-1").restart_policy(RestartPolicy::Permanent))
1474            .unwrap();
1475
1476        supervisor
1477            .handle_child_exit("child-1", ExitReason::Error("crash".to_string()))
1478            .await
1479            .unwrap();
1480        supervisor
1481            .handle_child_exit("child-1", ExitReason::Error("crash".to_string()))
1482            .await
1483            .unwrap();
1484
1485        let result = supervisor
1486            .handle_child_exit("child-1", ExitReason::Error("crash".to_string()))
1487            .await;
1488        assert!(result.is_err());
1489
1490        let child = supervisor.get_child("child-1").unwrap();
1491        matches!(child.state, ChildState::Stopped { .. });
1492    }
1493
1494    #[tokio::test]
1495    async fn test_supervisor_give_up_escalation() {
1496        let mut supervisor =
1497            Supervisor::new(SupervisionStrategy::one_for_one(1, Duration::from_secs(60)));
1498        supervisor.set_escalation_action(EscalationAction::GiveUp);
1499
1500        supervisor
1501            .start_child(ChildSpec::new("child-1").restart_policy(RestartPolicy::Permanent))
1502            .unwrap();
1503
1504        supervisor
1505            .handle_child_exit("child-1", ExitReason::Error("crash".to_string()))
1506            .await
1507            .unwrap();
1508
1509        let result = supervisor
1510            .handle_child_exit("child-1", ExitReason::Error("crash".to_string()))
1511            .await;
1512        assert!(result.is_ok());
1513    }
1514
1515    #[test]
1516    fn test_supervisor_tree_creation() {
1517        let tree = SupervisorTree::new(SupervisionStrategy::default());
1518
1519        assert!(tree.get_supervisor(&tree.root()).is_some());
1520        assert_eq!(tree.tree_depth(), 0);
1521    }
1522
1523    #[test]
1524    fn test_supervisor_tree_add_supervisor() {
1525        let mut tree = SupervisorTree::new(SupervisionStrategy::default());
1526        let root = tree.root();
1527
1528        let child_sup = tree
1529            .add_supervisor(
1530                root,
1531                SupervisionStrategy::one_for_all(3, Duration::from_secs(60)),
1532            )
1533            .unwrap();
1534
1535        assert!(tree.get_supervisor(&child_sup).is_some());
1536        assert_eq!(tree.tree_depth(), 1);
1537        assert_eq!(tree.tree_stats().supervisor_count, 2);
1538    }
1539
1540    #[test]
1541    fn test_supervisor_tree_start_child() {
1542        let mut tree = SupervisorTree::new(SupervisionStrategy::default());
1543        let root = tree.root();
1544
1545        let child_id = tree
1546            .start_child_under(root, ChildSpec::new("child-1"))
1547            .unwrap();
1548
1549        let supervisor = tree.get_supervisor(&root).unwrap();
1550        let child = supervisor.get_child("child-1").unwrap();
1551        assert_eq!(child.id, child_id);
1552    }
1553
1554    #[test]
1555    fn test_supervisor_tree_invalid_parent() {
1556        let mut tree = SupervisorTree::new(SupervisionStrategy::default());
1557
1558        let fake_id = ActorId::new();
1559        let result = tree.add_supervisor(fake_id, SupervisionStrategy::default());
1560
1561        assert!(result.is_err());
1562    }
1563
1564    #[tokio::test]
1565    async fn test_supervisor_tree_terminate() {
1566        let mut tree = SupervisorTree::new(SupervisionStrategy::default());
1567        let root = tree.root();
1568
1569        let sup1 = tree
1570            .add_supervisor(root, SupervisionStrategy::default())
1571            .unwrap();
1572        let sup2 = tree
1573            .add_supervisor(sup1, SupervisionStrategy::default())
1574            .unwrap();
1575
1576        tree.start_child_under(root, ChildSpec::new("child-1"))
1577            .unwrap();
1578        tree.start_child_under(sup1, ChildSpec::new("child-2"))
1579            .unwrap();
1580        tree.start_child_under(sup2, ChildSpec::new("child-3"))
1581            .unwrap();
1582
1583        assert_eq!(tree.tree_stats().supervisor_count, 3);
1584        assert_eq!(tree.tree_stats().total_children, 3);
1585
1586        tree.terminate_tree(&sup1).await.unwrap();
1587
1588        assert_eq!(tree.tree_stats().supervisor_count, 1);
1589        assert_eq!(tree.tree_stats().total_children, 1);
1590    }
1591
1592    #[test]
1593    fn test_supervisor_tree_stats() {
1594        let mut tree = SupervisorTree::new(SupervisionStrategy::default());
1595        let root = tree.root();
1596
1597        tree.start_child_under(root, ChildSpec::new("child-1"))
1598            .unwrap();
1599        tree.start_child_under(root, ChildSpec::new("child-2"))
1600            .unwrap();
1601
1602        let sup1 = tree
1603            .add_supervisor(root, SupervisionStrategy::default())
1604            .unwrap();
1605        tree.start_child_under(sup1, ChildSpec::new("child-3"))
1606            .unwrap();
1607
1608        let stats = tree.tree_stats();
1609        assert_eq!(stats.supervisor_count, 2);
1610        assert_eq!(stats.total_children, 3);
1611        assert_eq!(stats.running_children, 3);
1612    }
1613
1614    #[test]
1615    fn test_supervisor_handle() {
1616        let supervisor = Supervisor::new(SupervisionStrategy::default());
1617        let handle = SupervisorHandle::new(supervisor);
1618
1619        let id = handle.id();
1620        assert_ne!(id, ActorId::new());
1621
1622        handle.start_child(ChildSpec::new("child-1")).unwrap();
1623
1624        let stats = handle.count_children();
1625        assert_eq!(stats.total, 1);
1626
1627        let children = handle.which_children();
1628        assert_eq!(children.len(), 1);
1629    }
1630
1631    #[test]
1632    fn test_child_state_is_active() {
1633        assert!(ChildState::Running.is_active());
1634        assert!(ChildState::Restarting { attempt: 1 }.is_active());
1635        assert!(!ChildState::Stopping.is_active());
1636        assert!(!ChildState::Stopped {
1637            reason: ExitReason::Normal
1638        }
1639        .is_active());
1640    }
1641
1642    #[test]
1643    fn test_supervisor_remove_child() {
1644        let mut supervisor = Supervisor::new(SupervisionStrategy::default());
1645
1646        supervisor.start_child(ChildSpec::new("child-1")).unwrap();
1647        assert_eq!(supervisor.child_count(), 1);
1648
1649        let removed = supervisor.remove_child("child-1").unwrap();
1650        assert_eq!(removed.spec.name, "child-1");
1651        assert_eq!(supervisor.child_count(), 0);
1652
1653        let result = supervisor.remove_child("nonexistent");
1654        assert!(result.is_err());
1655    }
1656
1657    #[tokio::test]
1658    async fn test_one_for_all_strategy() {
1659        let mut supervisor =
1660            Supervisor::new(SupervisionStrategy::one_for_all(5, Duration::from_secs(60)));
1661
1662        supervisor.start_child(ChildSpec::new("child-1")).unwrap();
1663        supervisor.start_child(ChildSpec::new("child-2")).unwrap();
1664        supervisor.start_child(ChildSpec::new("child-3")).unwrap();
1665
1666        supervisor
1667            .handle_child_exit("child-2", ExitReason::Error("crashed".to_string()))
1668            .await
1669            .unwrap();
1670
1671        for name in ["child-1", "child-2", "child-3"] {
1672            let child = supervisor.get_child(name).unwrap();
1673            matches!(child.state, ChildState::Restarting { .. });
1674        }
1675    }
1676
1677    #[tokio::test]
1678    async fn test_rest_for_one_strategy() {
1679        let mut supervisor = Supervisor::new(SupervisionStrategy::rest_for_one(
1680            5,
1681            Duration::from_secs(60),
1682        ));
1683
1684        supervisor.start_child(ChildSpec::new("child-1")).unwrap();
1685        supervisor.start_child(ChildSpec::new("child-2")).unwrap();
1686        supervisor.start_child(ChildSpec::new("child-3")).unwrap();
1687
1688        supervisor
1689            .handle_child_exit("child-2", ExitReason::Error("crashed".to_string()))
1690            .await
1691            .unwrap();
1692
1693        let child1 = supervisor.get_child("child-1").unwrap();
1694        matches!(child1.state, ChildState::Running);
1695
1696        for name in ["child-2", "child-3"] {
1697            let child = supervisor.get_child(name).unwrap();
1698            matches!(child.state, ChildState::Restarting { .. });
1699        }
1700    }
1701
1702    #[tokio::test]
1703    async fn test_simple_one_for_one_strategy() {
1704        let mut supervisor = Supervisor::new(SupervisionStrategy::simple_one_for_one());
1705
1706        supervisor
1707            .start_child(ChildSpec::new("child-1").restart_policy(RestartPolicy::Permanent))
1708            .unwrap();
1709
1710        supervisor
1711            .handle_child_exit("child-1", ExitReason::Error("crashed".to_string()))
1712            .await
1713            .unwrap();
1714
1715        let child = supervisor.get_child("child-1").unwrap();
1716        match &child.state {
1717            ChildState::Stopped { .. } => {}
1718            _ => panic!("Expected stopped state for simple_one_for_one"),
1719        }
1720    }
1721
1722    // -----------------------------------------------------------------------
1723    // Graceful Degradation tests
1724    // -----------------------------------------------------------------------
1725
1726    #[test]
1727    fn test_pressure_level_ordering() {
1728        assert!(PressureLevel::Normal < PressureLevel::Elevated);
1729        assert!(PressureLevel::Elevated < PressureLevel::Critical);
1730    }
1731
1732    #[test]
1733    fn test_pressure_level_display() {
1734        assert_eq!(format!("{}", PressureLevel::Normal), "normal");
1735        assert_eq!(format!("{}", PressureLevel::Elevated), "elevated");
1736        assert_eq!(format!("{}", PressureLevel::Critical), "critical");
1737    }
1738
1739    #[test]
1740    fn test_static_resource_monitor() {
1741        let normal = StaticResourceMonitor::new(PressureLevel::Normal);
1742        assert_eq!(normal.current_pressure(), PressureLevel::Normal);
1743
1744        let critical = StaticResourceMonitor::new(PressureLevel::Critical);
1745        assert_eq!(critical.current_pressure(), PressureLevel::Critical);
1746    }
1747
1748    #[test]
1749    fn test_fn_resource_monitor() {
1750        let monitor = FnResourceMonitor::new(Box::new(|| true), Box::new(|| false));
1751        assert_eq!(monitor.current_pressure(), PressureLevel::Elevated);
1752
1753        let monitor2 = FnResourceMonitor::new(Box::new(|| true), Box::new(|| true));
1754        assert_eq!(monitor2.current_pressure(), PressureLevel::Critical);
1755
1756        let monitor3 = FnResourceMonitor::new(Box::new(|| false), Box::new(|| false));
1757        assert_eq!(monitor3.current_pressure(), PressureLevel::Normal);
1758    }
1759
1760    #[test]
1761    fn test_degradation_controller_admit_spawn() {
1762        let controller =
1763            DegradationController::new(StaticResourceMonitor::new(PressureLevel::Normal), 0.5);
1764        assert_eq!(controller.admit_spawn(), DegradationDecision::Allow);
1765
1766        let controller =
1767            DegradationController::new(StaticResourceMonitor::new(PressureLevel::Elevated), 0.5);
1768        assert_eq!(controller.admit_spawn(), DegradationDecision::Throttle);
1769
1770        let controller =
1771            DegradationController::new(StaticResourceMonitor::new(PressureLevel::Critical), 0.5);
1772        assert_eq!(controller.admit_spawn(), DegradationDecision::Reject);
1773    }
1774
1775    #[test]
1776    fn test_degradation_controller_batch_size() {
1777        let controller =
1778            DegradationController::new(StaticResourceMonitor::new(PressureLevel::Normal), 0.5);
1779        assert_eq!(controller.effective_batch_size(100), 100);
1780
1781        let controller =
1782            DegradationController::new(StaticResourceMonitor::new(PressureLevel::Elevated), 0.5);
1783        assert_eq!(controller.effective_batch_size(100), 50);
1784        assert_eq!(controller.effective_batch_size(1), 1);
1785
1786        let controller =
1787            DegradationController::new(StaticResourceMonitor::new(PressureLevel::Critical), 0.5);
1788        assert_eq!(controller.effective_batch_size(100), 1);
1789    }
1790
1791    #[test]
1792    fn test_degradation_controller_pause_non_essential() {
1793        let controller =
1794            DegradationController::new(StaticResourceMonitor::new(PressureLevel::Normal), 0.5);
1795        assert!(!controller.should_pause_non_essential());
1796
1797        let controller =
1798            DegradationController::new(StaticResourceMonitor::new(PressureLevel::Elevated), 0.5);
1799        assert!(controller.should_pause_non_essential());
1800
1801        let controller =
1802            DegradationController::new(StaticResourceMonitor::new(PressureLevel::Critical), 0.5);
1803        assert!(controller.should_pause_non_essential());
1804    }
1805
1806    #[test]
1807    fn test_degradation_decision_is_allowed() {
1808        assert!(DegradationDecision::Allow.is_allowed());
1809        assert!(DegradationDecision::Throttle.is_allowed());
1810        assert!(!DegradationDecision::Reject.is_allowed());
1811    }
1812}