Support for graceful exit from loop using shutdown flag

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<Duration>,
     static_directory: Option<PathBuf>,
+    pub shutdown: Arc<AtomicBool>,
 }
 
 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) => {

src/parsing.rs

@@ -95,7 +95,8 @@ pub fn try_parse_request(buffer: Vec<u8>) -> Result<ParseResult, httparse::Error
                 let method = RequestMethodIndices(method.0, method.1);
 
                 (r, method, proto, n)
-            }).map(|(r, method, proto, n)| {
+            })
+            .map(|(r, method, proto, n)| {
                 let headers = r
                     .headers
                     .iter()
@@ -109,7 +110,8 @@ pub fn try_parse_request(buffer: Vec<u8>) -> Result<ParseResult, httparse::Error
                                 value: slice_indices(&*buffer, value),
                             }
                         },
-                    ).collect::<Vec<_>>();
+                    )
+                    .collect::<Vec<_>>();
                 (method, proto, headers, n)
             })
     };