|
| 1 | +use jobserver_crate::{Client, HelperThread, Acquired}; |
| 2 | +use lazy_static::lazy_static; |
| 3 | +use std::sync::{Condvar, Arc, Mutex}; |
| 4 | +use std::mem; |
| 5 | + |
| 6 | +#[derive(Default)] |
| 7 | +pub struct LockedProxyData { |
| 8 | + /// The number of free thread tokens, this may include the implicit token given to the process |
| 9 | + free: usize, |
| 10 | + |
| 11 | + /// The number of threads waiting for a token |
| 12 | + waiters: usize, |
| 13 | + |
| 14 | + /// The number of tokens we requested from the server |
| 15 | + requested: usize, |
| 16 | + |
| 17 | + /// Stored tokens which will be dropped when we no longer need them |
| 18 | + tokens: Vec<Acquired>, |
| 19 | +} |
| 20 | + |
| 21 | +impl LockedProxyData { |
| 22 | + fn request_token(&mut self, thread: &Mutex<HelperThread>) { |
| 23 | + self.requested += 1; |
| 24 | + thread.lock().unwrap().request_token(); |
| 25 | + } |
| 26 | + |
| 27 | + fn release_token(&mut self, cond_var: &Condvar) { |
| 28 | + if self.waiters > 0 { |
| 29 | + self.free += 1; |
| 30 | + cond_var.notify_one(); |
| 31 | + } else { |
| 32 | + if self.tokens.is_empty() { |
| 33 | + // We are returning the implicit token |
| 34 | + self.free += 1; |
| 35 | + } else { |
| 36 | + // Return a real token to the server |
| 37 | + self.tokens.pop().unwrap(); |
| 38 | + } |
| 39 | + } |
| 40 | + } |
| 41 | + |
| 42 | + fn take_token(&mut self, thread: &Mutex<HelperThread>) -> bool { |
| 43 | + if self.free > 0 { |
| 44 | + self.free -= 1; |
| 45 | + self.waiters -= 1; |
| 46 | + |
| 47 | + // We stole some token reqested by someone else |
| 48 | + // Request another one |
| 49 | + if self.requested + self.free < self.waiters { |
| 50 | + self.request_token(thread); |
| 51 | + } |
| 52 | + |
| 53 | + true |
| 54 | + } else { |
| 55 | + false |
| 56 | + } |
| 57 | + } |
| 58 | + |
| 59 | + fn new_requested_token(&mut self, token: Acquired, cond_var: &Condvar) { |
| 60 | + self.requested -= 1; |
| 61 | + |
| 62 | + // Does anything need this token? |
| 63 | + if self.waiters > 0 { |
| 64 | + self.free += 1; |
| 65 | + self.tokens.push(token); |
| 66 | + cond_var.notify_one(); |
| 67 | + } else { |
| 68 | + // Otherwise we'll just drop it |
| 69 | + mem::drop(token); |
| 70 | + } |
| 71 | + } |
| 72 | +} |
| 73 | + |
| 74 | +#[derive(Default)] |
| 75 | +pub struct ProxyData { |
| 76 | + lock: Mutex<LockedProxyData>, |
| 77 | + cond_var: Condvar, |
| 78 | +} |
| 79 | + |
| 80 | +pub struct Proxy { |
| 81 | + thread: Mutex<HelperThread>, |
| 82 | + data: Arc<ProxyData>, |
| 83 | +} |
| 84 | + |
| 85 | +lazy_static! { |
| 86 | + // We can only call `from_env` once per process |
| 87 | + |
| 88 | + // Note that this is unsafe because it may misinterpret file descriptors |
| 89 | + // on Unix as jobserver file descriptors. We hopefully execute this near |
| 90 | + // the beginning of the process though to ensure we don't get false |
| 91 | + // positives, or in other words we try to execute this before we open |
| 92 | + // any file descriptors ourselves. |
| 93 | + // |
| 94 | + // Pick a "reasonable maximum" if we don't otherwise have |
| 95 | + // a jobserver in our environment, capping out at 32 so we |
| 96 | + // don't take everything down by hogging the process run queue. |
| 97 | + // The fixed number is used to have deterministic compilation |
| 98 | + // across machines. |
| 99 | + // |
| 100 | + // Also note that we stick this in a global because there could be |
| 101 | + // multiple rustc instances in this process, and the jobserver is |
| 102 | + // per-process. |
| 103 | + static ref GLOBAL_CLIENT: Client = unsafe { |
| 104 | + Client::from_env().unwrap_or_else(|| { |
| 105 | + Client::new(32).expect("failed to create jobserver") |
| 106 | + }) |
| 107 | + }; |
| 108 | + |
| 109 | + static ref GLOBAL_PROXY: Proxy = { |
| 110 | + let data = Arc::new(ProxyData::default()); |
| 111 | + |
| 112 | + Proxy { |
| 113 | + data: data.clone(), |
| 114 | + thread: Mutex::new(client().into_helper_thread(move |token| { |
| 115 | + data.lock.lock().unwrap().new_requested_token(token.unwrap(), &data.cond_var); |
| 116 | + }).unwrap()), |
| 117 | + } |
| 118 | + }; |
| 119 | +} |
| 120 | + |
| 121 | +pub fn client() -> Client { |
| 122 | + GLOBAL_CLIENT.clone() |
| 123 | +} |
| 124 | + |
| 125 | +pub fn acquire_thread() { |
| 126 | + GLOBAL_PROXY.acquire_token(); |
| 127 | +} |
| 128 | + |
| 129 | +pub fn release_thread() { |
| 130 | + GLOBAL_PROXY.release_token(); |
| 131 | +} |
| 132 | + |
| 133 | +impl Proxy { |
| 134 | + pub fn release_token(&self) { |
| 135 | + self.data.lock.lock().unwrap().release_token(&self.data.cond_var); |
| 136 | + } |
| 137 | + |
| 138 | + pub fn acquire_token(&self) { |
| 139 | + let mut data = self.data.lock.lock().unwrap(); |
| 140 | + data.waiters += 1; |
| 141 | + if data.take_token(&self.thread) { |
| 142 | + return; |
| 143 | + } |
| 144 | + // Request a token for us |
| 145 | + data.request_token(&self.thread); |
| 146 | + loop { |
| 147 | + data = self.data.cond_var.wait(data).unwrap(); |
| 148 | + if data.take_token(&self.thread) { |
| 149 | + return; |
| 150 | + } |
| 151 | + } |
| 152 | + } |
| 153 | +} |
0 commit comments