cuprated/
monitor.rs

1//! Task spawning and shutdown coordination.
2
3use std::{
4    future::Future,
5    sync::{
6        atomic::{AtomicBool, Ordering},
7        Arc,
8    },
9};
10
11use futures::FutureExt;
12use tokio::task::JoinHandle;
13use tokio_util::{sync::CancellationToken, task::TaskTracker};
14use tracing::{debug, error, info};
15
16/// An unexpected node-side failure that should trigger a shutdown.
17pub type FatalError = tower::BoxError;
18
19/// Why the node stopped.
20#[must_use]
21pub enum ShutdownReason {
22    /// Shutdown was requested.
23    Requested,
24
25    /// A critical task failed.
26    TaskFailed,
27}
28
29/// A handle for task spawning and shutdown coordination.
30#[derive(Clone, Default)]
31pub struct TaskExecutor {
32    token: CancellationToken,
33    tracker: TaskTracker,
34    failed: Arc<AtomicBool>,
35}
36
37impl TaskExecutor {
38    /// Create a new executor.
39    pub fn new() -> Self {
40        Self::default()
41    }
42
43    /// Spawn a tracked task.
44    pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
45    where
46        F: Future + Send + 'static,
47        F::Output: Send + 'static,
48    {
49        self.tracker.spawn(future)
50    }
51
52    /// Spawn a tracked task that triggers shutdown if the future returns early.
53    pub fn spawn_critical<F>(&self, name: &'static str, future: F) -> JoinHandle<()>
54    where
55        F: Future<Output = Result<(), FatalError>> + Send + 'static,
56    {
57        let executor = self.clone();
58        self.tracker.spawn(future.map(move |result| {
59            if executor.token.is_cancelled() {
60                // Node is shutting down, so an early exit or error is expected
61                if let Err(e) = result {
62                    debug!(subsystem = name, "{:#}", anyhow::Error::from_boxed(e));
63                }
64                return;
65            }
66            match result {
67                Ok(()) => error!(
68                    subsystem = name,
69                    "critical task exited before shutdown was requested"
70                ),
71                Err(e) => {
72                    error!(subsystem = name, "{:#}", anyhow::Error::from_boxed(e));
73                }
74            }
75
76            executor.failed.store(true, Ordering::Relaxed);
77            executor.trigger_shutdown();
78        }))
79    }
80
81    /// Get a clone of the cancellation token.
82    pub fn cancellation_token(&self) -> CancellationToken {
83        self.token.clone()
84    }
85
86    /// Trigger a graceful shutdown.
87    pub fn trigger_shutdown(&self) {
88        if !self.token.is_cancelled() {
89            info!("Shutting down...");
90        }
91        self.token.cancel();
92    }
93
94    /// Wait for shutdown to be triggered, then await all tracked tasks.
95    pub async fn wait_for_shutdown(&self) -> ShutdownReason {
96        self.token.cancelled().await;
97        self.tracker.close();
98        self.tracker.wait().await;
99        if self.failed.load(Ordering::Relaxed) {
100            ShutdownReason::TaskFailed
101        } else {
102            ShutdownReason::Requested
103        }
104    }
105}
cuprated/monitor.rs

cuprated/
monitor.rs
