From fca59c9713a091657d0863831603e90113a27256 Mon Sep 17 00:00:00 2001
From: Huan-Cheng Chang <huancheng.chang@trili.tech>
Date: Fri, 29 Nov 2024 15:44:10 +0000
Subject: [PATCH] feat(jstzd): listen to os signals

---
 crates/jstzd/src/lib.rs         |  9 +++++++-
 crates/jstzd/tests/main_test.rs | 38 +++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/crates/jstzd/src/lib.rs b/crates/jstzd/src/lib.rs
index 01e060c78..934b63c95 100644
--- a/crates/jstzd/src/lib.rs
+++ b/crates/jstzd/src/lib.rs
@@ -5,6 +5,7 @@ pub mod task;
 use crate::task::jstzd::{JstzdConfig, JstzdServer};
 pub use config::BOOTSTRAP_CONTRACT_NAMES;
 use std::process::exit;
+use tokio::signal::unix::{signal, SignalKind};
 
 pub const EXCHANGER_ADDRESS: &str = "KT1F3MuqvT9Yz57TgCS3EkDcKNZe9HpiavUJ";
 pub const JSTZ_ROLLUP_ADDRESS: &str = "sr1PuFMgaRUN12rKQ3J2ae5psNtwCxPNmGNK";
@@ -32,8 +33,14 @@ async fn run(port: u16, config: JstzdConfig) {
         exit(1);
     }
 
-    server.wait().await;
+    let mut sigterm = signal(SignalKind::terminate()).unwrap();
+    let mut sigint = signal(SignalKind::interrupt()).unwrap();
 
+    tokio::select! {
+        _ = server.wait() => (),
+        _ = sigterm.recv() => (),
+        _ = sigint.recv() => (),
+    };
     println!("Shutting down");
     server.stop().await.unwrap();
 }
diff --git a/crates/jstzd/tests/main_test.rs b/crates/jstzd/tests/main_test.rs
index bc57891e4..0f30c00b5 100644
--- a/crates/jstzd/tests/main_test.rs
+++ b/crates/jstzd/tests/main_test.rs
@@ -16,6 +16,25 @@ fn unknown_command() {
         .stderr(predicate::str::contains("unrecognized subcommand \'test\'"));
 }
 
+#[test]
+fn default_config() {
+    // Since the server's port number is unknown when jstzd runs on default config,
+    // here it's assumed that if the child process is still alive after 10 seconds,
+    // it means that jstzd successfully launched
+    let mut child = Command::cargo_bin("jstzd")
+        .unwrap()
+        .arg("run")
+        .spawn()
+        .unwrap();
+    thread::sleep(Duration::from_secs(10));
+    assert!(child.try_wait().unwrap().is_none());
+    Command::new("kill")
+        .args(["-s", "TERM", &child.id().to_string()])
+        .spawn()
+        .unwrap();
+    assert!(child.wait().is_ok());
+}
+
 #[test]
 fn valid_config_file() {
     let port = unused_port();
@@ -74,3 +93,22 @@ fn bad_config_file() {
             "should have at least one bootstrap account with at least 6000 tez",
         ));
 }
+
+#[test]
+fn terminate_with_sigint() {
+    // Since the server's port number is unknown when jstzd runs on default config,
+    // here it's assumed that if the child process is still alive after 10 seconds,
+    // it means that jstzd successfully launched
+    let mut child = Command::cargo_bin("jstzd")
+        .unwrap()
+        .arg("run")
+        .spawn()
+        .unwrap();
+    thread::sleep(Duration::from_secs(10));
+    assert!(child.try_wait().unwrap().is_none());
+    Command::new("kill")
+        .args(["-s", "INT", &child.id().to_string()])
+        .spawn()
+        .unwrap();
+    assert!(child.wait().is_ok());
+}