feat: add basis for signal handling

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"

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"

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<Mutex<server::ServerProcess>>, 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);
+        let clone = Arc::clone(&server);
-        std::thread::spawn(move || backups_thread(clone, cli.frequency));
+        std::thread::spawn(move || backups_thread(server, cli.frequency));
     }
 
-    let stdin = io::stdin();
+    // Spawn thread that handles the main stdin loop
-    let input = &mut String::new();
+    let clone = Arc::clone(&server);
+    std::thread::spawn(move || stdin::handle_stdin(server));
 
-    loop {
+    // Signal handler loop exits the process when necessary
-        input.clear();
+    signals::handle_signals(&mut signals, term, server)
-
-        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;
-        }
-    }
 }

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<()> {

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<AtomicBool>, 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<AtomicBool>, server: Arc<server::ServerProcess>) -> 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(())
+}

src/stdin.rs

@@ -0,0 +1,29 @@
+use std::sync::{Arc, Mutex};
+use std::io;
+
+use crate::server;
+
+pub fn handle_stdin(counter: Arc<Mutex<server::ServerProcess>>) {
+    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);
+        }
+    }
+}