Skip to main content

task_supervisor/supervisor/
builder.rs

1use std::{collections::HashMap, sync::Arc, time::Duration};
2
3use tokio::sync::mpsc;
4
5use crate::{task::TaskHandle, SupervisedTask, Supervisor};
6
7/// Builds a `Supervisor` instance with configurable parameters.
8///
9/// Configuration methods can be called in any order. Task restart settings
10/// are applied to all tasks uniformly at `build()` time, regardless of when
11/// `with_task()` was called.
12pub struct SupervisorBuilder {
13    tasks: HashMap<Arc<str>, TaskHandle>,
14    health_check_interval: Duration,
15    max_restart_attempts: Option<u32>,
16    base_restart_delay: Duration,
17    max_backoff_exponent: u32,
18    task_stable_after_delay: Duration,
19    max_dead_tasks_percentage_threshold: Option<f64>,
20}
21
22impl SupervisorBuilder {
23    /// Creates a new builder with default configuration values.
24    pub fn new() -> Self {
25        Self {
26            tasks: HashMap::new(),
27            health_check_interval: Duration::from_millis(200),
28            max_restart_attempts: Some(5),
29            base_restart_delay: Duration::from_secs(1),
30            max_backoff_exponent: 5,
31            task_stable_after_delay: Duration::from_secs(80),
32            max_dead_tasks_percentage_threshold: None,
33        }
34    }
35
36    /// Adds a task to the supervisor with the specified name.
37    ///
38    /// Task restart configuration is applied at `build()` time, so the
39    /// order of `with_task()` vs configuration methods does not matter.
40    pub fn with_task(mut self, name: &str, task: impl SupervisedTask + Clone) -> Self {
41        let handle = TaskHandle::from_task(task);
42        self.tasks.insert(Arc::from(name), handle);
43        self
44    }
45
46    /// Sets the maximum exponent used in exponential backoff.
47    ///
48    /// The restart delay is calculated as `base_restart_delay * 2^min(attempt, max_backoff_exponent)`.
49    /// For example, with a base delay of 1s and exponent cap of 5, the maximum delay is 32s.
50    pub fn with_max_backoff_exponent(mut self, exponent: u32) -> Self {
51        self.max_backoff_exponent = exponent;
52        self
53    }
54
55    /// Sets the interval between health checks.
56    pub fn with_health_check_interval(mut self, interval: Duration) -> Self {
57        self.health_check_interval = interval;
58        self
59    }
60
61    /// Sets the maximum number of restart attempts for tasks.
62    pub fn with_max_restart_attempts(mut self, attempts: u32) -> Self {
63        self.max_restart_attempts = Some(attempts);
64        self
65    }
66
67    /// Disables the restart limit, allowing tasks to restart indefinitely.
68    pub fn with_unlimited_restarts(mut self) -> Self {
69        self.max_restart_attempts = None;
70        self
71    }
72
73    /// Sets the base delay for task restarts, used in exponential backoff.
74    pub fn with_base_restart_delay(mut self, delay: Duration) -> Self {
75        self.base_restart_delay = delay;
76        self
77    }
78
79    /// Sets the delay after which a task is considered stable and healthy.
80    /// When a task is considered stable, its restarts are reset to zero.
81    pub fn with_task_being_stable_after(mut self, delay: Duration) -> Self {
82        self.task_stable_after_delay = delay;
83        self
84    }
85
86    /// Sets the threshold for the percentage of dead tasks that will trigger a supervisor shutdown.
87    ///
88    /// The `threshold_percentage` should be a value between 0.0 (0%) and 1.0 (100%).
89    /// If the percentage of dead tasks exceeds this value, the supervisor will shut down
90    /// and return an error.
91    pub fn with_dead_tasks_threshold(mut self, threshold_percentage: Option<f64>) -> Self {
92        self.max_dead_tasks_percentage_threshold = threshold_percentage.map(|t| t.clamp(0.0, 1.0));
93        self
94    }
95
96    /// Constructs the `Supervisor` with the configured settings.
97    ///
98    /// Applies restart configuration (max attempts, base delay, backoff exponent)
99    /// to all registered tasks.
100    pub fn build(mut self) -> Supervisor {
101        // Apply configuration to all tasks — this fixes the ordering issue where
102        // tasks added before config methods would get stale defaults.
103        for task_handle in self.tasks.values_mut() {
104            task_handle.max_restart_attempts = self.max_restart_attempts;
105            task_handle.base_restart_delay = self.base_restart_delay;
106            task_handle.max_backoff_exponent = self.max_backoff_exponent;
107        }
108
109        let (internal_tx, internal_rx) = mpsc::unbounded_channel();
110        let (user_tx, user_rx) = mpsc::unbounded_channel();
111        Supervisor {
112            tasks: self.tasks,
113            health_check_interval: self.health_check_interval,
114            base_restart_delay: self.base_restart_delay,
115            max_restart_attempts: self.max_restart_attempts,
116            max_backoff_exponent: self.max_backoff_exponent,
117            task_is_stable_after: self.task_stable_after_delay,
118            max_dead_tasks_percentage_threshold: self.max_dead_tasks_percentage_threshold,
119            internal_tx,
120            internal_rx,
121            external_tx: user_tx,
122            external_rx: user_rx,
123        }
124    }
125}
126
127impl Default for SupervisorBuilder {
128    fn default() -> Self {
129        Self::new()
130    }
131}