From 80cb8209a0b63f862dc13a51d7d74a9142269f1c Mon Sep 17 00:00:00 2001 From: Dhruv Manilawala Date: Mon, 30 Jun 2025 10:56:11 +0530 Subject: [PATCH 1/2] [ty] Add background request task support --- crates/ty_server/src/server/api.rs | 61 ++++++++++++++++--- .../src/server/api/requests/completion.rs | 10 ++- .../src/server/api/requests/diagnostic.rs | 6 +- .../api/requests/goto_type_definition.rs | 6 +- .../src/server/api/requests/hover.rs | 6 +- .../src/server/api/requests/inlay_hints.rs | 6 +- crates/ty_server/src/server/api/traits.rs | 36 +++++++---- crates/ty_server/src/session.rs | 36 +++++++++++ 8 files changed, 141 insertions(+), 26 deletions(-) diff --git a/crates/ty_server/src/server/api.rs b/crates/ty_server/src/server/api.rs index 438bc5f7540b4..e15d6b98a3443 100644 --- a/crates/ty_server/src/server/api.rs +++ b/crates/ty_server/src/server/api.rs @@ -28,23 +28,23 @@ pub(super) fn request(req: server::Request) -> Task { let id = req.id.clone(); match req.method.as_str() { - requests::DocumentDiagnosticRequestHandler::METHOD => background_request_task::< + requests::DocumentDiagnosticRequestHandler::METHOD => background_document_request_task::< requests::DocumentDiagnosticRequestHandler, >( req, BackgroundSchedule::Worker ), - requests::GotoTypeDefinitionRequestHandler::METHOD => background_request_task::< + requests::GotoTypeDefinitionRequestHandler::METHOD => background_document_request_task::< requests::GotoTypeDefinitionRequestHandler, >( req, BackgroundSchedule::Worker ), - requests::HoverRequestHandler::METHOD => background_request_task::< + requests::HoverRequestHandler::METHOD => background_document_request_task::< requests::HoverRequestHandler, >(req, BackgroundSchedule::Worker), - requests::InlayHintRequestHandler::METHOD => background_request_task::< + requests::InlayHintRequestHandler::METHOD => background_document_request_task::< requests::InlayHintRequestHandler, >(req, BackgroundSchedule::Worker), - requests::CompletionRequestHandler::METHOD => background_request_task::< + requests::CompletionRequestHandler::METHOD => background_document_request_task::< requests::CompletionRequestHandler, >( req, BackgroundSchedule::LatencySensitive @@ -135,7 +135,52 @@ where })) } -fn background_request_task( +#[expect(dead_code)] +fn background_request_task( + req: server::Request, + schedule: BackgroundSchedule, +) -> Result +where + <::RequestType as Request>::Params: UnwindSafe, +{ + let retry = R::RETRY_ON_CANCELLATION.then(|| req.clone()); + let (id, params) = cast_request::(req)?; + + Ok(Task::background(schedule, move |session: &Session| { + let cancellation_token = session + .request_queue() + .incoming() + .cancellation_token(&id) + .expect("request should have been tested for cancellation before scheduling"); + + let snapshot = session.take_workspace_snapshot(); + + Box::new(move |client| { + let _span = tracing::debug_span!("request", %id, method = R::METHOD).entered(); + + // Test again if the request was cancelled since it was scheduled on the background task + // and, if so, return early + if cancellation_token.is_cancelled() { + tracing::trace!( + "Ignoring request id={id} method={} because it was cancelled", + R::METHOD + ); + + // We don't need to send a response here because the `cancel` notification + // handler already responded with a message. + return; + } + + let result = ruff_db::panic::catch_unwind(|| R::run(snapshot, client, params)); + + if let Some(response) = request_result_to_response::(&id, client, result, retry) { + respond::(&id, response, client); + } + }) + })) +} + +fn background_document_request_task( req: server::Request, schedule: BackgroundSchedule, ) -> Result @@ -168,7 +213,7 @@ where }; let Some(snapshot) = session.take_snapshot(url) else { - tracing::warn!("Ignoring request because snapshot for path `{path:?}` doesn't exist."); + tracing::warn!("Ignoring request because snapshot for path `{path:?}` doesn't exist"); return Box::new(|_| {}); }; @@ -209,7 +254,7 @@ fn request_result_to_response( request: Option, ) -> Option::RequestType as Request>::Result>> where - R: traits::BackgroundDocumentRequestHandler, + R: traits::RetriableRequestHandler, { match result { Ok(response) => Some(response), diff --git a/crates/ty_server/src/server/api/requests/completion.rs b/crates/ty_server/src/server/api/requests/completion.rs index 5ce652a691165..e01424b8a867e 100644 --- a/crates/ty_server/src/server/api/requests/completion.rs +++ b/crates/ty_server/src/server/api/requests/completion.rs @@ -8,7 +8,9 @@ use ty_project::ProjectDatabase; use crate::DocumentSnapshot; use crate::document::PositionExt; -use crate::server::api::traits::{BackgroundDocumentRequestHandler, RequestHandler}; +use crate::server::api::traits::{ + BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler, +}; use crate::session::client::Client; pub(crate) struct CompletionRequestHandler; @@ -18,8 +20,6 @@ impl RequestHandler for CompletionRequestHandler { } impl BackgroundDocumentRequestHandler for CompletionRequestHandler { - const RETRY_ON_CANCELLATION: bool = true; - fn document_url(params: &CompletionParams) -> Cow { Cow::Borrowed(¶ms.text_document_position.text_document.uri) } @@ -65,3 +65,7 @@ impl BackgroundDocumentRequestHandler for CompletionRequestHandler { Ok(Some(response)) } } + +impl RetriableRequestHandler for CompletionRequestHandler { + const RETRY_ON_CANCELLATION: bool = true; +} diff --git a/crates/ty_server/src/server/api/requests/diagnostic.rs b/crates/ty_server/src/server/api/requests/diagnostic.rs index 115aa878e3098..8cead7f1356db 100644 --- a/crates/ty_server/src/server/api/requests/diagnostic.rs +++ b/crates/ty_server/src/server/api/requests/diagnostic.rs @@ -8,7 +8,9 @@ use lsp_types::{ use crate::server::Result; use crate::server::api::diagnostics::{Diagnostics, compute_diagnostics}; -use crate::server::api::traits::{BackgroundDocumentRequestHandler, RequestHandler}; +use crate::server::api::traits::{ + BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler, +}; use crate::session::DocumentSnapshot; use crate::session::client::Client; use ty_project::ProjectDatabase; @@ -43,7 +45,9 @@ impl BackgroundDocumentRequestHandler for DocumentDiagnosticRequestHandler { }), )) } +} +impl RetriableRequestHandler for DocumentDiagnosticRequestHandler { fn salsa_cancellation_error() -> lsp_server::ResponseError { lsp_server::ResponseError { code: lsp_server::ErrorCode::ServerCancelled as i32, diff --git a/crates/ty_server/src/server/api/requests/goto_type_definition.rs b/crates/ty_server/src/server/api/requests/goto_type_definition.rs index ead442c1a5907..59924a26b95d1 100644 --- a/crates/ty_server/src/server/api/requests/goto_type_definition.rs +++ b/crates/ty_server/src/server/api/requests/goto_type_definition.rs @@ -8,7 +8,9 @@ use ty_project::ProjectDatabase; use crate::DocumentSnapshot; use crate::document::{PositionExt, ToLink}; -use crate::server::api::traits::{BackgroundDocumentRequestHandler, RequestHandler}; +use crate::server::api::traits::{ + BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler, +}; use crate::session::client::Client; pub(crate) struct GotoTypeDefinitionRequestHandler; @@ -70,3 +72,5 @@ impl BackgroundDocumentRequestHandler for GotoTypeDefinitionRequestHandler { } } } + +impl RetriableRequestHandler for GotoTypeDefinitionRequestHandler {} diff --git a/crates/ty_server/src/server/api/requests/hover.rs b/crates/ty_server/src/server/api/requests/hover.rs index 6919e172372c0..bc65a583849fb 100644 --- a/crates/ty_server/src/server/api/requests/hover.rs +++ b/crates/ty_server/src/server/api/requests/hover.rs @@ -2,7 +2,9 @@ use std::borrow::Cow; use crate::DocumentSnapshot; use crate::document::{PositionExt, ToRangeExt}; -use crate::server::api::traits::{BackgroundDocumentRequestHandler, RequestHandler}; +use crate::server::api::traits::{ + BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler, +}; use crate::session::client::Client; use lsp_types::request::HoverRequest; use lsp_types::{HoverContents, HoverParams, MarkupContent, Url}; @@ -73,3 +75,5 @@ impl BackgroundDocumentRequestHandler for HoverRequestHandler { })) } } + +impl RetriableRequestHandler for HoverRequestHandler {} diff --git a/crates/ty_server/src/server/api/requests/inlay_hints.rs b/crates/ty_server/src/server/api/requests/inlay_hints.rs index 62ffe111a2436..aba2cd715124a 100644 --- a/crates/ty_server/src/server/api/requests/inlay_hints.rs +++ b/crates/ty_server/src/server/api/requests/inlay_hints.rs @@ -2,7 +2,9 @@ use std::borrow::Cow; use crate::DocumentSnapshot; use crate::document::{RangeExt, TextSizeExt}; -use crate::server::api::traits::{BackgroundDocumentRequestHandler, RequestHandler}; +use crate::server::api::traits::{ + BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler, +}; use crate::session::client::Client; use lsp_types::request::InlayHintRequest; use lsp_types::{InlayHintParams, Url}; @@ -64,3 +66,5 @@ impl BackgroundDocumentRequestHandler for InlayHintRequestHandler { Ok(Some(inlay_hints)) } } + +impl RetriableRequestHandler for InlayHintRequestHandler {} diff --git a/crates/ty_server/src/server/api/traits.rs b/crates/ty_server/src/server/api/traits.rs index 460da654dbffd..f96f5172c7cfb 100644 --- a/crates/ty_server/src/server/api/traits.rs +++ b/crates/ty_server/src/server/api/traits.rs @@ -1,7 +1,7 @@ //! A stateful LSP implementation that calls into the ty API. use crate::session::client::Client; -use crate::session::{DocumentSnapshot, Session}; +use crate::session::{DocumentSnapshot, Session, WorkspaceSnapshot}; use lsp_types::notification::Notification as LSPNotification; use lsp_types::request::Request; @@ -25,11 +25,24 @@ pub(super) trait SyncRequestHandler: RequestHandler { ) -> super::Result<<::RequestType as Request>::Result>; } -/// A request handler that can be run on a background thread. -pub(super) trait BackgroundDocumentRequestHandler: RequestHandler { - /// Whether this request be retried if it was cancelled due to a modification to the Salsa database. +pub(super) trait RetriableRequestHandler: RequestHandler { + /// Whether this request can be cancelled if the Salsa database is modified. const RETRY_ON_CANCELLATION: bool = false; + /// The error to return if the request was cancelled due to a modification to the Salsa database. + fn salsa_cancellation_error() -> lsp_server::ResponseError { + lsp_server::ResponseError { + code: lsp_server::ErrorCode::ContentModified as i32, + message: "content modified".to_string(), + data: None, + } + } +} + +/// A request handler that can be run on a background thread. +/// +/// This handler is specific to requests that operate on a single document. +pub(super) trait BackgroundDocumentRequestHandler: RetriableRequestHandler { fn document_url( params: &<::RequestType as Request>::Params, ) -> std::borrow::Cow; @@ -40,14 +53,15 @@ pub(super) trait BackgroundDocumentRequestHandler: RequestHandler { client: &Client, params: <::RequestType as Request>::Params, ) -> super::Result<<::RequestType as Request>::Result>; +} - fn salsa_cancellation_error() -> lsp_server::ResponseError { - lsp_server::ResponseError { - code: lsp_server::ErrorCode::ContentModified as i32, - message: "content modified".to_string(), - data: None, - } - } +/// A request handler that can be run on a background thread. +pub(super) trait BackgroundRequestHandler: RetriableRequestHandler { + fn run( + snapshot: WorkspaceSnapshot, + client: &Client, + params: <::RequestType as Request>::Params, + ) -> super::Result<<::RequestType as Request>::Result>; } /// A supertrait for any server notification handler. diff --git a/crates/ty_server/src/session.rs b/crates/ty_server/src/session.rs index b7c365240d831..5ebb9552c2ed9 100644 --- a/crates/ty_server/src/session.rs +++ b/crates/ty_server/src/session.rs @@ -2,6 +2,7 @@ use std::collections::{BTreeMap, VecDeque}; use std::ops::{Deref, DerefMut}; +use std::panic::AssertUnwindSafe; use std::path::{Path, PathBuf}; use std::sync::Arc; @@ -224,6 +225,19 @@ impl Session { self.index().key_from_url(url) } + pub(crate) fn take_workspace_snapshot(&self) -> WorkspaceSnapshot { + WorkspaceSnapshot { + projects: AssertUnwindSafe( + self.projects_by_workspace_folder + .values() + .cloned() + .collect(), + ), + index: self.index.clone().unwrap(), + position_encoding: self.position_encoding, + } + } + pub(crate) fn initialize_workspaces(&mut self, workspace_settings: Vec<(Url, ClientOptions)>) { assert!(!self.workspaces.all_initialized()); @@ -461,6 +475,28 @@ impl DocumentSnapshot { } } +/// An immutable snapshot of the current state of [`Session`]. +pub(crate) struct WorkspaceSnapshot { + projects: AssertUnwindSafe>, + index: Arc, + position_encoding: PositionEncoding, +} + +#[expect(dead_code)] +impl WorkspaceSnapshot { + pub(crate) fn projects(&self) -> &[ProjectDatabase] { + &self.projects + } + + pub(crate) fn index(&self) -> &index::Index { + &self.index + } + + pub(crate) fn position_encoding(&self) -> PositionEncoding { + self.position_encoding + } +} + #[derive(Debug, Default)] pub(crate) struct Workspaces { workspaces: BTreeMap, From a1fa874ee771b84a2bcfbd96230f130013f8c38c Mon Sep 17 00:00:00 2001 From: Dhruv Manilawala Date: Wed, 2 Jul 2025 14:47:25 +0530 Subject: [PATCH 2/2] Fix clippy --- crates/ty_server/src/session.rs | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/crates/ty_server/src/session.rs b/crates/ty_server/src/session.rs index 5ebb9552c2ed9..f45fdae504784 100644 --- a/crates/ty_server/src/session.rs +++ b/crates/ty_server/src/session.rs @@ -227,12 +227,7 @@ impl Session { pub(crate) fn take_workspace_snapshot(&self) -> WorkspaceSnapshot { WorkspaceSnapshot { - projects: AssertUnwindSafe( - self.projects_by_workspace_folder - .values() - .cloned() - .collect(), - ), + projects: AssertUnwindSafe(self.projects.values().cloned().collect()), index: self.index.clone().unwrap(), position_encoding: self.position_encoding, }