Add a jobserver proxy to ensure at least one token is always held

Author: Zoxc

compiler/rustc_data_structures/src/jobserver.rs

@@ -2,6 +2,7 @@ use std::sync::{LazyLock, OnceLock};
 
 pub use jobserver_crate::{Acquired, Client, HelperThread};
 use jobserver_crate::{FromEnv, FromEnvErrorKind};
+use parking_lot::{Condvar, Mutex};
 
 // We can only call `from_env_ext` once per process
 
@@ -52,8 +53,6 @@ fn default_client() -> Client {
     client
 }
 
-static GLOBAL_CLIENT_CHECKED: OnceLock<Client> = OnceLock::new();
-
 pub fn initialize_checked(report_warning: impl FnOnce(&'static str)) {
     let client_checked = match &*GLOBAL_CLIENT {
         Ok(client) => client.clone(),
@@ -62,19 +61,112 @@ pub fn initialize_checked(report_warning: impl FnOnce(&'static str)) {
             default_client()
         }
     };
-    GLOBAL_CLIENT_CHECKED.set(client_checked).ok();
+    let proxy = Proxy {
+        client: client_checked,
+        data: Mutex::new(ProxyData { total: 1, used: 1, needed: 0 }),
+        wake_needer: Condvar::new(),
+        wake_helper: Condvar::new(),
+    };
+    GLOBAL_PROXY.set(proxy).ok();
+
+    std::thread::spawn(|| {
+        GLOBAL_PROXY.get().unwrap().helper();
+    });
 }
 
 const ACCESS_ERROR: &str = "jobserver check should have been called earlier";
 
 pub fn client() -> Client {
-    GLOBAL_CLIENT_CHECKED.get().expect(ACCESS_ERROR).clone()
+    GLOBAL_PROXY.get().expect(ACCESS_ERROR).client.clone()
 }
 
 pub fn acquire_thread() {
-    GLOBAL_CLIENT_CHECKED.get().expect(ACCESS_ERROR).acquire_raw().ok();
+    GLOBAL_PROXY.get().expect(ACCESS_ERROR).acquire_thread();
 }
 
 pub fn release_thread() {
-    GLOBAL_CLIENT_CHECKED.get().expect(ACCESS_ERROR).release_raw().ok();
+    GLOBAL_PROXY.get().expect(ACCESS_ERROR).release_thread();
+}
+
+static GLOBAL_PROXY: OnceLock<Proxy> = OnceLock::new();
+
+struct ProxyData {
+    /// The number of tokens assigned to this process.
+    total: u16,
+
+    /// The number of tokens assigned to threads.
+    used: u16,
+
+    /// The number of threads requesting a token
+    needed: u16,
+}
+
+/// This is a jobserver proxy used to ensure that we hold on to at least one token.
+struct Proxy {
+    client: Client,
+    data: Mutex<ProxyData>,
+
+    /// Threads which are waiting on a token will wait on this.
+    wake_needer: Condvar,
+
+    /// This is used to wake the helper thread when tokens are needed.
+    wake_helper: Condvar,
+}
+
+impl Proxy {
+    fn helper(&self) {
+        let mut data = self.data.lock();
+        loop {
+            while data.needed > 0 {
+                drop(data);
+                self.client.acquire_raw().ok();
+                data = self.data.lock();
+                if data.needed > 0 {
+                    data.total += 1;
+                    data.used += 1;
+                    data.needed -= 1;
+                    self.wake_needer.notify_one();
+                } else {
+                    drop(data);
+                    self.client.release_raw().ok();
+                    data = self.data.lock();
+                }
+            }
+            self.wake_helper.wait(&mut data);
+        }
+    }
+
+    fn acquire_thread(&self) {
+        let mut data = self.data.lock();
+
+        if data.total > data.used {
+            assert_eq!(data.needed, 0);
+            data.used += 1;
+        } else {
+            if data.needed == 0 {
+                self.wake_helper.notify_one();
+            }
+            data.needed += 1;
+            self.wake_needer.wait(&mut data);
+        }
+    }
+
+    fn release_thread(&self) {
+        let mut data = self.data.lock();
+
+        if data.needed > 0 {
+            // Give the token to a waiting thread
+            data.needed -= 1;
+            self.wake_needer.notify_one();
+        } else {
+            data.used -= 1;
+
+            // Release the token unless it's the last one in the process
+            if data.total > 1 {
+                data.total -= 1;
+                drop(data);
+                self.client.release_raw().ok();
+            }
+        }
+    }
 }