Skip to main content

cuprated/
monitor.rs

1//! Task spawning and shutdown coordination.
2
3use std::future::Future;
4
5use tokio::task::JoinHandle;
6use tokio_util::{sync::CancellationToken, task::TaskTracker};
7use tracing::info;
8
9/// A handle for task spawning and shutdown coordination.
10#[derive(Clone, Default)]
11pub struct TaskExecutor {
12    token: CancellationToken,
13    tracker: TaskTracker,
14}
15
16impl TaskExecutor {
17    /// Create a new executor.
18    pub fn new() -> Self {
19        Self::default()
20    }
21
22    /// Spawn a tracked task.
23    pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
24    where
25        F: Future + Send + 'static,
26        F::Output: Send + 'static,
27    {
28        self.tracker.spawn(future)
29    }
30
31    /// Get a clone of the cancellation token.
32    pub fn cancellation_token(&self) -> CancellationToken {
33        self.token.clone()
34    }
35
36    /// Trigger a graceful shutdown.
37    pub fn trigger_shutdown(&self) {
38        if !self.token.is_cancelled() {
39            info!("Shutting down...");
40        }
41        self.token.cancel();
42    }
43
44    /// Wait for shutdown to be triggered, then await all tracked tasks.
45    pub async fn wait_for_shutdown(&self) {
46        self.token.cancelled().await;
47        self.tracker.close();
48        self.tracker.wait().await;
49    }
50}