[ty] Cancel background tasks when shutdown is requested (#20039)

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

@@ -3,6 +3,7 @@ use crate::server::api::traits::{RequestHandler, SyncRequestHandler};
 use crate::session::client::Client;
 
 use lsp_types::{WorkspaceDiagnosticReport, WorkspaceDiagnosticReportResult};
+use salsa::Database;
 
 pub(crate) struct ShutdownHandler;
 
@@ -28,6 +29,12 @@ impl SyncRequestHandler for ShutdownHandler {
 
         session.set_shutdown_requested(true);
 
+        // Trigger cancellation for every db to cancel any compute intensive background tasks
+        // (e.g. workspace diagnostics or workspace symbols).
+        for db in session.projects_mut() {
+            db.trigger_cancellation();
+        }
+
         Ok(())
     }
 }

crates/ty_server/src/session.rs

@@ -427,7 +427,7 @@ impl Session {
     /// Returns a mutable iterator over all project databases that have been initialized to this point.
     ///
     /// This iterator will only yield the default project database if it has been used.
-    fn projects_mut(&mut self) -> impl Iterator<Item = &'_ mut ProjectDatabase> + '_ {
+    pub(crate) fn projects_mut(&mut self) -> impl Iterator<Item = &'_ mut ProjectDatabase> + '_ {
         self.project_states_mut().map(|project| &mut project.db)
     }