From 8ecfafaf756461ed971a01850bde3b367f648a6b Mon Sep 17 00:00:00 2001
From: q1-silver <q1@silver.lan>
Date: Tue, 26 Oct 2021 01:23:12 +0300
Subject: [PATCH] add: basic signal handling

---
 Cargo.lock  | 20 ++++++++++++++++++++
 Cargo.toml  |  3 ++-
 src/main.rs | 22 ++++++++++++++++++----
 3 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 7e9afaa..5bf3d57 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -8,9 +8,29 @@ version = "0.2.105"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "869d572136620d55835903746bcb5cdc54cb2851fd0aeec53220b4bb65ef3013"
 
+[[package]]
+name = "signal-hook"
+version = "0.3.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9c98891d737e271a2954825ef19e46bd16bdb98e2746f2eec4f7a4ef7946efd1"
+dependencies = [
+ "libc",
+ "signal-hook-registry",
+]
+
+[[package]]
+name = "signal-hook-registry"
+version = "1.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e51e73328dc4ac0c7ccbda3a494dfa03df1de2f46018127f60c693f2648455b0"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "tepoll"
 version = "0.1.0"
 dependencies = [
  "libc",
+ "signal-hook",
 ]
diff --git a/Cargo.toml b/Cargo.toml
index 460733b..b63b0f5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -6,4 +6,5 @@ edition = "2018"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-libc = "*"
\ No newline at end of file
+libc = "*"
+signal-hook = "*"
diff --git a/src/main.rs b/src/main.rs
index fd87c69..2d2d125 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,6 +4,8 @@ use std::collections::HashMap;
 use std::io::{self, prelude::*};
 use std::net::{TcpListener, TcpStream};
 use std::os::unix::io::{AsRawFd, RawFd};
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
 
 #[allow(unused_macros)]
 macro_rules! syscall {
@@ -27,6 +29,9 @@ const READ_FLAGS: i32 = libc::EPOLLONESHOT | libc::EPOLLIN;
 const WRITE_FLAGS: i32 = libc::EPOLLONESHOT | libc::EPOLLOUT;
 
 fn main() -> io::Result<()> {
+    let term = Arc::new(AtomicBool::new(false));
+    signal_hook::flag::register(signal_hook::consts::SIGINT, Arc::clone(&term))?;
+
     let mut request_contexts: HashMap<u64, RequestContext> = HashMap::new();
     let mut events: Vec<libc::epoll_event> = Vec::with_capacity(1024);
     let mut key = 100;
@@ -39,8 +44,12 @@ fn main() -> io::Result<()> {
     add_interest(epoll_fd, listener_fd, listener_read_event(key))?;
 
     loop {
-        // TODO[0]: check signal and 0 contexts
-        println!("\nrequests in flight: {}", request_contexts.len());
+        // check signal and 0 contexts
+        if term.load(Ordering::Relaxed) && request_contexts.len() == 0 {
+            println!("received sigterm, queue is empty, shutting down..");
+            break;
+        }
+        //println!("\nrequests in flight: {}", request_contexts.len());
         events.clear();
         let res = match syscall!(epoll_wait(
             epoll_fd,
@@ -49,6 +58,7 @@ fn main() -> io::Result<()> {
             1000 as libc::c_int,
         )) {
             Ok(v) => v,
+            Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
             Err(e) => panic!("error during epoll wait: {}", e),
         };
 
@@ -68,8 +78,11 @@ fn main() -> io::Result<()> {
                         }
                         Err(e) => eprintln!("couldn't accept: {}", e),
                     };
-                    // TODO[0]: check signal
-                    modify_interest(epoll_fd, listener_fd, listener_read_event(100))?;
+                    if !term.load(Ordering::Relaxed) {
+                        modify_interest(epoll_fd, listener_fd, listener_read_event(100))?;
+                    } else {
+                        println!("received sigterm, not accepting new incoming connections");
+                    }
                 }
                 key => {
                     let mut to_delete = None;
@@ -93,6 +106,7 @@ fn main() -> io::Result<()> {
             }
         }
     }
+    Ok(())
 }
 
 fn epoll_create() -> io::Result<RawFd> {