Add a global cache for executed requests

Closes #219

Author: alexcrichton

src/bin/playbot.rs

@@ -10,27 +10,28 @@ extern crate url;
 use std::fs::File;
 use std::io::{self, Read};
 use std::str;
-use std::u16;
+use std::sync::Arc;
 use std::thread;
+use std::u16;
 
 use hyper::client::Client;
 use irc::client::prelude::*;
-use rust_playpen::ReleaseChannel;
+use rust_playpen::{ReleaseChannel, Cache};
 use rustc_serialize::json;
 use url::form_urlencoded;
 
 static DEFAULT_CHANNEL: ReleaseChannel = ReleaseChannel::Stable;
 
-fn get_rust_versions() -> Vec<String> {
+fn get_rust_versions(cache: &Cache) -> Vec<String> {
     let mut versions = Vec::new();
     // Note: Keep these in the same order as their discriminant values
     for channel in &[ReleaseChannel::Stable,
                      ReleaseChannel::Beta,
                      ReleaseChannel::Nightly] {
-        let (status, output) = rust_playpen::exec(*channel,
-                                                  "rustc",
-                                                  vec![String::from("-V")],
-                                                  String::new()).unwrap();
+        let (status, output) = cache.exec(*channel,
+                                          "rustc",
+                                          vec![String::from("-V")],
+                                          String::new()).unwrap();
         assert!(status.success(), "couldn't get version (this currently needs to run as root)");
         let version = str::from_utf8(&output).unwrap();
         // Strip the trailing newline
@@ -46,6 +47,7 @@ struct Playbot {
     conn: IrcServer,
     rust_versions: Vec<String>,
     shorten_key: String,
+    cache: Arc<Cache>,
 }
 
 impl Playbot {
@@ -77,10 +79,10 @@ impl Playbot {
                 full_code: &str,
                 channel: ReleaseChannel)
                 -> io::Result<String> {
-        let (_status, output) = try!(rust_playpen::exec(channel,
-                                                        "/usr/local/bin/evaluate.sh",
-                                                        Vec::new(),
-                                                        String::from(full_code)));
+        let (_status, output) = try!(self.cache.exec(channel,
+                                                     "/usr/local/bin/evaluate.sh",
+                                                     Vec::new(),
+                                                     String::from(full_code)));
 
         let output_merged = output.splitn(2, |b| *b == b'\xff')
                                   .map(|sub| String::from_utf8_lossy(sub).into_owned())
@@ -206,7 +208,8 @@ fn main() {{
 fn main() {
     env_logger::init().unwrap();
 
-    let rust_versions = get_rust_versions();
+    let cache = Arc::new(Cache::new());
+    let rust_versions = get_rust_versions(&cache);
 
     // FIXME All these unwraps are pretty bad UX, but they should only panic on misconfiguration
     let mut config = String::new();
@@ -235,17 +238,16 @@ fn main() {
                 ..Config::default()
             };
 
-            let rust_versions = rust_versions.clone();
-            let bitly_key = bitly_key.clone();
+            let server = IrcServer::from_config(conf).unwrap();
+            server.identify().unwrap();
+            let mut bot = Playbot {
+                conn: server,
+                rust_versions: rust_versions.clone(),
+                shorten_key: bitly_key.clone(),
+                cache: cache.clone(),
+            };
             threads.push(thread::spawn(move || {
-                let server = IrcServer::from_config(conf).unwrap();
-                server.identify().unwrap();
-
-                Playbot {
-                    conn: server,
-                    rust_versions: rust_versions,
-                    shorten_key: bitly_key,
-                }.main_loop();
+                bot.main_loop();
             }));
         }
     }

src/bin/playpen.rs

@@ -8,25 +8,26 @@ extern crate rustc_serialize;
 extern crate staticfile;
 extern crate unicase;
 
-use rust_playpen::*;
-
-use iron::prelude::*;
-use iron::status;
-use iron::headers;
-use iron::middleware::AfterMiddleware;
-use iron::modifiers::Header;
-use iron::method::Method;
-use hyper::header;
-use staticfile::Static;
-use router::Router;
-use unicase::UniCase;
-use rustc_serialize::json;
-
 use std::env;
 use std::fmt;
 use std::io::Read;
 use std::path::Path;
 use std::process::Command;
+use std::sync::Arc;
+
+use hyper::header;
+use iron::headers;
+use iron::method::Method;
+use iron::middleware::{BeforeMiddleware, AfterMiddleware};
+use iron::modifiers::Header;
+use iron::typemap;
+use iron::prelude::*;
+use iron::status;
+use router::Router;
+use rust_playpen::*;
+use rustc_serialize::json;
+use staticfile::Static;
+use unicase::UniCase;
 
 #[derive(Clone, Debug)]
 struct XXssProtection(bool);
@@ -100,10 +101,11 @@ fn evaluate(req: &mut Request) -> IronResult<Response> {
         args.push(String::from("--test"));
     }
 
-    let (_status, output) = itry!(rust_playpen::exec(version,
-                                                     "/usr/local/bin/evaluate.sh",
-                                                     args,
-                                                     data.code));
+    let cache = req.extensions.get::<AddCache>().unwrap();
+    let (_status, output) = itry!(cache.exec(version,
+                                             "/usr/local/bin/evaluate.sh",
+                                             args,
+                                             data.code));
 
     let mut obj = json::Object::new();
     if separate_output {
@@ -167,8 +169,11 @@ fn compile(req: &mut Request) -> IronResult<Response> {
         args.push(String::from("--color=always"));
     }
 
-    let (_status, output) = itry!(
-        rust_playpen::exec(version, "/usr/local/bin/compile.sh", args, data.code));
+    let cache = req.extensions.get::<AddCache>().unwrap();
+    let (_status, output) = itry!(cache.exec(version,
+                                             "/usr/local/bin/compile.sh",
+                                             args,
+                                             data.code));
     let mut split = output.splitn(2, |b| *b == b'\xff');
     let rustc = String::from_utf8(split.next().unwrap().into()).unwrap();
 
@@ -201,8 +206,11 @@ fn format(req: &mut Request) -> IronResult<Response> {
     let data: FormatReq = itry!(json::decode(&body));
     let version = itry!(data.version.map(|v| v.parse()).unwrap_or(Ok(ReleaseChannel::Stable)));
 
-    let (status, output) = itry!(
-        rust_playpen::exec(version, "rustfmt", Vec::new(), data.code));
+    let cache = req.extensions.get::<AddCache>().unwrap();
+    let (status, output) = itry!(cache.exec(version,
+                                            "rustfmt",
+                                            Vec::new(),
+                                            data.code));
     let output = String::from_utf8(output).unwrap();
     let mut response_obj = json::Object::new();
     if status.success() {
@@ -229,6 +237,19 @@ impl AfterMiddleware for EnablePostCors {
     }
 }
 
+struct AddCache {
+    cache: Arc<Cache>,
+}
+
+impl typemap::Key for AddCache { type Value = Arc<Cache>; }
+
+impl BeforeMiddleware for AddCache {
+    fn before(&self, req: &mut Request) -> IronResult<()> {
+        req.extensions.insert::<AddCache>(self.cache.clone());
+        Ok(())
+    }
+}
+
 fn main() {
     env_logger::init().unwrap();
 
@@ -244,6 +265,7 @@ fn main() {
 
     // Use our router as the middleware, and pass the generated response through `EnablePostCors`
     let mut chain = Chain::new(router);
+    chain.link_before(AddCache { cache: Arc::new(Cache::new()) });
     chain.link_after(EnablePostCors);
 
     let addr = env::args().skip(1).next().unwrap_or("127.0.0.1".to_string());

src/lib.rs

@@ -6,13 +6,13 @@ extern crate wait_timeout;
 
 use lru_cache::LruCache;
 
-use std::cell::RefCell;
 use std::error::Error;
 use std::fmt;
 use std::io::Write;
 use std::io;
 use std::process::{Command, ExitStatus, Stdio};
 use std::str::FromStr;
+use std::sync::Mutex;
 use std::time::Duration;
 
 use docker::Container;
@@ -60,57 +60,63 @@ impl FromStr for ReleaseChannel {
     }
 }
 
-/// Helper method for safely invoking a command inside a playpen
-pub fn exec(channel: ReleaseChannel,
-            cmd: &str,
-            args: Vec<String>,
-            input: String)
-            -> io::Result<(ExitStatus, Vec<u8>)> {
-    #[derive(PartialEq, Eq, Hash)]
-    struct CacheKey {
-        channel: ReleaseChannel,
-        cmd: String,
-        args: Vec<String>,
-        input: String,
-    }
+pub struct Cache {
+    cache: Mutex<LruCache<CacheKey, (ExitStatus, Vec<u8>)>>,
+}
 
-    thread_local! {
-        static CACHE: RefCell<LruCache<CacheKey, (ExitStatus, Vec<u8>)>> =
-            RefCell::new(LruCache::new(256))
-    }
+#[derive(PartialEq, Eq, Hash)]
+struct CacheKey {
+    channel: ReleaseChannel,
+    cmd: String,
+    args: Vec<String>,
+    input: String,
+}
 
-    // Build key to look up
-    let key = CacheKey {
-        channel: channel,
-        cmd: cmd.to_string(),
-        args: args,
-        input: input,
-    };
-    let prev = CACHE.with(|cache| {
-        cache.borrow_mut().get_mut(&key).map(|x| x.clone())
-    });
-    if let Some(prev) = prev {
-        return Ok(prev)
+impl Cache {
+    pub fn new() -> Cache {
+        Cache {
+            cache: Mutex::new(LruCache::new(256)),
+        }
     }
 
-    let chan = match channel {
-        ReleaseChannel::Stable => "stable",
-        ReleaseChannel::Beta => "beta",
-        ReleaseChannel::Nightly => "nightly",
-    };
-    let container = format!("rust-playpen-{}", chan);
+    /// Helper method for safely invoking a command inside a playpen
+    pub fn exec(&self,
+                channel: ReleaseChannel,
+                cmd: &str,
+                args: Vec<String>,
+                input: String)
+                -> io::Result<(ExitStatus, Vec<u8>)> {
+        // Build key to look up
+        let key = CacheKey {
+            channel: channel,
+            cmd: cmd.to_string(),
+            args: args,
+            input: input,
+        };
+        let mut cache = self.cache.lock().unwrap();
+        if let Some(prev) = cache.get_mut(&key) {
+            return Ok(prev.clone())
+        }
+        drop(cache);
 
-    let container = try!(Container::new(cmd, &key.args, &container));
+        let chan = match channel {
+            ReleaseChannel::Stable => "stable",
+            ReleaseChannel::Beta => "beta",
+            ReleaseChannel::Nightly => "nightly",
+        };
+        let container = format!("rust-playpen-{}", chan);
 
-    let tuple = try!(container.run(key.input.as_bytes(), Duration::new(5, 0)));
-    let (status, mut output, timeout) = tuple;
-    if timeout {
-        output.extend_from_slice(b"\ntimeout triggered!");
+        let container = try!(Container::new(cmd, &key.args, &container));
+
+        let tuple = try!(container.run(key.input.as_bytes(), Duration::new(5, 0)));
+        let (status, mut output, timeout) = tuple;
+        if timeout {
+            output.extend_from_slice(b"\ntimeout triggered!");
+        }
+        let mut cache = self.cache.lock().unwrap();
+        cache.insert(key, (status.clone(), output.clone()));
+        Ok((status, output))
     }
-    CACHE.with(|cache| {
-        cache.borrow_mut().insert(key, (status.clone(), output.clone()));
-    });
-    Ok((status, output))
 }
 
 pub enum AsmFlavor {