From 0faa6a85780b43214ad34ef8cc314b2026977e41 Mon Sep 17 00:00:00 2001 From: Chewing_Bever Date: Tue, 6 Jun 2023 16:57:59 +0200 Subject: [PATCH] feat: add basis for signal handling --- Cargo.lock | 20 +++++++++++++++ Cargo.toml | 1 + src/main.rs | 37 +++++++++------------------- src/server/process.rs | 4 +++ src/signals.rs | 57 +++++++++++++++++++++++++++++++++++++++++++ src/stdin.rs | 29 ++++++++++++++++++++++ 6 files changed, 123 insertions(+), 25 deletions(-) create mode 100644 src/signals.rs create mode 100644 src/stdin.rs diff --git a/Cargo.lock b/Cargo.lock index 5953204..5124e00 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -15,6 +15,7 @@ dependencies = [ "chrono", "clap", "flate2", + "signal-hook", "tar", ] @@ -383,6 +384,25 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "signal-hook" +version = "0.3.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "732768f1176d21d09e076c23a93123d40bba92d50c4058da34d45c8de8e682b9" +dependencies = [ + "libc", + "signal-hook-registry", +] + +[[package]] +name = "signal-hook-registry" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8229b473baa5980ac72ef434c4415e70c4b5e71b423043adb4ba059f89c99a1" +dependencies = [ + "libc", +] + [[package]] name = "strsim" version = "0.10.0" diff --git a/Cargo.toml b/Cargo.toml index dd2f79b..4c0ba69 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,7 @@ flate2 = "1.0.26" # Used for backup filenames chrono = "0.4.26" clap = { version = "4.3.1", features = ["derive", "env"] } +signal-hook = "0.3.15" [profile.release] lto = "fat" diff --git a/src/main.rs b/src/main.rs index 7837887..1c083fa 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,6 @@ mod server; +mod signals; +mod stdin; use clap::Parser; use server::ServerType; @@ -79,7 +81,8 @@ fn backups_thread(counter: Arc>, frequency: u64) { } } -fn main() { +fn main() -> io::Result<()> { + let (term, mut signals) = signals::install_signal_handlers()?; let cli = Cli::parse(); let cmd = server::ServerCommand::new(cli.type_, &cli.server_version) @@ -91,33 +94,17 @@ fn main() { .xms(cli.xms) .xmx(cli.xmx) .max_backups(cli.max_backups); - let counter = Arc::new(Mutex::new(cmd.spawn().expect("Failed to start server."))); + let server = Arc::new(cmd.spawn()?); if cli.frequency > 0 { - let clone = Arc::clone(&counter); - std::thread::spawn(move || backups_thread(clone, cli.frequency)); + let clone = Arc::clone(&server); + std::thread::spawn(move || backups_thread(server, cli.frequency)); } - let stdin = io::stdin(); - let input = &mut String::new(); + // Spawn thread that handles the main stdin loop + let clone = Arc::clone(&server); + std::thread::spawn(move || stdin::handle_stdin(server)); - loop { - input.clear(); - - if stdin.read_line(input).is_err() { - continue; - }; - - { - let mut server = counter.lock().unwrap(); - - if let Err(e) = server.send_command(input) { - println!("{}", e); - }; - } - - if input.trim() == "stop" { - break; - } - } + // Signal handler loop exits the process when necessary + signals::handle_signals(&mut signals, term, server) } diff --git a/src/server/process.rs b/src/server/process.rs index 5640b18..b5baae9 100644 --- a/src/server/process.rs +++ b/src/server/process.rs @@ -67,6 +67,10 @@ impl ServerProcess { Ok(()) } + pub fn kill(&mut self) -> std::io::Result<()> { + self.child.kill() + } + /// Perform a backup by disabling the server's save feature and flushing its data, before /// creating an archive file. pub fn backup(&mut self) -> std::io::Result<()> { diff --git a/src/signals.rs b/src/signals.rs new file mode 100644 index 0000000..a173047 --- /dev/null +++ b/src/signals.rs @@ -0,0 +1,57 @@ +use std::io; +use std::sync::{Arc, Mutex}; +use std::sync::atomic::{AtomicBool, Ordering}; + +use signal_hook::consts::TERM_SIGNALS; +use signal_hook::flag; +use signal_hook::iterator::{Signals, SignalsInfo}; + +use crate::server; + +/// Install the required signal handlers for terminating signals. +pub fn install_signal_handlers() -> io::Result<(Arc, SignalsInfo)> { + let term = Arc::new(AtomicBool::new(false)); + + // For each terminating signal, we register both a shutdown handler and a handler that sets an + // atomic bool. With this, the process will get killed immediately once it receives a second + // termination signal (e.g. a double ctrl-c). + // https://docs.rs/signal-hook/0.3.15/signal_hook/#a-complex-signal-handling-with-a-background-thread + // for sig in TERM_SIGNALS { + // // But this will "arm" the above for the second time, by setting it to true. + // // The order of registering these is important, if you put this one first, it will + // // first arm and then terminate ‒ all in the first round. + // flag::register(*sig, Arc::clone(&term))?; + // } + + let signals = TERM_SIGNALS; + + Ok((term, Signals::new(signals)?)) +} + +/// Loop that handles terminating signals as they come in. +pub fn handle_signals(signals: &mut SignalsInfo, term: Arc, server: Arc) -> io::Result<()> { + let mut force = false; + + // We only register terminating signals, so we don't need to differentiate between what kind of + // signal came in + for _ in signals { + // If term is already true, this is the second signal, meaning we kill the process + // immediately. + if force { + return server.kill() + } + // The stop command runs in a separate thread to avoid blocking the signal handling loop. + // After stopping the server, the thread terminates the process. + else { + let clone = Arc::clone(&server); + std::thread::spawn(move || { + let _ = clone.stop(); + std::process::exit(0); + }); + } + + force = true; + } + + Ok(()) +} diff --git a/src/stdin.rs b/src/stdin.rs new file mode 100644 index 0000000..98d8463 --- /dev/null +++ b/src/stdin.rs @@ -0,0 +1,29 @@ +use std::sync::{Arc, Mutex}; +use std::io; + +use crate::server; + +pub fn handle_stdin(counter: Arc>) { + let stdin = io::stdin(); + let input = &mut String::new(); + + loop { + input.clear(); + + if stdin.read_line(input).is_err() { + continue; + }; + + { + let mut server = counter.lock().unwrap(); + + if let Err(e) = server.send_command(input) { + println!("{}", e); + }; + } + + if input.trim() == "stop" { + std::process::exit(0); + } + } +}