Allow concurrent peer LSP requests via collab (#63736)

Before this change, two different guests requesting the same kind of
request would cancel each other's requests and only the latest one would
have arrived to the latest to query that.

Release Notes:

- Fixed collab LSP requests cancelling each other sometimes

crates/collab/src/rpc.rs

@@ -420,7 +420,7 @@ impl Server {
             .add_request_handler(forward_mutating_project_request::<proto::SaveBuffer>)
             .add_request_handler(forward_mutating_project_request::<proto::BlameBuffer>)
             .add_request_handler(lsp_query)
-            .add_message_handler(broadcast_project_message_from_host::<proto::LspQueryResponse>)
+            .add_message_handler(forward_lsp_query_response)
             .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>)
@@ -2473,6 +2473,26 @@ async fn lsp_query(
     }
 }
 
+async fn forward_lsp_query_response(
+    request: proto::LspQueryResponse,
+    session: MessageContext,
+) -> Result<()> {
+    let project_id = ProjectId::from_proto(request.project_id);
+    session
+        .db()
+        .await
+        .check_user_is_project_host(project_id, session.connection_id)
+        .await?;
+    if let Some(peer_id) = request.peer_id {
+        session
+            .peer
+            .forward_send(session.connection_id, peer_id.into(), request)?;
+        Ok(())
+    } else {
+        broadcast_project_message_from_host(request, session).await
+    }
+}
+
 /// Notify other participants that a new buffer has been created
 async fn create_buffer_for_peer(
     request: proto::CreateBufferForPeer,

crates/collab/tests/integration/integration_tests.rs

@@ -9,7 +9,10 @@ use client::{RECEIVE_TIMEOUT, User};
 use collab::rpc::{CLEANUP_TIMEOUT, RECONNECT_TIMEOUT};
 use collections::{BTreeMap, HashMap, HashSet};
 use fs::{FakeFs, Fs as _, RemoveOptions};
-use futures::{StreamExt as _, channel::mpsc};
+use futures::{
+    StreamExt as _,
+    channel::{mpsc, oneshot},
+};
 use git::{
     repository::repo_path,
     status::{FileStatus, StatusCode, TrackedStatus, UnmergedStatus, UnmergedStatusCode},
@@ -5489,6 +5492,107 @@ async fn test_references(
     });
 }
 
