task_supervisor/
handle.rs

1use std::{collections::HashMap, sync::Arc};
2
3use tokio::{
4    sync::{mpsc, oneshot, watch},
5    task::JoinHandle,
6};
7
8use crate::{
9    supervisor::{Outcome, SupervisorError},
10    task::{panic_message, DynTask, Reply, SupervisedTask, TaskStatus},
11};
12
13/// Why a [`SupervisorHandle`] request was refused.
14#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
15pub enum SupervisorHandleError {
16    /// The supervisor has exited or is shutting down.
17    #[error("supervisor is no longer running")]
18    Closed,
19    /// `add_task` was called with a name that is already registered.
20    #[error("task '{0}' already exists")]
21    TaskAlreadyExists(String),
22    /// No task is registered under this name.
23    #[error("task '{0}' not found")]
24    TaskNotFound(String),
25}
26
27pub(crate) enum Message {
28    AddTask {
29        name: Arc<str>,
30        task: Box<dyn DynTask>,
31        reply: Reply<()>,
32    },
33    RestartTask {
34        name: String,
35        reply: Reply<()>,
36    },
37    KillTask {
38        name: String,
39        reply: Reply<()>,
40    },
41    TaskStatus {
42        name: String,
43        reply: Reply<TaskStatus>,
44    },
45    TaskStatuses {
46        reply: oneshot::Sender<HashMap<Arc<str>, TaskStatus>>,
47    },
48    Shutdown,
49}
50
51/// Controls a running supervisor. Cheap to clone.
52///
53/// Dropping the last handle shuts the supervisor down gracefully.
54#[derive(Clone)]
55pub struct SupervisorHandle {
56    tx: mpsc::UnboundedSender<Message>,
57    outcome: watch::Receiver<Option<Outcome>>,
58}
59
60impl SupervisorHandle {
61    pub(crate) fn new(join: JoinHandle<Outcome>, tx: mpsc::UnboundedSender<Message>) -> Self {
62        let (outcome_tx, outcome) = watch::channel(None);
63        tokio::spawn(async move {
64            let outcome = match join.await {
65                Ok(outcome) => outcome,
66                Err(err) if err.is_panic() => {
67                    Err(SupervisorError::Panicked(panic_message(err.into_panic())))
68                }
69                Err(_) => Err(SupervisorError::Aborted),
70            };
71            let _ = outcome_tx.send(Some(outcome));
72        });
73        Self { tx, outcome }
74    }
75
76    /// Registers and starts a task. Returns once it is
77    /// [`Running`](TaskStatus::Running).
78    pub async fn add_task(
79        &self,
80        name: &str,
81        task: impl SupervisedTask,
82    ) -> Result<(), SupervisorHandleError> {
83        self.request(|reply| Message::AddTask {
84            name: Arc::from(name),
85            task: Box::new(task),
86            reply,
87        })
88        .await?
89    }
90
91    /// Restarts a task with a fresh restart budget and returns once the new
92    /// run is [`Running`](TaskStatus::Running).
93    ///
94    /// A running task is cancelled first and has `stop_timeout` to exit.
95    pub async fn restart_task(&self, name: &str) -> Result<(), SupervisorHandleError> {
96        self.request(|reply| Message::RestartTask {
97            name: name.to_owned(),
98            reply,
99        })
100        .await?
101    }
102
103    /// Stops a task permanently and returns once it is
104    /// [`Dead`](TaskStatus::Dead), which takes at most `stop_timeout`.
105    /// Idempotent.
106    pub async fn kill_task(&self, name: &str) -> Result<(), SupervisorHandleError> {
107        self.request(|reply| Message::KillTask {
108            name: name.to_owned(),
109            reply,
110        })
111        .await?
112    }
113
114    /// Current status of one task.
115    pub async fn task_status(&self, name: &str) -> Result<TaskStatus, SupervisorHandleError> {
116        self.request(|reply| Message::TaskStatus {
117            name: name.to_owned(),
118            reply,
119        })
120        .await?
121    }
122
123    /// Current status of every task, keyed by name.
124    pub async fn task_statuses(
125        &self,
126    ) -> Result<HashMap<Arc<str>, TaskStatus>, SupervisorHandleError> {
127        self.request(|reply| Message::TaskStatuses { reply }).await
128    }
129
130    /// Cancels every task, waits for them to exit (bounded by `stop_timeout`)
131    /// and returns the supervisor outcome. Same as [`wait`](Self::wait) if the
132    /// supervisor already exited.
133    pub async fn shutdown(&self) -> Outcome {
134        let _ = self.tx.send(Message::Shutdown);
135        self.wait().await
136    }
137
138    /// Waits for the supervisor to exit. Can be awaited any number of times,
139    /// from any clone.
140    pub async fn wait(&self) -> Outcome {
141        let mut outcome = self.outcome.clone();
142        loop {
143            if let Some(outcome) = outcome.borrow_and_update().clone() {
144                return outcome;
145            }
146            if outcome.changed().await.is_err() {
147                return Err(SupervisorError::Aborted);
148            }
149        }
150    }
151
152    async fn request<T>(
153        &self,
154        message: impl FnOnce(oneshot::Sender<T>) -> Message,
155    ) -> Result<T, SupervisorHandleError> {
156        let (reply, response) = oneshot::channel();
157        self.tx
158            .send(message(reply))
159            .map_err(|_| SupervisorHandleError::Closed)?;
160        response.await.map_err(|_| SupervisorHandleError::Closed)
161    }
162}
163
164impl std::fmt::Debug for SupervisorHandle {
165    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
166        f.debug_struct("SupervisorHandle")
167            .field("running", &!self.tx.is_closed())
168            .finish()
169    }
170}
task_supervisor/handle.rs

task_supervisor/
handle.rs
