Rename to SessionSnapshot, move unwind assertion closer (#19177)

This PR addresses the post-merge review comments from
#19041, specifically it:
- Rename `WorkspaceSnapshot` to `SessionSnapshot`
- Rename `take_workspace_snapshot` to `take_session_snapshot`
- Rename `take_snapshot` to `take_document_snapshot`
- Move `AssertUnwindSafe` closer to the `catch_unwind` call which
requires the assertion

Author: dhruvmanila

crates/ty_server/src/server/api.rs

@@ -6,7 +6,7 @@ use lsp_server as server;
 use lsp_server::RequestId;
 use lsp_types::notification::Notification;
 use lsp_types::request::Request;
-use std::panic::UnwindSafe;
+use std::panic::{AssertUnwindSafe, UnwindSafe};
 
 mod diagnostics;
 mod notifications;
@@ -157,7 +157,9 @@ where
             .cancellation_token(&id)
             .expect("request should have been tested for cancellation before scheduling");
 
-        let snapshot = session.take_workspace_snapshot();
+        // SAFETY: The `snapshot` is safe to move across the unwind boundary because it is not used
+        // after unwinding.
+        let snapshot = AssertUnwindSafe(session.take_session_snapshot());
 
         Box::new(move |client| {
             let _span = tracing::debug_span!("request", %id, method = R::METHOD).entered();
@@ -216,7 +218,7 @@ where
             AnySystemPath::SystemVirtual(_) => session.default_project_db().clone(),
         };
 
-        let Some(snapshot) = session.take_snapshot(url) else {
+        let Some(snapshot) = session.take_document_snapshot(url) else {
             tracing::warn!("Ignoring request because snapshot for path `{path:?}` doesn't exist");
             return Box::new(|_| {});
         };
@@ -317,7 +319,7 @@ where
     let (id, params) = cast_notification::<N>(req)?;
     Ok(Task::background(schedule, move |session: &Session| {
         let url = N::document_url(&params);
-        let Some(snapshot) = session.take_snapshot((*url).clone()) else {
+        let Some(snapshot) = session.take_document_snapshot((*url).clone()) else {
             tracing::debug!(
                 "Ignoring notification because snapshot for url `{url}` doesn't exist."
             );

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

@@ -80,7 +80,7 @@ pub(super) fn publish_diagnostics(
     let path = key.path();
 
     let snapshot = session
-        .take_snapshot(url.clone())
+        .take_document_snapshot(url.clone())
         .ok_or_else(|| anyhow::anyhow!("Unable to take snapshot for document with URL {url}"))
         .with_failure_code(lsp_server::ErrorCode::InternalError)?;
 

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

@@ -1,3 +1,5 @@
+use std::panic::AssertUnwindSafe;
+
 use lsp_types::request::WorkspaceDiagnosticRequest;
 use lsp_types::{
     FullDocumentDiagnosticReport, Url, WorkspaceDiagnosticParams, WorkspaceDiagnosticReport,
@@ -12,7 +14,7 @@ use crate::server::api::diagnostics::to_lsp_diagnostic;
 use crate::server::api::traits::{
     BackgroundRequestHandler, RequestHandler, RetriableRequestHandler,
 };
-use crate::session::WorkspaceSnapshot;
+use crate::session::SessionSnapshot;
 use crate::session::client::Client;
 use crate::system::file_to_url;
 
@@ -24,7 +26,7 @@ impl RequestHandler for WorkspaceDiagnosticRequestHandler {
 
 impl BackgroundRequestHandler for WorkspaceDiagnosticRequestHandler {
     fn run(
-        snapshot: WorkspaceSnapshot,
+        snapshot: AssertUnwindSafe<SessionSnapshot>,
         _client: &Client,
         _params: WorkspaceDiagnosticParams,
     ) -> Result<WorkspaceDiagnosticReportResult> {

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

@@ -1,7 +1,9 @@
 //! A stateful LSP implementation that calls into the ty API.
 
+use std::panic::AssertUnwindSafe;
+
 use crate::session::client::Client;
-use crate::session::{DocumentSnapshot, Session, WorkspaceSnapshot};
+use crate::session::{DocumentSnapshot, Session, SessionSnapshot};
 
 use lsp_types::notification::Notification as LSPNotification;
 use lsp_types::request::Request;
@@ -58,7 +60,7 @@ pub(super) trait BackgroundDocumentRequestHandler: RetriableRequestHandler {
 /// A request handler that can be run on a background thread.
 pub(super) trait BackgroundRequestHandler: RetriableRequestHandler {
     fn run(
-        snapshot: WorkspaceSnapshot,
+        snapshot: AssertUnwindSafe<SessionSnapshot>,
         client: &Client,
         params: <<Self as RequestHandler>::RequestType as Request>::Params,
     ) -> super::Result<<<Self as RequestHandler>::RequestType as Request>::Result>;

crates/ty_server/src/session.rs

@@ -2,7 +2,6 @@
 
 use std::collections::{BTreeMap, VecDeque};
 use std::ops::{Deref, DerefMut};
-use std::panic::AssertUnwindSafe;
 use std::sync::Arc;
 
 use anyhow::{Context, anyhow};
@@ -224,14 +223,6 @@ impl Session {
         self.index().key_from_url(url)
     }
 
-    pub(crate) fn take_workspace_snapshot(&self) -> WorkspaceSnapshot {
-        WorkspaceSnapshot {
-            projects: AssertUnwindSafe(self.projects.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());
 
@@ -289,7 +280,7 @@ impl Session {
     /// Creates a document snapshot with the URL referencing the document to snapshot.
     ///
     /// Returns `None` if the url can't be converted to a document key or if the document isn't open.
-    pub fn take_snapshot(&self, url: Url) -> Option<DocumentSnapshot> {
+    pub(crate) fn take_document_snapshot(&self, url: Url) -> Option<DocumentSnapshot> {
         let key = self.key_from_url(url).ok()?;
         Some(DocumentSnapshot {
             resolved_client_capabilities: self.resolved_client_capabilities.clone(),
@@ -299,6 +290,15 @@ impl Session {
         })
     }
 
+    /// Creates a snapshot of the current state of the [`Session`].
+    pub(crate) fn take_session_snapshot(&self) -> SessionSnapshot {
+        SessionSnapshot {
+            projects: self.projects.values().cloned().collect(),
+            index: self.index.clone().unwrap(),
+            position_encoding: self.position_encoding,
+        }
+    }
+
     /// Iterates over the document keys for all open text documents.
     pub(super) fn text_document_keys(&self) -> impl Iterator<Item = DocumentKey> + '_ {
         self.index()
@@ -467,13 +467,13 @@ impl DocumentSnapshot {
 }
 
 /// An immutable snapshot of the current state of [`Session`].
-pub(crate) struct WorkspaceSnapshot {
-    projects: AssertUnwindSafe<Vec<ProjectDatabase>>,
+pub(crate) struct SessionSnapshot {
+    projects: Vec<ProjectDatabase>,
     index: Arc<index::Index>,
     position_encoding: PositionEncoding,
 }
 
-impl WorkspaceSnapshot {
+impl SessionSnapshot {
     pub(crate) fn projects(&self) -> &[ProjectDatabase] {
         &self.projects
     }