diff --git a/src/lib.rs b/src/lib.rs index 383ae63..334ee41 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -46,6 +46,8 @@ use std::fs::File; use std::io::prelude::*; use std::net::{TcpListener, TcpStream}; use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; use std::time::Duration; use std::borrow::Borrow; @@ -69,6 +71,7 @@ pub struct Server { handler: Handler, timeout: Option, static_directory: Option, + pub shutdown: Arc, } impl fmt::Debug for Server { @@ -117,6 +120,7 @@ impl Server { handler: Box::new(handler), timeout: None, static_directory: Some(PathBuf::from("public")), + shutdown: Arc::new(AtomicBool::new(false)), } } @@ -157,6 +161,7 @@ impl Server { handler: Box::new(handler), timeout: Some(timeout), static_directory: Some(PathBuf::from("public")), + shutdown: Arc::new(AtomicBool::new(false)), } } @@ -207,7 +212,7 @@ impl Server { /// server.listen("127.0.0.1", "7979"); /// } /// ``` - pub fn listen(&self, host: &str, port: &str) -> ! { + pub fn listen(&self, host: &str, port: &str) { let listener = TcpListener::bind(format!("{}:{}", host, port)).expect("Error starting the server."); @@ -246,13 +251,17 @@ impl Server { /// server.listen_on_socket(listener); /// } /// ``` - pub fn listen_on_socket(&self, listener: TcpListener) -> ! { + pub fn listen_on_socket(&self, listener: TcpListener) { const READ_TIMEOUT_MS: u64 = 20; let num_threads = self.pool_size(); let mut pool = Pool::new(num_threads); let mut incoming = listener.incoming(); loop { + // check if shutdown received + if self.shutdown.load(Ordering::SeqCst) { + return; + } // Incoming is an endless iterator, so it's okay to unwrap on it. let stream = incoming.next().unwrap(); let stream = stream.expect("Error handling TCP stream."); @@ -341,7 +350,7 @@ impl Server { fn handle_connection(&self, mut stream: TcpStream) -> Result<(), Error> { let request = match request::read(&mut stream, self.timeout) { Err(Error::ConnectionClosed) | Err(Error::Timeout) | Err(Error::HttpParse(_)) => { - return Ok(()) + return Ok(()); } Err(Error::RequestTooLarge) => { diff --git a/src/parsing.rs b/src/parsing.rs index 3d5865c..f564a7f 100644 --- a/src/parsing.rs +++ b/src/parsing.rs @@ -95,7 +95,8 @@ pub fn try_parse_request(buffer: Vec) -> Result) -> Result>(); + ) + .collect::>(); (method, proto, headers, n) }) };