Skip to main content

task_supervisor/
builder.rs

1use std::{collections::HashMap, sync::Arc, time::Duration};
2
3use crate::{
4    supervisor::{Config, Supervisor},
5    task::{Slot, SupervisedTask},
6    SupervisorHandle,
7};
8
9/// Configures and spawns a supervisor.
10///
11/// | Setting | Default |
12/// | --- | --- |
13/// | `restart_limit` | 5 restarts in 60s |
14/// | `base_restart_delay` | 1s |
15/// | `max_restart_delay` | 30s |
16/// | `restart_jitter` | 0.1 |
17/// | `dead_tasks_threshold` | disabled |
18/// | `stop_timeout` | 5s |
19pub struct SupervisorBuilder {
20    tasks: HashMap<Arc<str>, Slot>,
21    config: Config,
22}
23
24impl SupervisorBuilder {
25    /// A builder with the defaults listed above and no tasks.
26    pub fn new() -> Self {
27        Self {
28            tasks: HashMap::new(),
29            config: Config {
30                max_restarts: Some(5),
31                restart_window: Duration::from_secs(60),
32                base_restart_delay: Duration::from_secs(1),
33                max_restart_delay: Duration::from_secs(30),
34                restart_jitter: 0.1,
35                dead_tasks_threshold: None,
36                stop_timeout: Duration::from_secs(5),
37            },
38        }
39    }
40
41    /// Registers a task to start with the supervisor. A later call with the
42    /// same name replaces the earlier task.
43    pub fn with_task(mut self, name: &str, task: impl SupervisedTask) -> Self {
44        let name: Arc<str> = Arc::from(name);
45        self.tasks
46            .insert(Arc::clone(&name), Slot::new(name, Box::new(task)));
47        self
48    }
49
50    /// A task that would need more than `max_restarts` restarts within any
51    /// `window` is declared [`Dead`](crate::TaskStatus::Dead) instead.
52    /// `max_restarts = 0` disables restarts. The backoff delay doubles with
53    /// each restart inside the window, so a task that stays up longer than
54    /// `window` starts again from `base_restart_delay`.
55    pub fn with_restart_limit(mut self, max_restarts: u32, window: Duration) -> Self {
56        self.config.max_restarts = Some(max_restarts);
57        self.config.restart_window = window;
58        self
59    }
60
61    /// Never give up restarting a failing task. The window still drives the
62    /// backoff delay.
63    pub fn with_unlimited_restarts(mut self) -> Self {
64        self.config.max_restarts = None;
65        self
66    }
67
68    /// Delay before the first restart in a window. Each further restart in
69    /// the window doubles it.
70    pub fn with_base_restart_delay(mut self, delay: Duration) -> Self {
71        self.config.base_restart_delay = delay;
72        self
73    }
74
75    /// Upper bound for the restart delay.
76    pub fn with_max_restart_delay(mut self, delay: Duration) -> Self {
77        self.config.max_restart_delay = delay;
78        self
79    }
80
81    /// Randomly shortens each restart delay by up to this fraction of itself
82    /// (`0.0..=1.0`), so tasks that failed together do not restart in
83    /// lockstep. `0.0` disables it.
84    pub fn with_restart_jitter(mut self, fraction: f64) -> Self {
85        self.config.restart_jitter = fraction.clamp(0.0, 1.0);
86        self
87    }
88
89    /// Shuts the supervisor down with
90    /// [`TooManyDeadTasks`](crate::SupervisorError::TooManyDeadTasks) once at
91    /// least one task is dead and `dead / total >= fraction`. Killed tasks
92    /// count as dead. `fraction` is clamped to `0.0..=1.0`: `0.0` means "any
93    /// dead task", `1.0` means "every task".
94    pub fn with_dead_tasks_threshold(mut self, fraction: f64) -> Self {
95        self.config.dead_tasks_threshold = Some(fraction.clamp(0.0, 1.0));
96        self
97    }
98
99    /// How long a cancelled run may keep going before its future is dropped.
100    /// Applies to kill, restart and shutdown.
101    pub fn with_stop_timeout(mut self, timeout: Duration) -> Self {
102        self.config.stop_timeout = timeout;
103        self
104    }
105
106    /// Spawns the supervisor and every registered task.
107    ///
108    /// # Panics
109    ///
110    /// Panics if called outside a Tokio runtime.
111    pub fn spawn(self) -> SupervisorHandle {
112        Supervisor::spawn(self.tasks, self.config)
113    }
114}
115
116impl Default for SupervisorBuilder {
117    fn default() -> Self {
118        Self::new()
119    }
120}