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

miden_node_utils/
shutdown.rs
