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/main.rs b/src/main.rs index a7775f47..9356593a 100644 --- a/src/main.rs +++ b/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; diff --git a/src/terminal/mod.rs b/src/terminal/mod.rs index 2ddfa3f4..61cf83d7 100644 --- a/src/terminal/mod.rs +++ b/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(¤t_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(""))?;