Added shutdown hook so we stop idle docker containers when stopping

the playground
Also containers are now stopped in parallel

Author: ffissore

ui/Cargo.lock

@@ -32,6 +32,7 @@ router = "0.6.0"
 openssl-probe = "0.1.2"
 dotenv = "0.13.0"
 snafu = "0.2.0"
+ctrlc = { version = "3.1.1", features = ["termination"] }
 
 [dependencies.playground-middleware]
 git = "https://github.com/integer32llc/playground-middleware"

ui/Cargo.toml

@@ -22,6 +22,7 @@ use std::{
     convert::{TryFrom, TryInto},
     env,
     path::PathBuf,
+    process,
     sync::{Arc, Mutex},
     time::{Duration, Instant},
 };
@@ -43,6 +44,13 @@ const ONE_YEAR_IN_SECONDS: u64 = 60 * 60 * 24 * 365;
 
 const SANDBOX_CACHE_TIME_TO_LIVE_IN_SECONDS: u64 = ONE_HOUR_IN_SECONDS as u64;
 
+fn set_graceful_shutdown_hook(containers: Arc<Mutex<DockerContainers>>) {
+    ctrlc::set_handler(move || {
+        containers.lock().unwrap().terminate();
+        process::exit(0);
+    }).expect("Error setting Ctrl-C handler");
+}
+
 fn main() {
     // Dotenv may be unable to load environment variables, but that's ok in production
     let _ = dotenv::dotenv();
@@ -58,7 +66,9 @@ fn main() {
     let cors_enabled = env::var_os("PLAYGROUND_CORS_ENABLED").is_some();
     let docker_containers_pool_size = env::var("DOCKER_CONTAINER_POOL_SIZE").ok().and_then(|v| v.parse().ok()).unwrap_or(DEFAULT_DOCKER_CONTAINER_POOL_SIZE);
 
-    let containers = Arc::new(DockerContainers::new(docker_containers_pool_size));
+    let containers = Arc::new(Mutex::new(DockerContainers::new(docker_containers_pool_size)));
+
+    set_graceful_shutdown_hook(containers.clone());
 
     let files = Staticfile::new(&root).expect("Unable to open root directory");
     let mut files = Chain::new(files);
@@ -142,21 +152,21 @@ impl iron::typemap::Key for GhToken {
     type Value = Self;
 }
 
-fn compile(req: &mut Request<'_, '_>, containers: &DockerContainers) -> IronResult<Response> {
+fn compile(req: &mut Request<'_, '_>, containers: &Mutex<DockerContainers>) -> IronResult<Response> {
     with_sandbox(req, |sandbox, req: CompileRequest| {
         let req: sandbox::CompileRequest = req.try_into()?;
-        let container = containers.pop(req.channel).unwrap();
+        let container = containers.lock().unwrap().pop(req.channel).unwrap();
         sandbox
             .compile(&req, &container)
             .map(CompileResponse::from)
             .eager_context(Compilation)
     })
 }
 
-fn handle(req: &mut Request<'_, '_>, containers: &DockerContainers) -> IronResult<Response> {
+fn handle(req: &mut Request<'_, '_>, containers: &Mutex<DockerContainers>) -> IronResult<Response> {
     with_sandbox(req, |sandbox, req: ExecuteRequest| {
         let req: sandbox::ExecuteRequest = req.try_into()?;
-        let container = containers.pop(req.channel).unwrap();
+        let container = containers.lock().unwrap().pop(req.channel).unwrap();
         sandbox
             .execute(&req, &container)
             .map(ExecuteResponse::from)
@@ -271,10 +281,10 @@ fn meta_gist_get(req: &mut Request<'_, '_>) -> IronResult<Response> {
 
 // This is a backwards compatibilty shim. The Rust homepage and the
 // documentation use this to run code in place.
-fn evaluate(req: &mut Request<'_, '_>, containers: &DockerContainers) -> IronResult<Response> {
+fn evaluate(req: &mut Request<'_, '_>, containers: &Mutex<DockerContainers>) -> IronResult<Response> {
     with_sandbox(req, |sandbox, req: EvaluateRequest| {
         let req = req.try_into()?;
-        let container = containers.pop(Channel::Stable).unwrap();
+        let container = containers.lock().unwrap().pop(Channel::Stable).unwrap();
         sandbox
             .execute(&req, &container)
             .map(EvaluateResponse::from)

ui/src/main.rs

@@ -12,7 +12,7 @@ use std::{
     sync::{Mutex, Arc},
     sync::mpsc::{self, Receiver, RecvError},
     time::Duration,
-    thread
+    thread::{self, JoinHandle}
 };
 use tempdir::TempDir;
 
@@ -396,17 +396,33 @@ pub struct DockerContainer {
     _temp_dir: TempDir, //we need to keep this, or it's dropped (= deleted)
     pub src_dir: PathBuf,
     pub output_dir: PathBuf,
+    terminated: Arc<Mutex<bool>>
 }
 
 impl DockerContainer {
 
-    pub fn terminate(&self) {
-        let mut cmd = Command::new("docker");
-        cmd
-            .arg("exec")
-            .arg(&self.id)
-            .args(&["pkill", "sleep"])
-            .status();
+    pub fn terminate(&mut self) -> JoinHandle<()> {
+        let mut terminated = self.terminated.lock().unwrap();
+        *terminated = true;
+
+        let container_id = self.id.clone();
+        thread::spawn(move || {
+            let mut cmd = Command::new("docker");
+            cmd
+                .arg("exec")
+                .arg(&container_id)
+                .args(&["pkill", "sleep"])
+                .status()
+                .expect(&format!("Unable to kill container with id {}", container_id));
+        })
+    }
+
+    // this was initially used to implement the Drop trait, keeping it for future reference 
+    #[allow(dead_code)]
+    fn is_terminated(&self) -> bool {
+        let terminated = &self.terminated.clone();
+        let terminated = terminated.lock().unwrap();
+        *terminated
     }
 
 }
@@ -431,7 +447,7 @@ impl DockerContainers {
             thread::spawn(move || {
                 let mut active = true;
                 while active {
-                    sender.send(DockerContainers::start_container(channel)).unwrap_or(());
+                    sender.send(Self::start_container(channel)).unwrap_or(());
                     active = *active_mutex.lock().unwrap();
                 }
             });
@@ -474,7 +490,7 @@ impl DockerContainers {
             Ok(child) => {
                 let id = Self::container_id_from_stdout(child.stdout.unwrap());
 
-                DockerContainer { id, _temp_dir: temp_dir, src_dir, output_dir }
+                DockerContainer { id, _temp_dir: temp_dir, src_dir, output_dir, terminated: Arc::new(Mutex::new(false)) }
             }
             Err(err) => panic!("Error starting container {}: {}", channel.container_name(), err)
         }
@@ -483,19 +499,28 @@ impl DockerContainers {
     pub fn pop(&self, channel: Channel) -> Result<DockerContainer, RecvError> {
         self.receivers.get(&channel).unwrap().lock().unwrap().recv()
     }
-}
-
-impl Drop for DockerContainers {
 
-    fn drop(&mut self) {
+    pub fn terminate(&self) {
+        log::info!("Shutting down docker containers pool...");
         let mut active = self.active.lock().unwrap();
         *active = false;
 
+        let mut terminate_handles = Vec::new();
         for (_, receiver) in self.receivers.iter() {
-            while let Ok(container) = receiver.lock().unwrap().recv_timeout(Duration::new(1, 0)) {
-                container.terminate();
+            while let Ok(mut container) = receiver.lock().unwrap().recv_timeout(Duration::new(1, 0)) {
+                terminate_handles.push(container.terminate());
             }
         }
+        for handle in terminate_handles.into_iter() {
+            handle.join().expect("error when joining");
+        }
+    }
+}
+
+impl Drop for DockerContainers {
+
+    fn drop(&mut self) {
+        self.terminate();
     }
 
 }
@@ -1447,8 +1472,8 @@ mod test {
         let containers = DockerContainers::new(0);
         thread::sleep(Duration::new(2, 0));
 
-        let container = containers.pop(Channel::Stable).unwrap();
+        let mut container = containers.pop(Channel::Stable).unwrap();
         assert_eq!(false, container.id.is_empty());
-        container.terminate();
+        container.terminate().join().expect("error when joining");
     }
 }