Fix/shutdown (#35)

* Update Cargo.toml, README, and CONTRIBUTING link for better documentation and accessibility 🦜🕺 (#32)

* Update the mod in src

* Add shutdown logger

* Add the handler for shutdown

* Use libc to handle unix

* Add shutdown.rs

* Update the mod in src

* Add the handler for shutdown

* Use libc to handle unix

* Fix shutdown module compilation error

* Fix shared reference in unix

* Fix shared reference in unix

* Access the global without shared reference

* Fix temporary value errors in shutdown.rs

* Simplified exit logic in unix
Add force quiting for repeatedly ctrl+c press

* Move close logic to main thread.

* Change signal handling and add safety checks to avoid UB

* Basic implementation to force the awake of listener thread.

this will attempt to connect to itself to wake the listener and check if
the flag is set to false

* Refactor to reach scalability and security in unix sign handler

* Add setup for graceful shutdown

---------

Co-authored-by: LoboGuardian 🐺 <[email protected]>
Co-authored-by: miky-rola <[email protected]>

Author: 3 people

Cargo.lock

@@ -10,7 +10,11 @@ documentation = "https://docs.rs/hteapot/"
 homepage = "https://github.com/az107/HTeaPot"
 repository = "https://github.com/az107/HTeaPot"
 keywords = ["http", "server", "web", "lightweight", "rust"]
-categories = ["network-programming", "web-programming", "command-line-utilities"]
+categories = [
+    "network-programming",
+    "web-programming",
+    "command-line-utilities",
+]
 exclude = ["config.toml", "demo/", "README.md"]
 
 [lib]
@@ -20,5 +24,9 @@ path = "src/hteapot/mod.rs"
 [[bin]]
 name = "hteapot"
 
+[dependencies]
+libc = "0.2.172"
+
+
 [package.metadata.docs.rs]
-no-readme = true
+no-readme = true

Cargo.toml

@@ -1,5 +1,5 @@
 // Written by Alberto Ruiz 2024-03-08
-// 
+//
 // This is the HTTP server module, it will handle the requests and responses
 // Also provides utilities to parse the requests and build the response
 
@@ -17,11 +17,11 @@
 //! ```
 
 /// Submodules for HTTP functionality.
-pub mod brew;        // HTTP client implementation
-mod methods;         // HTTP method and status enums
-mod request;         // Request parsing and builder
-mod response;        // Response types and streaming
-mod status;          // Status code mapping
+pub mod brew; // HTTP client implementation
+mod methods; // HTTP method and status enums
+mod request; // Request parsing and builder
+mod response; // Response types and streaming
+mod status; // Status code mapping
 
 // Internal types used for connection management
 use self::response::{EmptyHttpResponse, HttpResponseCommon, IterError};
@@ -37,6 +37,7 @@ pub use self::status::HttpStatus;
 use std::collections::VecDeque;
 use std::io::{self, Read, Write};
 use std::net::{Shutdown, TcpListener, TcpStream};
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Condvar, Mutex};
 use std::thread;
 use std::time::{Duration, Instant};
@@ -76,6 +77,8 @@ pub struct Hteapot {
     port: u16,
     address: String,
     threads: u16,
+    shutdown_signal: Option<Arc<AtomicBool>>,
+    shutdown_hooks: Vec<Arc<dyn Fn() + Send + Sync + 'static>>,
 }
 
 /// Represents the state of a connection's lifecycle.
@@ -95,12 +98,33 @@ struct SocketData {
 }
 
 impl Hteapot {
+    pub fn set_shutdown_signal(&mut self, signal: Arc<AtomicBool>) {
+        self.shutdown_signal = Some(signal);
+    }
+
+    pub fn get_shutdown_signal(&self) -> Option<Arc<AtomicBool>> {
+        self.shutdown_signal.clone()
+    }
+
+    pub fn add_shutdown_hook<F>(&mut self, hook: F)
+    where
+        F: Fn() + Send + Sync + 'static,
+    {
+        self.shutdown_hooks.push(Arc::new(hook));
+    }
+
+    pub fn get_addr(&self) -> (String, u16) {
+        return (self.address.clone(), self.port);
+    }
+
     // Constructor
     pub fn new(address: &str, port: u16) -> Self {
         Hteapot {
             port,
             address: address.to_string(),
             threads: 1,
+            shutdown_signal: None,
+            shutdown_hooks: Vec::new(),
         }
     }
 
@@ -109,6 +133,8 @@ impl Hteapot {
             port,
             address: address.to_string(),
             threads: if threads == 0 { 1 } else { threads },
+            shutdown_signal: None,
+            shutdown_hooks: Vec::new(),
         }
     }
 
@@ -132,23 +158,34 @@ impl Hteapot {
             Arc::new(Mutex::new(vec![0; self.threads as usize]));
         let arc_action = Arc::new(action);
 
+        // Clone shutdown_signal and share the shutdown_hooks via Arc
+        let shutdown_signal = self.shutdown_signal.clone();
+        let shutdown_hooks = Arc::new(self.shutdown_hooks.clone());
+
         for thread_index in 0..self.threads {
             let pool_clone = pool.clone();
             let action_clone = arc_action.clone();
             let priority_list_clone = priority_list.clone();
+            let shutdown_signal_clone = shutdown_signal.clone();
 
             thread::spawn(move || {
                 let mut streams_to_handle = Vec::new();
                 loop {
                     {
                         let (lock, cvar) = &*pool_clone;
                         let mut pool = lock.lock().expect("Error locking pool");
-
                         if streams_to_handle.is_empty() {
                             // Store the returned guard back into pool
-                            pool = cvar.wait_while(pool, |pool| pool.is_empty())
+                            pool = cvar
+                                .wait_while(pool, |pool| pool.is_empty())
                                 .expect("Error waiting on cvar");
                         }
+                        //TODO: move this to allow process the last request
+                        if let Some(signal) = &shutdown_signal_clone {
+                            if !signal.load(Ordering::SeqCst) {
+                                break; // Exit the server loop
+                            }
+                        }
 
                         while let Some(stream) = pool.pop_back() {
                             let socket_status = SocketStatus {
@@ -185,6 +222,17 @@ impl Hteapot {
         }
 
         loop {
+            if let Some(signal) = &shutdown_signal {
+                if !signal.load(Ordering::SeqCst) {
+                    let (lock, cvar) = &*pool;
+                    let _guard = lock.lock().unwrap();
+                    cvar.notify_all();
+                    for hook in shutdown_hooks.iter() {
+                        hook();
+                    }
+                    break;
+                }
+            }
             let stream = match listener.accept() {
                 Ok((stream, _)) => stream,
                 Err(_) => continue,
@@ -216,17 +264,14 @@ impl Hteapot {
     ) -> Option<()> {
         let status = socket_data.status.as_mut()?;
 
-        // Fix by miky-rola 2025-04-08
         // Check if the TTL (time-to-live) for the connection has expired.
-        // If the connection is idle for longer than `KEEP_ALIVE_TTL` and no data is being written,
-        // the connection is gracefully shut down to free resources.        
         if Instant::now().duration_since(status.ttl) > KEEP_ALIVE_TTL && !status.write {
             let _ = socket_data.stream.shutdown(Shutdown::Both);
             return None;
         }
+
         // If the request is not yet complete, read data from the stream into a buffer.
         // This ensures that the server can handle partial or chunked requests.
-
         if !status.request.done {
             let mut buffer = [0; BUFFER_SIZE];
             match socket_data.stream.read(&mut buffer) {
@@ -286,8 +331,7 @@ impl Hteapot {
             status.response = response;
         }
 
-        // Write the response to the client in chunks using the `peek` and `next` methods.
-        // This ensures that large responses are sent incrementally without blocking the server.
+        // Write the response to the client in chunks
         loop {
             match status.response.peek() {
                 Ok(n) => match socket_data.stream.write(&n) {
@@ -390,4 +434,4 @@ mod tests {
         assert!(response_str.contains("Server: HTeaPot/"));
         assert!(second_response_str.contains("Second Request"));
     }
-}
+}

src/hteapot/mod.rs

@@ -190,7 +190,6 @@ impl HttpRequestBuilder {
                 if parts.len() != 3 {
                     return Err("Invalid method + path + version request");
                 }
-
                 self.request.method = HttpMethod::from_str(parts[0]);
                 let path_parts: Vec<&str> = parts[1].split('?').collect();
                 self.request.path = path_parts[0].to_string();

src/hteapot/request.rs

@@ -1,9 +1,9 @@
+use std::fmt;
 use std::io::Write;
-use std::sync::mpsc::{channel, Sender};
+use std::sync::Arc;
+use std::sync::mpsc::{Sender, channel};
 use std::thread::{self, JoinHandle};
 use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
-use std::fmt;
-use std::sync::Arc;
 
 /// Differnt log levels
 #[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd, Copy)]
@@ -86,7 +86,15 @@ impl SimpleTime {
         // calculate millisecs from nanosecs
         let millis = nanos / 1_000_000;
 
-        (year, month as u32 + 1, day as u32, hour, minute, second, millis)
+        (
+            year,
+            month as u32 + 1,
+            day as u32,
+            hour,
+            minute,
+            second,
+            millis,
+        )
     }
 
     /// Returns a formatted timestamp string for the current system time.
@@ -133,7 +141,7 @@ impl Logger {
     pub fn new<W: Sized + Write + Send + Sync + 'static>(
         mut writer: W,
         min_level: LogLevel,
-        component: &str
+        component: &str,
     ) -> Logger {
         let (tx, rx) = channel::<LogMessage>();
         let thread = thread::spawn(move || {
@@ -151,7 +159,7 @@ impl Logger {
                             msg.timestamp, msg.level, msg.component, msg.content
                         );
                         buff.push(formatted);
-                    },
+                    }
                     Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {}
                     Err(_) => break,
                 }
@@ -238,8 +246,8 @@ impl Logger {
     pub fn fatal(&self, content: String) {
         self.log(LogLevel::FATAL, content);
     }
-    /// Log a message with TRACE level 
-    #[allow(dead_code)] 
+    /// Log a message with TRACE level
+    #[allow(dead_code)]
     pub fn trace(&self, content: String) {
         self.log(LogLevel::TRACE, content);
     }
@@ -255,20 +263,19 @@ impl Clone for Logger {
     }
 }
 
-
 #[cfg(test)]
 mod tests {
     use super::*;
     use std::io::stdout;
-    
+
     #[test]
     fn test_basic() {
         let logs = Logger::new(stdout(), LogLevel::DEBUG, "test");
         logs.info("test message".to_string());
         logs.debug("debug info".to_string());
-        
+
         // Create a sub-logger with a different component
         let sub_logger = logs.with_component("sub-component");
         sub_logger.warn("warning from sub-component".to_string());
     }
-}
+}