Add a field to MultiLSPQuery span showing the current request (#35372)

Mikayla Maki created 3 months ago

Release Notes:

- N/A

Change summary

crates/collab/src/rpc.rs  | 14 ++++++++++++--
crates/proto/src/proto.rs | 19 +++++++++++++++++++
2 files changed, 31 insertions(+), 2 deletions(-)

Detailed changes

crates/collab/src/rpc.rs

@@ -42,7 +42,7 @@ use collections::{HashMap, HashSet};
 pub use connection_pool::{ConnectionPool, ZedVersion};
 use core::fmt::{self, Debug, Formatter};
 use reqwest_client::ReqwestClient;
-use rpc::proto::split_repository_update;
+use rpc::proto::{MultiLspQuery, split_repository_update};
 use supermaven_api::{CreateExternalUserRequest, SupermavenAdminApi};
 
 use futures::{
@@ -374,7 +374,7 @@ impl Server {
             .add_request_handler(forward_mutating_project_request::<proto::OnTypeFormatting>)
             .add_request_handler(forward_mutating_project_request::<proto::SaveBuffer>)
             .add_request_handler(forward_mutating_project_request::<proto::BlameBuffer>)
-            .add_request_handler(forward_mutating_project_request::<proto::MultiLspQuery>)
+            .add_request_handler(multi_lsp_query)
             .add_request_handler(forward_mutating_project_request::<proto::RestartLanguageServers>)
             .add_request_handler(forward_mutating_project_request::<proto::StopLanguageServers>)
             .add_request_handler(forward_mutating_project_request::<proto::LinkedEditingRange>)
@@ -865,6 +865,7 @@ impl Server {
                                 user_id=field::Empty,
                                 login=field::Empty,
                                 impersonator=field::Empty,
+                                multi_lsp_query_request=field::Empty,
                             );
                             principal.update_span(&span);
                             let span_enter = span.enter();
@@ -2329,6 +2330,15 @@ where
     Ok(())
 }
 
+async fn multi_lsp_query(
+    request: MultiLspQuery,
+    response: Response<MultiLspQuery>,
+    session: Session,
+) -> Result<()> {
+    tracing::Span::current().record("multi_lsp_query_request", request.request_str());
+    forward_mutating_project_request(request, response, session).await
+}
+
 /// Notify other participants that a new buffer has been created
 async fn create_buffer_for_peer(
     request: proto::CreateBufferForPeer,

crates/proto/src/proto.rs

@@ -784,6 +784,25 @@ pub fn split_repository_update(
     }])
 }
 
+impl MultiLspQuery {
+    pub fn request_str(&self) -> &str {
+        match self.request {
+            Some(multi_lsp_query::Request::GetHover(_)) => "GetHover",
+            Some(multi_lsp_query::Request::GetCodeActions(_)) => "GetCodeActions",
+            Some(multi_lsp_query::Request::GetSignatureHelp(_)) => "GetSignatureHelp",
+            Some(multi_lsp_query::Request::GetCodeLens(_)) => "GetCodeLens",
+            Some(multi_lsp_query::Request::GetDocumentDiagnostics(_)) => "GetDocumentDiagnostics",
+            Some(multi_lsp_query::Request::GetDocumentColor(_)) => "GetDocumentColor",
+            Some(multi_lsp_query::Request::GetDefinition(_)) => "GetDefinition",
+            Some(multi_lsp_query::Request::GetDeclaration(_)) => "GetDeclaration",
+            Some(multi_lsp_query::Request::GetTypeDefinition(_)) => "GetTypeDefinition",
+            Some(multi_lsp_query::Request::GetImplementation(_)) => "GetImplementation",
+            Some(multi_lsp_query::Request::GetReferences(_)) => "GetReferences",
+            None => "<unknown>",
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;