miden_node_utils/
shutdown.rs

1use std::fmt::{self, Display, Formatter};
2use std::future::Future;
3use std::time::Duration;
4
5use anyhow::Context;
6pub use tokio_util::sync::CancellationToken;
7
8/// Time allowed for services to finish after a shutdown signal before the process exits.
9pub const GRACE_PERIOD: Duration = Duration::from_secs(10);
10
11/// Operating-system signal which requested service shutdown.
12#[derive(Clone, Copy, Debug, Eq, PartialEq)]
13pub enum ShutdownSignal {
14    Interrupt,
15    Terminate,
16}
17
18impl Display for ShutdownSignal {
19    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
20        match self {
21            Self::Interrupt => f.write_str("SIGINT"),
22            Self::Terminate => f.write_str("SIGTERM"),
23        }
24    }
25}
26
27/// Runs a service future until it completes or a shutdown signal is received.
28///
29/// On `SIGTERM` or Ctrl-C, the provided root cancellation token is cancelled and the service future
30/// is given [`GRACE_PERIOD`] to complete. If it does not, the process exits immediately so
31/// blocking work cannot hold the Tokio runtime alive indefinitely.
32pub async fn run_with_shutdown<F, Fut>(service_name: &'static str, run: F) -> anyhow::Result<()>
33where
34    F: FnOnce(CancellationToken) -> Fut,
35    Fut: Future<Output = anyhow::Result<()>>,
36{
37    run_with_shutdown_signal(service_name, run, shutdown_signal()).await
38}
39
40async fn run_with_shutdown_signal<F, Fut, Signal>(
41    service_name: &'static str,
42    run: F,
43    signal: Signal,
44) -> anyhow::Result<()>
45where
46    F: FnOnce(CancellationToken) -> Fut,
47    Fut: Future<Output = anyhow::Result<()>>,
48    Signal: Future<Output = anyhow::Result<ShutdownSignal>>,
49{
50    let token = CancellationToken::new();
51    let service = run(token.clone());
52    tokio::pin!(service);
53    tokio::pin!(signal);
54
55    tokio::select! {
56        result = &mut service => result,
57        result = &mut signal => {
58            let signal = result?;
59            tracing::info!(
60                service.name = service_name,
61                shutdown.signal = %signal,
62                "Shutdown requested",
63            );
64            token.cancel();
65
66            let Ok(result) = tokio::time::timeout(GRACE_PERIOD, &mut service).await else {
67                tracing::error!(
68                    service.name = service_name,
69                    grace_period = ?GRACE_PERIOD,
70                    "Graceful shutdown timed out; exiting process",
71                );
72                std::process::exit(1);
73            };
74
75            result?;
76            tracing::info!(service.name = service_name, "Shutdown complete");
77            Ok(())
78        },
79    }
80}
81
82/// Waits for SIGTERM or Ctrl-C.
83pub async fn shutdown_signal() -> anyhow::Result<ShutdownSignal> {
84    #[cfg(unix)]
85    {
86        let mut terminate =
87            tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
88                .context("failed to install SIGTERM handler")?;
89
90        tokio::select! {
91            _ = terminate.recv() => Ok(ShutdownSignal::Terminate),
92            result = tokio::signal::ctrl_c() => {
93                result
94                    .context("failed to install Ctrl-C handler")
95                    .map(|()| ShutdownSignal::Interrupt)
96            },
97        }
98    }
99
100    #[cfg(not(unix))]
101    {
102        tokio::signal::ctrl_c()
103            .await
104            .context("failed to install Ctrl-C handler")
105            .map(|()| ShutdownSignal::Interrupt)
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use std::sync::Arc;
112    use std::sync::atomic::{AtomicBool, Ordering};
113
114    use super::*;
115
116    #[tokio::test]
117    async fn signal_cancels_and_waits_for_service() {
118        let cancelled = Arc::new(AtomicBool::new(false));
119        let service_cancelled = Arc::clone(&cancelled);
120
121        run_with_shutdown_signal(
122            "test-service",
123            move |shutdown| async move {
124                shutdown.cancelled().await;
125                service_cancelled.store(true, Ordering::Relaxed);
126                Ok(())
127            },
128            std::future::ready(Ok(ShutdownSignal::Interrupt)),
129        )
130        .await
131        .expect("clean shutdown should succeed");
132
133        assert!(cancelled.load(Ordering::Relaxed));
134    }
135
136    #[tokio::test]
137    async fn signal_handler_error_is_propagated() {
138        let err = run_with_shutdown_signal(
139            "test-service",
140            |shutdown| async move {
141                shutdown.cancelled().await;
142                Ok(())
143            },
144            std::future::ready(Err(anyhow::anyhow!("signal handler failed"))),
145        )
146        .await
147        .expect_err("signal error should be returned");
148
149        assert_eq!(err.to_string(), "signal handler failed");
150    }
151}
miden_node_utils/shutdown.rs

miden_node_utils/
shutdown.rs
