Use a single queue for all TSP connections
Summary:
The TSP supports secondary ("extra") connections which Pylance uses to
communicate to Pyrefly from a background thread. Before this change, Pyrefly
dispatched requests received on extra connections directly in a separate thread.
With this change, the extra connection instead inserts the request into the
`lsp_queue`, so all requests are handled by the same thread.
Because we use a single loop for requests from different clients, it's important
that responses are routed to the correct one. To make this safer, this change
also refactors the structure of the TSP. Before, the event loop was owned by
`TspMainConnection` which also owns the `Sender` object, meaning all request
handlers had access to that sender from `self`. It seemed too risky to me that
code might accidentally send responses to the wrong client.
To address the risk, the event loop is now owned by `TspServer` and both the
main and extra connections provide a `Reply` object for each event: the server
owns no channel, and a handler can only reach the channel it was given.
This change also enables a stacked change -- all request handlers now share a
single `TransactionManager` object. The transaction manager is also used in the
LSP and helps deal with concurrent request handling, so a recheck and read-only
request can run in parallel and multiple read-only requests can build on top of
one another before a recheck commits. The stacked change uses that to optimize
the TSP.
Test Plan:
Replayed a captured Pylance session (pandas, 45251 TSP messages over both the
stdio and IPC connections) against `pyrefly tsp`. All 45102 requests were
answered on the connection that sent them and total request time is
unchanged.
diff --git a/pyrefly/lib/lsp/non_wasm/queue.rs b/pyrefly/lib/lsp/non_wasm/queue.rs
index 431a4a4..7d5323d 100644
--- a/pyrefly/lib/lsp/non_wasm/queue.rs
+++ b/pyrefly/lib/lsp/non_wasm/queue.rs
@@ -31,6 +31,7 @@
use tracing::debug;
use tracing::info;
+use crate::lsp::non_wasm::protocol::Message;
use crate::lsp::non_wasm::protocol::Request;
use crate::lsp::non_wasm::protocol::Response;
use crate::lsp::non_wasm::server::Server;
@@ -65,6 +66,13 @@
InvalidateConfigFind,
LspResponse(Response),
LspRequest(Request),
+ /// A request from an extra TSP connection. It shares the main connection's
+ /// queue so every TSP request is served by one loop, and carries the
+ /// channel its response goes back on.
+ TspExtraRequest {
+ request: Request,
+ response_sender: Sender<Message>,
+ },
Exit,
}
@@ -129,6 +137,9 @@
Self::DidSaveNotebookDocument(_) => "DidSaveNotebookDocument".to_owned(),
Self::LspResponse(_) => "LspResponse".to_owned(),
Self::LspRequest(request) => format!("LspRequest({})", request.method,),
+ Self::TspExtraRequest { request, .. } => {
+ format!("TspExtraRequest({})", request.method)
+ }
Self::Exit => "Exit".to_owned(),
}
}
@@ -173,7 +184,9 @@
| Self::DidChangeNotebookDocument(_)
| Self::InvalidateConfigFind
| Self::Exit => LspEventKind::Mutation,
- Self::LspResponse(_) | Self::LspRequest(_) => LspEventKind::Query,
+ Self::LspResponse(_) | Self::LspRequest(_) | Self::TspExtraRequest { .. } => {
+ LspEventKind::Query
+ }
}
}
}
diff --git a/pyrefly/lib/lsp/non_wasm/server.rs b/pyrefly/lib/lsp/non_wasm/server.rs
index 9b40b41..b42c35c 100644
--- a/pyrefly/lib/lsp/non_wasm/server.rs
+++ b/pyrefly/lib/lsp/non_wasm/server.rs
@@ -1797,6 +1797,9 @@
LspEvent::Exit => {
return Ok(ProcessEvent::Exit);
}
+ LspEvent::TspExtraRequest { .. } => {
+ unreachable!("extra TSP connection requests are answered by the TSP loop")
+ }
LspEvent::RecheckFinished => {
// We did a commit and want to get back to a stable state.
self.validate_in_memory_and_commit_if_possible(
diff --git a/pyrefly/lib/test/tsp/tsp_interaction.rs b/pyrefly/lib/test/tsp/tsp_interaction.rs
index 67cb452..3cc1244 100644
--- a/pyrefly/lib/test/tsp/tsp_interaction.rs
+++ b/pyrefly/lib/test/tsp/tsp_interaction.rs
@@ -13,5 +13,6 @@
pub mod get_type_queries;
pub mod notebook;
pub mod object_model;
+pub mod request_errors;
pub mod resolve_import;
pub mod snapshot_changed;
diff --git a/pyrefly/lib/test/tsp/tsp_interaction/request_errors.rs b/pyrefly/lib/test/tsp/tsp_interaction/request_errors.rs
new file mode 100644
index 0000000..8ff90b0
--- /dev/null
+++ b/pyrefly/lib/test/tsp/tsp_interaction/request_errors.rs
@@ -0,0 +1,72 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+//! A request the server cannot act on is answered with an error rather than
+//! dropped, so the client is never left waiting.
+
+use lsp_server::ErrorCode;
+use lsp_server::RequestId;
+use tempfile::TempDir;
+
+use crate::lsp::non_wasm::protocol::Message;
+use crate::lsp::non_wasm::protocol::Request;
+use crate::test::tsp::tsp_interaction::object_model::TspInteraction;
+
+#[test]
+fn test_tsp_unknown_method_returns_method_not_found() {
+ let temp_dir = TempDir::new().unwrap();
+
+ let mut tsp = TspInteraction::new();
+ tsp.set_root(temp_dir.path().to_path_buf());
+ tsp.initialize(Default::default());
+
+ tsp.server.send_message(Message::Request(Request {
+ id: RequestId::from(2),
+ method: "typeServer/thisMethodDoesNotExist".to_owned(),
+ params: serde_json::json!(null),
+ activity_key: None,
+ }));
+
+ let response = tsp.client.receive_response_skip_notifications();
+ assert_eq!(response.id, RequestId::from(2));
+ assert_eq!(response.result, None);
+ let error = response.error.expect("unknown method should be an error");
+ assert_eq!(error.code, ErrorCode::MethodNotFound as i32);
+ assert!(
+ error.message.contains("typeServer/thisMethodDoesNotExist"),
+ "error should name the unsupported method, got: {}",
+ error.message
+ );
+
+ tsp.shutdown();
+}
+
+#[test]
+fn test_tsp_malformed_params_return_invalid_params() {
+ let temp_dir = TempDir::new().unwrap();
+
+ let mut tsp = TspInteraction::new();
+ tsp.set_root(temp_dir.path().to_path_buf());
+ tsp.initialize(Default::default());
+
+ // A known method, but `arg` is missing the `uri` and `range` that
+ // `GetTypeParams` requires.
+ tsp.server.send_message(Message::Request(Request {
+ id: RequestId::from(2),
+ method: "typeServer/getComputedType".to_owned(),
+ params: serde_json::json!({ "arg": {}, "snapshot": 0 }),
+ activity_key: None,
+ }));
+
+ let response = tsp.client.receive_response_skip_notifications();
+ assert_eq!(response.id, RequestId::from(2));
+ assert_eq!(response.result, None);
+ let error = response.error.expect("malformed params should be an error");
+ assert_eq!(error.code, ErrorCode::InvalidParams as i32);
+
+ tsp.shutdown();
+}
diff --git a/pyrefly/lib/tsp/requests/get_computed_type.rs b/pyrefly/lib/tsp/requests/get_computed_type.rs
index d96e9a4..2145ff4 100644
--- a/pyrefly/lib/tsp/requests/get_computed_type.rs
+++ b/pyrefly/lib/tsp/requests/get_computed_type.rs
@@ -12,10 +12,10 @@
use tsp_types::Type;
use crate::lsp::non_wasm::server::TspInterface;
-use crate::tsp::server::TspConnection;
+use crate::tsp::server::TspServer;
use crate::tsp::validation::parse_uri;
-impl<T: TspInterface> TspConnection<T> {
+impl<T: TspInterface> TspServer<T> {
/// Return the computed (inferred) type at the given position.
///
/// The computed type reflects the type checker's analysis of the code
diff --git a/pyrefly/lib/tsp/requests/get_declared_type.rs b/pyrefly/lib/tsp/requests/get_declared_type.rs
index c2eaad2..96b46c4 100644
--- a/pyrefly/lib/tsp/requests/get_declared_type.rs
+++ b/pyrefly/lib/tsp/requests/get_declared_type.rs
@@ -12,10 +12,10 @@
use tsp_types::Type;
use crate::lsp::non_wasm::server::TspInterface;
-use crate::tsp::server::TspConnection;
+use crate::tsp::server::TspServer;
use crate::tsp::validation::parse_uri;
-impl<T: TspInterface> TspConnection<T> {
+impl<T: TspInterface> TspServer<T> {
/// Return the declared type at the given position.
///
/// The declared type is the annotation explicitly written by the user.
diff --git a/pyrefly/lib/tsp/requests/get_expected_type.rs b/pyrefly/lib/tsp/requests/get_expected_type.rs
index a229253..e993ea7 100644
--- a/pyrefly/lib/tsp/requests/get_expected_type.rs
+++ b/pyrefly/lib/tsp/requests/get_expected_type.rs
@@ -12,10 +12,10 @@
use tsp_types::Type;
use crate::lsp::non_wasm::server::TspInterface;
-use crate::tsp::server::TspConnection;
+use crate::tsp::server::TspServer;
use crate::tsp::validation::parse_uri;
-impl<T: TspInterface> TspConnection<T> {
+impl<T: TspInterface> TspServer<T> {
/// Return the expected type at the given position.
///
/// The expected type is the type that a surrounding context demands.
diff --git a/pyrefly/lib/tsp/requests/get_python_search_paths.rs b/pyrefly/lib/tsp/requests/get_python_search_paths.rs
index a4e94b2..c68ed34 100644
--- a/pyrefly/lib/tsp/requests/get_python_search_paths.rs
+++ b/pyrefly/lib/tsp/requests/get_python_search_paths.rs
@@ -16,11 +16,12 @@
use tsp_types::protocol::GetPythonSearchPathsParams;
use crate::lsp::non_wasm::server::TspInterface;
-use crate::tsp::server::TspConnection;
+use crate::tsp::server::Reply;
+use crate::tsp::server::TspServer;
use crate::tsp::validation::internal_error;
use crate::tsp::validation::parse_uri;
-impl<T: TspInterface> TspConnection<T> {
+impl<T: TspInterface> TspServer<T> {
/// Handle a `typeServer/getPythonSearchPaths` request.
///
/// Validates the snapshot, parses the `from_uri`, and delegates to
@@ -33,10 +34,11 @@
&self,
id: RequestId,
params: GetPythonSearchPathsParams,
+ reply: Reply,
) {
// --- 1. Validate snapshot ---
if let Err(err) = self.validate_snapshot(params.snapshot) {
- self.send_err(id, err);
+ reply.err(id, err);
return;
}
@@ -44,7 +46,7 @@
let url = match parse_uri(¶ms.from_uri) {
Ok(url) => url,
Err(err) => {
- self.send_err(id, err);
+ reply.err(id, err);
return;
}
};
@@ -60,7 +62,7 @@
Some(file_url) => file_url,
None => {
// Cannot resolve to a filesystem path — return empty list.
- self.send_ok::<Vec<String>>(id, vec![]);
+ reply.ok::<Vec<String>>(id, vec![]);
return;
}
}
@@ -69,8 +71,8 @@
};
match self.inner().get_python_search_paths(&resolved_url) {
- Ok(paths) => self.send_ok(id, paths),
- Err(detail) => self.send_err(id, internal_error(&detail)),
+ Ok(paths) => reply.ok(id, paths),
+ Err(detail) => reply.err(id, internal_error(&detail)),
}
}
}
diff --git a/pyrefly/lib/tsp/requests/get_snapshot.rs b/pyrefly/lib/tsp/requests/get_snapshot.rs
index d81752c..0b86409 100644
--- a/pyrefly/lib/tsp/requests/get_snapshot.rs
+++ b/pyrefly/lib/tsp/requests/get_snapshot.rs
@@ -8,22 +8,18 @@
//! Implementation of the getSnapshot TSP request
use crate::lsp::non_wasm::server::TspInterface;
-use crate::tsp::server::TspConnection;
+use crate::tsp::server::TspServer;
-impl<T: TspInterface> TspConnection<T> {
+impl<T: TspInterface> TspServer<T> {
/// Get the current snapshot version
///
/// The snapshot represents the current epoch of the global state.
/// It changes whenever files are modified, configuration changes,
/// or any other event that would trigger a recomputation.
pub fn get_snapshot(&self) -> i32 {
- *self
- .server
- .current_snapshot
- .lock()
- .unwrap_or_else(|poisoned| {
- eprintln!("TSP: Warning - snapshot mutex was poisoned, recovering");
- poisoned.into_inner()
- })
+ *self.current_snapshot.lock().unwrap_or_else(|poisoned| {
+ eprintln!("TSP: Warning - snapshot mutex was poisoned, recovering");
+ poisoned.into_inner()
+ })
}
}
diff --git a/pyrefly/lib/tsp/requests/get_supported_protocol_version.rs b/pyrefly/lib/tsp/requests/get_supported_protocol_version.rs
index d130753..3a63388 100644
--- a/pyrefly/lib/tsp/requests/get_supported_protocol_version.rs
+++ b/pyrefly/lib/tsp/requests/get_supported_protocol_version.rs
@@ -11,9 +11,9 @@
use tsp_types::protocol::TypeServerVersion;
use crate::lsp::non_wasm::server::TspInterface;
-use crate::tsp::server::TspConnection;
+use crate::tsp::server::TspServer;
-impl<T: TspInterface> TspConnection<T> {
+impl<T: TspInterface> TspServer<T> {
pub fn get_supported_protocol_version(&self) -> TypeServerVersion {
// Return the current protocol version from the generated enum
TSP_PROTOCOL_VERSION
diff --git a/pyrefly/lib/tsp/requests/resolve_import.rs b/pyrefly/lib/tsp/requests/resolve_import.rs
index a41f435..76eb00a 100644
--- a/pyrefly/lib/tsp/requests/resolve_import.rs
+++ b/pyrefly/lib/tsp/requests/resolve_import.rs
@@ -23,11 +23,12 @@
use crate::lsp::module_helpers::to_real_path;
use crate::lsp::non_wasm::server::TspInterface;
use crate::lsp::non_wasm::transaction_manager::TransactionManager;
-use crate::tsp::server::TspConnection;
+use crate::tsp::server::Reply;
+use crate::tsp::server::TspServer;
use crate::tsp::validation::invalid_params_error;
use crate::tsp::validation::parse_uri;
-impl<T: TspInterface> TspConnection<T> {
+impl<T: TspInterface> TspServer<T> {
/// Handle a `typeServer/resolveImport` request.
///
/// Converts the TSP [`ResolveImportParams`] into pyrefly's internal
@@ -39,10 +40,11 @@
id: RequestId,
params: ResolveImportParams,
ide_transaction_manager: &mut TransactionManager<'a>,
+ reply: Reply,
) {
// --- 1. Validate snapshot ---
if let Err(err) = self.validate_snapshot(params.snapshot) {
- self.send_err(id, err);
+ reply.err(id, err);
return;
}
@@ -50,7 +52,7 @@
let source_url = match parse_uri(¶ms.source_uri) {
Ok(url) => url,
Err(err) => {
- self.send_err(id, err);
+ reply.err(id, err);
return;
}
};
@@ -58,7 +60,7 @@
Some(p) => p,
None => {
// URI cannot be resolved to a filesystem path — return null.
- self.send_ok::<Option<String>>(id, None);
+ reply.ok::<Option<String>>(id, None);
return;
}
};
@@ -78,7 +80,7 @@
) {
Ok(name) => name,
Err(err) => {
- self.send_err(id, err);
+ reply.err(id, err);
return;
}
};
@@ -98,7 +100,7 @@
})
});
- self.send_ok(id, uri_string);
+ reply.ok(id, uri_string);
}
}
diff --git a/pyrefly/lib/tsp/server.rs b/pyrefly/lib/tsp/server.rs
index 3201383..c748b39 100644
--- a/pyrefly/lib/tsp/server.rs
+++ b/pyrefly/lib/tsp/server.rs
@@ -132,6 +132,22 @@
})
}
+ /// Convenience accessor for the inner LSP server.
+ pub(crate) fn inner(&self) -> &T {
+ &self.inner
+ }
+
+ /// Validate that the client-supplied snapshot matches the server's current
+ /// snapshot. Returns `Ok(())` on match or `Err(ResponseError)` on mismatch.
+ pub(crate) fn validate_snapshot(&self, client_snapshot: i32) -> Result<(), ResponseError> {
+ let current = self.get_snapshot();
+ if client_snapshot != current {
+ Err(snapshot_outdated_error(client_snapshot, current))
+ } else {
+ Ok(())
+ }
+ }
+
/// Send a `snapshotChanged` notification to the main connection.
fn broadcast_snapshot_changed(
&self,
@@ -162,84 +178,84 @@
response_sender,
}
}
+}
- /// Convenience accessor for the inner LSP server.
- pub(crate) fn inner(&self) -> &T {
- &self.server.inner
- }
+/// Where one request's response is written.
+///
+/// A handler is handed the reply channel of the connection that asked, so it
+/// cannot answer a different client, and the server itself owns no channel to
+/// answer through.
+pub(crate) struct Reply<'a>(pub(crate) &'a crossbeam_channel::Sender<Message>);
- fn send_response(&self, response: Response) {
- if let Err(error) = self.response_sender.send(Message::Response(response)) {
+impl Reply<'_> {
+ fn send(&self, response: Response) {
+ if let Err(error) = self.0.send(Message::Response(response)) {
warn!("Failed to send TSP response: {error}");
}
}
/// Send a successful JSON-RPC response for `id` with `result`.
- pub(crate) fn send_ok<R: Serialize>(&self, id: RequestId, result: R) {
- self.send_response(new_response(id, Ok(result)));
+ pub(crate) fn ok<R: Serialize>(&self, id: RequestId, result: R) {
+ self.send(new_response(id, Ok(result)));
}
/// Send a JSON-RPC error response for `id`.
- pub(crate) fn send_err(&self, id: RequestId, error: ResponseError) {
- self.send_response(Response {
+ pub(crate) fn err(&self, id: RequestId, error: ResponseError) {
+ self.send(Response {
id,
result: None,
error: Some(error),
});
}
+}
- /// Validate that the client-supplied snapshot matches the server's current
- /// snapshot. Returns `Ok(())` on match or `Err(ResponseError)` on mismatch.
- pub(crate) fn validate_snapshot(&self, client_snapshot: i32) -> Result<(), ResponseError> {
- let current = self.get_snapshot();
- if client_snapshot != current {
- Err(snapshot_outdated_error(client_snapshot, current))
- } else {
- Ok(())
- }
- }
-
+impl<T: TspInterface> TspServer<T> {
+ /// Handle one TSP request, answering on `reply` -- the channel belonging to
+ /// the connection that sent it.
+ ///
+ /// Infallible: a request that cannot be served is reported to its own
+ /// client as an error response. Every connection shares one event loop, so
+ /// returning an error here would take down every other client with it.
fn dispatch_tsp_request<'a>(
&'a self,
ide_transaction_manager: &mut TransactionManager<'a>,
+ reply: Reply,
request: &Request,
msg: TSPRequests,
- ) -> anyhow::Result<bool> {
+ ) {
match msg {
TSPRequests::GetSupportedProtocolVersionRequest { .. } => {
- self.send_ok(request.id.clone(), self.get_supported_protocol_version());
- Ok(true)
+ reply.ok(request.id.clone(), self.get_supported_protocol_version());
}
TSPRequests::GetSnapshotRequest { .. } => {
// Get snapshot doesn't need a transaction since it just returns the cached value
- self.send_ok(request.id.clone(), self.get_snapshot());
- Ok(true)
+ reply.ok(request.id.clone(), self.get_snapshot());
}
TSPRequests::ResolveImportRequest { params, .. } => {
- self.handle_resolve_import(request.id.clone(), params, ide_transaction_manager);
- Ok(true)
+ self.handle_resolve_import(
+ request.id.clone(),
+ params,
+ ide_transaction_manager,
+ reply,
+ );
}
TSPRequests::GetPythonSearchPathsRequest { params, .. } => {
- self.handle_get_python_search_paths(request.id.clone(), params);
- Ok(true)
+ self.handle_get_python_search_paths(request.id.clone(), params, reply);
}
TSPRequests::GetDeclaredTypeRequest { params, .. } => {
- self.dispatch_get_type_request(request.id.clone(), params, |s, p| {
- s.handle_get_declared_type(p)
+ self.dispatch_get_type_request(request.id.clone(), params, reply, |p| {
+ self.handle_get_declared_type(p)
});
- Ok(true)
}
TSPRequests::GetComputedTypeRequest { params, .. } => {
- self.dispatch_get_type_request(request.id.clone(), params, |s, p| {
- s.handle_get_computed_type(p)
+ self.dispatch_get_type_request(request.id.clone(), params, reply, |p| {
+ self.handle_get_computed_type(p)
});
- Ok(true)
}
TSPRequests::GetExpectedTypeRequest { params, .. } => {
- self.dispatch_get_type_request(request.id.clone(), params, |s, p| {
- s.handle_get_expected_type(p)
+ self.dispatch_get_type_request(request.id.clone(), params, reply, |p| {
+ self.handle_get_expected_type(p)
});
- Ok(true)
}
TSPRequests::ConnectionRequest { .. } => {
// Multi-connection management is handled at the transport layer,
@@ -256,49 +272,38 @@
&self,
id: RequestId,
raw_params: serde_json::Value,
+ reply: Reply,
handler: impl FnOnce(
- &Self,
GetTypeParams,
) -> Result<Option<tsp_types::Type>, lsp_server::ResponseError>,
) {
let params: GetTypeParams = match serde_json::from_value::<GetTypeParams>(raw_params) {
Ok(p) => p,
Err(e) => {
- self.send_err(id, invalid_params_error(&e.to_string()));
+ reply.err(id, invalid_params_error(&e.to_string()));
return;
}
};
- match handler(self, params) {
+ match handler(params) {
Ok(result) => {
- self.send_ok(id, result);
+ reply.ok(id, result);
}
Err(err) => {
- self.send_err(id, err);
+ reply.err(id, err);
}
}
}
}
-/// The main (stdio) connection. Only this type can manage extra connections
-/// and trigger `snapshotChanged` notifications.
-pub struct TspMainConnection<T: TspInterface>(TspConnection<T>);
-
-impl<T: TspInterface> TspMainConnection<T> {
- fn new(server: Arc<TspServer<T>>, response_sender: crossbeam_channel::Sender<Message>) -> Self {
- Self(TspConnection::new(server, response_sender))
- }
-}
-
-impl<T: TspInterface> Deref for TspMainConnection<T> {
- type Target = TspConnection<T>;
- fn deref(&self) -> &Self::Target {
- &self.0
- }
-}
-impl<T: TspInterface> TspMainConnection<T> {
- /// Process a single event on the main connection.
+impl<T: TspInterface> TspServer<T> {
+ /// Process a single event.
fn process_event<'a>(
- &'a self,
+ self: &'a Arc<Self>,
+ // The channel of the connection that sent this event.
+ reply: Reply,
+ // The main connection, where broadcasts and connection management go
+ // regardless of who asked.
+ main_reply: Reply,
ide_transaction_manager: &mut TransactionManager<'a>,
canceled_requests: &mut HashSet<RequestId>,
telemetry: &'a impl Telemetry,
@@ -319,16 +324,21 @@
};
// For TSP requests, handle them specially
- if let LspEvent::LspRequest(request) = event.event() {
+ let tsp_request = match event.event() {
+ LspEvent::LspRequest(request) => Some(request),
+ LspEvent::TspExtraRequest { request, .. } => Some(request),
+ _ => None,
+ };
+ if let Some(request) = tsp_request {
match parse_tsp_request(request) {
Some(TSPRequests::ConnectionRequest { params, .. }) => {
- self.handle_connection_request(request.id.clone(), params);
+ self.handle_connection_request(request.id.clone(), params, reply);
}
Some(msg) => {
- self.dispatch_tsp_request(ide_transaction_manager, request, msg)?;
+ self.dispatch_tsp_request(ide_transaction_manager, reply, request, msg);
}
None => {
- self.send_response(Response::new_err(
+ reply.send(Response::new_err(
request.id.clone(),
ErrorCode::MethodNotFound as i32,
format!("TSP server does not support LSP method: {}", request.method),
@@ -338,7 +348,7 @@
return Ok(ProcessEvent::Continue);
}
- let result = self.inner().process_event(
+ let result = self.inner.process_event(
ide_transaction_manager,
canceled_requests,
telemetry,
@@ -350,7 +360,6 @@
// Increment snapshot after the inner server has processed the event
if should_increment_snapshot {
let mut current = self
- .server
.current_snapshot
.lock()
.expect("current_snapshot mutex poisoned");
@@ -358,17 +367,18 @@
*current += 1;
let new_snapshot = *current;
drop(current);
- self.server.broadcast_snapshot_changed(
- &self.0.response_sender,
- old_snapshot,
- new_snapshot,
- );
+ self.broadcast_snapshot_changed(main_reply.0, old_snapshot, new_snapshot);
}
Ok(result)
}
- fn handle_connection_request(&self, id: RequestId, params: ConnectionRequestParams) {
+ fn handle_connection_request(
+ self: &Arc<Self>,
+ id: RequestId,
+ params: ConnectionRequestParams,
+ reply: Reply,
+ ) {
let result = match params.type_.as_str() {
"open" => self.open_extra_connection(params),
"close" => self.close_extra_connection(params),
@@ -378,20 +388,19 @@
};
match result {
- Ok(connection_result) => self.send_ok(id, connection_result),
- Err(error) => self.send_err(id, error),
+ Ok(connection_result) => reply.ok(id, connection_result),
+ Err(error) => reply.err(id, error),
}
}
fn open_extra_connection(
- &self,
+ self: &Arc<Self>,
params: ConnectionRequestParams,
) -> Result<ConnectionRequestResult, ResponseError> {
let transport = IpcTransportNames::from_connection_request(¶ms)?;
let description = transport.description();
let mut extra_connections = self
- .server
.extra_connections
.lock()
.map_err(|_| internal_error("extra connection state was poisoned"))?;
@@ -420,7 +429,7 @@
};
let extra_sender = ipc_connection.sender.clone();
- let extra_conn = TspExtraConnection::new(self.server.clone(), extra_sender.clone());
+ let extra_conn = TspExtraConnection::new(self.clone(), extra_sender.clone());
let (close_tx, close_rx) = crossbeam_channel::bounded::<()>(1);
extra_connections.insert(transport.clone(), ExtraConnectionHandle { close_tx });
@@ -443,7 +452,6 @@
let description = transport.description();
let handle = self
- .server
.extra_connections
.lock()
.expect("extra_connections mutex poisoned")
@@ -486,21 +494,20 @@
close_rx: crossbeam_channel::Receiver<()>,
transport: IpcTransportNames,
) {
- let (message_tx, message_rx) = crossbeam_channel::unbounded();
+ // The reader runs on its own thread because `recv` blocks and this loop has
+ // to stay responsive to `close_rx`. A rendezvous rather than a buffer: every
+ // message is forwarded to the unbounded main queue immediately, so buffering
+ // here would only duplicate that queue.
+ let (message_tx, message_rx) = crossbeam_channel::bounded(0);
+ std::thread::spawn(move || {
+ while let Some(message) = reader.recv() {
+ if message_tx.send(message).is_err() {
+ break;
+ }
+ }
+ });
std::thread::spawn(move || {
- std::thread::spawn(move || {
- while let Some(message) = reader.recv() {
- let (processed_tx, processed_rx) = crossbeam_channel::bounded(1);
- if message_tx.send((message, processed_tx)).is_err() {
- break;
- }
- if processed_rx.recv().is_err() {
- break;
- }
- }
- });
-
let mut selector = crossbeam_channel::Select::new();
let close_index = selector.recv(&close_rx);
let message_index = selector.recv(&message_rx);
@@ -512,16 +519,15 @@
break;
}
i if i == message_index => {
- let Ok((message, processed_tx)) = selected.recv(&message_rx) else {
+ let Ok(message) = selected.recv(&message_rx) else {
break;
};
match message {
Message::Request(request) => {
- let mut tm = TransactionManager::default();
match parse_tsp_request(&request) {
Some(TSPRequests::ConnectionRequest { .. }) => {
- self.send_err(
+ Reply(&self.response_sender).err(
request.id,
ResponseError {
code: ErrorCode::InvalidRequest as i32,
@@ -533,30 +539,27 @@
},
);
}
- Some(msg) => {
- if let Err(error) =
- self.dispatch_tsp_request(&mut tm, &request, msg)
+ // Everything else joins the main queue,
+ // stamped with this connection's channel so
+ // the loop answers the right client.
+ _ => {
+ if self
+ .server
+ .inner
+ .lsp_queue()
+ .send(LspEvent::TspExtraRequest {
+ request,
+ response_sender: self.response_sender.clone(),
+ })
+ .is_err()
{
- warn!("Extra TSP connection error: {error}");
break;
}
}
- None => {
- self.send_response(Response::new_err(
- request.id,
- ErrorCode::MethodNotFound as i32,
- format!(
- "Extra TSP connection does not support method: {}",
- request.method
- ),
- ));
- }
}
}
Message::Notification(_) | Message::Response(_) => {}
}
-
- let _ = processed_tx.send(());
}
_ => unreachable!(),
}
@@ -610,7 +613,7 @@
telemetry: &impl Telemetry,
) -> anyhow::Result<()> {
let server = TspServer::new(lsp_server);
- let main_conn = TspMainConnection::new(server.clone(), server.inner.sender().clone());
+ let main_sender = server.inner.sender().clone();
std::thread::scope(|scope| {
scope.spawn(|| server.inner.run_recheck_queue(telemetry));
@@ -636,7 +639,17 @@
);
let event_description = event.describe();
- let result = main_conn.process_event(
+ // Answer on the channel of whichever connection sent this request.
+ let reply_sender = match event.event() {
+ LspEvent::TspExtraRequest {
+ response_sender, ..
+ } => response_sender.clone(),
+ _ => main_sender.clone(),
+ };
+
+ let result = server.process_event(
+ Reply(&reply_sender),
+ Reply(&main_sender),
&mut ide_transaction_manager,
&mut canceled_requests,
telemetry,
