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}