Prepare functions to run on release/acquire of a token

Author: Mark-Simulacrum

src/librustc_jobserver/lib.rs

@@ -1,5 +1,4 @@
-pub use jobserver_crate::Acquired;
-use jobserver_crate::{HelperThread, Client};
+use jobserver_crate::Client;
 use lazy_static::lazy_static;
 
 lazy_static! {
@@ -30,19 +29,61 @@ lazy_static! {
     };
 }
 
+// Unlike releasing tokens, there's not really a "one size fits all" approach, as we have two
+// primary ways of acquiring a token: via the helper thread, and via the acquire_thread function.
+//
+// That makes this function necessary unlike in the release case where everything is piped through
+// `release_thread`.
+fn notify_acquiring_token() {
+    // this function does nothing for now but will be wired up to send a message to Cargo
+}
+
 pub fn initialize() {
     lazy_static::initialize(&GLOBAL_CLIENT)
 }
 
+pub struct HelperThread {
+    helper: jobserver_crate::HelperThread,
+}
+
+impl HelperThread {
+    // This API does not block, but is shimmed so that we can inform Cargo of our interest here.
+    pub fn request_token(&self) {
+        notify_acquiring_token();
+        self.helper.request_token();
+    }
+}
+
+pub struct Acquired(());
+
+impl Drop for Acquired {
+    fn drop(&mut self) {
+        release_thread();
+    }
+}
+
 pub fn helper_thread<F>(mut cb: F) -> HelperThread
     where F: FnMut(Acquired) + Send + 'static,
 {
-    GLOBAL_CLIENT.clone().into_helper_thread(move |token| {
-        cb(token.expect("acquire token"))
-    }).expect("failed to spawn helper thread")
+    let thread = GLOBAL_CLIENT.clone().into_helper_thread(move |token| {
+        // we've acquired a token, but we need to not use it as we have our own
+        // custom release-on-drop struct since we'll want different logic than
+        // just normally releasing the token in this case.
+        //
+        // On unix this unfortunately means that we lose the specific byte that
+        // was in the pipe (i.e., we just write back the same byte all the time)
+        // but that's not expected to be a problem.
+        std::mem::forget(token.expect("acquire token"));
+        cb(Acquired(()))
+    }).expect("failed to spawn helper thread");
+
+    HelperThread {
+        helper: thread,
+    }
 }
 
 pub fn acquire_thread() {
+    notify_acquiring_token();
     GLOBAL_CLIENT.acquire_raw().ok();
 }