task_supervisor/supervisor/
builder.rs

1use std::{collections::HashMap, time::Duration};
2
3use tokio::sync::mpsc;
4
5use crate::{
6    task::{CloneableSupervisedTask, TaskHandle},
7    Supervisor,
8};
9
10/// Builds a `Supervisor` instance with configurable parameters.
11///
12/// Allows customization of task timeout, heartbeat interval, health check timing,
13/// and per-task restart settings.
14pub struct SupervisorBuilder {
15    tasks: HashMap<String, TaskHandle>,
16    health_check_interval: Duration,
17    max_restart_attempts: u32,
18    base_restart_delay: Duration,
19    max_backoff_exponent: u32,
20    task_stable_after_delay: Duration,
21    max_dead_tasks_percentage_threshold: Option<f64>,
22}
23
24impl SupervisorBuilder {
25    /// Creates a new builder with default configuration values.
26    pub fn new() -> Self {
27        Self {
28            tasks: HashMap::new(),
29            health_check_interval: Duration::from_millis(200),
30            max_restart_attempts: 5,
31            base_restart_delay: Duration::from_secs(1),
32            max_backoff_exponent: 5,
33            task_stable_after_delay: Duration::from_secs(80),
34            max_dead_tasks_percentage_threshold: None,
35        }
36    }
37
38    /// Adds a task to the supervisor with the specified name.
39    pub fn with_task(mut self, name: &str, task: impl CloneableSupervisedTask) -> Self {
40        let handle = TaskHandle::from_task(
41            task,
42            self.max_restart_attempts,
43            self.base_restart_delay,
44            self.max_backoff_exponent,
45        );
46        self.tasks.insert(name.into(), handle);
47        self
48    }
49
50    /// Sets the maximum exponent used in exponential backoff.
51    ///
52    /// The restart delay is calculated as `base_restart_delay * 2^min(attempt, max_backoff_exponent)`.
53    /// For example, with a base delay of 1s and exponent cap of 5, the maximum delay is 32s.
54    pub fn with_max_backoff_exponent(mut self, exponent: u32) -> Self {
55        self.max_backoff_exponent = exponent;
56        self
57    }
58
59    /// Sets the interval between health checks.
60    pub fn with_health_check_interval(mut self, interval: Duration) -> Self {
61        self.health_check_interval = interval;
62        self
63    }
64
65    /// Sets the maximum number of restart attempts for tasks.
66    pub fn with_max_restart_attempts(mut self, attempts: u32) -> Self {
67        self.max_restart_attempts = attempts;
68        self
69    }
70
71    /// Sets the base delay for task restarts, used in exponential backoff.
72    pub fn with_base_restart_delay(mut self, delay: Duration) -> Self {
73        self.base_restart_delay = delay;
74        self
75    }
76
77    /// Sets the delay after which a task is considered stable and healthy.
78    /// When a task is considered stable, its restarts are reset to zero.
79    pub fn with_task_being_stable_after(mut self, delay: Duration) -> Self {
80        self.task_stable_after_delay = delay;
81        self
82    }
83
84    /// Sets the threshold for the percentage of dead tasks that will trigger a supervisor shutdown.
85    ///
86    /// The `threshold_percentage` should be a value between 0.0 (0%) and 1.0 (100%).
87    /// If the percentage of dead tasks exceeds this value, the supervisor will shut down
88    /// and return an error.
89    pub fn with_dead_tasks_threshold(mut self, threshold_percentage: Option<f64>) -> Self {
90        self.max_dead_tasks_percentage_threshold = threshold_percentage.map(|t| t.clamp(0.0, 1.0));
91        self
92    }
93
94    /// Constructs the `Supervisor` with the configured settings.
95    pub fn build(self) -> Supervisor {
96        let (internal_tx, internal_rx) = mpsc::unbounded_channel();
97        let (user_tx, user_rx) = mpsc::unbounded_channel();
98        Supervisor {
99            tasks: self.tasks,
100            health_check_interval: self.health_check_interval,
101            base_restart_delay: self.base_restart_delay,
102            max_restart_attempts: self.max_restart_attempts,
103            max_backoff_exponent: self.max_backoff_exponent,
104            task_is_stable_after: self.task_stable_after_delay,
105            max_dead_tasks_percentage_threshold: self.max_dead_tasks_percentage_threshold,
106            internal_tx,
107            internal_rx,
108            external_tx: user_tx,
109            external_rx: user_rx,
110        }
111    }
112}
113
114impl Default for SupervisorBuilder {
115    fn default() -> Self {
116        Self::new()
117    }
118}