perf(workspace): one memory budget for the caches of the open documents

crates/pgls_workspace/src/workspace/server.rs

@@ -18,7 +18,8 @@ use connection_manager::ConnectionManager;
 #[cfg(feature = "db")]
 use document::TypecheckDiagnosticsMapper;
 use document::{DefaultMapper, Document, ExecuteStatementMapper};
-use document_cache::CacheKey;
+use document_cache::{CacheBudget, CacheKey};
+pub use document_cache::{DOCUMENT_CACHE_ENTRIES_MAX, DocumentCacheStats};
 #[cfg(feature = "db")]
 use futures::{StreamExt, TryStreamExt, stream};
 use pgls_analyse::{
@@ -109,6 +110,10 @@ pub struct WorkspaceServer {
     /// Incremented whenever the settings change, see [`CacheKey`].
     settings_version: AtomicU64,
 
+    /// What the caches of the open documents keep together, see
+    /// [`DOCUMENT_CACHE_ENTRIES_MAX`].
+    document_caches: CacheBudget<PgLSPath>,
+
     #[cfg(feature = "db")]
     connection: ConnectionManager,
 
@@ -142,6 +147,7 @@ impl WorkspaceServer {
             snapshot: SnapshotManager::new(),
             catalog_base: Mutex::default(),
             settings_version: AtomicU64::default(),
+            document_caches: CacheBudget::new(DOCUMENT_CACHE_ENTRIES_MAX),
             connection: ConnectionManager::new(),
             #[cfg(test)]
             walked_statements: Default::default(),
@@ -157,6 +163,7 @@ impl WorkspaceServer {
             snapshot: SnapshotManager::new(),
             catalog_base: Mutex::default(),
             settings_version: AtomicU64::default(),
+            document_caches: CacheBudget::new(DOCUMENT_CACHE_ENTRIES_MAX),
             #[cfg(test)]
             walked_statements: Default::default(),
         }
@@ -175,6 +182,33 @@ impl WorkspaceServer {
         handle
     }
 
+    /// What the caches of the open document at `path` hold, for tests of their memory.
+    #[doc(hidden)]
+    pub fn document_cache_stats(&self, path: &PgLSPath) -> Option<DocumentCacheStats> {
+        let documents = self.documents.read().unwrap();
+        documents.get(path).map(Document::cache_stats)
+    }
+
+    /// Records what the caches of the document at `path` keep after a request, and drops the
+    /// caches of the least recently used other documents to stay within
+    /// [`DOCUMENT_CACHE_ENTRIES_MAX`]. The document of the request is never dropped, so a
+    /// single file larger than the budget stays cached while it is used.
+    fn charge_document_caches(&self, documents: &HashMap<PgLSPath, Document>, path: &PgLSPath) {
+        let Some(doc) = documents.get(path) else {
+            return;
+        };
+        for evicted in self.document_caches.record(path, doc.cache().usage()) {
+            if let Some(doc) = documents.get(&evicted) {
+                tracing::debug!(
+                    "Dropping the caches of {} to keep at most {} cache entries",
+                    evicted.display(),
+                    self.document_caches.max()
+                );
+                doc.cache().clear();
+            }
+        }
+    }
+
     /// The database snapshot of the database in `settings`, if it is reachable or a snapshot was
     /// loaded from JSON.
     fn database_snapshot(&self, settings: &Settings) -> Option<Arc<Snapshot>> {
@@ -662,6 +696,7 @@ impl Workspace for WorkspaceServer {
         documents
             .remove(&params.path)
             .ok_or_else(WorkspaceError::not_found)?;
+        self.document_caches.remove(&params.path);
 
         Ok(())
     }
@@ -784,6 +819,7 @@ impl Workspace for WorkspaceServer {
             disabled_reason: invalidate_disabled_reason,
         });
 
+        self.charge_document_caches(&documents, &params.path);
         Ok(CodeActionsResult { actions })
     }
 
@@ -885,6 +921,8 @@ impl Workspace for WorkspaceServer {
             .ok_or(WorkspaceError::not_found())?;
 
         let stmt = parser.find(params.statement_id, ExecuteStatementMapper);
+        // Finding the statement parsed the statements before it.
+        self.charge_document_caches(&documents, &params.path);
 
         if stmt.is_none() {
             return Ok(ExecuteStatementResult {
@@ -1110,6 +1148,7 @@ impl Workspace for WorkspaceServer {
         diagnostics.retain(|d| !suppressions.is_suppressed(d));
         diagnostics.extend(suppression_errors.into_iter().map(SDiagnostic::new));
 
+        self.charge_document_caches(&documents, &params.path);
         info!("Pulled {:?} diagnostic(s)", diagnostics.len());
         Ok(PullDiagnosticsResult {
             diagnostics,
@@ -1209,12 +1248,14 @@ impl Workspace for WorkspaceServer {
             .get(&params.path)
             .ok_or(WorkspaceError::not_found())?;
 
-        Ok(formatting::format_document(
+        let result = formatting::format_document(
             doc,
             &settings.formatter,
             params.range,
             &params.path.as_path().display().to_string(),
-        ))
+        );
+        self.charge_document_caches(&documents, &params.path);
+        Ok(result)
     }
 
     #[cfg(feature = "db")]
@@ -1317,6 +1358,7 @@ impl Workspace for WorkspaceServer {
                     text: stmt.id.content().to_string(),
                 });
 
+                self.charge_document_caches(&documents, &params.path);
                 Ok(CompletionsResult::in_document(items, &stmt))
             }
         }
@@ -1371,6 +1413,7 @@ impl Workspace for WorkspaceServer {
                 })
                 .unwrap_or_default();
 
+                self.charge_document_caches(&documents, &params.path);
                 Ok(OnHoverResult { markdown_blocks })
             }
             None => Ok(OnHoverResult::default()),
@@ -1439,6 +1482,7 @@ impl Workspace for WorkspaceServer {
         }
 
         hints.sort_by_key(|hint| hint.offset);
+        self.charge_document_caches(&documents, &params.path);
         Ok(PullInlayHintsResult { hints })
     }
 
@@ -1516,6 +1560,7 @@ impl Workspace for WorkspaceServer {
             }
             keep
         });
+        self.charge_document_caches(&documents, &params.path);
         Ok(PullSemanticTokensResult { tokens })
     }
 }

