add worker abstraction

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

@@ -21,7 +21,7 @@ use crate::{
             serialize_span_data_index_map,
         },
     },
-    process::ProcMacroServerProcess,
+    process::{ProcMacroWorker, SynIO},
     transport::codec::postcard::PostcardProtocol,
     version,
 };
@@ -84,7 +84,7 @@ fn wrap_decode(err: io::Error) -> ServerError {
 }
 
 pub(crate) fn version_check(
-    srv: &ProcMacroServerProcess,
+    srv: &dyn ProcMacroWorker,
     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: &ProcMacroServerProcess,
+    srv: &dyn ProcMacroWorker,
     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: &ProcMacroServerProcess,
+    srv: &dyn ProcMacroWorker,
     dylib_path: &AbsPath,
     callback: SubCallback<'_>,
 ) -> Result<Result<Vec<(String, ProcMacroKind)>, String>, ServerError> {
@@ -175,7 +175,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(proc_macro.process.as_ref(), task, callback)?;
 
     match response_payload {
         BidirectionalMessage::Response(Response::ExpandMacro(it)) => Ok(it
@@ -205,7 +205,7 @@ pub(crate) fn expand(
 }
 
 fn run_request(
-    srv: &ProcMacroServerProcess,
+    srv: &dyn ProcMacroWorker,
     msg: BidirectionalMessage,
     callback: SubCallback<'_>,
 ) -> Result<BidirectionalMessage, ServerError> {

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

@@ -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(proc_macro.process.as_ref(), Request::ExpandMacro(Box::new(task)))?;
 
     match response {
         Response::ExpandMacro(it) => Ok(it
@@ -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: &ProcMacroServerProcess, req: Request) -> Result<Response, ServerError> {
+fn send_task(srv: &dyn ProcMacroWorker, req: Request) -> Result<Response, ServerError> {
     if let Some(server_error) = srv.exited() {
         return Err(server_error.clone());
     }

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

@@ -27,7 +27,10 @@ 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,
+    process::{ProcMacroServerProcess, ProcMacroWorker},
+};
 
 /// The versions of the server protocol
 pub mod version {
@@ -85,7 +88,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>,
+    worker: Arc<dyn ProcMacroWorker>,
     path: AbsPathBuf,
 }
 
@@ -107,7 +110,7 @@ impl MacroDylib {
 /// we share a single expander process for all macros within a workspace.
 #[derive(Debug, Clone)]
 pub struct ProcMacro {
-    process: Arc<ProcMacroServerProcess>,
+    process: Arc<dyn ProcMacroWorker>,
     dylib_path: Arc<AbsPathBuf>,
     name: Box<str>,
     kind: ProcMacroKind,
@@ -171,7 +174,7 @@ impl ProcMacroClient {
         version: Option<&Version>,
     ) -> io::Result<ProcMacroClient> {
         let process = ProcMacroServerProcess::run(spawn, version, || "<unknown>".to_owned())?;
-        Ok(ProcMacroClient { process: Arc::new(process), path: process_path.to_owned() })
+        Ok(ProcMacroClient { worker: Arc::new(process), path: process_path.to_owned() })
     }
 
     /// Returns the absolute path to the proc-macro server.
@@ -186,7 +189,7 @@ impl ProcMacroClient {
         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 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())
@@ -196,7 +199,7 @@ impl ProcMacroClient {
             Ok(macros) => Ok(macros
                 .into_iter()
                 .map(|(name, kind)| ProcMacro {
-                    process: self.process.clone(),
+                    process: self.worker.clone(),
                     name: name.into(),
                     kind,
                     dylib_path: dylib_path.clone(),
@@ -209,7 +212,7 @@ impl ProcMacroClient {
 
     /// Checks if the proc-macro server has exited.
     pub fn exited(&self) -> Option<&ServerError> {
-        self.process.exited()
+        self.worker.exited()
     }
 }
 

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

@@ -1,8 +1,9 @@
 //! Handle process life-time and message passing for proc-macro client
 
 use std::{
+    fmt::Debug,
     io::{self, BufRead, BufReader, Read, Write},
-    panic::AssertUnwindSafe,
+    panic::{AssertUnwindSafe, RefUnwindSafe},
     process::{Child, ChildStdin, ChildStdout, Command, Stdio},
     sync::{Arc, Mutex, OnceLock},
 };
@@ -74,12 +75,79 @@ impl ProcessExit for Process {
 }
 
 /// Maintains the state of the proc-macro server process.
-struct ProcessSrvState {
+pub(crate) struct ProcessSrvState {
     process: Box<dyn ProcessExit>,
     stdin: Box<dyn Write + Send + Sync>,
     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
+    }
+}
+
 impl ProcMacroServerProcess {
     /// Starts the proc-macro server and performs a version check
     pub(crate) fn spawn<'a>(
@@ -291,9 +359,13 @@ impl ProcMacroServerProcess {
             ),
         }
     }
+}
 
+pub(crate) struct SynIO;
+
+impl SynIO {
     pub(crate) fn send_task<Request, Response, C: Codec>(
-        &self,
+        proc_macro_worker: &dyn ProcMacroWorker,
         send: impl FnOnce(
             &mut dyn Write,
             &mut dyn BufRead,
@@ -302,7 +374,7 @@ impl ProcMacroServerProcess {
         ) -> Result<Option<Response>, ServerError>,
         req: Request,
     ) -> Result<Response, ServerError> {
-        self.with_locked_io::<C, _>(|writer, reader, buf| {
+        SynIO::with_locked_io::<C, _>(proc_macro_worker, |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();
@@ -319,10 +391,10 @@ impl ProcMacroServerProcess {
     }
 
     pub(crate) fn with_locked_io<C: Codec, R>(
-        &self,
+        proc_macro_worker: &dyn ProcMacroWorker,
         f: impl FnOnce(&mut dyn Write, &mut dyn BufRead, &mut C::Buf) -> Result<R, ServerError>,
     ) -> Result<R, ServerError> {
-        let state = &mut *self.state.lock().unwrap();
+        let state = &mut *proc_macro_worker.state().lock().unwrap();
         let mut buf = C::Buf::default();
 
         f(&mut state.stdin, &mut state.stdout, &mut buf).map_err(|e| {
@@ -330,7 +402,11 @@ impl ProcMacroServerProcess {
                 match state.process.exit_err() {
                     None => e,
                     Some(server_error) => {
-                        self.exited.get_or_init(|| AssertUnwindSafe(server_error)).0.clone()
+                        proc_macro_worker
+                            .get_exited()
+                            .get_or_init(|| AssertUnwindSafe(server_error))
+                            .0
+                            .clone()
                     }
                 }
             } else {
@@ -340,11 +416,11 @@ impl ProcMacroServerProcess {
     }
 
     pub(crate) fn run_bidirectional<C: Codec>(
-        &self,
+        proc_macro_worker: &dyn ProcMacroWorker,
         initial: BidirectionalMessage,
         callback: SubCallback<'_>,
     ) -> Result<BidirectionalMessage, ServerError> {
-        self.with_locked_io::<C, _>(|writer, reader, buf| {
+        SynIO::with_locked_io::<C, _>(proc_macro_worker, |writer, reader, buf| {
             bidirectional_protocol::run_conversation::<C>(writer, reader, buf, initial, callback)
         })
     }