miden_node_utils/
shutdown.rs

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

miden_node_utils/
shutdown.rs