crates/pgls_workspace/src/workspace/server.tests.rs

@@ -2594,6 +2594,132 @@ fn test_no_concurrently_fix_in_a_transaction() {
     );
 }
 
+/// When the open documents together keep more than the budget, the caches of the least recently
+/// used ones are dropped, and recomputed when they are used again.
+#[test]
+fn document_caches_stay_within_the_budget() {
+    let (mut workspace, project_key) =
+        get_test_workspace(None).expect("Unable to create test workspace");
+    // Each document keeps a parse result for each of its six statements.
+    workspace.document_caches = super::document_cache::CacheBudget::new(8);
+
+    let content = "drop table a;\ndrop table b;\ndrop table c;\nselect 1;\nselect 2;\nselect 3;\n";
+    let paths = [PgLSPath::new("a.sql"), PgLSPath::new("b.sql")];
+    for path in &paths {
+        workspace
+            .open_file(OpenFileParams {
+                project_key,
+                path: path.clone(),
+                content: content.into(),
+                version: 1,
+            })
+            .expect("Unable to open test file");
+    }
+    let diagnostics = |path: &PgLSPath| {
+        workspace
+            .pull_file_diagnostics(crate::workspace::PullFileDiagnosticsParams {
+                project_key,
+                path: path.clone(),
+                categories: RuleCategories::all(),
+                max_diagnostics: 100,
+                only: vec![],
+                skip: vec![],
+            })
+            .expect("Unable to pull diagnostics")
+            .diagnostics
+            .len()
+    };
+    let cached = |path: &PgLSPath| {
+        let stats = workspace.document_cache_stats(path).unwrap();
+        (stats.parse_results, stats.analysis)
+    };
+
+    let expected = diagnostics(&paths[0]);
+    assert_eq!(expected, 3);
+    assert_eq!(cached(&paths[0]), (6, true));
+
+    // Both documents together keep 12 parse results, so the first one's caches are dropped.
+    assert_eq!(diagnostics(&paths[1]), expected);
+    assert_eq!(cached(&paths[0]), (0, false));
+    assert_eq!(cached(&paths[1]), (6, true));
+
+    // And recomputed when it is used again.
+    assert_eq!(diagnostics(&paths[0]), expected);
+    assert_eq!(cached(&paths[0]), (6, true));
+    assert_eq!(cached(&paths[1]), (0, false));
+
+    workspace
+        .close_file(crate::workspace::CloseFileParams {
+            project_key,
+            path: paths[0].clone(),
+        })
+        .expect("Unable to close test file");
+    assert_eq!(workspace.document_caches.total(), 0);
+}
+
+/// Executing a statement parses the statements before it to find it, and charges them to the
+/// budget, also when it returns early, like without a database connection.
+#[test]
+fn executing_a_statement_charges_the_document_caches() {
+    let (mut workspace, project_key) =
+        get_test_workspace(None).expect("Unable to create test workspace");
+    workspace.document_caches = super::document_cache::CacheBudget::new(8);
+
+    let content = "drop table a;\ndrop table b;\ndrop table c;\nselect 1;\nselect 2;\nselect 3;\n";
+    let paths = [PgLSPath::new("a.sql"), PgLSPath::new("b.sql")];
+    for path in &paths {
+        workspace
+            .open_file(OpenFileParams {
+                project_key,
+                path: path.clone(),
+                content: content.into(),
+                version: 1,
+            })
+            .expect("Unable to open test file");
+    }
+    workspace
+        .pull_file_diagnostics(crate::workspace::PullFileDiagnosticsParams {
+            project_key,
+            path: paths[0].clone(),
+            categories: RuleCategories::all(),
+            max_diagnostics: 100,
+            only: vec![],
+            skip: vec![],
+        })
+        .expect("Unable to pull diagnostics");
+    assert_eq!(
+        workspace
+            .document_cache_stats(&paths[0])
+            .unwrap()
+            .parse_results,
+        6
+    );
+
+    workspace
+        .execute_statement(ExecuteStatementParams {
+            project_key,
+            path: paths[1].clone(),
+            statement_id: StatementId::new("select 3;"),
+        })
+        .expect("Unable to execute statement");
+
+    // Both documents together keep 12 parse results, so the first one's caches are dropped.
+    assert_eq!(
+        workspace
+            .document_cache_stats(&paths[1])
+            .unwrap()
+            .parse_results,
+        6
+    );
+    assert_eq!(
+        workspace
+            .document_cache_stats(&paths[0])
+            .unwrap()
+            .parse_results,
+        0
+    );
+}
+
 /// Code actions are requested on every cursor move. They reuse the analysis that the
 /// diagnostics made until the document or the settings change.
 #[test]
