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    timeout_threshold: Duration,
17    heartbeat_interval: Duration,
18    health_check_initial_delay: Duration,
19    health_check_interval: Duration,
20    max_restart_attempts: u32,
21    base_restart_delay: Duration,
22    task_stable_after_delay: Duration,
23}
24
25impl SupervisorBuilder {
26    /// Creates a new builder with default configuration values.
27    pub fn new() -> Self {
28        Self {
29            tasks: HashMap::new(),
30            timeout_threshold: Duration::from_secs(2),
31            heartbeat_interval: Duration::from_millis(100),
32            health_check_initial_delay: Duration::from_secs(3),
33            health_check_interval: Duration::from_millis(200),
34            max_restart_attempts: 5,
35            base_restart_delay: Duration::from_secs(1),
36            task_stable_after_delay: Duration::from_secs(80),
37        }
38    }
39
40    /// Adds a task to the supervisor with the specified name.
41    pub fn with_task(mut self, name: &str, task: impl CloneableSupervisedTask) -> Self {
42        let handle =
43            TaskHandle::from_task(task, self.max_restart_attempts, self.base_restart_delay);
44        self.tasks.insert(name.into(), handle);
45        self
46    }
47
48    /// Sets the timeout threshold for detecting task crashes.
49    pub fn with_timeout_threshold(mut self, threshold: Duration) -> Self {
50        self.timeout_threshold = threshold;
51        self
52    }
53
54    /// Sets the interval at which tasks send heartbeats.
55    pub fn with_heartbeat_interval(mut self, interval: Duration) -> Self {
56        self.heartbeat_interval = interval;
57        self
58    }
59
60    /// Sets the initial delay before health checks begin.
61    pub fn with_health_check_initial_delay(mut self, delay: Duration) -> Self {
62        self.health_check_initial_delay = delay;
63        self
64    }
65
66    /// Sets the interval between health checks.
67    pub fn with_health_check_interval(mut self, interval: Duration) -> Self {
68        self.health_check_interval = interval;
69        self
70    }
71
72    /// Sets the maximum number of restart attempts for tasks.
73    pub fn with_max_restart_attempts(mut self, attempts: u32) -> Self {
74        self.max_restart_attempts = attempts;
75        self
76    }
77
78    /// Sets the base delay for task restarts, used in exponential backoff.
79    pub fn with_base_restart_delay(mut self, delay: Duration) -> Self {
80        self.base_restart_delay = delay;
81        self
82    }
83
84    /// Sets the delay after which a task is considered stable and healthy.
85    /// When a task is considered stable, its restarts are reset to zero.
86    pub fn with_task_being_stable_after(mut self, delay: Duration) -> Self {
87        self.task_stable_after_delay = delay;
88        self
89    }
90
91    /// Constructs the `Supervisor` with the configured settings.
92    pub fn build(self) -> Supervisor {
93        let (internal_tx, internal_rx) = mpsc::unbounded_channel();
94        let (user_tx, user_rx) = mpsc::unbounded_channel();
95        Supervisor {
96            tasks: self.tasks,
97            timeout_threshold: self.timeout_threshold,
98            heartbeat_interval: self.heartbeat_interval,
99            health_check_initial_delay: self.health_check_initial_delay,
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            task_is_stable_after: self.task_stable_after_delay,
104            internal_tx,
105            internal_rx,
106            external_tx: user_tx,
107            external_rx: user_rx,
108        }
109    }
110}
111
112impl Default for SupervisorBuilder {
113    fn default() -> Self {
114        Self::new()
115    }
116}