add request id to requests

Author: Shourya742

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

@@ -30,7 +30,7 @@ use crate::{
 pub mod msg;
 
 pub trait ClientCallbacks {
-    fn handle_sub_request(&mut self, id: u64, req: SubRequest) -> Result<SubResponse, ServerError>;
+    fn handle_sub_request(&mut self, req: SubRequest) -> Result<SubResponse, ServerError>;
 }
 
 pub fn run_conversation<C: Codec>(
@@ -65,7 +65,7 @@ pub fn run_conversation<C: Codec>(
 
         match (msg.kind, msg.payload) {
             (Kind::SubRequest, Payload::SubRequest(sr)) => {
-                let resp = callbacks.handle_sub_request(id, sr)?;
+                let resp = callbacks.handle_sub_request(sr)?;
                 let reply =
                     Envelope { id, kind: Kind::SubResponse, payload: Payload::SubResponse(resp) };
                 let encoded = C::encode(&reply).map_err(wrap_encode)?;
@@ -104,19 +104,15 @@ pub(crate) fn version_check(srv: &ProcMacroServerProcess) -> Result<u32, ServerE
 
     struct NoCallbacks;
     impl ClientCallbacks for NoCallbacks {
-        fn handle_sub_request(
-            &mut self,
-            _id: u64,
-            _req: SubRequest,
-        ) -> Result<SubResponse, ServerError> {
+        fn handle_sub_request(&mut self, _req: SubRequest) -> Result<SubResponse, ServerError> {
             Err(ServerError { message: "sub-request not supported here".into(), io: None })
         }
     }
 
     let mut callbacks = NoCallbacks;
 
     let response_payload =
-        run_bidirectional(srv, (0, Kind::Request, request).into(), &mut callbacks)?;
+        run_bidirectional(srv, (srv.request_id(), Kind::Request, request).into(), &mut callbacks)?;
 
     match response_payload {
         Payload::Response(Response::ApiVersionCheck(version)) => Ok(version),
@@ -135,19 +131,15 @@ pub(crate) fn enable_rust_analyzer_spans(
 
     struct NoCallbacks;
     impl ClientCallbacks for NoCallbacks {
-        fn handle_sub_request(
-            &mut self,
-            _id: u64,
-            _req: SubRequest,
-        ) -> Result<SubResponse, ServerError> {
+        fn handle_sub_request(&mut self, _req: SubRequest) -> Result<SubResponse, ServerError> {
             Err(ServerError { message: "sub-request not supported here".into(), io: None })
         }
     }
 
     let mut callbacks = NoCallbacks;
 
     let response_payload =
-        run_bidirectional(srv, (0, Kind::Request, request).into(), &mut callbacks)?;
+        run_bidirectional(srv, (srv.request_id(), Kind::Request, request).into(), &mut callbacks)?;
 
     match response_payload {
         Payload::Response(Response::SetConfig(ServerConfig { span_mode })) => Ok(span_mode),
@@ -165,19 +157,15 @@ pub(crate) fn find_proc_macros(
 
     struct NoCallbacks;
     impl ClientCallbacks for NoCallbacks {
-        fn handle_sub_request(
-            &mut self,
-            _id: u64,
-            _req: SubRequest,
-        ) -> Result<SubResponse, ServerError> {
+        fn handle_sub_request(&mut self, _req: SubRequest) -> Result<SubResponse, ServerError> {
             Err(ServerError { message: "sub-request not supported here".into(), io: None })
         }
     }
 
     let mut callbacks = NoCallbacks;
 
     let response_payload =
-        run_bidirectional(srv, (0, Kind::Request, request).into(), &mut callbacks)?;
+        run_bidirectional(srv, (srv.request_id(), Kind::Request, request).into(), &mut callbacks)?;
 
     match response_payload {
         Payload::Response(Response::ListMacros(it)) => Ok(it),
@@ -225,15 +213,13 @@ pub(crate) fn expand(
         current_dir: Some(current_dir),
     })));
 
+    let request_id = proc_macro.process.request_id();
+
     struct Callbacks<'de> {
         db: &'de dyn SourceDatabase,
     }
     impl<'db> ClientCallbacks for Callbacks<'db> {
-        fn handle_sub_request(
-            &mut self,
-            _id: u64,
-            req: SubRequest,
-        ) -> Result<SubResponse, ServerError> {
+        fn handle_sub_request(&mut self, req: SubRequest) -> Result<SubResponse, ServerError> {
             match req {
                 SubRequest::SourceText { file_id, start, end } => {
                     let file = FileId::from_raw(file_id);
@@ -249,8 +235,11 @@ pub(crate) fn expand(
 
     let mut callbacks = Callbacks { db };
 
-    let response_payload =
-        run_bidirectional(&proc_macro.process, (0, Kind::Request, task).into(), &mut callbacks)?;
+    let response_payload = run_bidirectional(
+        &proc_macro.process,
+        (request_id, Kind::Request, task).into(),
+        &mut callbacks,
+    )?;
 
     match response_payload {
         Payload::Response(Response::ExpandMacro(it)) => Ok(it

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

@@ -4,7 +4,10 @@ use std::{
     io::{self, BufRead, BufReader, Read, Write},
     panic::AssertUnwindSafe,
     process::{Child, ChildStdin, ChildStdout, Command, Stdio},
-    sync::{Arc, Mutex, OnceLock},
+    sync::{
+        Arc, Mutex, OnceLock,
+        atomic::{AtomicU64, Ordering},
+    },
 };
 
 use base_db::SourceDatabase;
@@ -33,6 +36,7 @@ pub(crate) struct ProcMacroServerProcess {
     protocol: Protocol,
     /// Populated when the server exits.
     exited: OnceLock<AssertUnwindSafe<ServerError>>,
+    next_request_id: AtomicU64,
 }
 
 #[derive(Debug, Clone)]
@@ -90,6 +94,7 @@ impl ProcMacroServerProcess {
                     version: 0,
                     protocol: protocol.clone(),
                     exited: OnceLock::new(),
+                    next_request_id: AtomicU64::new(1),
                 })
             };
             let mut srv = create_srv()?;
@@ -312,6 +317,10 @@ impl ProcMacroServerProcess {
             )
         })
     }
+
+    pub(crate) fn request_id(&self) -> RequestId {
+        self.next_request_id.fetch_add(1, Ordering::Relaxed)
+    }
 }
 
 /// Manages the execution of the proc-macro server process.

crates/proc-macro-srv/src/server_impl/rust_analyzer_span.rs

@@ -163,12 +163,9 @@ impl server::Span for RaSpanServer {
                 start,
                 end,
             });
-            self.cli_to_server
-                .clone()
-                .unwrap()
-                .recv()
-                .and_then(|SubResponse::SourceTextResult { text }| Ok(text))
-                .expect("REASON")
+            match self.cli_to_server.as_ref()?.recv().ok()? {
+                SubResponse::SourceTextResult { text } => text,
+            }
         } else {
             None
         }