diff --git a/src/tools/rust-analyzer/crates/proc-macro-api/src/bidirectional_protocol.rs b/src/tools/rust-analyzer/crates/proc-macro-api/src/bidirectional_protocol.rs index 643ba98f5176..137f2dafc0de 100644 --- a/src/tools/rust-analyzer/crates/proc-macro-api/src/bidirectional_protocol.rs +++ b/src/tools/rust-analyzer/crates/proc-macro-api/src/bidirectional_protocol.rs @@ -21,7 +21,7 @@ use crate::{ serialize_span_data_index_map, }, }, - process::{ProcMacroWorker, SynIO}, + process::ProcMacroServerProcess, transport::codec::postcard::PostcardProtocol, version, }; @@ -84,7 +84,7 @@ fn wrap_decode(err: io::Error) -> ServerError { } pub(crate) fn version_check( - srv: &dyn ProcMacroWorker, + srv: &ProcMacroServerProcess, callback: SubCallback<'_>, ) -> Result { let request = BidirectionalMessage::Request(Request::ApiVersionCheck {}); @@ -101,7 +101,7 @@ pub(crate) fn version_check( /// Enable support for rust-analyzer span mode if the server supports it. pub(crate) fn enable_rust_analyzer_spans( - srv: &dyn ProcMacroWorker, + srv: &ProcMacroServerProcess, callback: SubCallback<'_>, ) -> Result { let request = BidirectionalMessage::Request(Request::SetConfig(ServerConfig { @@ -120,7 +120,7 @@ pub(crate) fn enable_rust_analyzer_spans( /// Finds proc-macros in a given dynamic library. pub(crate) fn find_proc_macros( - srv: &dyn ProcMacroWorker, + srv: &ProcMacroServerProcess, dylib_path: &AbsPath, callback: SubCallback<'_>, ) -> Result, String>, ServerError> { @@ -205,7 +205,7 @@ pub(crate) fn expand( } fn run_request( - srv: &dyn ProcMacroWorker, + srv: &ProcMacroServerProcess, msg: BidirectionalMessage, callback: SubCallback<'_>, ) -> Result { diff --git a/src/tools/rust-analyzer/crates/proc-macro-api/src/legacy_protocol.rs b/src/tools/rust-analyzer/crates/proc-macro-api/src/legacy_protocol.rs index 56bf863a88e9..7b546cf7aef6 100644 --- a/src/tools/rust-analyzer/crates/proc-macro-api/src/legacy_protocol.rs +++ b/src/tools/rust-analyzer/crates/proc-macro-api/src/legacy_protocol.rs @@ -143,7 +143,7 @@ pub(crate) fn expand( } /// Sends a request to the proc-macro server and waits for a response. -fn send_task(srv: &dyn ProcMacroWorker, req: Request) -> Result { +fn send_task(srv: &ProcMacroServerProcess, req: Request) -> Result { if let Some(server_error) = srv.exited() { return Err(server_error.clone()); } diff --git a/src/tools/rust-analyzer/crates/proc-macro-api/src/lib.rs b/src/tools/rust-analyzer/crates/proc-macro-api/src/lib.rs index dd0c89103a2b..ffae28f92c05 100644 --- a/src/tools/rust-analyzer/crates/proc-macro-api/src/lib.rs +++ b/src/tools/rust-analyzer/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; @@ -29,7 +30,8 @@ use std::{fmt, io, sync::Arc, time::SystemTime}; pub use crate::transport::codec::Codec; use crate::{ bidirectional_protocol::SubCallback, - process::{ProcMacroServerProcess, ProcMacroWorker}, + pool::{ProcMacroServerPool, default_pool_size}, + process::ProcMacroServerProcess, }; /// The versions of the server protocol @@ -88,7 +90,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. - worker: Arc, + pool: Arc, path: AbsPathBuf, } @@ -110,7 +112,7 @@ impl MacroDylib { /// we share a single expander process for all macros within a workspace. #[derive(Debug, Clone)] pub struct ProcMacro { - process: Arc, + process: Arc, dylib_path: Arc, name: Box, kind: ProcMacroKind, @@ -188,31 +190,12 @@ impl ProcMacroClient { dylib: MacroDylib, callback: Option>, ) -> Result, ServerError> { - let _p = tracing::info_span!("ProcMacroServer::load_dylib").entered(); - let macros = self.worker.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.worker.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.worker.exited() + self.pool.exited() } } diff --git a/src/tools/rust-analyzer/crates/proc-macro-api/src/pool.rs b/src/tools/rust-analyzer/crates/proc-macro-api/src/pool.rs new file mode 100644 index 000000000000..685bc05be62a --- /dev/null +++ b/src/tools/rust-analyzer/crates/proc-macro-api/src/pool.rs @@ -0,0 +1,61 @@ +use std::sync::Arc; + +use crate::{ + MacroDylib, ProcMacro, ServerError, bidirectional_protocol::SubCallback, + process::ProcMacroServerProcess, +}; + +#[derive(Debug)] +pub(crate) struct ProcMacroServerPool { + workers: Vec>, +} + +impl ProcMacroServerPool { + pub(crate) fn new(workers: Vec>) -> Self { + Self { workers } + } +} + +impl ProcMacroServerPool { + pub(crate) fn exited(&self) -> Option<&ServerError> { + for worker in &self.workers { + if let Some(e) = worker.exited() { + return Some(e); + } + } + None + } + + pub(crate) fn load_dylib( + &self, + dylib: &MacroDylib, + _callback: Option>, + ) -> Result, 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.load_dylib(&dylib.path, None)?; + + for (name, kind) in macros { + all_macros.push(ProcMacro { + process: worker.clone(), + name: name.into(), + kind, + dylib_path: Arc::new(dylib.path.clone()), + dylib_last_modified, + }); + } + } + + Ok(all_macros) + } +} + +pub(crate) fn default_pool_size() -> usize { + std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1).min(4) +} diff --git a/src/tools/rust-analyzer/crates/proc-macro-api/src/process.rs b/src/tools/rust-analyzer/crates/proc-macro-api/src/process.rs index efb8e0e84a1f..30877c5cf491 100644 --- a/src/tools/rust-analyzer/crates/proc-macro-api/src/process.rs +++ b/src/tools/rust-analyzer/crates/proc-macro-api/src/process.rs @@ -3,7 +3,7 @@ use std::{ fmt::Debug, io::{self, BufRead, BufReader, Read, Write}, - panic::{AssertUnwindSafe, RefUnwindSafe}, + panic::AssertUnwindSafe, process::{Child, ChildStdin, ChildStdout, Command, Stdio}, sync::{Arc, Mutex, OnceLock}, }; @@ -82,77 +82,6 @@ pub(crate) struct ProcessSrvState { stdout: Box, } -impl ProcMacroWorker for ProcMacroServerProcess { - fn find_proc_macros( - &self, - dylib_path: &AbsPath, - callback: Option>, - ) -> Result, String>, ServerError> { - ProcMacroServerProcess::find_proc_macros(self, dylib_path, callback) - } - - fn expand( - &self, - proc_macro: &ProcMacro, - subtree: tt::SubtreeView<'_>, - attr: Option>, - env: Vec<(String, String)>, - def_site: Span, - call_site: Span, - mixed_site: Span, - current_dir: String, - callback: Option>, - ) -> Result, ServerError> { - ProcMacroServerProcess::expand( - self, - proc_macro, - subtree, - attr, - env, - def_site, - call_site, - mixed_site, - current_dir, - callback, - ) - } - - fn exited(&self) -> Option<&ServerError> { - ProcMacroServerProcess::exited(self) - } - - fn version(&self) -> u32 { - ProcMacroServerProcess::version(self) - } - - fn rust_analyzer_spans(&self) -> bool { - ProcMacroServerProcess::rust_analyzer_spans(self) - } - - fn enable_rust_analyzer_spans( - &self, - callback: Option>, - ) -> Result { - ProcMacroServerProcess::enable_rust_analyzer_spans(self, callback) - } - - fn use_postcard(&self) -> bool { - ProcMacroServerProcess::use_postcard(self) - } - - fn state(&self) -> &Mutex { - &self.state - } - - fn get_exited(&self) -> &OnceLock> { - &self.exited - } - - fn is_reusable(&self) -> bool { - !self.single_use - } -} - impl ProcMacroServerProcess { /// Starts the proc-macro server and performs a version check pub(crate) fn spawn<'a>( @@ -220,7 +149,11 @@ impl ProcMacroServerProcess { let (process, stdin, stdout) = spawn(format)?; io::Result::Ok(ProcMacroServerProcess { - state: Mutex::new(ProcessSrvState { process, stdin, stdout }), + state: Mutex::new(ProcessSrvState { + process, + stdin, + stdout, + }), version: 0, protocol: match format { Some(ProtocolFormat::BidirectionalPostcardPrototype) => { @@ -271,6 +204,37 @@ impl ProcMacroServerProcess { Err(err.unwrap()) } + pub(crate) fn load_dylib( + &self, + dylib_path: &AbsPath, + callback: Option>, + ) -> Result, ServerError> { + let _state = self.state.lock().unwrap(); + + // if state.loaded_dylibs.contains(dylib_path) { + // // Already loaded in this worker + // return Ok(Vec::new()); + // } + + let result = match self.protocol { + Protocol::LegacyJson { .. } | Protocol::LegacyPostcard { .. } => { + legacy_protocol::find_proc_macros(self, dylib_path)? + } + Protocol::BidirectionalPostcardPrototype { .. } => { + let cb = callback.expect("callback required"); + bidirectional_protocol::find_proc_macros(self, dylib_path, cb)? + } + }; + + match result { + Ok(macros) => { + // state.loaded_dylibs.insert(dylib_path.to_owned()); + Ok(macros) + } + Err(message) => Err(ServerError { message, io: None }), + } + } + /// Returns the server error if the process has exited. pub(crate) fn exited(&self) -> Option<&ServerError> { self.exited.get().map(|it| &it.0) @@ -314,21 +278,6 @@ impl ProcMacroServerProcess { } } - /// Finds proc-macros in a given dynamic library. - pub(crate) fn find_proc_macros( - &self, - dylib_path: &AbsPath, - callback: Option>, - ) -> Result, String>, ServerError> { - match self.protocol { - Protocol::LegacyJson { .. } => legacy_protocol::find_proc_macros(self, dylib_path), - Protocol::BidirectionalPostcardPrototype { .. } => { - let cb = callback.expect("callback required for bidirectional protocol"); - bidirectional_protocol::find_proc_macros(self, dylib_path, cb) - } - } - } - pub(crate) fn expand( &self, proc_macro: &ProcMacro, @@ -365,25 +314,25 @@ impl ProcMacroServerProcess { ), }; - if self.is_reusable() { + if !self.is_reusable() { self.terminate(); } result } + fn is_reusable(&self) -> bool { + self.single_use + } + fn terminate(&self) { if let Ok(mut state) = self.state.lock() { let _ = state.process.child.kill(); } } -} -pub(crate) struct SynIO; - -impl SynIO { pub(crate) fn send_task( - proc_macro_worker: &dyn ProcMacroWorker, + &self, send: impl FnOnce( &mut dyn Write, &mut dyn BufRead, @@ -392,7 +341,7 @@ impl SynIO { ) -> Result, ServerError>, req: Request, ) -> Result { - SynIO::with_locked_io::(proc_macro_worker, |writer, reader, buf| { + self.with_locked_io::(|writer, reader, buf| { send(writer, reader, req, buf).and_then(|res| { res.ok_or_else(|| { let message = "proc-macro server did not respond with data".to_owned(); @@ -409,10 +358,10 @@ impl SynIO { } pub(crate) fn with_locked_io( - proc_macro_worker: &dyn ProcMacroWorker, + &self, f: impl FnOnce(&mut dyn Write, &mut dyn BufRead, &mut C::Buf) -> Result, ) -> Result { - let state = &mut *proc_macro_worker.state().lock().unwrap(); + let state = &mut *self.state.lock().unwrap(); let mut buf = C::Buf::default(); f(&mut state.stdin, &mut state.stdout, &mut buf).map_err(|e| { @@ -434,11 +383,11 @@ impl SynIO { } pub(crate) fn run_bidirectional( - proc_macro_worker: &dyn ProcMacroWorker, + &self, initial: BidirectionalMessage, callback: SubCallback<'_>, ) -> Result { - SynIO::with_locked_io::(proc_macro_worker, |writer, reader, buf| { + self.with_locked_io::(|writer, reader, buf| { bidirectional_protocol::run_conversation::(writer, reader, buf, initial, callback) }) }