refactor: serve read requests from session snapshots (#673)

Author: joshuadavidthomas

crates/djls-server/Cargo.toml

@@ -19,14 +19,14 @@ percent-encoding = { workspace = true }
 rustc-hash = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
+salsa = { workspace = true }
 tokio = { workspace = true }
 tower-lsp-server = { workspace = true }
 tracing = { workspace = true }
 tracing-appender = { workspace = true }
 tracing-subscriber = { workspace = true }
 
 [dev-dependencies]
-salsa = { workspace = true }
 tempfile = { workspace = true }
 
 [build-dependencies]

crates/djls-server/src/server.rs

@@ -1,4 +1,5 @@
 use std::future::Future;
+use std::panic::AssertUnwindSafe;
 use std::sync::Arc;
 
 use djls_project::Db as ProjectDb;
@@ -17,6 +18,9 @@ use crate::ext::UriExt;
 use crate::logging::LoggingGuard;
 use crate::queue::Queue;
 use crate::session::Session;
+use crate::session::SessionSnapshot;
+
+const SNAPSHOT_CANCEL_RETRIES: usize = 2;
 
 pub(crate) struct DjangoLanguageServer {
     client: Client,
@@ -52,6 +56,47 @@ impl DjangoLanguageServer {
         f(&mut session)
     }
 
+    /// Capture a snapshot under a brief lock, then compute on the blocking
+    /// pool so the single-threaded event loop stays responsive.
+    pub(crate) async fn with_snapshot<F, R>(&self, f: F) -> R
+    where
+        F: Fn(&SessionSnapshot) -> R + Send + Sync + 'static,
+        R: Default + Send + 'static,
+    {
+        let f = Arc::new(f);
+
+        for attempt in 0..=SNAPSHOT_CANCEL_RETRIES {
+            let snapshot = { self.session.lock().await.snapshot() };
+            let f = Arc::clone(&f);
+            let result = tokio::task::spawn_blocking(move || {
+                salsa::Cancelled::catch(AssertUnwindSafe(|| f(&snapshot)))
+            })
+            .await
+            .expect("snapshot task must not panic");
+
+            match result {
+                Ok(result) => return result,
+                Err(cancelled) if attempt < SNAPSHOT_CANCEL_RETRIES => {
+                    tracing::debug!(
+                        ?cancelled,
+                        attempt = attempt + 1,
+                        "Snapshot request cancelled; retrying with fresh snapshot"
+                    );
+                }
+                Err(cancelled) => {
+                    tracing::debug!(
+                        ?cancelled,
+                        retries = SNAPSHOT_CANCEL_RETRIES,
+                        "Snapshot request cancelled; returning fallback"
+                    );
+                    return R::default();
+                }
+            }
+        }
+
+        unreachable!("snapshot retry loop must return")
+    }
+
     pub(crate) async fn with_session_mut_task<F, Fut>(
         &self,
         f: F,
@@ -85,14 +130,15 @@ impl DjangoLanguageServer {
     }
 
     async fn maybe_push_diagnostics(&self, document: &TextDocument) {
+        let file = document.file();
         let Some(diagnostics) = self
-            .with_session(|session| {
-                if session.client_info().supports_pull_diagnostics() {
+            .with_snapshot(move |snapshot| {
+                if snapshot.client_info().supports_pull_diagnostics() {
                     tracing::debug!("Client supports pull diagnostics, skipping push");
                     return None;
                 }
 
-                djls_ide::collect_diagnostics(session.db(), document.file())
+                djls_ide::collect_diagnostics(snapshot.db(), file)
             })
             .await
         else {
@@ -260,13 +306,13 @@ impl LanguageServer for DjangoLanguageServer {
         params: ls_types::CompletionParams,
     ) -> LspResult<Option<ls_types::CompletionResponse>> {
         let response = self
-            .with_session(|session| {
-                let (file, offset) = session.position_for_document_request(
+            .with_snapshot(move |snapshot| {
+                let (file, offset) = snapshot.position_for_document_request(
                     &params.text_document_position.text_document,
                     params.text_document_position.position,
                     "completion",
                 )?;
-                let db = session.db();
+                let db = snapshot.db();
 
                 if *file.source(db).kind() != FileKind::Template {
                     return None;
@@ -276,8 +322,8 @@ impl LanguageServer for DjangoLanguageServer {
                     db,
                     file,
                     offset,
-                    session.client_info().position_encoding(),
-                    session.client_info().supports_snippets(),
+                    snapshot.client_info().position_encoding(),
+                    snapshot.client_info().supports_snippets(),
                 )
             })
             .await;
@@ -287,13 +333,13 @@ impl LanguageServer for DjangoLanguageServer {
 
     async fn hover(&self, params: ls_types::HoverParams) -> LspResult<Option<ls_types::Hover>> {
         let response = self
-            .with_session(|session| {
-                let (file, offset) = session.position_for_document_request(
+            .with_snapshot(move |snapshot| {
+                let (file, offset) = snapshot.position_for_document_request(
                     &params.text_document_position_params.text_document,
                     params.text_document_position_params.position,
                     "hover",
                 )?;
-                let db = session.db();
+                let db = snapshot.db();
 
                 if *file.source(db).kind() != FileKind::Template {
                     return None;
@@ -316,14 +362,14 @@ impl LanguageServer for DjangoLanguageServer {
         );
 
         let diagnostics = self
-            .with_session(|session| {
+            .with_snapshot(move |snapshot| {
                 let Some(file) =
-                    session.file_for_document_request(&params.text_document, "diagnostic")
+                    snapshot.file_for_document_request(&params.text_document, "diagnostic")
                 else {
                     return Vec::new();
                 };
 
-                djls_ide::collect_diagnostics(session.db(), file).unwrap_or_default()
+                djls_ide::collect_diagnostics(snapshot.db(), file).unwrap_or_default()
             })
             .await;
 
@@ -345,13 +391,13 @@ impl LanguageServer for DjangoLanguageServer {
         params: ls_types::FoldingRangeParams,
     ) -> LspResult<Option<Vec<ls_types::FoldingRange>>> {
         let ranges = self
-            .with_session(|session| {
+            .with_snapshot(move |snapshot| {
                 let Some(file) =
-                    session.file_for_document_request(&params.text_document, "folding")
+                    snapshot.file_for_document_request(&params.text_document, "folding")
                 else {
                     return Vec::new();
                 };
-                let db = session.db();
+                let db = snapshot.db();
 
                 if *file.source(db).kind() != FileKind::Template {
                     return Vec::new();
@@ -369,13 +415,13 @@ impl LanguageServer for DjangoLanguageServer {
         params: ls_types::DocumentSymbolParams,
     ) -> LspResult<Option<ls_types::DocumentSymbolResponse>> {
         let symbols = self
-            .with_session(|session| {
+            .with_snapshot(move |snapshot| {
                 let Some(file) =
-                    session.file_for_document_request(&params.text_document, "document symbol")
+                    snapshot.file_for_document_request(&params.text_document, "document symbol")
                 else {
                     return Vec::new();
                 };
-                let db = session.db();
+                let db = snapshot.db();
 
                 if *file.source(db).kind() != FileKind::Template {
                     return Vec::new();
@@ -393,13 +439,13 @@ impl LanguageServer for DjangoLanguageServer {
         params: ls_types::GotoDefinitionParams,
     ) -> LspResult<Option<ls_types::GotoDefinitionResponse>> {
         let response = self
-            .with_session(|session| {
-                let (file, offset) = session.position_for_document_request(
+            .with_snapshot(move |snapshot| {
+                let (file, offset) = snapshot.position_for_document_request(
                     &params.text_document_position_params.text_document,
                     params.text_document_position_params.position,
                     "goto definition",
                 )?;
-                let db = session.db();
+                let db = snapshot.db();
 
                 if *file.source(db).kind() != FileKind::Template {
                     return None;
@@ -417,13 +463,13 @@ impl LanguageServer for DjangoLanguageServer {
         params: ls_types::ReferenceParams,
     ) -> LspResult<Option<Vec<ls_types::Location>>> {
         let response = self
-            .with_session(|session| {
-                let (file, offset) = session.position_for_document_request(
+            .with_snapshot(move |snapshot| {
+                let (file, offset) = snapshot.position_for_document_request(
                     &params.text_document_position.text_document,
                     params.text_document_position.position,
                     "references",
                 )?;
-                let db = session.db();
+                let db = snapshot.db();
 
                 if *file.source(db).kind() != FileKind::Template {
                     return None;
@@ -441,13 +487,13 @@ impl LanguageServer for DjangoLanguageServer {
         params: ls_types::DocumentFormattingParams,
     ) -> LspResult<Option<Vec<ls_types::TextEdit>>> {
         let edits = self
-            .with_session(|session| {
+            .with_snapshot(move |snapshot| {
                 let Some(file) =
-                    session.file_for_document_request(&params.text_document, "formatting")
+                    snapshot.file_for_document_request(&params.text_document, "formatting")
                 else {
                     return Vec::new();
                 };
-                let db = session.db();
+                let db = snapshot.db();
                 let format_config = db.settings().format().clone();
 
                 if !format_config.enabled() {
@@ -462,7 +508,7 @@ impl LanguageServer for DjangoLanguageServer {
                 djls_ide::format_document(
                     db,
                     file,
-                    session.client_info().position_encoding(),
+                    snapshot.client_info().position_encoding(),
                     format_config.backend(),
                     &params.options,
                 )

crates/djls-server/src/session.rs

@@ -86,8 +86,7 @@ impl Session {
         }
     }
 
-    #[cfg(test)]
-    fn snapshot(&self) -> SessionSnapshot {
+    pub(crate) fn snapshot(&self) -> SessionSnapshot {
         SessionSnapshot::new(self.db.clone(), self.client_info.clone())
     }
 
@@ -196,40 +195,26 @@ impl Session {
     /// Open editor buffers are exposed to Salsa through the workspace overlay,
     /// so feature code should read current text through [`File::source`]
     /// instead of reaching back into [`TextDocument`] state.
+    #[allow(dead_code)]
     pub(crate) fn file_for_document_request(
         &self,
         text_document: &ls_types::TextDocumentIdentifier,
         request: &str,
     ) -> Option<File> {
-        let Some(path) = text_document.uri.to_utf8_path_buf() else {
-            tracing::debug!(
-                "Skipping non-file URI in {} request: {}",
-                request,
-                text_document.uri.as_str()
-            );
-            return None;
-        };
-
-        Some(self.db.get_or_create_file(&path))
+        self.snapshot()
+            .file_for_document_request(text_document, request)
     }
 
     /// Resolve an LSP positioned document request to a tracked file and byte offset.
+    #[allow(dead_code)]
     pub(crate) fn position_for_document_request(
         &self,
         text_document: &ls_types::TextDocumentIdentifier,
         position: ls_types::Position,
         request: &str,
     ) -> Option<(File, Offset)> {
-        let file = self.file_for_document_request(text_document, request)?;
-        let source = file.source(&self.db);
-        let line_index = file.line_index(&self.db);
-        let offset = position.to_offset(
-            source.as_str(),
-            line_index,
-            self.client_info.position_encoding(),
-        );
-
-        Some((file, offset))
+        self.snapshot()
+            .position_for_document_request(text_document, position, request)
     }
 
     /// Get all currently open documents.
@@ -248,27 +233,66 @@ impl Default for Session {
     }
 }
 
-/// Immutable snapshot of session state for tests.
-#[cfg(test)]
+/// Immutable snapshot of session state.
 #[derive(Clone)]
-struct SessionSnapshot {
+pub(crate) struct SessionSnapshot {
     db: DjangoDatabase,
     client_info: ClientInfo,
 }
 
-#[cfg(test)]
 impl SessionSnapshot {
     fn new(db: DjangoDatabase, client_info: ClientInfo) -> Self {
         Self { db, client_info }
     }
 
-    fn db(&self) -> &DjangoDatabase {
+    pub(crate) fn db(&self) -> &DjangoDatabase {
         &self.db
     }
 
-    fn client_info(&self) -> &ClientInfo {
+    pub(crate) fn client_info(&self) -> &ClientInfo {
         &self.client_info
     }
+
+    /// Resolve an LSP document request to the tracked file for that URI.
+    ///
+    /// Open editor buffers are exposed to Salsa through the workspace overlay,
+    /// so feature code should read current text through [`File::source`]
+    /// instead of reaching back into [`TextDocument`] state.
+    pub(crate) fn file_for_document_request(
+        &self,
+        text_document: &ls_types::TextDocumentIdentifier,
+        request: &str,
+    ) -> Option<File> {
+        let Some(path) = text_document.uri.to_utf8_path_buf() else {
+            tracing::debug!(
+                "Skipping non-file URI in {} request: {}",
+                request,
+                text_document.uri.as_str()
+            );
+            return None;
+        };
+
+        Some(self.db.get_or_create_file(&path))
+    }
+
+    /// Resolve an LSP positioned document request to a tracked file and byte offset.
+    pub(crate) fn position_for_document_request(
+        &self,
+        text_document: &ls_types::TextDocumentIdentifier,
+        position: ls_types::Position,
+        request: &str,
+    ) -> Option<(File, Offset)> {
+        let file = self.file_for_document_request(text_document, request)?;
+        let source = file.source(&self.db);
+        let line_index = file.line_index(&self.db);
+        let offset = position.to_offset(
+            source.as_str(),
+            line_index,
+            self.client_info.position_encoding(),
+        );
+
+        Some((file, offset))
+    }
 }
 
 #[cfg(test)]