diff --git a/bin/katana/src/cli/node.rs b/bin/katana/src/cli/node.rs index 2d46261dbf..c20cf2df8e 100644 --- a/bin/katana/src/cli/node.rs +++ b/bin/katana/src/cli/node.rs @@ -34,12 +34,11 @@ use katana_primitives::genesis::constant::DEFAULT_PREFUNDED_ACCOUNT_BALANCE; use katana_primitives::genesis::Genesis; use katana_rpc::config::ServerConfig; use katana_rpc_api::ApiKind; -use tokio::signal::ctrl_c; use tracing::{info, Subscriber}; use tracing_subscriber::{fmt, EnvFilter}; use url::Url; -use crate::utils::{parse_genesis, parse_seed}; +use crate::utils::{parse_genesis, parse_seed, wait_signal}; #[derive(Parser, Debug)] pub struct NodeArgs { @@ -234,9 +233,9 @@ impl NodeArgs { print_intro(&self, genesis, node.rpc.addr); } - // Wait until ctrl-c signal is received or TaskManager signals shutdown + // Wait until an OS signal is received or TaskManager shutdown tokio::select! { - _ = ctrl_c() => {}, + _ = wait_signal() => {}, _ = node.task_manager.wait_for_shutdown() => {} } diff --git a/bin/katana/src/utils.rs b/bin/katana/src/utils.rs index f05f5a7634..92a63e3a2e 100644 --- a/bin/katana/src/utils.rs +++ b/bin/katana/src/utils.rs @@ -22,6 +22,29 @@ pub fn parse_genesis(value: &str) -> Result { Ok(genesis) } +pub async fn wait_signal() { + use tokio::signal::ctrl_c; + + #[cfg(unix)] + tokio::select! { + _ = ctrl_c() => {}, + _ = sigterm() => {}, + } + + #[cfg(not(unix))] + tokio::select! { + _ = ctrl_c() => {}, + } +} + +/// Returns a future that can be awaited to wait for the SIGTERM signal. +#[cfg(unix)] +async fn sigterm() -> std::io::Result<()> { + use tokio::signal::unix::{signal, SignalKind}; + signal(SignalKind::terminate())?.recv().await; + Ok(()) +} + #[cfg(test)] mod tests { use super::*;