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