Merge pull request #111 from uqbar-dao/hf/intercept-sigs

gracefully exit in case of many SIG*s

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",

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"] }

src/main.rs

@@ -472,6 +472,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!(
@@ -521,6 +522,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;

src/terminal/mod.rs

@@ -14,6 +14,7 @@ use crossterm::{
 use futures::{future::FutureExt, StreamExt};
 use std::fs::{read_to_string, OpenOptions};
 use std::io::{stdout, BufWriter, Write};
+use tokio::signal::unix::{signal, SignalKind};
 
 mod utils;
 
@@ -129,6 +130,20 @@ pub async fn terminal(
         .unwrap();
     let mut log_writer = BufWriter::new(log_handle);
 
+    // use to trigger cleanup if receive signal to kill process
+    let mut sigalrm = signal(SignalKind::alarm()).expect("uqbar: failed to set up SIGALRM handler");
+    let mut sighup = signal(SignalKind::hangup()).expect("uqbar: failed to set up SIGHUP handler");
+    let mut sigint =
+        signal(SignalKind::interrupt()).expect("uqbar: failed to set up SIGINT handler");
+    let mut sigpipe = signal(SignalKind::pipe()).expect("uqbar: failed to set up SIGPIPE handler");
+    let mut sigquit = signal(SignalKind::quit()).expect("uqbar: failed to set up SIGQUIT handler");
+    let mut sigterm =
+        signal(SignalKind::terminate()).expect("uqbar: failed to set up SIGTERM handler");
+    let mut sigusr1 =
+        signal(SignalKind::user_defined1()).expect("uqbar: failed to set up SIGUSR1 handler");
+    let mut sigusr2 =
+        signal(SignalKind::user_defined2()).expect("uqbar: failed to set up SIGUSR2 handler");
+
     loop {
         let event = reader.next().fuse();
 
@@ -172,7 +187,7 @@ pub async fn terminal(
                     Print(utils::truncate_in_place(&current_line, prompt_len, win_cols, (line_col, cursor_col))),
                     cursor::MoveTo(cursor_col, win_rows),
                 )?;
-            },
+            }
             Some(Ok(event)) = event => {
                 let mut stdout = stdout.lock();
                 match event {
@@ -599,7 +614,15 @@ pub async fn terminal(
                     },
                     _ => {},
                 }
-                }
+            }
+            _ = sigalrm.recv() => return Err(anyhow::anyhow!("exiting due to SIGALRM")),
+            _ = sighup.recv() =>  return Err(anyhow::anyhow!("exiting due to SIGHUP")),
+            _ = sigint.recv() =>  return Err(anyhow::anyhow!("exiting due to SIGINT")),
+            _ = sigpipe.recv() => return Err(anyhow::anyhow!("exiting due to SIGPIPE")),
+            _ = sigquit.recv() => return Err(anyhow::anyhow!("exiting due to SIGQUIT")),
+            _ = sigterm.recv() => return Err(anyhow::anyhow!("exiting due to SIGTERM")),
+            _ = sigusr1.recv() => return Err(anyhow::anyhow!("exiting due to SIGUSR1")),
+            _ = sigusr2.recv() => return Err(anyhow::anyhow!("exiting due to SIGUSR2")),
         }
     }
     execute!(stdout.lock(), DisableBracketedPaste, terminal::SetTitle(""))?;