Add support for using a jobserver with Rayon

3 files changed, 158 insertions, 2 deletions

diff --git a/src/librustc_data_structures/Cargo.toml b/src/librustc_data_structures/Cargo.toml
index 6aa262715ec..6002bf69b70 100644
--- a/src/librustc_data_structures/Cargo.toml
+++ b/src/librustc_data_structures/Cargo.toml
@@ -12,13 +12,15 @@ crate-type = ["dylib"]
 [dependencies]
 ena = "0.11"
 log = "0.4"
+jobserver_crate = { version = "0.1", package = "jobserver" }
+lazy_static = "1"
 rustc_cratesio_shim = { path = "../librustc_cratesio_shim" }
 serialize = { path = "../libserialize" }
 graphviz = { path = "../libgraphviz" }
 cfg-if = "0.1.2"
 stable_deref_trait = "1.0.0"
-rayon = { version = "0.1.1", package = "rustc-rayon" }
-rayon-core = { version = "0.1.1", package = "rustc-rayon-core" }
+rayon = { version = "0.1.2", package = "rustc-rayon" }
+rayon-core = { version = "0.1.2", package = "rustc-rayon-core" }
 rustc-hash = "1.0.1"
 smallvec = { version = "0.6.7", features = ["union", "may_dangle"] }
 
diff --git a/src/librustc_data_structures/jobserver.rs b/src/librustc_data_structures/jobserver.rs
new file mode 100644
index 00000000000..c85cdfbdcd8
--- /dev/null
+++ b/src/librustc_data_structures/jobserver.rs
@@ -0,0 +1,153 @@
+use jobserver_crate::{Client, HelperThread, Acquired};
+use lazy_static::lazy_static;
+use std::sync::{Condvar, Arc, Mutex};
+use std::mem;
+
+#[derive(Default)]
+pub struct LockedProxyData {
+    /// The number of free thread tokens, this may include the implicit token given to the process
+    free: usize,
+
+    /// The number of threads waiting for a token
+    waiters: usize,
+
+    /// The number of tokens we requested from the server
+    requested: usize,
+
+    /// Stored tokens which will be dropped when we no longer need them
+    tokens: Vec<Acquired>,
+}
+
+impl LockedProxyData {
+    fn request_token(&mut self, thread: &Mutex<HelperThread>) {
+        self.requested += 1;
+        thread.lock().unwrap().request_token();
+    }
+
+    fn release_token(&mut self, cond_var: &Condvar) {
+        if self.waiters > 0 {
+            self.free += 1;
+            cond_var.notify_one();
+        } else {
+            if self.tokens.is_empty() {
+                // We are returning the implicit token
+                self.free += 1;
+            } else {
+                // Return a real token to the server
+                self.tokens.pop().unwrap();
+            }
+        }
+    }
+
+    fn take_token(&mut self, thread: &Mutex<HelperThread>) -> bool {
+        if self.free > 0 {
+            self.free -= 1;
+            self.waiters -= 1;
+
+            // We stole some token reqested by someone else
+            // Request another one
+            if self.requested + self.free < self.waiters {
+                self.request_token(thread);
+            }
+
+            true
+        } else {
+            false
+        }
+    }
+
+    fn new_requested_token(&mut self, token: Acquired, cond_var: &Condvar) {
+        self.requested -= 1;
+
+        // Does anything need this token?
+        if self.waiters > 0 {
+            self.free += 1;
+            self.tokens.push(token);
+            cond_var.notify_one();
+        } else {
+            // Otherwise we'll just drop it
+            mem::drop(token);
+        }
+    }
+}
+
+#[derive(Default)]
+pub struct ProxyData {
+    lock: Mutex<LockedProxyData>,
+    cond_var: Condvar,
+}
+
+pub struct Proxy {
+    thread: Mutex<HelperThread>,
+    data: Arc<ProxyData>,
+}
+
+lazy_static! {
+    // We can only call `from_env` once per process
+
+    // Note that this is unsafe because it may misinterpret file descriptors
+    // on Unix as jobserver file descriptors. We hopefully execute this near
+    // the beginning of the process though to ensure we don't get false
+    // positives, or in other words we try to execute this before we open
+    // any file descriptors ourselves.
+    //
+    // Pick a "reasonable maximum" if we don't otherwise have
+    // a jobserver in our environment, capping out at 32 so we
+    // don't take everything down by hogging the process run queue.
+    // The fixed number is used to have deterministic compilation
+    // across machines.
+    //
+    // Also note that we stick this in a global because there could be
+    // multiple rustc instances in this process, and the jobserver is
+    // per-process.
+    static ref GLOBAL_CLIENT: Client = unsafe {
+        Client::from_env().unwrap_or_else(|| {
+            Client::new(32).expect("failed to create jobserver")
+        })
+    };
+
+    static ref GLOBAL_PROXY: Proxy = {
+        let data = Arc::new(ProxyData::default());
+
+        Proxy {
+            data: data.clone(),
+            thread: Mutex::new(client().into_helper_thread(move |token| {
+                data.lock.lock().unwrap().new_requested_token(token.unwrap(), &data.cond_var);
+            }).unwrap()),
+        }
+    };
+}
+
+pub fn client() -> Client {
+    GLOBAL_CLIENT.clone()
+}
+
+pub fn acquire_thread() {
+    GLOBAL_PROXY.acquire_token();
+}
+
+pub fn release_thread() {
+    GLOBAL_PROXY.release_token();
+}
+
+impl Proxy {
+    pub fn release_token(&self) {
+        self.data.lock.lock().unwrap().release_token(&self.data.cond_var);
+    }
+
+    pub fn acquire_token(&self) {
+        let mut data = self.data.lock.lock().unwrap();
+        data.waiters += 1;
+        if data.take_token(&self.thread) {
+            return;
+        }
+        // Request a token for us
+        data.request_token(&self.thread);
+        loop {
+            data = self.data.cond_var.wait(data).unwrap();
+            if data.take_token(&self.thread) {
+                return;
+            }
+        }
+    }
+}
diff --git a/src/librustc_data_structures/lib.rs b/src/librustc_data_structures/lib.rs
index 2bfb1b24a81..09482340b1a 100644
--- a/src/librustc_data_structures/lib.rs
+++ b/src/librustc_data_structures/lib.rs
@@ -77,6 +77,7 @@ pub mod fx;
 pub mod graph;
 pub mod indexed_vec;
 pub mod interner;
+pub mod jobserver;
 pub mod obligation_forest;
 pub mod owning_ref;
 pub mod ptr_key;