@@ -4871,8 +4997,13 @@ async fn walks_resume_from_checkpoints_before_the_edit(test_db: PgPool) {
     let (count, _) = walked(&|| hover(&marked, "checkpoint_early;"));
     assert_eq!(count, 2);
 
-    let documents = workspace.documents.read().unwrap();
-    assert!(documents.get(&path).unwrap().cache().walk_size().1 <= 4);
+    assert!(
+        workspace
+            .document_cache_stats(&path)
+            .unwrap()
+            .walk_checkpoints
+            <= 4
+    );
 }
 
 /// After edits anywhere in the file, hover and inlay hints that resume from checkpoints give

crates/pgls_workspace/src/workspace/server/document.rs

@@ -8,7 +8,7 @@ use pgls_text_size::{TextRange, TextSize};
 use rustc_hash::FxHashSet;
 
 use super::{
-    document_cache::{ContentChange, DocumentCache, StaleEntries},
+    document_cache::{ContentChange, DocumentCache, DocumentCacheStats, StaleEntries},
     statement_identifier::StatementId,
 };
 
@@ -95,6 +95,16 @@ impl Document {
         &self.cache
     }
 
+    /// What the caches of the document hold, for tests of their memory. It doesn't fill them.
+    pub(crate) fn cache_stats(&self) -> DocumentCacheStats {
+        let statements: FxHashSet<StatementId> = self
+            .ranges
+            .iter()
+            .map(|range| StatementId::new(&self.content[*range]))
+            .collect();
+        self.cache.stats(&statements)
+    }
+
     pub fn suppressions(&self) -> &Suppressions {
         &self.suppressions
     }