rust/src/librustc_data_structures/jobserver.rs

use jobserver_crate::{Client, HelperThread, Acquired};
use lazy_static::lazy_static;
use std::sync::{Condvar, Arc, Mutex};
use std::mem;

#[derive(Default)]
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)]
struct ProxyData {
    lock: Mutex<LockedProxyData>,
    cond_var: Condvar,
}

/// A helper type which makes managing jobserver tokens easier.
/// It also allows you to treat the implicit token given to the process
/// in the same manner as requested tokens.
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 {
    fn release_token(&self) {
        self.data.lock.lock().unwrap().release_token(&self.data.cond_var);
    }

    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;
            }
        }
    }
}