From cf0c10a32e452a8e877ea30a3647b4e4dbfc11c0 Mon Sep 17 00:00:00 2001
From: Ammar Arif <evergreenkary@gmail.com>
Date: Mon, 16 Sep 2024 13:05:19 -0400
Subject: [PATCH] graceful-shutdown-on-sigterm

---
 bin/katana/src/cli/node.rs      |  4 ++--
 bin/katana/src/utils.rs         | 23 -----------------------
 bin/saya/src/main.rs            |  9 ++++-----
 bin/torii/src/main.rs           |  1 +
 crates/dojo-utils/Cargo.toml    |  4 ++--
 crates/dojo-utils/src/lib.rs    |  2 ++
 crates/dojo-utils/src/signal.rs | 27 +++++++++++++++++++++++++++
 7 files changed, 38 insertions(+), 32 deletions(-)
 create mode 100644 crates/dojo-utils/src/signal.rs

diff --git a/bin/katana/src/cli/node.rs b/bin/katana/src/cli/node.rs
index c20cf2df8e..f452bcce00 100644
--- a/bin/katana/src/cli/node.rs
+++ b/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() => {}
         }
 
diff --git a/bin/katana/src/utils.rs b/bin/katana/src/utils.rs
index 92a63e3a2e..f05f5a7634 100644
--- a/bin/katana/src/utils.rs
+++ b/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::*;
diff --git a/bin/saya/src/main.rs b/bin/saya/src/main.rs
index 376a337b21..07f91bc289 100644
--- a/bin/saya/src/main.rs
+++ b/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(())
 }
diff --git a/bin/torii/src/main.rs b/bin/torii/src/main.rs
index 40a3514fd9..94b6e63d70 100644
--- a/bin/torii/src/main.rs
+++ b/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(())
diff --git a/crates/dojo-utils/Cargo.toml b/crates/dojo-utils/Cargo.toml
index ae9889a37f..798b732482 100644
--- a/crates/dojo-utils/Cargo.toml
+++ b/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
diff --git a/crates/dojo-utils/src/lib.rs b/crates/dojo-utils/src/lib.rs
index eace35faf7..519a7027c4 100644
--- a/crates/dojo-utils/src/lib.rs
+++ b/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;
diff --git a/crates/dojo-utils/src/signal.rs b/crates/dojo-utils/src/signal.rs
new file mode 100644
index 0000000000..d5cb8f0569
--- /dev/null
+++ b/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(())
+}