Handle signals and shutdown gracefully

Author: prasmussen

Cargo.lock

@@ -16,3 +16,4 @@ log = "0.4.11"
 env_logger = "0.7.1"
 tiny_http = { git = "https://github.com/prasmussen/tiny-http" }
 ascii = "1.0.0"
+signal-hook = "0.1.16"

Cargo.toml

@@ -16,54 +16,78 @@ pub struct Config<C, H> {
     pub handler: H,
 }
 
-pub fn start<C, H>(config: Config<C, H>) -> Result<(), Error>
-    where
-        C: Send + Clone + 'static,
-        H: Send + Copy + 'static,
-        H: FnOnce(&C, tiny_http::Request) {
-
-    let server = tiny_http::Server::new(config.listen_addr)
-        .map_err(Error::Bind)?;
-
-    let mut handles = Vec::new();
-    let request_handler = config.handler;
-
-    for n in 0..config.worker_threads {
-        let handler_config = config.handler_config.clone();
-        let server = server.try_clone()
-            .map_err(|err| Error::CloneServer(err, n))?;
-
-        handles.push(thread::spawn(move || {
-            loop {
-                match server.accept() {
-                    Ok(client) => {
-                        for request in client {
-                            request_handler(&handler_config, request);
+pub struct Server {
+    server: tiny_http::Server,
+}
+
+impl Server {
+    pub fn new(listen_addr: String) -> Result<Server, io::Error> {
+        let server = tiny_http::Server::new(listen_addr)?;
+
+        Ok(Server{
+            server,
+        })
+    }
+
+    pub fn start<C, H>(&self, config: Config<C, H>) -> Result<Workers, Error>
+        where
+            C: Send + Clone + 'static,
+            H: Send + Copy + 'static,
+            H: FnOnce(&C, tiny_http::Request) {
+
+        let mut handles = Vec::new();
+        let request_handler = config.handler;
+
+        for n in 0..config.worker_threads {
+            let handler_config = config.handler_config.clone();
+            let server = self.server.try_clone()
+                .map_err(|err| Error::CloneServer(err, n))?;
+
+            handles.push(thread::spawn(move || {
+                loop {
+                    match server.accept() {
+                        Ok(client) => {
+                            for request in client {
+                                request_handler(&handler_config, request);
+                            }
                         }
-                    }
 
-                    Err(tiny_http::AcceptError::Accept(err)) => {
-                        log::error!("Accept error on thread {}: {:?}", n, err);
-                        break;
-                    }
+                        Err(tiny_http::AcceptError::Accept(err)) => {
+                            log::error!("Accept error on thread {}: {:?}", n, err);
+                            break;
+                        }
 
-                    Err(tiny_http::AcceptError::ShuttingDown()) => {
-                        log::info!("Thread {} shutting down", n);
-                        break;
+                        Err(tiny_http::AcceptError::ShuttingDown()) => {
+                            log::info!("Thread {} shutting down", n);
+                            break;
+                        }
                     }
                 }
-            }
-        }))
-    }
+            }))
+        }
 
-    // Wait for threads to complete, in practice this will block forever unless there is a panic
-    for handle in handles {
-        handle.join().unwrap();
+        Ok(Workers{
+            handles
+        })
     }
+}
 
-    Ok(())
+pub struct Workers {
+    handles: Vec<thread::JoinHandle<()>>
 }
 
+impl Workers {
+    pub fn wait(self) {
+        // Wait for threads to complete, in practice this will block forever unless:
+        // - The server is shutdown
+        // - One of the threads panics
+        for handle in self.handles {
+            handle.join().unwrap();
+        }
+    }
+}
+
+
 #[derive(Debug, Clone)]
 pub struct ApiConfig {
     pub access_token: ascii::AsciiString,
@@ -181,17 +205,12 @@ pub fn error_response(request: tiny_http::Request, error: ErrorResponse) -> Resu
 
 
 pub enum Error {
-    Bind(io::Error),
     CloneServer(io::Error, u16),
 }
 
 impl fmt::Display for Error {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
-            Error::Bind(err) => {
-                write!(f, "Failed to bind: {}", err)
-            }
-
             Error::CloneServer(err, n) => {
                 write!(f, "Failed to clone server (n = {}): {}", n, err)
             }

src/docker_run/api/mod.rs

@@ -1,8 +1,13 @@
 mod docker_run;
 
+use std::io;
+use std::fmt;
+use std::thread;
 use std::process;
 use std::time::Duration;
 
+use signal_hook::iterator::Signals;
+
 use docker_run::config;
 use docker_run::environment;
 use docker_run::unix_stream;
@@ -14,38 +19,72 @@ fn main() {
     env_logger::init();
 
     match start() {
-        Ok(()) => {}
-
-        Err(Error::BuildConfig(err)) => {
-            log::error!("Failed to build config: {}", err);
-            process::exit(1)
+        Ok(()) => {
+            log::info!("Exiting gracefully")
         }
 
-        Err(Error::StartServer(err)) => {
-            log::error!("Failed to start api server: {}", err);
+        Err(err) => {
+            log::error!("{}", err);
             process::exit(1)
         }
     }
 }
 
 enum Error {
     BuildConfig(environment::Error),
+    CreateServer(io::Error),
     StartServer(api::Error),
+    Signal(io::Error),
+}
+
+impl fmt::Display for Error {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Error::BuildConfig(err) => {
+                write!(f, "Failed to build config: {}", err)
+            }
+
+            Error::CreateServer(err) => {
+                write!(f, "Failed to create api server: {}", err)
+            }
+
+            Error::StartServer(err) => {
+                write!(f, "Failed to start api server: {}", err)
+            }
+
+            Error::Signal(err) => {
+                write!(f, "Failed to register signal handler: {}", err)
+            }
+        }
+    }
 }
 
+
 fn start() -> Result<(), Error> {
     let env = environment::get_environment();
     let config = build_config(&env)
         .map_err(Error::BuildConfig)?;
 
     log::info!("Listening on {} with {} worker threads", config.server.listen_addr_with_port(), config.server.worker_threads);
 
-    api::start(api::Config{
+    let server = api::Server::new(config.server.listen_addr_with_port())
+        .map_err(Error::CreateServer)?;
+
+    let workers = server.start(api::Config{
         listen_addr: config.server.listen_addr_with_port(),
         worker_threads: config.server.worker_threads,
         handler_config: config,
         handler: handle_request,
-    }).map_err(Error::StartServer)
+    }).map_err(Error::StartServer)?;
+
+    // Handle OS signals
+    handle_signals(server)
+        .map_err(Error::Signal)?;
+
+    // Wait for workers
+    workers.wait();
+
+    Ok(())
 }
 
 
@@ -92,6 +131,39 @@ fn router(config: &config::Config, request: &mut tiny_http::Request) -> Result<a
 }
 
 
+fn handle_signals(server: api::Server) -> Result<(), io::Error> {
+    let signals = Signals::new(&[
+        signal_hook::SIGTERM,
+        signal_hook::SIGINT,
+    ])?;
+
+    thread::spawn(move || {
+        for signal in signals.forever() {
+            match signal {
+                signal_hook::SIGTERM => {
+                    log::info!("Caught SIGTERM signal");
+                    break
+                }
+
+                signal_hook::SIGINT => {
+                    log::info!("Caught SIGINT signal");
+                    break
+                }
+
+                _ => {
+                    log::info!("Ignoring signal {}", signal);
+                }
+            }
+        }
+
+        log::info!("Shutting down server...");
+        drop(server)
+    });
+
+    Ok(())
+}
+
+
 fn build_config(env: &environment::Environment) -> Result<config::Config, environment::Error> {
     let server = build_server_config(env)?;
     let api = build_api_config(env)?;