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    task_stable_after_delay: Duration,
20    max_dead_tasks_percentage_threshold: Option<f64>,
21}
22
23impl SupervisorBuilder {
24    /// Creates a new builder with default configuration values.
25    pub fn new() -> Self {
26        Self {
27            tasks: HashMap::new(),
28            health_check_interval: Duration::from_millis(200),
29            max_restart_attempts: 5,
30            base_restart_delay: Duration::from_secs(1),
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    pub fn with_task(mut self, name: &str, task: impl CloneableSupervisedTask) -> Self {
38        let handle =
39            TaskHandle::from_task(task, self.max_restart_attempts, self.base_restart_delay);
40        self.tasks.insert(name.into(), handle);
41        self
42    }
43
44    /// Sets the interval between health checks.
45    pub fn with_health_check_interval(mut self, interval: Duration) -> Self {
46        self.health_check_interval = interval;
47        self
48    }
49
50    /// Sets the maximum number of restart attempts for tasks.
51    pub fn with_max_restart_attempts(mut self, attempts: u32) -> Self {
52        self.max_restart_attempts = attempts;
53        self
54    }
55
56    /// Sets the base delay for task restarts, used in exponential backoff.
57    pub fn with_base_restart_delay(mut self, delay: Duration) -> Self {
58        self.base_restart_delay = delay;
59        self
60    }
61
62    /// Sets the delay after which a task is considered stable and healthy.
63    /// When a task is considered stable, its restarts are reset to zero.
64    pub fn with_task_being_stable_after(mut self, delay: Duration) -> Self {
65        self.task_stable_after_delay = delay;
66        self
67    }
68
69    /// Sets the threshold for the percentage of dead tasks that will trigger a supervisor shutdown.
70    ///
71    /// The `threshold_percentage` should be a value between 0.0 (0%) and 1.0 (100%).
72    /// If the percentage of dead tasks exceeds this value, the supervisor will shut down
73    /// and return an error.
74    pub fn with_dead_tasks_threshold(mut self, threshold_percentage: Option<f64>) -> Self {
75        self.max_dead_tasks_percentage_threshold = threshold_percentage.map(|t| t.clamp(0.0, 1.0));
76        self
77    }
78
79    /// Constructs the `Supervisor` with the configured settings.
80    pub fn build(self) -> Supervisor {
81        let (internal_tx, internal_rx) = mpsc::unbounded_channel();
82        let (user_tx, user_rx) = mpsc::unbounded_channel();
83        Supervisor {
84            tasks: self.tasks,
85            health_check_interval: self.health_check_interval,
86            base_restart_delay: self.base_restart_delay,
87            max_restart_attempts: self.max_restart_attempts,
88            task_is_stable_after: self.task_stable_after_delay,
89            max_dead_tasks_percentage_threshold: self.max_dead_tasks_percentage_threshold,
90            internal_tx,
91            internal_rx,
92            external_tx: user_tx,
93            external_rx: user_rx,
94        }
95    }
96}
97
98impl Default for SupervisorBuilder {
99    fn default() -> Self {
100        Self::new()
101    }
102}