graceful-shutdown-on-sigterm

bin/katana/src/cli/node.rs

@@ -38,7 +38,7 @@ use tracing::{info, Subscriber};
 use tracing_subscriber::{fmt, EnvFilter};
 use url::Url;
 
-use crate::utils::{parse_genesis, parse_seed, wait_signal};
+use crate::utils::{parse_genesis, parse_seed};
 
 #[derive(Parser, Debug)]
 pub struct NodeArgs {
@@ -235,7 +235,7 @@ impl NodeArgs {
 
         // Wait until an OS signal is received or TaskManager shutdown
         tokio::select! {
-            _ = wait_signal() => {},
+            _ = dojo_utils::signal::wait_signals() => {},
             _ = node.task_manager.wait_for_shutdown() => {}
         }
 

bin/katana/src/utils.rs

@@ -22,29 +22,6 @@ pub fn parse_genesis(value: &str) -> Result<Genesis, anyhow::Error> {
     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::*;

bin/saya/src/main.rs

@@ -2,7 +2,6 @@
 use clap::Parser;
 use console::Style;
 use saya_core::{Saya, SayaConfig};
-use tokio::signal::ctrl_c;
 
 mod args;
 
@@ -20,11 +19,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     print_intro(&config);
 
     let mut saya = Saya::new(config).await?;
-    saya.start().await?;
 
-    // Wait until Ctrl + C is pressed, then shutdown
-    ctrl_c().await?;
-    // handle.stop()?;
+    tokio::select! {
+        res = saya.start() => res?,
+        _ = dojo_utils::signal::wait_signals() => {}
+    }
 
     Ok(())
 }

bin/torii/src/main.rs

@@ -270,6 +270,7 @@ async fn main() -> anyhow::Result<()> {
         _ = graphql_server => {},
         _ = grpc_server => {},
         _ = libp2p_relay_server.run() => {},
+        _ = dojo_utils::signal::wait_signals() => {},
     };
 
     Ok(())

crates/dojo-utils/Cargo.toml

@@ -9,10 +9,10 @@ version.workspace = true
 anyhow.workspace = true
 futures.workspace = true
 reqwest.workspace = true
+rpassword.workspace = true
 starknet.workspace = true
 thiserror.workspace = true
-rpassword.workspace = true
-tokio = { version = "1", features = [ "time" ], default-features = false }
+tokio = { version = "1", features = [ "signal", "time" ], default-features = false }
 
 [dev-dependencies]
 assert_matches.workspace = true

crates/dojo-utils/src/lib.rs

@@ -8,3 +8,5 @@ pub use tx::{TransactionExt, TxnAction, TxnConfig};
 
 pub mod env;
 pub mod keystore;
+
+pub mod signal;

crates/dojo-utils/src/signal.rs

@@ -0,0 +1,27 @@
+use std::io;
+
+use tokio::signal::ctrl_c;
+
+/// Returns a future for awaiting on OS signals to be received - SIGTERM, SIGINT.
+///
+/// Used for handling graceful shutdowns.
+pub async fn wait_signals() {
+    #[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() -> io::Result<()> {
+    use tokio::signal::unix::{signal, SignalKind};
+    signal(SignalKind::terminate())?.recv().await;
+    Ok(())
+}