proc-macro-api: Implement client side timeout handling

crates/load-cargo/src/lib.rs

@@ -116,6 +116,7 @@ pub fn load_workspace_into_db(
                     extra_env,
                     ws.toolchain.as_ref(),
                     load_config.proc_macro_processes,
+                    Some(proc_macro_api::DEFAULT_EXPANSION_TIMEOUT),
                 )
                 .map_err(Into::into)
             })
@@ -127,6 +128,7 @@ pub fn load_workspace_into_db(
                 extra_env,
                 ws.toolchain.as_ref(),
                 load_config.proc_macro_processes,
+                Some(proc_macro_api::DEFAULT_EXPANSION_TIMEOUT),
             )
             .map_err(|e| ProcMacroLoadingError::ProcMacroSrvError(e.to_string().into_boxed_str())),
         ),

crates/proc-macro-api/Cargo.toml

@@ -19,6 +19,7 @@ serde_json = { workspace = true, features = ["unbounded_depth"] }
 tracing.workspace = true
 rustc-hash.workspace = true
 indexmap.workspace = true
+parking_lot = "0.12.4"
 
 # local deps
 paths = { workspace = true, features = ["serde1"] }

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

@@ -100,7 +100,7 @@ pub(crate) fn version_check(
 ) -> Result<u32, ServerError> {
     let request = BidirectionalMessage::Request(Request::ApiVersionCheck(ApiVersionCheck {}));
 
-    let response_payload = run_request(srv, request, callback)?;
+    let response_payload = run_request(srv, request, callback, None)?;
 
     match response_payload {
         BidirectionalMessage::Response(Response::ApiVersionCheck(version)) => Ok(version),
@@ -119,7 +119,7 @@ pub(crate) fn enable_rust_analyzer_spans(
         span_mode: SpanMode::RustAnalyzer,
     }));
 
-    let response_payload = run_request(srv, request, callback)?;
+    let response_payload = run_request(srv, request, callback, None)?;
 
     match response_payload {
         BidirectionalMessage::Response(Response::SetConfig(ServerConfig { span_mode })) => {
@@ -139,7 +139,7 @@ pub(crate) fn find_proc_macros(
         dylib_path: dylib_path.to_path_buf().into(),
     }));
 
-    let response_payload = run_request(srv, request, callback)?;
+    let response_payload = run_request(srv, request, callback, None)?;
 
     match response_payload {
         BidirectionalMessage::Response(Response::ListMacros(it)) => Ok(it),
@@ -182,7 +182,7 @@ pub(crate) fn expand(
         current_dir: Some(current_dir),
     })));
 
