aster_forge_runtime/
shutdown.rs

1//! Termination signal helpers.
2//!
3//! This module only waits for process termination signals. Product crates remain
4//! responsible for recording audit events, stopping background tasks, flushing
5//! buffers, and closing database or network handles in their preferred order.
6
7use std::future::Future;
8use std::pin::Pin;
9use std::time::{Duration, Instant};
10
11use tokio::task::JoinHandle;
12use tokio_util::sync::CancellationToken;
13
14type ShutdownFuture = Pin<Box<dyn Future<Output = Result<(), String>> + Send>>;
15type ShutdownPhaseFn = dyn FnMut() -> ShutdownFuture + Send;
16
17/// Termination signal observed by the runtime.
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum TerminationSignal {
20    /// Unix SIGINT or cross-platform Ctrl+C.
21    Interrupt,
22    /// Unix SIGTERM.
23    Terminate,
24}
25
26impl TerminationSignal {
27    /// Returns a stable label for logging and tests.
28    pub const fn as_str(self) -> &'static str {
29        match self {
30            Self::Interrupt => "SIGINT",
31            Self::Terminate => "SIGTERM",
32        }
33    }
34}
35
36/// Errors returned while installing signal listeners.
37#[derive(Debug, thiserror::Error)]
38pub enum RuntimeSignalError {
39    /// Failed to install or await a process signal handler.
40    #[error("failed to install termination signal handler: {0}")]
41    Install(String),
42}
43
44/// Final status for one shutdown phase.
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub enum ShutdownPhaseStatus {
47    /// The phase completed successfully.
48    Succeeded,
49    /// The phase returned an error string.
50    Failed(String),
51    /// The phase exceeded its timeout.
52    TimedOut,
53}
54
55impl ShutdownPhaseStatus {
56    /// Returns whether this phase did not complete successfully.
57    pub const fn is_failure(&self) -> bool {
58        !matches!(self, Self::Succeeded)
59    }
60}
61
62/// Report for one executed shutdown phase.
63#[derive(Debug, Clone, PartialEq, Eq)]
64pub struct ShutdownPhaseReport {
65    /// Stable phase name.
66    pub name: &'static str,
67    /// Phase result.
68    pub status: ShutdownPhaseStatus,
69    /// Execution duration.
70    pub duration: Duration,
71}
72
73/// Aggregate report for a shutdown run.
74#[derive(Debug, Clone, PartialEq, Eq)]
75pub struct ShutdownReport {
76    /// Phase reports in execution order.
77    pub phases: Vec<ShutdownPhaseReport>,
78}
79
80impl ShutdownReport {
81    /// Returns a report from phase entries.
82    pub fn new(phases: Vec<ShutdownPhaseReport>) -> Self {
83        Self { phases }
84    }
85
86    /// Returns whether any phase failed or timed out.
87    pub fn has_failures(&self) -> bool {
88        self.phases.iter().any(|phase| phase.status.is_failure())
89    }
90}
91
92/// Logs the aggregate result of a shutdown run.
93pub fn log_shutdown_report(report: &ShutdownReport) {
94    if report.has_failures() {
95        tracing::warn!("shutdown completed with one or more failed phases");
96    } else {
97        tracing::info!("shutdown complete");
98    }
99}
100
101struct RegisteredShutdownPhase {
102    name: &'static str,
103    timeout: Option<Duration>,
104    phase: Box<ShutdownPhaseFn>,
105}
106
107/// Ordered shutdown phase coordinator.
108///
109/// The coordinator owns phase ordering, timeout handling, duration collection,
110/// and error aggregation. Product crates provide the actual phase closures.
111/// Phases are `FnMut` so shutdown code can move owned handles into the
112/// coordinator and consume them exactly once during the shutdown run.
113#[derive(Default)]
114pub struct ShutdownCoordinator {
115    phases: Vec<RegisteredShutdownPhase>,
116}
117
118impl ShutdownCoordinator {
119    /// Creates an empty shutdown coordinator.
120    pub fn new() -> Self {
121        Self::default()
122    }
123
124    /// Registers a shutdown phase.
125    pub fn phase<F, Fut>(
126        &mut self,
127        name: &'static str,
128        timeout: Option<Duration>,
129        mut phase: F,
130    ) -> &mut Self
131    where
132        F: FnMut() -> Fut + Send + 'static,
133        Fut: Future<Output = Result<(), String>> + Send + 'static,
134    {
135        self.phases.push(RegisteredShutdownPhase {
136            name,
137            timeout,
138            phase: Box::new(move || Box::pin(phase())),
139        });
140        self
141    }
142
143    /// Runs phases sequentially and returns a report.
144    ///
145    /// Later phases still run when an earlier phase fails. This lets product
146    /// shutdown code make best-effort progress through independent resources.
147    pub async fn run(&mut self) -> ShutdownReport {
148        let mut reports = Vec::with_capacity(self.phases.len());
149
150        for phase in &mut self.phases {
151            tracing::info!(phase = phase.name, "starting shutdown phase");
152            let started_at = Instant::now();
153            let future = (phase.phase)();
154            let status = match phase.timeout {
155                Some(timeout) => match tokio::time::timeout(timeout, future).await {
156                    Ok(Ok(())) => ShutdownPhaseStatus::Succeeded,
157                    Ok(Err(error)) => ShutdownPhaseStatus::Failed(error),
158                    Err(_) => ShutdownPhaseStatus::TimedOut,
159                },
160                None => match future.await {
161                    Ok(()) => ShutdownPhaseStatus::Succeeded,
162                    Err(error) => ShutdownPhaseStatus::Failed(error),
163                },
164            };
165            let duration = started_at.elapsed();
166            match &status {
167                ShutdownPhaseStatus::Succeeded => {
168                    tracing::info!(phase = phase.name, ?duration, "shutdown phase completed");
169                }
170                ShutdownPhaseStatus::Failed(error) => {
171                    tracing::error!(phase = phase.name, ?duration, %error, "shutdown phase failed");
172                }
173                ShutdownPhaseStatus::TimedOut => {
174                    tracing::error!(phase = phase.name, ?duration, "shutdown phase timed out");
175                }
176            }
177            reports.push(ShutdownPhaseReport {
178                name: phase.name,
179                status,
180                duration,
181            });
182        }
183
184        ShutdownReport::new(reports)
185    }
186
187    /// Returns how many phases are registered.
188    pub fn len(&self) -> usize {
189        self.phases.len()
190    }
191
192    /// Returns whether no phases are registered.
193    pub fn is_empty(&self) -> bool {
194        self.phases.is_empty()
195    }
196}
197
198/// Waits until the process receives a termination signal.
199pub async fn wait_for_termination_signal() -> Result<TerminationSignal, RuntimeSignalError> {
200    let signal = wait_for_signal_impl().await?;
201    tracing::info!(
202        signal = signal.as_str(),
203        "received termination signal, shutting down gracefully..."
204    );
205    Ok(signal)
206}
207
208/// Spawns a task that waits for a termination signal, cancels `shutdown_token`,
209/// and then runs `on_signal`.
210///
211/// This keeps product entrypoints from duplicating the same signal-listener
212/// task while leaving the actual server stop primitive product-specific.
213pub fn spawn_termination_signal_handler<F, Fut>(
214    shutdown_token: CancellationToken,
215    on_signal: F,
216) -> JoinHandle<()>
217where
218    F: FnOnce() -> Fut + Send + 'static,
219    Fut: Future<Output = ()> + Send + 'static,
220{
221    tokio::spawn(async move {
222        if let Err(error) = wait_for_termination_signal().await {
223            tracing::error!(%error, "shutdown signal listener failed");
224        }
225        shutdown_token.cancel();
226        on_signal().await;
227    })
228}
229
230#[cfg(unix)]
231async fn wait_for_signal_impl() -> Result<TerminationSignal, RuntimeSignalError> {
232    use tokio::signal::unix::{SignalKind, signal};
233
234    let mut sigint = signal(SignalKind::interrupt())
235        .map_err(|error| RuntimeSignalError::Install(error.to_string()))?;
236    let mut sigterm = signal(SignalKind::terminate())
237        .map_err(|error| RuntimeSignalError::Install(error.to_string()))?;
238
239    tokio::select! {
240        _ = sigint.recv() => Ok(TerminationSignal::Interrupt),
241        _ = sigterm.recv() => Ok(TerminationSignal::Terminate),
242    }
243}
244
245#[cfg(not(unix))]
246async fn wait_for_signal_impl() -> Result<TerminationSignal, RuntimeSignalError> {
247    tokio::signal::ctrl_c()
248        .await
249        .map_err(|error| RuntimeSignalError::Install(error.to_string()))?;
250    Ok(TerminationSignal::Interrupt)
251}
252
253#[cfg(test)]
254mod tests {
255    use super::{ShutdownCoordinator, ShutdownPhaseStatus, TerminationSignal};
256    use std::time::Duration;
257
258    #[test]
259    fn termination_signal_reports_stable_labels() {
260        assert_eq!(TerminationSignal::Interrupt.as_str(), "SIGINT");
261        assert_eq!(TerminationSignal::Terminate.as_str(), "SIGTERM");
262    }
263
264    #[tokio::test]
265    async fn shutdown_coordinator_runs_all_phases_in_order() {
266        let mut coordinator = ShutdownCoordinator::new();
267        coordinator
268            .phase("tasks", None, || async { Ok(()) })
269            .phase("audit", None, || async { Err("flush failed".to_string()) })
270            .phase("db", None, || async { Ok(()) });
271
272        let report = coordinator.run().await;
273
274        assert_eq!(coordinator.len(), 3);
275        assert!(report.has_failures());
276        assert_eq!(report.phases[0].name, "tasks");
277        assert_eq!(report.phases[0].status, ShutdownPhaseStatus::Succeeded);
278        assert_eq!(
279            report.phases[1].status,
280            ShutdownPhaseStatus::Failed("flush failed".to_string())
281        );
282        assert_eq!(report.phases[2].status, ShutdownPhaseStatus::Succeeded);
283    }
284
285    #[test]
286    fn shutdown_report_logger_accepts_success_and_failure_reports() {
287        super::log_shutdown_report(&super::ShutdownReport::new(vec![
288            super::ShutdownPhaseReport {
289                name: "tasks",
290                status: ShutdownPhaseStatus::Succeeded,
291                duration: Duration::from_millis(1),
292            },
293        ]));
294
295        super::log_shutdown_report(&super::ShutdownReport::new(vec![
296            super::ShutdownPhaseReport {
297                name: "database",
298                status: ShutdownPhaseStatus::Failed("close failed".to_string()),
299                duration: Duration::from_millis(1),
300            },
301        ]));
302    }
303
304    #[tokio::test]
305    async fn shutdown_coordinator_reports_timeouts() {
306        let mut coordinator = ShutdownCoordinator::new();
307        coordinator.phase("slow", Some(Duration::from_millis(1)), || async {
308            tokio::time::sleep(Duration::from_millis(50)).await;
309            Ok(())
310        });
311
312        let report = coordinator.run().await;
313
314        assert!(report.has_failures());
315        assert_eq!(report.phases[0].status, ShutdownPhaseStatus::TimedOut);
316    }
317
318    #[tokio::test]
319    async fn shutdown_coordinator_supports_consumed_phase_handles() {
320        let mut coordinator = ShutdownCoordinator::new();
321        let mut owned_handle = Some("resource");
322        coordinator.phase("owned", None, move || {
323            let handle = owned_handle.take();
324            async move {
325                if handle == Some("resource") {
326                    Ok(())
327                } else {
328                    Err("resource already consumed".to_string())
329                }
330            }
331        });
332
333        let report = coordinator.run().await;
334
335        assert!(!report.has_failures());
336        assert_eq!(report.phases[0].status, ShutdownPhaseStatus::Succeeded);
337    }
338}
aster_forge_runtime/shutdown.rs

aster_forge_runtime/
shutdown.rs
