add pool of processes

This commit is contained in:
bit-aloo 2026-01-01 16:05:33 +05:30
parent 98d4496a1e
commit 9fb5d34626
No known key found for this signature in database
5 changed files with 122 additions and 129 deletions

View file

@ -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<u32, ServerError> {
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<SpanMode, ServerError> {
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<Result<Vec<(String, ProcMacroKind)>, String>, ServerError> {
@ -205,7 +205,7 @@ pub(crate) fn expand(
}
fn run_request(
srv: &dyn ProcMacroWorker,
srv: &ProcMacroServerProcess,
msg: BidirectionalMessage,
callback: SubCallback<'_>,
) -> Result<BidirectionalMessage, ServerError> {

View file

@ -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<Response, ServerError> {
fn send_task(srv: &ProcMacroServerProcess, req: Request) -> Result<Response, ServerError> {
if let Some(server_error) = srv.exited() {
return Err(server_error.clone());
}

View file

@ -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<dyn ProcMacroWorker>,
pool: Arc<ProcMacroServerPool>,
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<dyn ProcMacroWorker>,
process: Arc<ProcMacroServerProcess>,
dylib_path: Arc<AbsPathBuf>,
name: Box<str>,
kind: ProcMacroKind,
@ -188,31 +190,12 @@ impl ProcMacroClient {
dylib: MacroDylib,
callback: Option<SubCallback<'_>>,
) -> Result<Vec<ProcMacro>, 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()
}
}

View file

@ -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<Arc<ProcMacroServerProcess>>,
}
impl ProcMacroServerPool {
pub(crate) fn new(workers: Vec<Arc<ProcMacroServerProcess>>) -> 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<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.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)
}

View file

@ -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<dyn BufRead + Send + Sync>,
}
impl ProcMacroWorker for ProcMacroServerProcess {
fn find_proc_macros(
&self,
dylib_path: &AbsPath,
callback: Option<SubCallback<'_>>,
) -> Result<Result<Vec<(String, ProcMacroKind)>, String>, ServerError> {
ProcMacroServerProcess::find_proc_macros(self, dylib_path, callback)
}
fn expand(
&self,
proc_macro: &ProcMacro,
subtree: tt::SubtreeView<'_>,
attr: Option<tt::SubtreeView<'_>>,
env: Vec<(String, String)>,
def_site: Span,
call_site: Span,
mixed_site: Span,
current_dir: String,
callback: Option<SubCallback<'_>>,
) -> Result<Result<tt::TopSubtree, String>, 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<SubCallback<'_>>,
) -> Result<SpanMode, ServerError> {
ProcMacroServerProcess::enable_rust_analyzer_spans(self, callback)
}
fn use_postcard(&self) -> bool {
ProcMacroServerProcess::use_postcard(self)
}
fn state(&self) -> &Mutex<ProcessSrvState> {
&self.state
}
fn get_exited(&self) -> &OnceLock<AssertUnwindSafe<ServerError>> {
&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<SubCallback<'_>>,
) -> Result<Vec<(String, ProcMacroKind)>, 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<SubCallback<'_>>,
) -> Result<Result<Vec<(String, ProcMacroKind)>, 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<Request, Response, C: Codec>(
proc_macro_worker: &dyn ProcMacroWorker,
&self,
send: impl FnOnce(
&mut dyn Write,
&mut dyn BufRead,
@ -392,7 +341,7 @@ impl SynIO {
) -> Result<Option<Response>, ServerError>,
req: Request,
) -> Result<Response, ServerError> {
SynIO::with_locked_io::<C, _>(proc_macro_worker, |writer, reader, buf| {
self.with_locked_io::<C, _>(|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<C: Codec, R>(
proc_macro_worker: &dyn ProcMacroWorker,
&self,
f: impl FnOnce(&mut dyn Write, &mut dyn BufRead, &mut C::Buf) -> Result<R, ServerError>,
) -> Result<R, ServerError> {
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<C: Codec>(
proc_macro_worker: &dyn ProcMacroWorker,
&self,
initial: BidirectionalMessage,
callback: SubCallback<'_>,
) -> Result<BidirectionalMessage, ServerError> {
SynIO::with_locked_io::<C, _>(proc_macro_worker, |writer, reader, buf| {
self.with_locked_io::<C, _>(|writer, reader, buf| {
bidirectional_protocol::run_conversation::<C>(writer, reader, buf, initial, callback)
})
}