task_supervisor/
task.rs

1use std::{collections::VecDeque, future::Future, pin::Pin, sync::Arc};
2
3use tokio::sync::oneshot;
4use tokio_util::sync::CancellationToken;
5
6use crate::handle::SupervisorHandleError;
7
8/// Error type returned by [`SupervisedTask::run`].
9///
10/// Any `std::error::Error + Send + Sync` converts into it with `?`, including
11/// `anyhow::Error`.
12pub type TaskError = Box<dyn std::error::Error + Send + Sync + 'static>;
13
14/// Outcome of one run of a [`SupervisedTask`].
15pub type TaskResult = Result<(), TaskError>;
16
17/// A long-lived unit of work managed by the supervisor.
18///
19/// Every start or restart clones the registered instance and passes the clone
20/// to `run`. Owned fields reset on every run; `Arc` fields are shared.
21///
22/// | `run` outcome | Supervisor reaction |
23/// | --- | --- |
24/// | `Ok(())` | [`Completed`](TaskStatus::Completed); not restarted |
25/// | `Err(_)` or panic | Restarted with backoff, [`Dead`](TaskStatus::Dead) once the restart limit is reached |
26///
27/// `cancel` fires on kill, restart and shutdown. The run then has
28/// `stop_timeout` to return before its future is dropped. Ignore it, wrap the
29/// work in `cancel.run_until_cancelled(..)`, or `select!` on
30/// `cancel.cancelled()` when there is cleanup to do.
31///
32/// # Example
33///
34/// ```rust
35/// use std::sync::Arc;
36/// use std::sync::atomic::{AtomicUsize, Ordering};
37/// use std::time::Duration;
38/// use task_supervisor::{CancellationToken, SupervisedTask, TaskResult};
39///
40/// #[derive(Clone)]
41/// struct Worker {
42///     /// Reset on every run (owned, cloned from the original).
43///     polls: u64,
44///     /// Shared across runs (`Arc`, cloned by reference).
45///     total_polls: Arc<AtomicUsize>,
46///     db: Arc<Db>,
47/// }
48///
49/// impl SupervisedTask for Worker {
50///     async fn run(mut self, cancel: CancellationToken) -> TaskResult {
51///         loop {
52///             tokio::select! {
53///                 _ = cancel.cancelled() => break,
54///                 _ = tokio::time::sleep(Duration::from_secs(1)) => {
55///                     self.polls += 1;
56///                     self.total_polls.fetch_add(1, Ordering::Relaxed);
57///                 }
58///             }
59///         }
60///         self.db.flush().await?; // cleanup after cancellation
61///         Ok(())
62///     }
63/// }
64/// # struct Db;
65/// # impl Db { async fn flush(&self) -> Result<(), std::io::Error> { Ok(()) } }
66/// ```
67pub trait SupervisedTask: Clone + Send + 'static {
68    /// Runs one instance of the task until it completes, fails, or is cancelled.
69    fn run(self, cancel: CancellationToken) -> impl Future<Output = TaskResult> + Send;
70}
71
72/// Object-safe view of [`SupervisedTask`].
73pub(crate) trait DynTask: Send + 'static {
74    fn run(
75        self: Box<Self>,
76        cancel: CancellationToken,
77    ) -> Pin<Box<dyn Future<Output = TaskResult> + Send>>;
78    fn clone_box(&self) -> Box<dyn DynTask>;
79}
80
81impl<T: SupervisedTask> DynTask for T {
82    fn run(
83        self: Box<Self>,
84        cancel: CancellationToken,
85    ) -> Pin<Box<dyn Future<Output = TaskResult> + Send>> {
86        Box::pin(SupervisedTask::run(*self, cancel))
87    }
88
89    fn clone_box(&self) -> Box<dyn DynTask> {
90        Box::new(self.clone())
91    }
92}
93
94/// Observable state of a supervised task.
95#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
96pub enum TaskStatus {
97    /// `run` is executing.
98    Running,
99    /// `run` failed; a restart is scheduled after the backoff delay.
100    Restarting,
101    /// Cancellation was requested; waiting for `run` to return
102    /// (bounded by `stop_timeout`).
103    Stopping,
104    /// `run` returned `Ok(())`. Terminal unless restarted manually.
105    Completed,
106    /// Killed, or restart limit reached. Terminal unless restarted manually.
107    Dead,
108}
109
110impl std::fmt::Display for TaskStatus {
111    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112        f.write_str(match self {
113            Self::Running => "running",
114            Self::Restarting => "restarting",
115            Self::Stopping => "stopping",
116            Self::Completed => "completed",
117            Self::Dead => "dead",
118        })
119    }
120}
121
122pub(crate) type Reply<T> = oneshot::Sender<Result<T, SupervisorHandleError>>;
123
124/// Supervisor-side bookkeeping for one task. At most one run is alive per slot.
125pub(crate) struct Slot {
126    pub(crate) name: Arc<str>,
127    pub(crate) status: TaskStatus,
128    pub(crate) task: Box<dyn DynTask>,
129    pub(crate) cancel: Option<CancellationToken>,
130    /// Incremented on every start; stale backoff timers carry an older value.
131    pub(crate) generation: u64,
132    /// Restart instants inside the current window, oldest first.
133    pub(crate) restarts: VecDeque<tokio::time::Instant>,
134    /// What to do once a `Stopping` run exits.
135    pub(crate) restart_after_stop: bool,
136    /// `kill_task` callers, answered once the run has exited.
137    pub(crate) kill_waiters: Vec<Reply<()>>,
138    /// `restart_task` callers, answered once the new run has started.
139    pub(crate) restart_waiters: Vec<Reply<()>>,
140}
141
142impl Slot {
143    pub(crate) fn new(name: Arc<str>, task: Box<dyn DynTask>) -> Self {
144        Self {
145            name,
146            status: TaskStatus::Dead,
147            task,
148            cancel: None,
149            generation: 0,
150            restarts: VecDeque::new(),
151            restart_after_stop: false,
152            kill_waiters: Vec::new(),
153            restart_waiters: Vec::new(),
154        }
155    }
156
157    pub(crate) fn request_stop(&mut self, restart_after: bool) {
158        if let Some(cancel) = &self.cancel {
159            cancel.cancel();
160        }
161        self.status = TaskStatus::Stopping;
162        self.restart_after_stop = restart_after;
163    }
164}
165
166pub(crate) fn panic_message(payload: Box<dyn std::any::Any + Send>) -> String {
167    payload
168        .downcast_ref::<&str>()
169        .map(|s| (*s).to_owned())
170        .or_else(|| payload.downcast_ref::<String>().cloned())
171        .unwrap_or_else(|| "non-string panic payload".to_owned())
172}
task_supervisor/task.rs

task_supervisor/
task.rs
