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 an empty restart history and returns once the
92    /// new run is [`Running`](TaskStatus::Running).
93    ///
94    /// A running task is cancelled first and has `stop_timeout` to exit.
95    /// Returns [`Closed`](SupervisorHandleError::Closed) if the supervisor
96    /// shuts down before the new run starts.
97    pub async fn restart_task(&self, name: &str) -> Result<(), SupervisorHandleError> {
98        self.request(|reply| Message::RestartTask {
99            name: name.to_owned(),
100            reply,
101        })
102        .await?
103    }
104
105    /// Stops a task permanently and returns once it is
106    /// [`Dead`](TaskStatus::Dead), which takes at most `stop_timeout`.
107    /// Idempotent.
108    pub async fn kill_task(&self, name: &str) -> Result<(), SupervisorHandleError> {
109        self.request(|reply| Message::KillTask {
110            name: name.to_owned(),
111            reply,
112        })
113        .await?
114    }
115
116    /// Current status of one task.
117    pub async fn task_status(&self, name: &str) -> Result<TaskStatus, SupervisorHandleError> {
118        self.request(|reply| Message::TaskStatus {
119            name: name.to_owned(),
120            reply,
121        })
122        .await?
123    }
124
125    /// Current status of every task, keyed by name.
126    pub async fn task_statuses(
127        &self,
128    ) -> Result<HashMap<Arc<str>, TaskStatus>, SupervisorHandleError> {
129        self.request(|reply| Message::TaskStatuses { reply }).await
130    }
131
132    /// Cancels every task, waits for them to exit (bounded by `stop_timeout`)
133    /// and returns the supervisor outcome. Same as [`wait`](Self::wait) if the
134    /// supervisor already exited.
135    pub async fn shutdown(&self) -> Outcome {
136        let _ = self.tx.send(Message::Shutdown);
137        self.wait().await
138    }
139
140    /// Waits for the supervisor to exit. Can be awaited any number of times,
141    /// from any clone.
142    pub async fn wait(&self) -> Outcome {
143        let mut outcome = self.outcome.clone();
144        loop {
145            if let Some(outcome) = outcome.borrow_and_update().clone() {
146                return outcome;
147            }
148            if outcome.changed().await.is_err() {
149                return Err(SupervisorError::Aborted);
150            }
151        }
152    }
153
154    async fn request<T>(
155        &self,
156        message: impl FnOnce(oneshot::Sender<T>) -> Message,
157    ) -> Result<T, SupervisorHandleError> {
158        let (reply, response) = oneshot::channel();
159        self.tx
160            .send(message(reply))
161            .map_err(|_| SupervisorHandleError::Closed)?;
162        response.await.map_err(|_| SupervisorHandleError::Closed)
163    }
164}
165
166impl std::fmt::Debug for SupervisorHandle {
167    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
168        f.debug_struct("SupervisorHandle")
169            .field("running", &!self.tx.is_closed())
170            .finish()
171    }
172}
task_supervisor/handle.rs

task_supervisor/
handle.rs
