Unnamed repository; edit this file 'description' to name the repository.
| -rw-r--r-- | crates/proc-macro-api/src/pool.rs | 50 | ||||
| -rw-r--r-- | crates/proc-macro-api/src/process.rs | 14 |
2 files changed, 33 insertions, 31 deletions
diff --git a/crates/proc-macro-api/src/pool.rs b/crates/proc-macro-api/src/pool.rs index 4639374f3e..4215b0f2c0 100644 --- a/crates/proc-macro-api/src/pool.rs +++ b/crates/proc-macro-api/src/pool.rs @@ -27,12 +27,10 @@ impl ProcMacroServerPool { } fn pick_process(&self) -> &ProcMacroServerProcess { - for workers in &*self.workers { - if workers.can_use() { - return workers; - } - } - &self.workers[0] + self.workers + .iter() + .min_by_key(|w| w.number_of_active_req()) + .expect("worker pool must not be empty") } pub(crate) fn load_dylib( @@ -41,27 +39,31 @@ impl ProcMacroServerPool { _callback: Option<SubCallback<'_>>, ) -> Result<Vec<ProcMacro>, ServerError> { let _p = tracing::info_span!("ProcMacroServer::load_dylib").entered(); - let mut all_macros = Vec::new(); - for worker in &*self.workers { - let dylib_path = Arc::new(dylib.path.clone()); - let dylib_last_modified = std::fs::metadata(dylib_path.as_path()) - .ok() - .and_then(|metadata| metadata.modified().ok()); - let macros = worker.find_proc_macros(&dylib.path, None)?.unwrap(); - - for (name, kind) in macros { - all_macros.push(ProcMacro { - process: self.clone(), - name: name.into(), - kind, - dylib_path: Arc::new(dylib.path.clone()), - dylib_last_modified, - }); - } + let dylib_path = Arc::new(dylib.path.clone()); + let dylib_last_modified = std::fs::metadata(dylib_path.as_path()) + .ok() + .and_then(|metadata| metadata.modified().ok()); + + let first = &self.workers[0]; + let macros = first.find_proc_macros(&dylib.path, None)?.unwrap(); + + for worker in &self.workers[1..] { + let _ = worker.find_proc_macros(&dylib.path, None)?; } - Ok(all_macros) + let result = macros + .into_iter() + .map(|(name, kind)| ProcMacro { + process: self.clone(), + name: name.into(), + kind, + dylib_path: dylib_path.clone(), + dylib_last_modified, + }) + .collect(); + + Ok(result) } pub(crate) fn expand( diff --git a/crates/proc-macro-api/src/process.rs b/crates/proc-macro-api/src/process.rs index a41bb58e74..775d59174f 100644 --- a/crates/proc-macro-api/src/process.rs +++ b/crates/proc-macro-api/src/process.rs @@ -7,7 +7,7 @@ use std::{ process::{Child, ChildStdin, ChildStdout, Command, Stdio}, sync::{ Arc, Mutex, OnceLock, - atomic::{AtomicBool, Ordering}, + atomic::{AtomicU32, Ordering}, }, }; @@ -32,7 +32,7 @@ pub(crate) struct ProcMacroServerProcess { protocol: Protocol, /// Populated when the server exits. exited: OnceLock<AssertUnwindSafe<ServerError>>, - can_use: AtomicBool, + active: AtomicU32, } impl std::fmt::Debug for ProcMacroServerProcess { @@ -163,7 +163,7 @@ impl ProcMacroServerProcess { } }, exited: OnceLock::new(), - can_use: AtomicBool::new(true), + active: AtomicU32::new(0), }) }; let mut srv = create_srv()?; @@ -275,7 +275,7 @@ impl ProcMacroServerProcess { current_dir: String, callback: Option<SubCallback<'_>>, ) -> Result<Result<tt::TopSubtree, String>, ServerError> { - self.can_use.store(false, Ordering::Release); + self.active.fetch_add(1, Ordering::AcqRel); let result = match self.protocol { Protocol::LegacyJson { .. } => legacy_protocol::expand( proc_macro, @@ -302,7 +302,7 @@ impl ProcMacroServerProcess { ), }; - self.can_use.store(true, Ordering::Release); + self.active.fetch_sub(1, Ordering::AcqRel); result } @@ -367,8 +367,8 @@ impl ProcMacroServerProcess { }) } - pub(crate) fn can_use(&self) -> bool { - self.can_use.load(Ordering::Acquire) + pub(crate) fn number_of_active_req(&self) -> u32 { + self.active.load(Ordering::Acquire) } } |