Parallelize proc macro expansion

crates/load-cargo/src/lib.rs

@@ -45,6 +45,7 @@ pub struct LoadCargoConfig {
     pub load_out_dirs_from_check: bool,
     pub with_proc_macro_server: ProcMacroServerChoice,
     pub prefill_caches: bool,
+    pub proc_macro_processes: usize,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -113,15 +114,25 @@ pub fn load_workspace_into_db(
     let proc_macro_server = match &load_config.with_proc_macro_server {
         ProcMacroServerChoice::Sysroot => ws.find_sysroot_proc_macro_srv().map(|it| {
             it.and_then(|it| {
-                ProcMacroClient::spawn(&it, extra_env, ws.toolchain.as_ref()).map_err(Into::into)
+                ProcMacroClient::spawn(
+                    &it,
+                    extra_env,
+                    ws.toolchain.as_ref(),
+                    load_config.proc_macro_processes,
+                )
+                .map_err(Into::into)
             })
             .map_err(|e| ProcMacroLoadingError::ProcMacroSrvError(e.to_string().into_boxed_str()))
         }),
-        ProcMacroServerChoice::Explicit(path) => {
-            Some(ProcMacroClient::spawn(path, extra_env, ws.toolchain.as_ref()).map_err(|e| {
-                ProcMacroLoadingError::ProcMacroSrvError(e.to_string().into_boxed_str())
-            }))
-        }
+        ProcMacroServerChoice::Explicit(path) => Some(
+            ProcMacroClient::spawn(
+                path,
+                extra_env,
+                ws.toolchain.as_ref(),
+                load_config.proc_macro_processes,
+            )
+            .map_err(|e| ProcMacroLoadingError::ProcMacroSrvError(e.to_string().into_boxed_str())),
+        ),
         ProcMacroServerChoice::None => Some(Err(ProcMacroLoadingError::Disabled)),
     };
     match &proc_macro_server {
@@ -435,7 +446,7 @@ pub fn load_proc_macro(
 ) -> ProcMacroLoadResult {
     let res: Result<Vec<_>, _> = (|| {
         let dylib = MacroDylib::new(path.to_path_buf());
-        let vec = server.load_dylib(dylib, Some(&mut reject_subrequests)).map_err(|e| {
+        let vec = server.load_dylib(dylib, Some(&reject_subrequests)).map_err(|e| {
             ProcMacroLoadingError::ProcMacroSrvError(format!("{e}").into_boxed_str())
         })?;
         if vec.is_empty() {
@@ -541,7 +552,7 @@ impl ProcMacroExpander for Expander {
         mixed_site: Span,
         current_dir: String,
     ) -> Result<tt::TopSubtree, ProcMacroExpansionError> {
-        let mut cb = |req| match req {
+        let cb = |req| match req {
             SubRequest::LocalFilePath { file_id } => {
                 let file_id = FileId::from_raw(file_id);
                 let source_root_id = db.file_source_root(file_id).source_root_id(db);
@@ -600,7 +611,7 @@ impl ProcMacroExpander for Expander {
             call_site,
             mixed_site,
             current_dir,
-            Some(&mut cb),
+            Some(&cb),
         ) {
             Ok(Ok(subtree)) => Ok(subtree),
             Ok(Err(err)) => Err(ProcMacroExpansionError::Panic(err)),
@@ -644,6 +655,7 @@ mod tests {
             load_out_dirs_from_check: false,
             with_proc_macro_server: ProcMacroServerChoice::None,
             prefill_caches: false,
+            proc_macro_processes: 1,
         };
         let (db, _vfs, _proc_macro) =
             load_workspace_at(path, &cargo_config, &load_cargo_config, &|_| {}).unwrap();

crates/proc-macro-api/src/bidirectional_protocol.rs

@@ -28,7 +28,7 @@ use crate::{
 
 pub mod msg;
 
-pub type SubCallback<'a> = &'a mut dyn FnMut(SubRequest) -> Result<SubResponse, ServerError>;
+pub type SubCallback<'a> = &'a dyn Fn(SubRequest) -> Result<SubResponse, ServerError>;
 
 pub fn run_conversation<C: Codec>(
     writer: &mut dyn Write,
@@ -138,6 +138,7 @@ pub(crate) fn find_proc_macros(
 
 pub(crate) fn expand(
     proc_macro: &ProcMacro,
+    process: &ProcMacroServerProcess,
     subtree: tt::SubtreeView<'_>,
     attr: Option<tt::SubtreeView<'_>>,
     env: Vec<(String, String)>,
@@ -147,7 +148,7 @@ pub(crate) fn expand(
     current_dir: String,
     callback: SubCallback<'_>,
 ) -> Result<Result<tt::TopSubtree, String>, crate::ServerError> {
-    let version = proc_macro.process.version();
+    let version = process.version();
     let mut span_data_table = SpanDataIndexMap::default();
     let def_site = span_data_table.insert_full(def_site).0;
     let call_site = span_data_table.insert_full(call_site).0;
@@ -164,7 +165,7 @@ pub(crate) fn expand(
                 call_site,
                 mixed_site,
             },
-            span_data_table: if proc_macro.process.rust_analyzer_spans() {
+            span_data_table: if process.rust_analyzer_spans() {
                 serialize_span_data_index_map(&span_data_table)
             } else {
                 Vec::new()
@@ -175,7 +176,7 @@ pub(crate) fn expand(
         current_dir: Some(current_dir),
     })));
 
-    let response_payload = run_request(&proc_macro.process, task, callback)?;
+    let response_payload = run_request(process, task, callback)?;
 
     match response_payload {
         BidirectionalMessage::Response(Response::ExpandMacro(it)) => Ok(it

crates/proc-macro-api/src/legacy_protocol.rs

@@ -18,8 +18,7 @@ use crate::{
         flat::serialize_span_data_index_map,
     },
     process::ProcMacroServerProcess,
-    transport::codec::Codec,
-    transport::codec::{json::JsonProtocol, postcard::PostcardProtocol},
+    transport::codec::{Codec, json::JsonProtocol, postcard::PostcardProtocol},
     version,
 };
 
@@ -77,6 +76,7 @@ pub(crate) fn find_proc_macros(
 
 pub(crate) fn expand(
     proc_macro: &ProcMacro,
+    process: &ProcMacroServerProcess,
     subtree: tt::SubtreeView<'_>,
     attr: Option<tt::SubtreeView<'_>>,
     env: Vec<(String, String)>,
@@ -85,7 +85,7 @@ pub(crate) fn expand(
     mixed_site: Span,
     current_dir: String,
 ) -> Result<Result<tt::TopSubtree, String>, crate::ServerError> {
-    let version = proc_macro.process.version();
+    let version = process.version();
     let mut span_data_table = SpanDataIndexMap::default();
     let def_site = span_data_table.insert_full(def_site).0;
     let call_site = span_data_table.insert_full(call_site).0;
@@ -102,7 +102,7 @@ pub(crate) fn expand(
                 call_site,
                 mixed_site,
             },
-            span_data_table: if proc_macro.process.rust_analyzer_spans() {
+            span_data_table: if process.rust_analyzer_spans() {
                 serialize_span_data_index_map(&span_data_table)
             } else {
                 Vec::new()
@@ -113,7 +113,7 @@ pub(crate) fn expand(
         current_dir: Some(current_dir),
     };
 
-    let response = send_task(&proc_macro.process, Request::ExpandMacro(Box::new(task)))?;
+    let response = send_task(process, Request::ExpandMacro(Box::new(task)))?;
 
     match response {
         Response::ExpandMacro(it) => Ok(it

crates/proc-macro-api/src/lib.rs

@@ -18,6 +18,7 @@ extern crate rustc_driver as _;
 
 pub mod bidirectional_protocol;
 pub mod legacy_protocol;
+pub mod pool;
 pub mod process;
 pub mod transport;
 
@@ -27,7 +28,9 @@ use span::{ErasedFileAstId, FIXUP_ERASED_FILE_AST_ID_MARKER, Span};
 use std::{fmt, io, sync::Arc, time::SystemTime};
 
 pub use crate::transport::codec::Codec;
-use crate::{bidirectional_protocol::SubCallback, process::ProcMacroServerProcess};
+use crate::{
+    bidirectional_protocol::SubCallback, pool::ProcMacroServerPool, process::ProcMacroServerProcess,
+};
 
 /// The versions of the server protocol
 pub mod version {
@@ -84,7 +87,7 @@ pub struct ProcMacroClient {
     ///
     /// That means that concurrent salsa requests may block each other when expanding proc macros,
     /// which is unfortunate, but simple and good enough for the time being.
-    process: Arc<ProcMacroServerProcess>,
+    pool: Arc<ProcMacroServerPool>,
     path: AbsPathBuf,
 }
 
@@ -106,7 +109,7 @@ impl MacroDylib {
 /// we share a single expander process for all macros within a workspace.
 #[derive(Debug, Clone)]
 pub struct ProcMacro {
-    process: Arc<ProcMacroServerProcess>,
+    pool: ProcMacroServerPool,
     dylib_path: Arc<AbsPathBuf>,
     name: Box<str>,
     kind: ProcMacroKind,
@@ -120,7 +123,6 @@ impl PartialEq for ProcMacro {
             && self.kind == other.kind
             && self.dylib_path == other.dylib_path
             && self.dylib_last_modified == other.dylib_last_modified
-            && Arc::ptr_eq(&self.process, &other.process)
     }
 }
 
@@ -150,9 +152,17 @@ impl ProcMacroClient {
             Item = (impl AsRef<std::ffi::OsStr>, &'a Option<impl 'a + AsRef<std::ffi::OsStr>>),
         > + Clone,
         version: Option<&Version>,
+        num_process: usize,
     ) -> io::Result<ProcMacroClient> {
-        let process = ProcMacroServerProcess::spawn(process_path, env, version)?;
-        Ok(ProcMacroClient { process: Arc::new(process), path: process_path.to_owned() })
+        let pool_size = num_process;
+        let mut workers = Vec::with_capacity(pool_size);
+        for _ in 0..pool_size {
+            let worker = ProcMacroServerProcess::spawn(process_path, env.clone(), version)?;
+            workers.push(worker);
+        }
+
+        let pool = ProcMacroServerPool::new(workers);
+        Ok(ProcMacroClient { pool: Arc::new(pool), path: process_path.to_owned() })
     }
 
     /// Invokes `spawn` and returns a client connected to the resulting read and write handles.
@@ -166,11 +176,20 @@ impl ProcMacroClient {
             Box<dyn process::ProcessExit>,
             Box<dyn io::Write + Send + Sync>,
             Box<dyn io::BufRead + Send + Sync>,
-        )>,
+        )> + Clone,
         version: Option<&Version>,
+        num_process: usize,
     ) -> io::Result<ProcMacroClient> {
-        let process = ProcMacroServerProcess::run(spawn, version, || "<unknown>".to_owned())?;
-        Ok(ProcMacroClient { process: Arc::new(process), path: process_path.to_owned() })
+        let pool_size = num_process;
+        let mut workers = Vec::with_capacity(pool_size);
+        for _ in 0..pool_size {
+            let worker =
+                ProcMacroServerProcess::run(spawn.clone(), version, || "<unknown>".to_owned())?;
+            workers.push(worker);
+        }
+
+        let pool = ProcMacroServerPool::new(workers);
+        Ok(ProcMacroClient { pool: Arc::new(pool), path: process_path.to_owned() })
     }
 
     /// Returns the absolute path to the proc-macro server.
@@ -184,31 +203,12 @@ impl ProcMacroClient {
         dylib: MacroDylib,
         callback: Option<SubCallback<'_>>,
     ) -> Result<Vec<ProcMacro>, ServerError> {
-        let _p = tracing::info_span!("ProcMacroServer::load_dylib").entered();
-        let macros = self.process.find_proc_macros(&dylib.path, callback)?;
-
-        let dylib_path = Arc::new(dylib.path);
-        let dylib_last_modified = std::fs::metadata(dylib_path.as_path())
-            .ok()
-            .and_then(|metadata| metadata.modified().ok());
-        match macros {
-            Ok(macros) => Ok(macros
-                .into_iter()
-                .map(|(name, kind)| ProcMacro {
-                    process: self.process.clone(),
-                    name: name.into(),
-                    kind,
-                    dylib_path: dylib_path.clone(),
-                    dylib_last_modified,
-                })
-                .collect()),
-            Err(message) => Err(ServerError { message, io: None }),
-        }
+        self.pool.load_dylib(&dylib, callback)
     }
 
     /// Checks if the proc-macro server has exited.
     pub fn exited(&self) -> Option<&ServerError> {
-        self.process.exited()
+        self.pool.exited()
     }
 }
 
@@ -224,7 +224,7 @@ impl ProcMacro {
     }
 
     fn needs_fixup_change(&self) -> bool {
-        let version = self.process.version();
+        let version = self.pool.version();
         (version::RUST_ANALYZER_SPAN_SUPPORT..version::HASHED_AST_ID).contains(&version)
     }
 
@@ -268,7 +268,7 @@ impl ProcMacro {
             }
         }
 
-        self.process.expand(
+        self.pool.pick_process()?.expand(
             self,
             subtree,
             attr,

crates/proc-macro-api/src/pool.rs

@@ -0,0 +1,91 @@
+//! A pool of proc-macro server processes
+use std::sync::Arc;
+
+use crate::{
+    MacroDylib, ProcMacro, ServerError, bidirectional_protocol::SubCallback,
+    process::ProcMacroServerProcess,
+};
+
+#[derive(Debug, Clone)]
+pub(crate) struct ProcMacroServerPool {
+    workers: Arc<[ProcMacroServerProcess]>,
+    version: u32,
+}
+
+impl ProcMacroServerPool {
+    pub(crate) fn new(workers: Vec<ProcMacroServerProcess>) -> Self {
+        let version = workers[0].version();
+        Self { workers: workers.into(), version }
+    }
+}
+
+impl ProcMacroServerPool {
+    pub(crate) fn exited(&self) -> Option<&ServerError> {
+        for worker in &*self.workers {
+            worker.exited()?;
+        }
+        self.workers[0].exited()
+    }
+
+    pub(crate) fn pick_process(&self) -> Result<&ProcMacroServerProcess, ServerError> {
+        let mut best: Option<&ProcMacroServerProcess> = None;
+        let mut best_load = u32::MAX;
+
+        for w in self.workers.iter().filter(|w| w.exited().is_none()) {
+            let load = w.number_of_active_req();
+
+            if load == 0 {
+                return Ok(w);
+            }
+
+            if load < best_load {
+                best = Some(w);
+                best_load = load;
+            }
+        }
+
+        best.ok_or_else(|| ServerError {
+            message: "all proc-macro server workers have exited".into(),
+            io: None,
+        })
+    }
+
+    pub(crate) fn load_dylib(
+        &self,
+        dylib: &MacroDylib,
+        callback: Option<SubCallback<'_>>,
+    ) -> Result<Vec<ProcMacro>, ServerError> {
+        let _span = tracing::info_span!("ProcMacroServer::load_dylib").entered();
+
+        let dylib_path = Arc::new(dylib.path.clone());
+        let dylib_last_modified =
+            std::fs::metadata(dylib_path.as_path()).ok().and_then(|m| m.modified().ok());
+
+        let (first, rest) = self.workers.split_first().expect("worker pool must not be empty");
+
+        let macros = first
+            .find_proc_macros(&dylib.path, callback)?
+            .map_err(|e| ServerError { message: e, io: None })?;
+
+        for worker in rest {
+            worker
+                .find_proc_macros(&dylib.path, callback)?
+                .map_err(|e| ServerError { message: e, io: None })?;
+        }
+
+        Ok(macros
+            .into_iter()
+            .map(|(name, kind)| ProcMacro {
+                pool: self.clone(),
+                name: name.into(),
+                kind,
+                dylib_path: dylib_path.clone(),
+                dylib_last_modified,
+            })
+            .collect())
+    }
+
+    pub(crate) fn version(&self) -> u32 {
+        self.version
+    }
+}