From 62b2bf08cb4df8e62a962c2861ea4c8b13dc60e6 Mon Sep 17 00:00:00 2001
From: hosted-fornet <hosted-fornet@protonmail.com>
Date: Thu, 21 Dec 2023 16:11:44 -0800
Subject: [PATCH] gracefully exit in case of many SIG*s

---
 Cargo.lock            |  1 +
 Cargo.toml            |  6 +++---
 src/intercept_sigs.rs | 25 +++++++++++++++++++++++++
 src/main.rs           | 12 ++++++++++++
 4 files changed, 41 insertions(+), 3 deletions(-)
 create mode 100644 src/intercept_sigs.rs

diff --git a/Cargo.lock b/Cargo.lock
index b64c11ca..8bb5b2d6 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4698,6 +4698,7 @@ dependencies = [
  "mio",
  "num_cpus",
  "pin-project-lite",
+ "signal-hook-registry",
  "socket2 0.5.4",
  "tokio-macros",
  "windows-sys 0.48.0",
diff --git a/Cargo.toml b/Cargo.toml
index 828fdecd..e349c1ea 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -55,7 +55,9 @@ rand = "0.8.4"
 reqwest = "0.11.18"
 ring = "0.16.20"
 rmp-serde = "1.1.2"
+rocksdb = { version = "0.21.0", features = ["multi-threaded-cf"] }
 route-recognizer = "0.3.1"
+rusqlite = { version = "0.30.0", features = ["bundled"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 serde_urlencoded = "0.7"
@@ -63,7 +65,7 @@ sha2 = "0.10"
 snow = { version = "0.9.3", features = ["ring-resolver"] }
 static_dir = "0.2.0"
 thiserror = "1.0"
-tokio = { version = "1.28", features = ["fs", "macros", "rt-multi-thread", "sync"] }
+tokio = { version = "1.28", features = ["fs", "macros", "rt-multi-thread", "signal", "sync"] }
 tokio-tungstenite = "0.20.1"
 url = "2.4.1"
 uuid = { version = "1.1.2", features = ["serde", "v4"] }
@@ -71,5 +73,3 @@ warp = "0.3.5"
 wasmtime = "15.0.1"
 wasmtime-wasi = "15.0.1"
 zip = "0.6"
-rocksdb = { version = "0.21.0", features = ["multi-threaded-cf"] }
-rusqlite = { version = "0.30.0", features = ["bundled"] }
diff --git a/src/intercept_sigs.rs b/src/intercept_sigs.rs
new file mode 100644
index 00000000..de632256
--- /dev/null
+++ b/src/intercept_sigs.rs
@@ -0,0 +1,25 @@
+use tokio::signal::unix::{signal, SignalKind};
+
+/// trigger cleanup if receive signal to kill process
+pub async fn intercept_sigs() -> String {
+    let mut sigalrm = signal(SignalKind::alarm()).expect("uqdev run-tests: failed to set up SIGALRM handler");
+    let mut sighup = signal(SignalKind::hangup()).expect("uqdev run-tests: failed to set up SIGHUP handler");
+    let mut sigint = signal(SignalKind::interrupt()).expect("uqdev run-tests: failed to set up SIGINT handler");
+    let mut sigpipe = signal(SignalKind::pipe()).expect("uqdev run-tests: failed to set up SIGPIPE handler");
+    let mut sigquit = signal(SignalKind::quit()).expect("uqdev run-tests: failed to set up SIGQUIT handler");
+    let mut sigterm = signal(SignalKind::terminate()).expect("uqdev run-tests: failed to set up SIGTERM handler");
+    let mut sigusr1 = signal(SignalKind::user_defined1()).expect("uqdev run-tests: failed to set up SIGUSR1 handler");
+    let mut sigusr2 = signal(SignalKind::user_defined2()).expect("uqdev run-tests: failed to set up SIGUSR2 handler");
+
+    tokio::select! {
+        _ = sigalrm.recv() => "exiting due to SIGALRM",
+        _ = sighup.recv() =>  "exiting due to SIGHUP",
+        _ = sigint.recv() =>  "exiting due to SIGINT",
+        _ = sigpipe.recv() => "exiting due to SIGPIPE",
+        _ = sigquit.recv() => "exiting due to SIGQUIT",
+        _ = sigterm.recv() => "exiting due to SIGTERM",
+        _ = sigusr1.recv() => "exiting due to SIGUSR1",
+        _ = sigusr2.recv() => "exiting due to SIGUSR2",
+    }.into()
+}
+
diff --git a/src/main.rs b/src/main.rs
index a7775f47..64221f67 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -13,6 +13,7 @@ use ring::{rand::SystemRandom, signature, signature::KeyPair};
 
 mod eth_rpc;
 mod http;
+mod intercept_sigs;
 mod kernel;
 mod keygen;
 mod kv;
@@ -472,6 +473,7 @@ async fn main() {
     ));
     // if a runtime task exits, try to recover it,
     // unless it was terminal signaling a quit
+    // or a SIG* was intercepted
     let quit_msg: String = tokio::select! {
         Some(Ok(res)) = tasks.join_next() => {
             format!(
@@ -493,6 +495,9 @@ async fn main() {
                 Err(e) => e.to_string(),
             }
         }
+        s = intercept_sigs::intercept_sigs() => {
+            s.into()
+        }
     };
 
     // gracefully abort all running processes in kernel
@@ -521,6 +526,13 @@ async fn main() {
 
     // abort all remaining tasks
     tasks.shutdown().await;
+    let stdout = std::io::stdout();
+    let mut stdout = stdout.lock();
+    let _ = crossterm::execute!(
+        stdout,
+        crossterm::event::DisableBracketedPaste,
+        crossterm::terminal::SetTitle(""),
+    );
     let _ = crossterm::terminal::disable_raw_mode();
     println!("\r\n\x1b[38;5;196m{}\x1b[0m", quit_msg);
     return;