[ty] Fix server panic in workspace diagnostics request handler when typing (#19631)

crates/ty_server/src/session.rs

@@ -82,7 +82,6 @@ pub(crate) struct Session {
 
 /// LSP State for a Project
 pub(crate) struct ProjectState {
-    pub(crate) db: ProjectDatabase,
     /// Files that we have outstanding otherwise-untracked pushed diagnostics for.
     ///
     /// In `CheckMode::OpenFiles` we still read some files that the client hasn't
@@ -97,6 +96,14 @@ pub(crate) struct ProjectState {
     /// for so that we can clear the diagnostics for all of them before we go
     /// to update any of them.
     pub(crate) untracked_files_with_pushed_diagnostics: Vec<Url>,
+
+    // Note: This field should be last to ensure the `db` gets dropped last.
+    // The db drop order matters because we call `Arc::into_inner` on some Arc's
+    // and we use Salsa's cancellation to guarantee that there's only a single reference to the `Arc`.
+    // However, this requires that the db drops last.
+    // This shouldn't matter here because the db's stored in the session are the
+    // only reference we want to hold on, but better be safe than sorry ;).
+    pub(crate) db: ProjectDatabase,
 }
 
 impl Session {
@@ -639,10 +646,20 @@ impl DocumentSnapshot {
 
 /// An immutable snapshot of the current state of [`Session`].
 pub(crate) struct SessionSnapshot {
-    projects: Vec<ProjectDatabase>,
     index: Arc<Index>,
     position_encoding: PositionEncoding,
     resolved_client_capabilities: ResolvedClientCapabilities,
+
+    /// IMPORTANT: It's important that the databases come last, or at least,
+    /// after any `Arc` that we try to extract or mutate in-place using `Arc::into_inner`
+    /// and that relies on Salsa's cancellation to guarantee that there's now only a
+    /// single reference to it (e.g. see [`Session::index_mut`]).
+    ///
+    /// Making this field come last guarantees that the db's `Drop` handler is
+    /// dropped after all other fields, which ensures that
+    /// Salsa's cancellation blocks until all fields are dropped (and not only
+    /// waits for the db to be dropped while we still hold on to the `Index`).
+    projects: Vec<ProjectDatabase>,
 }
 
 impl SessionSnapshot {