hbakd: graceful shutdown on SIGINT, SIGTERM and SIGHUP

5 files changed, 95 insertions, 17 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 4267893..2af4447 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -115,6 +115,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
 
 [[package]]
+name = "bitflags"
+version = "2.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ed570934406eb16438a4e976b1b4500774099c13b8cb96eec99f620f05090ddf"
+
+[[package]]
 name = "blake2"
 version = "0.10.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -276,6 +282,16 @@ dependencies = [
 ]
 
 [[package]]
+name = "ctrlc"
+version = "3.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b467862cc8610ca6fc9a1532d7777cee0804e678ab45410897b9396495994a0b"
+dependencies = [
+ "nix",
+ "windows-sys 0.52.0",
+]
+
+[[package]]
 name = "digest"
 version = "0.10.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -363,8 +379,10 @@ name = "hbakd"
 version = "0.1.0-dev"
 dependencies = [
  "clap",
+ "ctrlc",
  "fork",
  "hbak_common",
+ "thiserror",
 ]
 
 [[package]]
@@ -458,6 +476,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "523dc4f511e55ab87b694dc30d0f820d60906ef06413f93d4d7a1385599cc149"
 
 [[package]]
+name = "nix"
+version = "0.27.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2eb04e9c688eff1c89d72b407f168cf79bb9e867a9d3323ed6c01519eb9cc053"
+dependencies = [
+ "bitflags 2.4.2",
+ "cfg-if",
+ "libc",
+]
+
+[[package]]
 name = "num-traits"
 version = "0.2.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -672,7 +701,7 @@ version = "2.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ab6b8b397ed1fa2364625189a004d532b5d58f772943c57aa144feb1966142bd"
 dependencies = [
- "bitflags",
+ "bitflags 1.3.2",
  "libc",
  "smart-default",
  "thiserror",
diff --git a/hbak_common/src/conn.rs b/hbak_common/src/conn.rs
index dca5769..e50d868 100644
--- a/hbak_common/src/conn.rs
+++ b/hbak_common/src/conn.rs
@@ -24,8 +24,9 @@ pub const DEFAULT_PORT: u16 = 20406;
 
 /// TCP connect timeout. Connection attempt is aborted if remote doesn't respond.
 pub const CONNECT_TIMEOUT: Duration = Duration::from_secs(30);
-/// TCP read timeout. Internal value for [`StreamConn::data_sync`] receive thread cancellation.
-const READ_TIMEOUT: Duration = Duration::from_millis(200);
+/// TCP read timeout. Used for cancellation of [`StreamConn::data_sync`] receive thread
+/// and `hbakd` TCP accept loop.
+pub const READ_TIMEOUT: Duration = Duration::from_millis(200);
 
 mod private {
     pub trait Sealed {}
diff --git a/hbakd/Cargo.toml b/hbakd/Cargo.toml
index 1038ceb..0e360e7 100644
--- a/hbakd/Cargo.toml
+++ b/hbakd/Cargo.toml
@@ -7,5 +7,7 @@ edition = "2021"
 
 [dependencies]
 clap = { version = "4.4.18", features = ["derive"] }
+ctrlc = { version = "3.4.2", features = ["termination"] }
 fork = "0.1.22"
 hbak_common = { version = "0.1.0-dev", path = "../hbak_common" }
+thiserror = "1.0"
diff --git a/hbakd/src/error.rs b/hbakd/src/error.rs
new file mode 100644
index 0000000..2ee386a
--- /dev/null
+++ b/hbakd/src/error.rs
@@ -0,0 +1,19 @@
+use std::io;
+
+use thiserror::Error;
+
+#[derive(Debug, Error)]
+pub enum Error {
+    #[error("An error occured on the local node: {0}")]
+    HbakLocalNode(#[from] hbak_common::LocalNodeError),
+    #[error("A network error occured: {0}")]
+    HbakNetwork(#[from] hbak_common::NetworkError),
+
+    #[error("IO error: {0}")]
+    Io(#[from] io::Error),
+
+    #[error("Unable to set signal handler: {0}")]
+    Ctrlc(#[from] ctrlc::Error),
+}
+
+pub type Result<T> = std::result::Result<T, Error>;
diff --git a/hbakd/src/main.rs b/hbakd/src/main.rs
index c1427fc..cc9ce3f 100644
--- a/hbakd/src/main.rs
+++ b/hbakd/src/main.rs
@@ -1,12 +1,16 @@
-use hbak_common::conn::{AuthServ, DEFAULT_PORT};
+mod error;
+use error::*;
+
+use hbak_common::conn::{AuthServ, DEFAULT_PORT, READ_TIMEOUT};
 use hbak_common::message::SyncInfo;
 use hbak_common::proto::{LocalNode, Mode, Node, Snapshot};
-use hbak_common::{LocalNodeError, NetworkError, RemoteError};
+use hbak_common::RemoteError;
 
 use std::collections::HashMap;
 use std::fs::{self, File};
 use std::io;
 use std::net::{IpAddr, Ipv6Addr, SocketAddr, TcpListener, TcpStream};
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use std::thread;
 
@@ -42,7 +46,15 @@ fn main() {
     }
 }
 
-fn serve() -> Result<(), LocalNodeError> {
+fn serve() -> Result<()> {
+    let should_exit = Arc::new(AtomicBool::new(false));
+    let should_exit2 = Arc::clone(&should_exit);
+
+    ctrlc::set_handler(move || {
+        println!("[info] Caught SIGINT, SIGTERM or SIGHUP, exiting");
+        should_exit2.store(true, Ordering::SeqCst);
+    })?;
+
     let local_node = Arc::new(LocalNode::new(Mode::Server)?);
 
     let bind_addr = local_node.config().bind_addr.unwrap_or(SocketAddr::new(
@@ -52,23 +64,36 @@ fn serve() -> Result<(), LocalNodeError> {
 
     let listener = TcpListener::bind(bind_addr)?;
 
+    listener.set_nonblocking(true)?;
+
     println!("[info] <{}> Listening", bind_addr);
 
     for stream in listener.incoming() {
-        let stream = stream?;
-        let peer_addr = stream.peer_addr()?;
-
-        let local_node = Arc::clone(&local_node);
-        thread::spawn(move || match handle_client(&local_node, stream) {
-            Ok(_) => println!("[info] <{}> Disconnected", peer_addr),
-            Err(e) => eprintln!("[warn] <{}> Cannot handle client: {}", peer_addr, e),
-        });
+        match stream {
+            Ok(stream) => {
+                let peer_addr = stream.peer_addr()?;
+
+                let local_node = Arc::clone(&local_node);
+                thread::spawn(move || match handle_client(&local_node, stream) {
+                    Ok(_) => println!("[info] <{}> Disconnected", peer_addr),
+                    Err(e) => eprintln!("[warn] <{}> Cannot handle client: {}", peer_addr, e),
+                });
+            }
+            Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
+                if should_exit.load(Ordering::SeqCst) {
+                    break;
+                } else {
+                    thread::sleep(READ_TIMEOUT);
+                }
+            }
+            Err(e) => return Err(e.into()),
+        }
     }
 
-    unreachable!()
+    Ok(())
 }
 
-fn handle_client(local_node: &LocalNode, stream: TcpStream) -> Result<(), NetworkError> {
+fn handle_client(local_node: &LocalNode, stream: TcpStream) -> Result<()> {
     let peer_addr = stream.peer_addr()?;
 
     let auth_serv = AuthServ::from(stream);
@@ -171,5 +196,7 @@ fn handle_client(local_node: &LocalNode, stream: TcpStream) -> Result<(), Networ
         Ok(())
     };
 
-    stream_conn.data_sync(tx, rx_setup, rx_finish)
+    stream_conn.data_sync(tx, rx_setup, rx_finish)?;
+
+    Ok(())
 }