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` with configurable parameters.
8///
9/// `with_task()` and configuration methods can be called in any order;
10/// task restart settings are applied uniformly at `build()` time.
11pub struct SupervisorBuilder {
12    tasks: HashMap<Arc<str>, TaskHandle>,
13    health_check_interval: Duration,
14    max_restart_attempts: Option<u32>,
15    base_restart_delay: Duration,
16    max_backoff_exponent: u32,
17    task_stable_after_delay: Duration,
18    max_dead_tasks_percentage_threshold: Option<f64>,
19}
20
21impl SupervisorBuilder {
22    pub fn new() -> Self {
23        Self {
24            tasks: HashMap::new(),
25            health_check_interval: Duration::from_millis(200),
26            max_restart_attempts: Some(5),
27            base_restart_delay: Duration::from_secs(1),
28            max_backoff_exponent: 5,
29            task_stable_after_delay: Duration::from_secs(80),
30            max_dead_tasks_percentage_threshold: None,
31        }
32    }
33
34    pub fn with_task(mut self, name: &str, task: impl SupervisedTask + Clone) -> Self {
35        let handle = TaskHandle::from_task(task);
36        self.tasks.insert(Arc::from(name), handle);
37        self
38    }
39
40    /// Caps the exponent in the backoff formula:
41    /// `delay = base_restart_delay * 2^min(attempt, max_backoff_exponent)`.
42    pub fn with_max_backoff_exponent(mut self, exponent: u32) -> Self {
43        self.max_backoff_exponent = exponent;
44        self
45    }
46
47    pub fn with_health_check_interval(mut self, interval: Duration) -> Self {
48        self.health_check_interval = interval;
49        self
50    }
51
52    pub fn with_max_restart_attempts(mut self, attempts: u32) -> Self {
53        self.max_restart_attempts = Some(attempts);
54        self
55    }
56
57    pub fn with_unlimited_restarts(mut self) -> Self {
58        self.max_restart_attempts = None;
59        self
60    }
61
62    pub fn with_base_restart_delay(mut self, delay: Duration) -> Self {
63        self.base_restart_delay = delay;
64        self
65    }
66
67    /// Once a task has been healthy for this long, its restart counter resets to zero.
68    pub fn with_task_being_stable_after(mut self, delay: Duration) -> Self {
69        self.task_stable_after_delay = delay;
70        self
71    }
72
73    /// Shuts down the supervisor if the fraction of dead tasks exceeds
74    /// `threshold_percentage` (0.0 – 1.0). `None` disables the check.
75    pub fn with_dead_tasks_threshold(mut self, threshold_percentage: Option<f64>) -> Self {
76        self.max_dead_tasks_percentage_threshold = threshold_percentage.map(|t| t.clamp(0.0, 1.0));
77        self
78    }
79
80    pub fn build(mut self) -> Supervisor {
81        for task_handle in self.tasks.values_mut() {
82            task_handle.max_restart_attempts = self.max_restart_attempts;
83            task_handle.base_restart_delay = self.base_restart_delay;
84            task_handle.max_backoff_exponent = self.max_backoff_exponent;
85        }
86
87        let (internal_tx, internal_rx) = mpsc::unbounded_channel();
88        let (user_tx, user_rx) = mpsc::unbounded_channel();
89        Supervisor {
90            tasks: self.tasks,
91            health_check_interval: self.health_check_interval,
92            base_restart_delay: self.base_restart_delay,
93            max_restart_attempts: self.max_restart_attempts,
94            max_backoff_exponent: self.max_backoff_exponent,
95            task_is_stable_after: self.task_stable_after_delay,
96            max_dead_tasks_percentage_threshold: self.max_dead_tasks_percentage_threshold,
97            internal_tx,
98            internal_rx,
99            external_tx: user_tx,
100            external_rx: user_rx,
101        }
102    }
103}
104
105impl Default for SupervisorBuilder {
106    fn default() -> Self {
107        Self::new()
108    }
109}
task_supervisor/supervisor/builder.rs

task_supervisor/supervisor/
builder.rs