+#[gpui::test]
+async fn test_concurrent_guest_lsp_requests_do_not_cancel_each_other(
+    executor: BackgroundExecutor,
+    cx_a: &mut TestAppContext,
+    cx_b: &mut TestAppContext,
+    cx_c: &mut TestAppContext,
+) {
+    let mut server = TestServer::start(executor.clone()).await;
+    let client_a = server.create_client(cx_a, "user_a").await;
+    let client_b = server.create_client(cx_b, "user_b").await;
+    let client_c = server.create_client(cx_c, "user_c").await;
+    server
+        .create_room(&mut [(&client_a, cx_a), (&client_b, cx_b), (&client_c, cx_c)])
+        .await;
+    let active_call_a = cx_a.read(ActiveCall::global);
+
+    let capabilities = lsp::ServerCapabilities {
+        references_provider: Some(lsp::OneOf::Left(true)),
+        ..lsp::ServerCapabilities::default()
+    };
+    client_a.language_registry().add(rust_lang());
+    let mut fake_language_servers = client_a.language_registry().register_fake_lsp(
+        "Rust",
+        FakeLspAdapter {
+            name: "my-fake-lsp-adapter",
+            capabilities: capabilities.clone(),
+            ..FakeLspAdapter::default()
+        },
+    );
+    for client in [&client_b, &client_c] {
+        client.language_registry().add(rust_lang());
+        client.language_registry().register_fake_lsp_adapter(
+            "Rust",
+            FakeLspAdapter {
+                name: "my-fake-lsp-adapter",
+                capabilities: capabilities.clone(),
+                ..FakeLspAdapter::default()
+            },
+        );
+    }
+
+    client_a
+        .fs()
+        .insert_tree(
+            path!("/root"),
+            json!({
+                "main.rs": "const ONE: usize = 1;"
+            }),
+        )
+        .await;
+    let (project_a, worktree_id) = client_a.build_local_project(path!("/root"), cx_a).await;
+    let project_id = active_call_a
+        .update(cx_a, |call, cx| call.share_project(project_a.clone(), cx))
+        .await
+        .unwrap();
+    let project_b = client_b.join_remote_project(project_id, cx_b).await;
+    let project_c = client_c.join_remote_project(project_id, cx_c).await;
+    let (buffer_b, _handle_b) = project_b
+        .update(cx_b, |project, cx| {
+            project.open_buffer_with_lsp((worktree_id, rel_path("main.rs")), cx)
+        })
+        .await
+        .unwrap();
+    let (buffer_c, _handle_c) = project_c
+        .update(cx_c, |project, cx| {
+            project.open_buffer_with_lsp((worktree_id, rel_path("main.rs")), cx)
+        })
+        .await
+        .unwrap();
+    let fake_language_server = fake_language_servers.next().await.unwrap();
+    cx_a.run_until_parked();
+    cx_b.run_until_parked();
+    cx_c.run_until_parked();
+
+    let (request_started_tx, mut request_started_rx) = mpsc::unbounded();
+    let _reference_requests = fake_language_server
+        .set_request_handler::<lsp::request::References, _, _>(move |params, _| {
+            let (response_tx, response_rx) = oneshot::channel();
+            request_started_tx
+                .unbounded_send((params.text_document_position.position, response_tx))
+                .unwrap();
+            async move { response_rx.await.unwrap() }
+        });
+
+    let references_b = project_b.update(cx_b, |project, cx| project.references(&buffer_b, 1, cx));
+    let (position_b, response_b) = request_started_rx.next().await.unwrap();
+    assert_eq!(position_b, lsp::Position::new(0, 1));
+
+    let references_c = project_c.update(cx_c, |project, cx| project.references(&buffer_c, 7, cx));
+    let (position_c, response_c) = request_started_rx.next().await.unwrap();
+    assert_eq!(position_c, lsp::Position::new(0, 7));
+    assert!(!response_b.is_canceled());
+    assert!(!response_c.is_canceled());
+
+    response_b.send(Ok(Some(Vec::new()))).unwrap();
+    response_c.send(Ok(Some(Vec::new()))).unwrap();
+
+    assert_eq!(references_b.await.unwrap().unwrap(), Vec::new());
+    assert_eq!(references_c.await.unwrap().unwrap(), Vec::new());
+}
+
 #[gpui::test(iterations = 10)]
 async fn test_project_search(
     executor: BackgroundExecutor,

crates/project/src/lsp_store.rs

@@ -4527,6 +4527,7 @@ pub struct BufferLspData {
 struct LspKey {
     request_type: TypeId,
     server_queried: Option<LanguageServerId>,
+    sender_id: proto::PeerId,
 }
 
 impl BufferLspData {
@@ -10840,6 +10841,7 @@ impl LspStore {
                 Self::deduplicate_range_based_lsp_requests::<InlayHints>(
                     &lsp_store,
                     server_id,
+                    sender_id,
                     lsp_request_id,
                     &inlay_hints,
                     query_start..query_end,
@@ -10921,6 +10923,7 @@ impl LspStore {
                     let key = LspKey {
                         request_type: TypeId::of::<GetDocumentDiagnostics>(),
                         server_queried: server_id,
+                        sender_id,
                     };
                     if <GetDocumentDiagnostics as LspCommand>::ProtoRequest::stop_previous_requests(
                     ) {
@@ -14407,6 +14410,7 @@ impl LspStore {
     async fn deduplicate_range_based_lsp_requests<T>(
         lsp_store: &Entity<Self>,
         server_id: Option<LanguageServerId>,
+        sender_id: proto::PeerId,
         lsp_request_id: LspRequestId,
         proto_request: &T::ProtoRequest,
         range: Range<Anchor>,
@@ -14436,6 +14440,7 @@ impl LspStore {
                     let key = LspKey {
                         request_type: TypeId::of::<T>(),
                         server_queried: server_id,
+                        sender_id,
                     };
                     let previous_request = lsp_data
                         .chunk_lsp_requests
@@ -14481,6 +14486,7 @@ impl LspStore {
         let key = LspKey {
             request_type: TypeId::of::<T>(),
             server_queried: for_server_id,
+            sender_id,
         };
         lsp_store.update(cx, |lsp_store, cx| {
             let request_task = match for_server_id {
@@ -14540,6 +14546,7 @@ impl LspStore {
                                     .collect::<HashMap<_, _>>();
                                 match client.send_lsp_response::<T::ProtoRequest>(
                                     project_id,
+                                    sender_id,
                                     lsp_request_id,
                                     response,
                                 ) {
@@ -14579,6 +14586,7 @@ impl LspStore {
         let key = LspKey {
             request_type: TypeId::of::<T>(),
             server_queried: server_id,
+            sender_id,
         };
         if T::ProtoRequest::stop_previous_requests() {
             if let Some(lsp_requests) = lsp_data.lsp_requests.get_mut(&key) {
@@ -14609,6 +14617,7 @@ impl LspStore {
                             .collect::<HashMap<_, _>>();
                         if let Err(e) = client.send_lsp_response::<T::ProtoRequest>(
                             project_id,
+                            sender_id,
                             lsp_request_id,
                             response,
                         ) {

crates/proto/proto/lsp.proto

@@ -2,6 +2,7 @@ syntax = "proto3";
 package zed.messages;
 
 import "buffer.proto";
+import "core.proto";
 
 message GetDefinition {
   uint64 project_id = 1;
@@ -1035,6 +1036,7 @@ message LspQueryResponse {
   uint64 project_id = 1;
   uint64 lsp_request_id = 2;
   repeated LspResponse responses = 3;
+  PeerId peer_id = 4;
 }
 
 message LspResponse {

crates/rpc/src/proto_client.rs

@@ -339,11 +339,13 @@ impl AnyProtoClient {
     pub fn send_lsp_response<T: LspRequestMessage>(
         &self,
         project_id: u64,
+        peer_id: proto::PeerId,
         lsp_request_id: LspRequestId,
         server_responses: HashMap<u64, T::Response>,
     ) -> Result<()> {
         self.send(proto::LspQueryResponse {
             project_id,
+            peer_id: Some(peer_id),
             lsp_request_id: lsp_request_id.0,
             responses: server_responses
                 .into_iter()