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

aster_forge_runtime/
shutdown.rs
