cc/parallel/
job_token.rs

Help
1use std::marker::PhantomData;
2
3use crate::{utilities::OnceLock, Error};
4
5pub(crate) struct JobToken(PhantomData<()>);
6
7impl JobToken {
8    fn new() -> Self {
9        Self(PhantomData)
10    }
11}
12
13impl Drop for JobToken {
14    fn drop(&mut self) {
15        match JobTokenServer::new() {
16            JobTokenServer::Inherited(jobserver) => jobserver.release_token_raw(),
17            JobTokenServer::InProcess(jobserver) => jobserver.release_token_raw(),
18        }
19    }
20}
21
22enum JobTokenServer {
23    Inherited(inherited_jobserver::JobServer),
24    InProcess(inprocess_jobserver::JobServer),
25}
26
27impl JobTokenServer {
28    /// This function returns a static reference to the jobserver because
29    ///  - creating a jobserver from env is a bit fd-unsafe (e.g. the fd might
30    ///    be closed by other jobserver users in the process) and better do it
31    ///    at the start of the program.
32    ///  - in case a jobserver cannot be created from env (e.g. it's not
33    ///    present), we will create a global in-process only jobserver
34    ///    that has to be static so that it will be shared by all cc
35    ///    compilation.
36    fn new() -> &'static Self {
37        // TODO: Replace with a OnceLock once MSRV is 1.70
38        static JOBSERVER: OnceLock<JobTokenServer> = OnceLock::new();
39
40        JOBSERVER.get_or_init(|| {
41            unsafe { inherited_jobserver::JobServer::from_env() }
42                .map(Self::Inherited)
43                .unwrap_or_else(|| Self::InProcess(inprocess_jobserver::JobServer::new()))
44        })
45    }
46}
47
48pub(crate) enum ActiveJobTokenServer {
49    Inherited(inherited_jobserver::ActiveJobServer<'static>),
50    InProcess(&'static inprocess_jobserver::JobServer),
51}
52
53impl ActiveJobTokenServer {
54    pub(crate) fn new() -> Self {
55        match JobTokenServer::new() {
56            JobTokenServer::Inherited(inherited_jobserver) => {
57                Self::Inherited(inherited_jobserver.enter_active())
58            }
59            JobTokenServer::InProcess(inprocess_jobserver) => Self::InProcess(inprocess_jobserver),
60        }
61    }
62
63    pub(crate) async fn acquire(&mut self) -> Result<JobToken, Error> {
64        match self {
65            Self::Inherited(jobserver) => jobserver.acquire().await,
66            Self::InProcess(jobserver) => Ok(jobserver.acquire().await),
67        }
68    }
69}
70
71mod inherited_jobserver {
72    use super::JobToken;
73
74    use crate::{parallel::async_executor::YieldOnce, Error, ErrorKind};
75
76    use std::{
77        io, mem,
78        sync::{mpsc, Mutex, MutexGuard, PoisonError},
79    };
80
81    pub(super) struct JobServer {
82        /// Implicit token for this process which is obtained and will be
83        /// released in parent. Since JobTokens only give back what they got,
84        /// there should be at most one global implicit token in the wild.
85        ///
86        /// Since Rust does not execute any `Drop` for global variables,
87        /// we can't just put it back to jobserver and then re-acquire it at
88        /// the end of the process.
89        ///
90        /// Use `Mutex` to avoid race between acquire and release.
91        /// If an `AtomicBool` is used, then it's possible for:
92        ///  - `release_token_raw`: Tries to set `global_implicit_token` to true, but it is already
93        ///    set  to `true`, continue to release it to jobserver
94        ///  - `acquire` takes the global implicit token, set `global_implicit_token` to false
95        ///  - `release_token_raw` now writes the token back into the jobserver, while
96        ///    `global_implicit_token` is `false`
97        ///
98        /// If the program exits here, then cc effectively increases parallelism by one, which is
99        /// incorrect, hence we use a `Mutex` here.
100        global_implicit_token: Mutex<bool>,
101        inner: jobserver::Client,
102    }
103
104    impl JobServer {
105        pub(super) unsafe fn from_env() -> Option<Self> {
106            jobserver::Client::from_env().map(|inner| Self {
107                inner,
108                global_implicit_token: Mutex::new(true),
109            })
110        }
111
112        fn get_global_implicit_token(&self) -> MutexGuard<'_, bool> {
113            self.global_implicit_token
114                .lock()
115                .unwrap_or_else(PoisonError::into_inner)
116        }
117
118        /// All tokens except for the global implicit token will be put back into the jobserver
119        /// immediately and they cannot be cached, since Rust does not call `Drop::drop` on
120        /// global variables.
121        pub(super) fn release_token_raw(&self) {
122            let mut global_implicit_token = self.get_global_implicit_token();
123
124            if *global_implicit_token {
125                // There's already a global implicit token, so this token must
126                // be released back into jobserver.
127                //
128                // `release_raw` should not block
129                let _ = self.inner.release_raw();
130            } else {
131                *global_implicit_token = true;
132            }
133        }
134
135        pub(super) fn enter_active(&self) -> ActiveJobServer<'_> {
136            ActiveJobServer {
137                jobserver: self,
138                helper_thread: None,
139            }
140        }
141    }
142
143    struct HelperThread {
144        inner: jobserver::HelperThread,
145        /// When rx is dropped, all the token stored within it will be dropped.
146        rx: mpsc::Receiver<io::Result<jobserver::Acquired>>,
147    }
148
149    impl HelperThread {
150        fn new(jobserver: &JobServer) -> Result<Self, Error> {
151            let (tx, rx) = mpsc::channel();
152
153            Ok(Self {
154                rx,
155                inner: jobserver.inner.clone().into_helper_thread(move |res| {
156                    let _ = tx.send(res);
157                })?,
158            })
159        }
160    }
161
162    pub(crate) struct ActiveJobServer<'a> {
163        jobserver: &'a JobServer,
164        helper_thread: Option<HelperThread>,
165    }
166
167    impl<'a> ActiveJobServer<'a> {
168        pub(super) async fn acquire(&mut self) -> Result<JobToken, Error> {
169            let mut has_requested_token = false;
170
171            loop {
172                // Fast path
173                if mem::replace(&mut *self.jobserver.get_global_implicit_token(), false) {
174                    break Ok(JobToken::new());
175                }
176
177                match self.jobserver.inner.try_acquire() {
178                    Ok(Some(acquired)) => {
179                        acquired.drop_without_releasing();
180                        break Ok(JobToken::new());
181                    }
182                    Ok(None) => YieldOnce::default().await,
183                    Err(err) if err.kind() == io::ErrorKind::Unsupported => {
184                        // Fallback to creating a help thread with blocking acquire
185                        let helper_thread = if let Some(thread) = self.helper_thread.as_ref() {
186                            thread
187                        } else {
188                            self.helper_thread
189                                .insert(HelperThread::new(self.jobserver)?)
190                        };
191
192                        match helper_thread.rx.try_recv() {
193                            Ok(res) => {
194                                let acquired = res?;
195                                acquired.drop_without_releasing();
196                                break Ok(JobToken::new());
197                            }
198                            Err(mpsc::TryRecvError::Disconnected) => break Err(Error::new(
199                                ErrorKind::JobserverHelpThreadError,
200                                "jobserver help thread has returned before ActiveJobServer is dropped",
201                            )),
202                            Err(mpsc::TryRecvError::Empty) => {
203                                if !has_requested_token {
204                                    helper_thread.inner.request_token();
205                                    has_requested_token = true;
206                                }
207                                YieldOnce::default().await
208                            }
209                        }
210                    }
211                    Err(err) => break Err(err.into()),
212                }
213            }
214        }
215    }
216}
217
218mod inprocess_jobserver {
219    use super::JobToken;
220
221    use crate::parallel::async_executor::YieldOnce;
222
223    use std::{
224        env::var,
225        sync::atomic::{
226            AtomicU32,
227            Ordering::{AcqRel, Acquire},
228        },
229    };
230
231    pub(crate) struct JobServer(AtomicU32);
232
233    impl JobServer {
234        pub(super) fn new() -> Self {
235            // Use `NUM_JOBS` if set (it's configured by Cargo) and otherwise
236            // just fall back to a semi-reasonable number.
237            //
238            // Note that we could use `num_cpus` here but it's an extra
239            // dependency that will almost never be used, so
240            // it's generally not too worth it.
241            let mut parallelism = 4;
242            // TODO: Use std::thread::available_parallelism as an upper bound
243            // when MSRV is bumped.
244            if let Ok(amt) = var("NUM_JOBS") {
245                if let Ok(amt) = amt.parse() {
246                    parallelism = amt;
247                }
248            }
249
250            Self(AtomicU32::new(parallelism))
251        }
252
253        pub(super) async fn acquire(&self) -> JobToken {
254            loop {
255                let res = self
256                    .0
257                    .fetch_update(AcqRel, Acquire, |tokens| tokens.checked_sub(1));
258
259                if res.is_ok() {
260                    break JobToken::new();
261                }
262
263                YieldOnce::default().await
264            }
265        }
266
267        pub(super) fn release_token_raw(&self) {
268            self.0.fetch_add(1, AcqRel);
269        }
270    }
271}
cc/parallel/job_token.rs

cc/parallel/
job_token.rs
