[ty] Small LSP cleanups (#18201)

crates/ty_project/src/db.rs

@@ -96,7 +96,8 @@ impl ProjectDatabase {
         // https://salsa.zulipchat.com/#narrow/stream/333573-salsa-3.2E0/topic/Expose.20an.20API.20to.20cancel.20other.20queries
         let _ = self.zalsa_mut();
 
-        Arc::get_mut(&mut self.system).unwrap()
+        Arc::get_mut(&mut self.system)
+            .expect("ref count should be 1 because `zalsa_mut` drops all other DB references.")
     }
 
     pub(crate) fn with_db<F, T>(&self, f: F) -> Result<T, Cancelled>

crates/ty_server/src/server/api.rs

@@ -1,6 +1,7 @@
 use crate::server::schedule::Task;
 use crate::session::Session;
 use crate::system::{AnySystemPath, url_to_any_system_path};
+use anyhow::anyhow;
 use lsp_server as server;
 use lsp_types::notification::Notification;
 
@@ -44,7 +45,11 @@ pub(super) fn request<'a>(req: server::Request) -> Task<'a> {
 
         method => {
             tracing::warn!("Received request {method} which does not have a handler");
-            return Task::nothing();
+            let result: Result<()> = Err(Error::new(
+                anyhow!("Unknown request"),
+                server::ErrorCode::MethodNotFound,
+            ));
+            return Task::immediate(id, result);
         }
     }
     .unwrap_or_else(|err| {
@@ -120,9 +125,10 @@ fn background_request_task<'a, R: traits::BackgroundDocumentRequestHandler>(
         let url = R::document_url(&params).into_owned();
 
         let Ok(path) = url_to_any_system_path(&url) else {
+            tracing::warn!("Ignoring request for invalid `{url}`");
             return Box::new(|_, _| {});
         };
-        let db = match path {
+        let db = match &path {
             AnySystemPath::System(path) => match session.project_db_for_path(path.as_std_path()) {
                 Some(db) => db.clone(),
                 None => session.default_project_db().clone(),
@@ -131,12 +137,13 @@ fn background_request_task<'a, R: traits::BackgroundDocumentRequestHandler>(
         };
 
         let Some(snapshot) = session.take_snapshot(url) else {
+            tracing::warn!("Ignoring request because snapshot for path `{path:?}` doesn't exist.");
             return Box::new(|_, _| {});
         };
 
         Box::new(move |notifier, responder| {
             let _span = tracing::trace_span!("request", %id, method = R::METHOD).entered();
-            let result = R::run_with_snapshot(snapshot, db, notifier, params);
+            let result = R::run_with_snapshot(&db, snapshot, notifier, params);
             respond::<R>(id, result, &responder);
         })
     }))
@@ -162,8 +169,11 @@ fn background_notification_thread<'a, N: traits::BackgroundDocumentNotificationH
 ) -> super::Result<Task<'a>> {
     let (id, params) = cast_notification::<N>(req)?;
     Ok(Task::background(schedule, move |session: &Session| {
-        // TODO(jane): we should log an error if we can't take a snapshot.
+        let url = N::document_url(&params);
-        let Some(snapshot) = session.take_snapshot(N::document_url(&params).into_owned()) else {
+        let Some(snapshot) = session.take_snapshot((*url).clone()) else {
+            tracing::debug!(
+                "Ignoring notification because snapshot for url `{url}` doesn't exist."
+            );
             return Box::new(|_, _| {});
         };
         Box::new(move |notifier, _| {

crates/ty_server/src/server/api/requests/completion.rs

@@ -23,24 +23,24 @@ impl BackgroundDocumentRequestHandler for CompletionRequestHandler {
     }
 
     fn run_with_snapshot(
+        db: &ProjectDatabase,
         snapshot: DocumentSnapshot,
-        db: ProjectDatabase,
         _notifier: Notifier,
         params: CompletionParams,
     ) -> crate::server::Result<Option<CompletionResponse>> {
-        let Some(file) = snapshot.file(&db) else {
+        let Some(file) = snapshot.file(db) else {
             tracing::debug!("Failed to resolve file for {:?}", params);
             return Ok(None);
         };
 
-        let source = source_text(&db, file);
+        let source = source_text(db, file);
-        let line_index = line_index(&db, file);
+        let line_index = line_index(db, file);
         let offset = params.text_document_position.position.to_text_size(
             &source,
             &line_index,
             snapshot.encoding(),
         );
-        let completions = completion(&db, file, offset);
+        let completions = completion(db, file, offset);
         if completions.is_empty() {
             return Ok(None);
         }

crates/ty_server/src/server/api/requests/diagnostic.rs

@@ -29,12 +29,12 @@ impl BackgroundDocumentRequestHandler for DocumentDiagnosticRequestHandler {
     }
 
     fn run_with_snapshot(
+        db: &ProjectDatabase,
         snapshot: DocumentSnapshot,
-        db: ProjectDatabase,
         _notifier: Notifier,
         _params: DocumentDiagnosticParams,
     ) -> Result<DocumentDiagnosticReportResult> {
-        let diagnostics = compute_diagnostics(&snapshot, &db);
+        let diagnostics = compute_diagnostics(&snapshot, db);
 
         Ok(DocumentDiagnosticReportResult::Report(
             DocumentDiagnosticReport::Full(RelatedFullDocumentDiagnosticReport {

crates/ty_server/src/server/api/requests/goto_type_definition.rs

@@ -23,25 +23,25 @@ impl BackgroundDocumentRequestHandler for GotoTypeDefinitionRequestHandler {
     }
 
     fn run_with_snapshot(
+        db: &ProjectDatabase,
         snapshot: DocumentSnapshot,
-        db: ProjectDatabase,
         _notifier: Notifier,
         params: GotoTypeDefinitionParams,
     ) -> crate::server::Result<Option<GotoDefinitionResponse>> {
-        let Some(file) = snapshot.file(&db) else {
+        let Some(file) = snapshot.file(db) else {
             tracing::debug!("Failed to resolve file for {:?}", params);
             return Ok(None);
         };
 
-        let source = source_text(&db, file);
+        let source = source_text(db, file);
-        let line_index = line_index(&db, file);
+        let line_index = line_index(db, file);
         let offset = params.text_document_position_params.position.to_text_size(
             &source,
             &line_index,
             snapshot.encoding(),
         );
 
-        let Some(ranged) = goto_type_definition(&db, file, offset) else {
+        let Some(ranged) = goto_type_definition(db, file, offset) else {
             return Ok(None);
         };
 
@@ -52,14 +52,14 @@ impl BackgroundDocumentRequestHandler for GotoTypeDefinitionRequestHandler {
             let src = Some(ranged.range);
             let links: Vec<_> = ranged
                 .into_iter()
-                .filter_map(|target| target.to_link(&db, src, snapshot.encoding()))
+                .filter_map(|target| target.to_link(db, src, snapshot.encoding()))
                 .collect();
 
             Ok(Some(GotoDefinitionResponse::Link(links)))
         } else {
             let locations: Vec<_> = ranged
                 .into_iter()
-                .filter_map(|target| target.to_location(&db, snapshot.encoding()))
+                .filter_map(|target| target.to_location(db, snapshot.encoding()))
                 .collect();
 
             Ok(Some(GotoDefinitionResponse::Array(locations)))

crates/ty_server/src/server/api/requests/hover.rs

@@ -23,25 +23,25 @@ impl BackgroundDocumentRequestHandler for HoverRequestHandler {
     }
 
     fn run_with_snapshot(
+        db: &ProjectDatabase,
         snapshot: DocumentSnapshot,
-        db: ProjectDatabase,
         _notifier: Notifier,
         params: HoverParams,
     ) -> crate::server::Result<Option<lsp_types::Hover>> {
-        let Some(file) = snapshot.file(&db) else {
+        let Some(file) = snapshot.file(db) else {
             tracing::debug!("Failed to resolve file for {:?}", params);
             return Ok(None);
         };
 
-        let source = source_text(&db, file);
+        let source = source_text(db, file);
-        let line_index = line_index(&db, file);
+        let line_index = line_index(db, file);
         let offset = params.text_document_position_params.position.to_text_size(
             &source,
             &line_index,
             snapshot.encoding(),
         );
 
-        let Some(range_info) = hover(&db, file, offset) else {
+        let Some(range_info) = hover(db, file, offset) else {
             return Ok(None);
         };
 
@@ -54,7 +54,7 @@ impl BackgroundDocumentRequestHandler for HoverRequestHandler {
             (MarkupKind::PlainText, lsp_types::MarkupKind::PlainText)
         };
 
-        let contents = range_info.display(&db, markup_kind).to_string();
+        let contents = range_info.display(db, markup_kind).to_string();
 
         Ok(Some(lsp_types::Hover {
             contents: HoverContents::Markup(MarkupContent {

crates/ty_server/src/server/api/requests/inlay_hints.rs

@@ -22,24 +22,24 @@ impl BackgroundDocumentRequestHandler for InlayHintRequestHandler {
     }
 
     fn run_with_snapshot(
+        db: &ProjectDatabase,
         snapshot: DocumentSnapshot,
-        db: ProjectDatabase,
         _notifier: Notifier,
         params: InlayHintParams,
     ) -> crate::server::Result<Option<Vec<lsp_types::InlayHint>>> {
-        let Some(file) = snapshot.file(&db) else {
+        let Some(file) = snapshot.file(db) else {
             tracing::debug!("Failed to resolve file for {:?}", params);
             return Ok(None);
         };
 
-        let index = line_index(&db, file);
+        let index = line_index(db, file);
-        let source = source_text(&db, file);
+        let source = source_text(db, file);
 
         let range = params
             .range
             .to_text_range(&source, &index, snapshot.encoding());
 
-        let inlay_hints = inlay_hints(&db, file, range);
+        let inlay_hints = inlay_hints(db, file, range);
 
         let inlay_hints = inlay_hints
             .into_iter()
@@ -47,7 +47,7 @@ impl BackgroundDocumentRequestHandler for InlayHintRequestHandler {
                 position: hint
                     .position
                     .to_position(&source, &index, snapshot.encoding()),
-                label: lsp_types::InlayHintLabel::String(hint.display(&db).to_string()),
+                label: lsp_types::InlayHintLabel::String(hint.display(db).to_string()),
                 kind: Some(lsp_types::InlayHintKind::TYPE),
                 tooltip: None,
                 padding_left: None,

crates/ty_server/src/server/api/traits.rs

@@ -34,8 +34,8 @@ pub(super) trait BackgroundDocumentRequestHandler: RequestHandler {
     ) -> std::borrow::Cow<lsp_types::Url>;
 
     fn run_with_snapshot(
+        db: &ProjectDatabase,
         snapshot: DocumentSnapshot,
-        db: ProjectDatabase,
         notifier: Notifier,
         params: <<Self as RequestHandler>::RequestType as Request>::Params,
     ) -> super::Result<<<Self as RequestHandler>::RequestType as Request>::Result>;