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: Option<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: Some(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 = Some(attempts);
68        self
69    }
70
71    /// Disables the restart limit, allowing tasks to restart indefinitely.
72    pub fn with_unlimited_restarts(mut self) -> Self {
73        self.max_restart_attempts = None;
74        self
75    }
76
77    /// Sets the base delay for task restarts, used in exponential backoff.
78    pub fn with_base_restart_delay(mut self, delay: Duration) -> Self {
79        self.base_restart_delay = delay;
80        self
81    }
82
83    /// Sets the delay after which a task is considered stable and healthy.
84    /// When a task is considered stable, its restarts are reset to zero.
85    pub fn with_task_being_stable_after(mut self, delay: Duration) -> Self {
86        self.task_stable_after_delay = delay;
87        self
88    }
89
90    /// Sets the threshold for the percentage of dead tasks that will trigger a supervisor shutdown.
91    ///
92    /// The `threshold_percentage` should be a value between 0.0 (0%) and 1.0 (100%).
93    /// If the percentage of dead tasks exceeds this value, the supervisor will shut down
94    /// and return an error.
95    pub fn with_dead_tasks_threshold(mut self, threshold_percentage: Option<f64>) -> Self {
96        self.max_dead_tasks_percentage_threshold = threshold_percentage.map(|t| t.clamp(0.0, 1.0));
97        self
98    }
99
100    /// Constructs the `Supervisor` with the configured settings.
101    pub fn build(self) -> Supervisor {
102        let (internal_tx, internal_rx) = mpsc::unbounded_channel();
103        let (user_tx, user_rx) = mpsc::unbounded_channel();
104        Supervisor {
105            tasks: self.tasks,
106            health_check_interval: self.health_check_interval,
107            base_restart_delay: self.base_restart_delay,
108            max_restart_attempts: self.max_restart_attempts,
109            max_backoff_exponent: self.max_backoff_exponent,
110            task_is_stable_after: self.task_stable_after_delay,
111            max_dead_tasks_percentage_threshold: self.max_dead_tasks_percentage_threshold,
112            internal_tx,
113            internal_rx,
114            external_tx: user_tx,
115            external_rx: user_rx,
116        }
117    }
118}
119
120impl Default for SupervisorBuilder {
121    fn default() -> Self {
122        Self::new()
123    }
124}
task_supervisor/supervisor/builder.rs

task_supervisor/supervisor/
builder.rs
