Skip to main content

miden_node_utils/
shutdown.rs

1use std::future::Future;
2use std::time::Duration;
3
4use anyhow::Context;
5pub use tokio_util::sync::CancellationToken;
6
7/// Time allowed for services to finish after a shutdown signal before the process exits.
8pub const GRACE_PERIOD: Duration = Duration::from_secs(10);
9
10/// Runs a service future until it completes or a shutdown signal is received.
11///
12/// On `SIGTERM` or Ctrl-C, the provided root cancellation token is cancelled and the service future
13/// is given [`GRACE_PERIOD`] to complete. If it does not, the process exits immediately so
14/// blocking work cannot hold the Tokio runtime alive indefinitely.
15pub async fn run_with_shutdown<F, Fut>(run: F) -> anyhow::Result<()>
16where
17    F: FnOnce(CancellationToken) -> Fut,
18    Fut: Future<Output = anyhow::Result<()>>,
19{
20    let token = CancellationToken::new();
21    let service = run(token.clone());
22    tokio::pin!(service);
23
24    tokio::select! {
25        result = &mut service => result,
26        result = shutdown_signal() => {
27            result?;
28            tracing::info!("Shutdown signal received; cancelling service tasks");
29            token.cancel();
30
31            let Ok(result) = tokio::time::timeout(GRACE_PERIOD, &mut service).await else {
32                tracing::error!(
33                    grace_period = ?GRACE_PERIOD,
34                    "Graceful shutdown timed out; exiting process",
35                );
36                std::process::exit(0);
37            };
38
39            result
40        },
41    }
42}
43
44/// Waits for SIGTERM or Ctrl-C.
45pub async fn shutdown_signal() -> anyhow::Result<()> {
46    #[cfg(unix)]
47    {
48        let mut terminate =
49            tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
50                .context("failed to install SIGTERM handler")?;
51
52        tokio::select! {
53            _ = terminate.recv() => Ok(()),
54            result = tokio::signal::ctrl_c() => {
55                result.context("failed to install Ctrl-C handler")
56            },
57        }
58    }
59
60    #[cfg(not(unix))]
61    {
62        tokio::signal::ctrl_c().await.context("failed to install Ctrl-C handler")
63    }
64}