-    let response_payload = run_request(process, task, callback)?;
+    let response_payload = run_request(process, task, callback, Some(proc_macro.name()))?;
 
     match response_payload {
         BidirectionalMessage::Response(Response::ExpandMacro(it)) => Ok(it
@@ -202,11 +202,12 @@ fn run_request(
     srv: &ProcMacroServerProcess,
     msg: BidirectionalMessage,
     callback: SubCallback<'_>,
+    macro_name: Option<&str>,
 ) -> Result<BidirectionalMessage, ServerError> {
     if let Some(err) = srv.exited() {
         return Err(err.clone());
     }
-    srv.run_bidirectional(msg, callback)
+    srv.run_bidirectional(msg, callback, macro_name)
 }
 
 pub fn reject_subrequests(req: SubRequest) -> Result<SubResponse, ServerError> {

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

@@ -37,7 +37,7 @@ impl std::fmt::Debug for SpanId {
 
 pub(crate) fn version_check(srv: &ProcMacroServerProcess) -> Result<u32, ServerError> {
     let request = Request::ApiVersionCheck {};
-    let response = send_task(srv, request)?;
+    let response = send_task(srv, request, None)?;
 
     match response {
         Response::ApiVersionCheck(version) => Ok(version),
@@ -50,7 +50,7 @@ pub(crate) fn enable_rust_analyzer_spans(
     srv: &ProcMacroServerProcess,
 ) -> Result<SpanMode, ServerError> {
     let request = Request::SetConfig(ServerConfig { span_mode: SpanMode::RustAnalyzer });
-    let response = send_task(srv, request)?;
+    let response = send_task(srv, request, None)?;
 
     match response {
         Response::SetConfig(ServerConfig { span_mode }) => Ok(span_mode),
@@ -65,7 +65,7 @@ pub(crate) fn find_proc_macros(
 ) -> Result<Result<Vec<(String, ProcMacroKind)>, String>, ServerError> {
     let request = Request::ListMacros { dylib_path: dylib_path.to_path_buf().into() };
 
-    let response = send_task(srv, request)?;
+    let response = send_task(srv, request, None)?;
 
     match response {
         Response::ListMacros(it) => Ok(it),
@@ -112,7 +112,8 @@ pub(crate) fn expand(
         current_dir: Some(current_dir),
     };
 
-    let response = send_task(process, Request::ExpandMacro(Box::new(task)))?;
+    let response =
+        send_task(process, Request::ExpandMacro(Box::new(task)), Some(proc_macro.name()))?;
 
     match response {
         Response::ExpandMacro(it) => Ok(it
@@ -142,12 +143,16 @@ 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: &ProcMacroServerProcess,
+    req: Request,
+    macro_name: Option<&str>,
+) -> Result<Response, ServerError> {
     if let Some(server_error) = srv.exited() {
         return Err(server_error.clone());
     }
 
-    srv.send_task_legacy::<_, _>(send_request, req)
+    srv.send_task_legacy::<_, _>(send_request, req, macro_name)
 }
 
 /// Sends a request to the server and reads the response.

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

@@ -24,12 +24,23 @@ pub mod transport;
 use paths::{AbsPath, AbsPathBuf};
 use semver::Version;
 use span::{ErasedFileAstId, FIXUP_ERASED_FILE_AST_ID_MARKER, Span};
-use std::{fmt, io, sync::Arc, time::SystemTime};
+use std::{
+    ffi::OsString,
+    fmt, io,
+    sync::Arc,
+    time::{Duration, SystemTime},
+};
 
 use crate::{
-    bidirectional_protocol::SubCallback, pool::ProcMacroServerPool, process::ProcMacroServerProcess,
+    bidirectional_protocol::SubCallback,
+    pool::{ProcMacroServerPool, ProcessFactory},
+    process::{ProcMacroServerProcess, ProcessWatchdog},
 };
 
+/// How long a single proc-macro expansion may take before the server process is killed
+/// and restarted.
+pub const DEFAULT_EXPANSION_TIMEOUT: Duration = Duration::from_secs(30);
+
 /// The versions of the server protocol
 pub mod version {
     pub const NO_VERSION_CHECK_VERSION: u32 = 0;
@@ -108,7 +119,7 @@ impl MacroDylib {
 /// we share a single expander process for all macros within a workspace.
 #[derive(Debug, Clone)]
 pub struct ProcMacro {
-    pool: ProcMacroServerPool,
+    pool: Arc<ProcMacroServerPool>,
     dylib_path: Arc<AbsPathBuf>,
     name: Box<str>,
     kind: ProcMacroKind,
@@ -149,19 +160,40 @@ impl ProcMacroClient {
         process_path: &AbsPath,
         env: impl IntoIterator<
             Item = (impl AsRef<std::ffi::OsStr>, &'a Option<impl 'a + AsRef<std::ffi::OsStr>>),
-        > + Clone,
+        >,
         version: Option<&Version>,
         num_process: usize,
+        expansion_timeout: Option<Duration>,
     ) -> io::Result<ProcMacroClient> {
-        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 process_path = process_path.to_owned();
+        let env: Arc<[(OsString, Option<OsString>)]> = env
+            .into_iter()
+            .map(|(key, value)| {
+                (
+                    key.as_ref().to_os_string(),
+                    value.as_ref().map(|value| value.as_ref().to_os_string()),
+                )
+            })
+            .collect();
+        let version = version.cloned();
+        let watchdog = expansion_timeout
+            .map(|timeout| io::Result::Ok((ProcessWatchdog::spawn()?, timeout)))
+            .transpose()?;
+        let spawn: ProcessFactory = Box::new({
+            let process_path = process_path.clone();
+            move || {
+                ProcMacroServerProcess::spawn(
+                    &process_path,
+                    env.iter().map(|(key, value)| (key, value)),
+                    version.as_ref(),
+                    watchdog.clone(),
+                )
+            }
+        });
+        let workers = (0..num_process).map(|_| spawn()).collect::<io::Result<Vec<_>>>()?;
 
-        let pool = ProcMacroServerPool::new(workers);
-        Ok(ProcMacroClient { pool: Arc::new(pool), path: process_path.to_owned() })
+        let pool = Arc::new(ProcMacroServerPool::new(workers, spawn));
+        Ok(ProcMacroClient { pool, path: process_path })
     }
 
     /// Invokes `spawn` and returns a client connected to the resulting read and write handles.
@@ -175,20 +207,30 @@ impl ProcMacroClient {
             Box<dyn process::ProcessExit>,
             Box<dyn io::Write + Send + Sync>,
             Box<dyn io::BufRead + Send + Sync>,
-        )> + Clone,
+        )> + Clone
+        + Send
+        + Sync
+        + 'static,
         version: Option<&Version>,
         num_process: usize,
+        expansion_timeout: Option<Duration>,
     ) -> io::Result<ProcMacroClient> {
-        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 version = version.cloned();
+        let watchdog = expansion_timeout
+            .map(|timeout| io::Result::Ok((ProcessWatchdog::spawn()?, timeout)))
+            .transpose()?;
+        let spawn: ProcessFactory = Box::new(move || {
+            ProcMacroServerProcess::run(
+                spawn.clone(),
+                version.as_ref(),
+                || "<unknown>".to_owned(),
+                watchdog.clone(),
+            )
+        });
+        let workers = (0..num_process).map(|_| spawn()).collect::<io::Result<Vec<_>>>()?;
 
-        let pool = ProcMacroServerPool::new(workers);
-        Ok(ProcMacroClient { pool: Arc::new(pool), path: process_path.to_owned() })
+        let pool = Arc::new(ProcMacroServerPool::new(workers, spawn));
+        Ok(ProcMacroClient { pool, path: process_path.to_owned() })
     }
 
     /// Returns the absolute path to the proc-macro server.
@@ -202,7 +244,7 @@ impl ProcMacroClient {
     }
 
     /// Checks if the proc-macro server has exited.
-    pub fn exited(&self) -> Option<&ServerError> {
+    pub fn exited(&self) -> Option<ServerError> {
         self.pool.exited()
     }
 }
@@ -263,7 +305,8 @@ impl ProcMacro {
             }
         }
 
-        self.pool.pick_process()?.expand(
+        let process = self.pool.pick_process()?;
+        let result = process.expand(
             self,
             subtree,
             attr,
@@ -273,6 +316,10 @@ impl ProcMacro {
             mixed_site,
             current_dir,
             callback,
-        )
+        );
+        if process.timed_out() {
+            self.pool.replace_timed_out_process_in_background(process);
+        }
+        result
     }
 }