From d79d10f52d1554d881ce28f7d491e2e3e83c71e1 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Mon, 26 Oct 2026 17:08:58 +0000 Subject: [PATCH 1/4] perf(workspace): one memory budget for the caches of the open documents A CacheBudget across documents counts what each document's cache keeps (parse results, trees, resolved statements and walk checkpoints) after each request, and drops the caches of the least recently used documents when all of them keep more than DOCUMENT_CACHE_ENTRIES_MAX. Closing a document drops its caches. Two tests count the memory: open documents stay within the budget and nothing is left after closing them, and a long editing session doesn't grow the caches. --- crates/pgls_workspace/src/workspace/server.rs | 49 ++- .../src/workspace/server.tests.rs | 72 +++- .../src/workspace/server/document.rs | 12 +- .../src/workspace/server/document_cache.rs | 80 +++- .../workspace/server/document_cache/budget.rs | 128 +++++++ .../workspace/server/document_cache/parse.rs | 21 ++ .../workspace/server/document_cache/tree.rs | 10 + .../workspace/server/document_cache/walk.rs | 1 - crates/pgls_workspace/tests/common/mod.rs | 70 ++++ crates/pgls_workspace/tests/edit_churn.rs | 352 ++++++++++++++++++ crates/pgls_workspace/tests/memory_bounds.rs | 269 +++++++++++++ 11 files changed, 1051 insertions(+), 13 deletions(-) create mode 100644 crates/pgls_workspace/src/workspace/server/document_cache/budget.rs create mode 100644 crates/pgls_workspace/tests/common/mod.rs create mode 100644 crates/pgls_workspace/tests/edit_churn.rs create mode 100644 crates/pgls_workspace/tests/memory_bounds.rs diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index 190e48972..999e1d510 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/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, + #[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 { + 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, 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> { @@ -662,6 +696,7 @@ impl Workspace for WorkspaceServer { documents .remove(¶ms.path) .ok_or_else(WorkspaceError::not_found)?; + self.document_caches.remove(¶ms.path); Ok(()) } @@ -783,6 +818,7 @@ impl Workspace for WorkspaceServer { disabled_reason: invalidate_disabled_reason, }); + self.charge_document_caches(&documents, ¶ms.path); Ok(CodeActionsResult { actions }) } @@ -1109,6 +1145,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, ¶ms.path); info!("Pulled {:?} diagnostic(s)", diagnostics.len()); Ok(PullDiagnosticsResult { diagnostics, @@ -1208,12 +1245,14 @@ impl Workspace for WorkspaceServer { .get(¶ms.path) .ok_or(WorkspaceError::not_found())?; - Ok(formatting::format_document( + let result = formatting::format_document( doc, &settings.formatter, params.range, ¶ms.path.as_path().display().to_string(), - )) + ); + self.charge_document_caches(&documents, ¶ms.path); + Ok(result) } #[cfg(feature = "db")] @@ -1315,6 +1354,7 @@ impl Workspace for WorkspaceServer { text: stmt.id.content().to_string(), }); + self.charge_document_caches(&documents, ¶ms.path); Ok(CompletionsResult::in_document(items, &stmt)) } } @@ -1369,6 +1409,7 @@ impl Workspace for WorkspaceServer { }) .unwrap_or_default(); + self.charge_document_caches(&documents, ¶ms.path); Ok(OnHoverResult { markdown_blocks }) } None => Ok(OnHoverResult::default()), @@ -1437,6 +1478,7 @@ impl Workspace for WorkspaceServer { } hints.sort_by_key(|hint| hint.offset); + self.charge_document_caches(&documents, ¶ms.path); Ok(PullInlayHintsResult { hints }) } @@ -1514,6 +1556,7 @@ impl Workspace for WorkspaceServer { } keep }); + self.charge_document_caches(&documents, ¶ms.path); Ok(PullSemanticTokensResult { tokens }) } } diff --git a/crates/pgls_workspace/src/workspace/server.tests.rs b/crates/pgls_workspace/src/workspace/server.tests.rs index 08bcadf67..36ed94bc3 100644 --- a/crates/pgls_workspace/src/workspace/server.tests.rs +++ b/crates/pgls_workspace/src/workspace/server.tests.rs @@ -2528,6 +2528,69 @@ 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); +} + /// Code actions are requested on every cursor move. They reuse the analysis that the /// diagnostics made until the document or the settings change. #[test] @@ -4805,8 +4868,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 diff --git a/crates/pgls_workspace/src/workspace/server/document.rs b/crates/pgls_workspace/src/workspace/server/document.rs index 26340d39f..8d007ec63 100644 --- a/crates/pgls_workspace/src/workspace/server/document.rs +++ b/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}, + document_cache::{ContentChange, DocumentCache, DocumentCacheStats}, statement_identifier::StatementId, }; @@ -93,6 +93,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 = self + .ranges + .iter() + .map(|range| StatementId::new(&self.content[*range])) + .collect(); + self.cache.stats(&statements) + } + pub fn suppressions(&self) -> &Suppressions { &self.suppressions } diff --git a/crates/pgls_workspace/src/workspace/server/document_cache.rs b/crates/pgls_workspace/src/workspace/server/document_cache.rs index a1054f054..fa4aa19e6 100644 --- a/crates/pgls_workspace/src/workspace/server/document_cache.rs +++ b/crates/pgls_workspace/src/workspace/server/document_cache.rs @@ -4,24 +4,63 @@ //! was cached for the statements it changed. //! - Per [`CacheKey`]: what also depends on the settings and the database, like the analysis of //! the document. It is dropped when the key changes, and the analysis when the content does. +//! +//! All of it can be recomputed. [`DocumentCache::usage`] says how much a document keeps, and +//! the workspace drops the caches of the least recently used documents when all of them keep more +//! than [`DOCUMENT_CACHE_ENTRIES_MAX`], see [`CacheBudget`]. use std::sync::{Arc, Mutex}; use pgls_catalog::CatalogBase; use pgls_query_ext::{diagnostics::SyntaxDiagnostic, function_body::EmbeddedStatement}; +use rustc_hash::FxHashSet; + use super::{StatementId, analysis::DocumentAnalysis}; use crate::workspace::ProjectKey; +mod budget; mod parse; mod tree; mod walk; +pub(crate) use budget::CacheBudget; pub(crate) use parse::ParseResult; use parse::PgQueryStore; use tree::TreeSitterStore; pub(crate) use walk::{ContentChange, ResolvedStatement, WalkCache}; +/// The number of cache entries that all open documents keep together, beyond the document in +/// use, see [`DocumentCache::usage`]. When they keep more, the caches of the least recently used +/// documents are dropped, and recomputed when they are used again. +/// +/// An entry is a statement's parse result or tree, or the references of a statement, and takes +/// about 10 KB on average (measured on migration files with tables, indexes, SQL functions and +/// queries), so this is about 200 MB. That keeps a few files of several thousand statements +/// fully cached, so that switching between them doesn't parse them again. Without a limit, 20 +/// open files of 5,000 statements each kept 2.7 GB. +pub const DOCUMENT_CACHE_ENTRIES_MAX: usize = 20_000; + +/// What a walk checkpoint counts as in [`DocumentCache::usage`]. It copies the maps of the +/// objects that the file changed before it, about 100 bytes per object, so a checkpoint after +/// the 5,000 tables of a 20,000-statement file takes about as much as 50 entries. +const CHECKPOINT_ENTRIES: usize = 50; + +/// What the caches of an open document hold, for tests of their memory. +#[doc(hidden)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DocumentCacheStats { + /// The distinct statements of the document, including the bodies of SQL functions. The + /// parse results and trees are keyed by the text of the statement. + pub statements: usize, + pub parse_results: usize, + pub trees: usize, + pub analysis: bool, + /// The statements whose references are kept, see [`WalkCache`]. + pub resolved_statements: usize, + pub walk_checkpoints: usize, +} + /// What the caches of a document depend on besides its text. #[derive(Clone)] pub(crate) struct CacheKey { @@ -101,18 +140,47 @@ impl DocumentCache { analysis } + /// How much the document keeps, in entries: one per parse result, tree and resolved + /// statement, see [`DOCUMENT_CACHE_ENTRIES_MAX`], and [`CHECKPOINT_ENTRIES`] per walk + /// checkpoint. + pub(crate) fn usage(&self) -> usize { + let (resolved, checkpoints) = self.keyed.lock().unwrap().walk.len(); + self.parses.len() + self.trees.len() + resolved + checkpoints * CHECKPOINT_ENTRIES + } + + /// Drops everything that is cached. It is recomputed when it is needed again. + pub(crate) fn clear(&self) { + self.parses.clear(); + self.trees.clear(); + *self.keyed.lock().unwrap() = KeyedCaches::default(); + } + + /// What the caches hold for a document with the top-level statements `statements`. The + /// statements of a function body are counted if they are cached. + pub(crate) fn stats(&self, statements: &FxHashSet) -> DocumentCacheStats { + let embedded: usize = statements + .iter() + .filter_map(|id| self.parses.cached_embedded_len(id)) + .sum(); + let statements = statements.len() + embedded; + let keyed = self.keyed.lock().unwrap(); + let (resolved_statements, walk_checkpoints) = keyed.walk.len(); + DocumentCacheStats { + statements, + parse_results: self.parses.len(), + trees: self.trees.len(), + analysis: keyed.analysis.is_some(), + resolved_statements, + walk_checkpoints, + } + } + /// The cached analysis, whatever its key. #[cfg(test)] pub(crate) fn cached_analysis(&self) -> Option> { self.keyed.lock().unwrap().analysis.clone() } - /// The number of resolved statements and walk checkpoints, whatever their key. - #[cfg(test)] - pub(crate) fn walk_size(&self) -> (usize, usize) { - self.keyed.lock().unwrap().walk.len() - } - /// The caches for `key`, emptied first if they were for another key. fn keyed(&self, key: &CacheKey) -> std::sync::MutexGuard<'_, KeyedCaches> { let mut keyed = self.keyed.lock().unwrap(); diff --git a/crates/pgls_workspace/src/workspace/server/document_cache/budget.rs b/crates/pgls_workspace/src/workspace/server/document_cache/budget.rs new file mode 100644 index 000000000..da0c990c7 --- /dev/null +++ b/crates/pgls_workspace/src/workspace/server/document_cache/budget.rs @@ -0,0 +1,128 @@ +//! A budget for caches of data that can be recomputed, kept per key, like the caches of each +//! open document. + +use std::hash::Hash; +use std::sync::Mutex; + +use rustc_hash::FxHashMap; + +/// Tracks how much each key's caches use, in a unit of the caller's choice, and which keys to +/// evict to stay within `max`: the least recently used first. +/// +/// The key that was just used is never evicted, so the caches of a single key can exceed the +/// budget while it is in use. The caller owns the caches: it reports their usage with +/// [`CacheBudget::record`] after using them, and drops the caches of the keys it returns. +pub(crate) struct CacheBudget { + max: usize, + state: Mutex>, +} + +struct State { + entries: FxHashMap, + total: usize, + /// Incremented on every use, to order the entries by their last use. + clock: u64, +} + +struct Entry { + usage: usize, + last_used: u64, +} + +impl CacheBudget { + pub(crate) fn new(max: usize) -> Self { + Self { + max, + state: Mutex::new(State { + entries: FxHashMap::default(), + total: 0, + clock: 0, + }), + } + } + + pub(crate) fn max(&self) -> usize { + self.max + } + + /// Records that the caches of `key` were just used and now use `usage`. Returns the keys + /// whose caches the caller has to drop to get back within the budget, least recently used + /// first. They are no longer tracked until they are recorded again. + pub(crate) fn record(&self, key: &K, usage: usize) -> Vec { + let mut state = self.state.lock().unwrap(); + state.clock += 1; + let clock = state.clock; + let previous = state + .entries + .insert( + key.clone(), + Entry { + usage, + last_used: clock, + }, + ) + .map_or(0, |entry| entry.usage); + state.total = state.total - previous + usage; + if state.total <= self.max { + return Vec::new(); + } + + let mut candidates: Vec<(u64, K)> = state + .entries + .iter() + .filter(|(candidate, _)| *candidate != key) + .map(|(candidate, entry)| (entry.last_used, candidate.clone())) + .collect(); + candidates.sort_unstable_by_key(|(last_used, _)| *last_used); + + let mut evicted = Vec::new(); + for (_, candidate) in candidates { + if state.total <= self.max { + break; + } + if let Some(entry) = state.entries.remove(&candidate) { + state.total -= entry.usage; + evicted.push(candidate); + } + } + evicted + } + + /// Stops tracking `key`, e.g. when its caches are dropped with the document. + pub(crate) fn remove(&self, key: &K) { + let mut state = self.state.lock().unwrap(); + if let Some(entry) = state.entries.remove(key) { + state.total -= entry.usage; + } + } + + #[cfg(test)] + pub(crate) fn total(&self) -> usize { + self.state.lock().unwrap().total + } +} + +#[cfg(test)] +mod tests { + use super::CacheBudget; + + #[test] + fn evicts_the_least_recently_used_keys_first() { + let budget = CacheBudget::new(10); + assert!(budget.record(&"a", 4).is_empty()); + assert!(budget.record(&"b", 4).is_empty()); + // `a` was used last, so `b` goes first. + assert!(budget.record(&"a", 4).is_empty()); + assert_eq!(budget.record(&"c", 4), ["b"]); + assert_eq!(budget.total(), 8); + + // The key in use is never evicted, even alone over the budget. + assert_eq!(budget.record(&"c", 20), ["a"]); + assert_eq!(budget.total(), 20); + assert!(budget.record(&"c", 20).is_empty()); + + budget.remove(&"c"); + assert_eq!(budget.total(), 0); + assert!(budget.record(&"a", 10).is_empty()); + } +} diff --git a/crates/pgls_workspace/src/workspace/server/document_cache/parse.rs b/crates/pgls_workspace/src/workspace/server/document_cache/parse.rs index bf30cfe62..13cc6b927 100644 --- a/crates/pgls_workspace/src/workspace/server/document_cache/parse.rs +++ b/crates/pgls_workspace/src/workspace/server/document_cache/parse.rs @@ -32,6 +32,27 @@ impl PgQueryStore { } } + /// The number of statements embedded in `statement`, if they are cached. + pub fn cached_embedded_len(&self, statement: &StatementId) -> Option { + self.embedded_db + .lock() + .unwrap() + .get(statement) + .map(|embedded| embedded.len()) + } + + /// The number of statements with a cached parse result. + pub fn len(&self) -> usize { + self.ast_db.lock().unwrap().len() + } + + /// Drops every cached result. + pub fn clear(&self) { + self.ast_db.lock().unwrap().clear(); + self.plpgsql_db.lock().unwrap().clear(); + self.embedded_db.lock().unwrap().clear(); + } + /// Drops the results of the statements for which `keep` returns `false`. pub fn retain(&self, keep: impl Fn(&StatementId) -> bool) { self.ast_db.lock().unwrap().retain(|id, _| keep(id)); diff --git a/crates/pgls_workspace/src/workspace/server/document_cache/tree.rs b/crates/pgls_workspace/src/workspace/server/document_cache/tree.rs index c699d2f67..eb8b4dafe 100644 --- a/crates/pgls_workspace/src/workspace/server/document_cache/tree.rs +++ b/crates/pgls_workspace/src/workspace/server/document_cache/tree.rs @@ -24,6 +24,16 @@ impl TreeSitterStore { } } + /// The number of statements with a cached tree. + pub fn len(&self) -> usize { + self.db.lock().unwrap().len() + } + + /// Drops every cached tree. + pub fn clear(&self) { + self.db.lock().unwrap().clear(); + } + /// Drops the trees of the statements for which `keep` returns `false`. pub fn retain(&self, keep: impl Fn(&StatementId) -> bool) { self.db.lock().unwrap().retain(|id, _| keep(id)); diff --git a/crates/pgls_workspace/src/workspace/server/document_cache/walk.rs b/crates/pgls_workspace/src/workspace/server/document_cache/walk.rs index 48eef4e99..a7367091c 100644 --- a/crates/pgls_workspace/src/workspace/server/document_cache/walk.rs +++ b/crates/pgls_workspace/src/workspace/server/document_cache/walk.rs @@ -132,7 +132,6 @@ impl WalkCache { } /// The number of resolved statements and checkpoints. - #[cfg(test)] pub(crate) fn len(&self) -> (usize, usize) { ( self.resolved diff --git a/crates/pgls_workspace/tests/common/mod.rs b/crates/pgls_workspace/tests/common/mod.rs new file mode 100644 index 000000000..61ea080f1 --- /dev/null +++ b/crates/pgls_workspace/tests/common/mod.rs @@ -0,0 +1,70 @@ +//! A global allocator that counts the bytes allocated by the test binary, for the memory tests. +//! Each test binary that uses it runs a single test, so that nothing else allocates at the same +//! time. + +// Each test binary uses a different part of it. +#![allow(dead_code)] + +use std::alloc::{GlobalAlloc, Layout, System}; +use std::sync::atomic::{AtomicUsize, Ordering}; + +struct Counting; + +pub static CURRENT: AtomicUsize = AtomicUsize::new(0); +pub static PEAK: AtomicUsize = AtomicUsize::new(0); + +unsafe impl GlobalAlloc for Counting { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { System.alloc(layout) }; + if !ptr.is_null() { + let current = CURRENT.fetch_add(layout.size(), Ordering::Relaxed) + layout.size(); + PEAK.fetch_max(current, Ordering::Relaxed); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { System.dealloc(ptr, layout) }; + CURRENT.fetch_sub(layout.size(), Ordering::Relaxed); + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + let new = unsafe { System.realloc(ptr, layout, new_size) }; + if !new.is_null() { + if new_size >= layout.size() { + let grown = new_size - layout.size(); + let current = CURRENT.fetch_add(grown, Ordering::Relaxed) + grown; + PEAK.fetch_max(current, Ordering::Relaxed); + } else { + CURRENT.fetch_sub(layout.size() - new_size, Ordering::Relaxed); + } + } + new + } +} + +#[global_allocator] +static ALLOCATOR: Counting = Counting; + +pub const MIB: usize = 1024 * 1024; + +pub fn mib(bytes: usize) -> f64 { + bytes as f64 / MIB as f64 +} + +/// Resets the peak to the current usage and returns it. +pub fn reset_peak() -> usize { + let current = CURRENT.load(Ordering::Relaxed); + PEAK.store(current, Ordering::Relaxed); + current +} + +/// The bytes allocated right now. +pub fn current() -> usize { + CURRENT.load(Ordering::Relaxed) +} + +/// The most bytes allocated at once since the last [`reset_peak`]. +pub fn peak() -> usize { + PEAK.load(Ordering::Relaxed) +} diff --git a/crates/pgls_workspace/tests/edit_churn.rs b/crates/pgls_workspace/tests/edit_churn.rs new file mode 100644 index 000000000..bea6ada0d --- /dev/null +++ b/crates/pgls_workspace/tests/edit_churn.rs @@ -0,0 +1,352 @@ +//! The caches of a document during a long editing session: the parse results, the analysis and +//! what the walk resolves follow the statements that the document has now, and don't grow with +//! the number of edits. +//! +//! Counts every allocation of the test binary with the allocator in `common`, so the test runs +//! on its own in this file. `CHURN_STATEMENTS` and `CHURN_EDITS` change the size of the workload. + +mod common; + +use common::{current, mib}; +use pgls_analyse::RuleCategories; +use pgls_configuration::{Merge, PartialConfiguration, database::PartialDatabaseConfiguration}; +use pgls_fs::PgLSPath; +use pgls_text_size::{TextRange, TextSize}; +use pgls_workspace::{ + Workspace, + features::{ + code_actions::CodeActionsParams, diagnostics::PullFileDiagnosticsParams, + inlay_hints::PullInlayHintsParams, on_hover::OnHoverParams, + semantic_tokens::PullSemanticTokensParams, + }, + workspace::{ + ChangeFileParams, CloseFileParams, OpenFileParams, ProjectKey, RegisterProjectFolderParams, + UpdateSettingsParams, server::WorkspaceServer, + }, +}; +use sqlx::PgPool; + +/// The statement edits applied by each change of the file. +const EDITS_PER_CHANGE: usize = 4; + +fn env(name: &str, default: usize) -> usize { + std::env::var(name).map_or(default, |value| value.parse().unwrap()) +} + +/// A small deterministic random number generator (xorshift), so that a failure reproduces. +struct Rng(u64); + +impl Rng { + fn below(&mut self, n: usize) -> usize { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + (self.0 % n as u64) as usize + } +} + +#[derive(Clone, Copy)] +enum Kind { + Table, + Index, + Function, + Query, +} + +/// A statement of the document. `id` makes its text unique, and `version` changes the body of a +/// SQL function. +#[derive(Clone, Copy)] +struct Statement { + kind: Kind, + id: usize, + version: usize, +} + +impl Statement { + fn sql(&self) -> String { + let Self { id, version, .. } = self; + match self.kind { + Kind::Table => format!( + "create table t{id} (id bigint primary key, name text not null, created_at timestamp);" + ), + Kind::Index => format!("create index t{id}_name_idx on t{id} (name);"), + Kind::Function => format!( + "create function f{id}(a int, b text) returns int language sql stable as $$ select count(*)::int + {version} from pg_class where relname = b and relnatts = a $$;" + ), + Kind::Query => format!( + "select c.relname, f{id}(1, 'x') from pg_class c where c.relnatts = {id} order by c.relname;" + ), + } + } +} + +struct Document { + statements: Vec, + next_id: usize, +} + +impl Document { + fn new(len: usize) -> Self { + let mut document = Self { + statements: Vec::new(), + next_id: 0, + }; + for _ in 0..len { + let statement = document.new_statement(document.statements.len()); + document.statements.push(statement); + } + document + } + + fn new_statement(&mut self, seed: usize) -> Statement { + self.next_id += 1; + let kind = match seed % 4 { + 0 => Kind::Table, + 1 => Kind::Index, + 2 => Kind::Function, + _ => Kind::Query, + }; + Statement { + kind, + id: self.next_id, + version: 0, + } + } + + /// The text, and the offset of each statement in it. + fn render(&self) -> (String, Vec) { + let mut sql = String::new(); + let mut offsets = Vec::new(); + for statement in &self.statements { + offsets.push(sql.len()); + sql.push_str(&statement.sql()); + sql.push('\n'); + } + (sql, offsets) + } + + /// Rewrites, inserts or deletes a statement, or changes the body of a SQL function, and keeps + /// the length of the document within 10% of `len`. + fn edit(&mut self, rng: &mut Rng, len: usize) { + let at = rng.below(self.statements.len()); + let operation = if self.statements.len() < len * 9 / 10 { + 1 + } else if self.statements.len() > len * 11 / 10 { + 2 + } else { + rng.below(4) + }; + match operation { + 0 => self.statements[at] = self.new_statement(rng.below(4)), + 1 => { + let statement = self.new_statement(rng.below(4)); + self.statements.insert(at, statement); + } + 2 => { + self.statements.remove(at); + } + _ => { + let function = (at..self.statements.len()) + .chain(0..at) + .find(|i| matches!(self.statements[*i].kind, Kind::Function)); + if let Some(function) = function { + self.statements[function].version += 1; + } + } + } + } +} + +fn workspace(test_db: &PgPool) -> (WorkspaceServer, ProjectKey) { + let mut conf = PartialConfiguration::init(); + conf.merge_with(PartialConfiguration { + db: Some(PartialDatabaseConfiguration { + database: Some( + test_db + .connect_options() + .get_database() + .unwrap() + .to_string(), + ), + ..Default::default() + }), + ..Default::default() + }); + let workspace = WorkspaceServer::new(); + let project_key = workspace + .register_project_folder(RegisterProjectFolderParams { path: None }) + .unwrap(); + workspace + .update_settings(UpdateSettingsParams { + project_key, + configuration: conf, + gitignore_matches: vec![], + vcs_base_path: None, + workspace_directory: None, + }) + .unwrap(); + (workspace, project_key) +} + +fn diagnostics_and_code_actions( + (workspace, project_key): &(WorkspaceServer, ProjectKey), + path: &PgLSPath, + at: TextSize, +) { + workspace + .pull_file_diagnostics(PullFileDiagnosticsParams { + project_key: *project_key, + path: path.clone(), + categories: RuleCategories::all(), + max_diagnostics: u32::MAX, + only: vec![], + skip: vec![], + }) + .unwrap(); + workspace + .pull_code_actions(CodeActionsParams { + project_key: *project_key, + path: path.clone(), + cursor_position: at, + only: vec![], + skip: vec![], + }) + .unwrap(); +} + +/// Hover on the statement at `offset`, and inlay hints around it. +fn hover_and_inlay_hints( + (workspace, project_key): &(WorkspaceServer, ProjectKey), + path: &PgLSPath, + sql: &str, + offset: usize, +) { + let at = TextSize::try_from(offset + 9).unwrap(); + workspace + .on_hover(OnHoverParams { + project_key: *project_key, + path: path.clone(), + position: at, + }) + .unwrap(); + let end = TextSize::try_from((offset + 2_000).min(sql.len())).unwrap(); + workspace + .pull_inlay_hints(PullInlayHintsParams { + project_key: *project_key, + path: path.clone(), + range: TextRange::new(TextSize::try_from(offset).unwrap(), end), + }) + .unwrap(); +} + +/// Semantic tokens for the whole file, which resolves every statement. +fn all_tokens((workspace, project_key): &(WorkspaceServer, ProjectKey), path: &PgLSPath) { + workspace + .pull_semantic_tokens(PullSemanticTokensParams { + project_key: *project_key, + path: path.clone(), + range: None, + }) + .unwrap(); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn caches_follow_the_document_through_many_edits(test_db: PgPool) { + let len = env("CHURN_STATEMENTS", 1_000); + let edits = env("CHURN_EDITS", 2_000); + let workspace = workspace(&test_db); + let project_key = workspace.1; + let path = PgLSPath::new("churn.sql"); + let mut rng = Rng(0x2545_f491_4f6c_dd1d); + let mut document = Document::new(len); + + // A full pass over the document as it is opened. + let baseline = current(); + let (sql, offsets) = document.render(); + workspace + .0 + .open_file(OpenFileParams { + project_key, + path: path.clone(), + content: sql.clone(), + version: 1, + }) + .unwrap(); + let last = *offsets.last().unwrap(); + diagnostics_and_code_actions(&workspace, &path, TextSize::try_from(last + 9).unwrap()); + hover_and_inlay_hints(&workspace, &path, &sql, last); + all_tokens(&workspace, &path); + let first_pass = current() - baseline; + + // Each change of the file applies a few edits, like typing a few characters between two + // requests: every change re-splits the whole document, which takes most of the time of a + // debug build. + let changes = edits / EDITS_PER_CHANGE; + for change in 0..changes { + for _ in 0..EDITS_PER_CHANGE { + document.edit(&mut rng, len); + } + let (sql, offsets) = document.render(); + workspace + .0 + .change_file(ChangeFileParams { + project_key, + path: path.clone(), + content: sql.clone(), + version: i32::try_from(change).unwrap() + 2, + }) + .unwrap(); + let offset = offsets[rng.below(offsets.len())]; + if change % 5 == 0 { + hover_and_inlay_hints(&workspace, &path, &sql, offset); + } + if change % 50 == 0 { + diagnostics_and_code_actions( + &workspace, + &path, + TextSize::try_from(offset + 9).unwrap(), + ); + all_tokens(&workspace, &path); + } + // The walk keeps a few checkpoints, and drops the references of the statements an edit + // changed. + let stats = workspace.0.document_cache_stats(&path).unwrap(); + assert!(stats.walk_checkpoints <= 4, "{stats:?}"); + assert!(stats.resolved_statements <= stats.statements, "{stats:?}"); + } + + // The same full pass at the end. + let (sql, offsets) = document.render(); + let last = *offsets.last().unwrap(); + diagnostics_and_code_actions(&workspace, &path, TextSize::try_from(last + 9).unwrap()); + hover_and_inlay_hints(&workspace, &path, &sql, last); + all_tokens(&workspace, &path); + let after_edits = current() - baseline; + + let stats = workspace.0.document_cache_stats(&path).unwrap(); + println!( + "{len} statements, {edits} edits in {changes} changes: retained {:.1} MiB after the first pass, {:.1} MiB after \ + the edits; {stats:?}", + mib(first_pass), + mib(after_edits) + ); + + // Every statement of the document, including the function bodies, has a parse result, a + // tree, and its references, and nothing else is cached. + assert_eq!(stats.parse_results, stats.statements, "{stats:?}"); + assert_eq!(stats.trees, stats.statements, "{stats:?}"); + assert_eq!(stats.resolved_statements, stats.statements, "{stats:?}"); + assert!(stats.analysis); + // The document stays within 10% of its length, so the memory doesn't grow with the edits. + assert!( + after_edits <= first_pass * 3 / 2, + "{:.1} MiB retained after {edits} edits, {:.1} MiB after the first pass", + mib(after_edits), + mib(first_pass) + ); + + workspace + .0 + .close_file(CloseFileParams { project_key, path }) + .unwrap(); +} diff --git a/crates/pgls_workspace/tests/memory_bounds.rs b/crates/pgls_workspace/tests/memory_bounds.rs new file mode 100644 index 000000000..4c067d67e --- /dev/null +++ b/crates/pgls_workspace/tests/memory_bounds.rs @@ -0,0 +1,269 @@ +//! The memory that the workspace keeps for open documents. +//! +//! Counts every allocation of the test binary with the allocator in `common`, so the test runs +//! on its own in this file. `MEMORY_DOCUMENTS` and `MEMORY_GROUPS` (four statements and one SQL function +//! body each) change the size of the workload, e.g. `MEMORY_DOCUMENTS=20 MEMORY_GROUPS=1250` for +//! 20 documents with 5,000 statements each. Run it in a release build for numbers that match +//! what users see. + +mod common; + +use common::{MIB, current, mib, peak, reset_peak}; +use pgls_analyse::RuleCategories; +use pgls_configuration::{Merge, PartialConfiguration, database::PartialDatabaseConfiguration}; +use pgls_fs::PgLSPath; +use pgls_text_size::{TextRange, TextSize}; +use pgls_workspace::{ + Workspace, + features::{ + code_actions::CodeActionsParams, diagnostics::PullFileDiagnosticsParams, + inlay_hints::PullInlayHintsParams, on_hover::OnHoverParams, + }, + workspace::server::DOCUMENT_CACHE_ENTRIES_MAX, + workspace::{ + CloseFileParams, OpenFileParams, ProjectKey, RegisterProjectFolderParams, + UpdateSettingsParams, + }, +}; +use sqlx::PgPool; + +fn env(name: &str, default: usize) -> usize { + std::env::var(name).map_or(default, |value| value.parse().unwrap()) +} + +/// A migration-style file: per group a table, an index, a SQL function and a query. +fn migration_file(document: usize, groups: usize) -> String { + let mut sql = String::new(); + for i in 0..groups { + sql.push_str(&format!( + "create table t{document}_{i} (id bigint primary key, name text not null, status text, created_at timestamp default now());\n\ + create index t{document}_{i}_name_idx on t{document}_{i} (name);\n\ + create function f{document}_{i}(a int, b text) returns int language sql stable as $$ select count(*)::int from t{document}_{i} where id = a and name = b $$;\n\ + select t.id, t.name, f{document}_{i}(1, 'x') from t{document}_{i} t where t.status = 'active' order by t.created_at;\n" + )); + } + sql +} + +fn workspace(test_db: &PgPool) -> (Box, ProjectKey) { + let mut conf = PartialConfiguration::init(); + conf.merge_with(PartialConfiguration { + db: Some(PartialDatabaseConfiguration { + database: Some( + test_db + .connect_options() + .get_database() + .unwrap() + .to_string(), + ), + ..Default::default() + }), + ..Default::default() + }); + let workspace = pgls_workspace::workspace::server(); + let project_key = workspace + .register_project_folder(RegisterProjectFolderParams { path: None }) + .unwrap(); + workspace + .update_settings(UpdateSettingsParams { + project_key, + configuration: conf, + gitignore_matches: vec![], + vcs_base_path: None, + workspace_directory: None, + }) + .unwrap(); + (workspace, project_key) +} + +/// What an editor asks for on a document: diagnostics, then hover, inlay hints and code +/// actions at the end of the file. +fn use_document( + workspace: &dyn Workspace, + project_key: ProjectKey, + path: &PgLSPath, + content: &str, +) { + workspace + .pull_file_diagnostics(PullFileDiagnosticsParams { + project_key, + path: path.clone(), + categories: RuleCategories::all(), + max_diagnostics: u32::MAX, + only: vec![], + skip: vec![], + }) + .unwrap(); + let last = TextSize::try_from(content.rfind("select t.id").unwrap()).unwrap(); + workspace + .on_hover(OnHoverParams { + project_key, + path: path.clone(), + position: last + TextSize::from(9), + }) + .unwrap(); + let len = TextSize::of(content); + let start = TextSize::try_from(content.len().saturating_sub(4_000)).unwrap(); + workspace + .pull_inlay_hints(PullInlayHintsParams { + project_key, + path: path.clone(), + range: TextRange::new(start, len), + }) + .unwrap(); + workspace + .pull_code_actions(CodeActionsParams { + project_key, + path: path.clone(), + cursor_position: last + TextSize::from(9), + only: vec![], + skip: vec![], + }) + .unwrap(); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn open_documents_stay_within_the_budget(test_db: PgPool) { + let documents = env("MEMORY_DOCUMENTS", 20); + let groups = env("MEMORY_GROUPS", 200); + let (workspace, project_key) = workspace(&test_db); + + // Loads the database snapshot, which stays for the lifetime of the workspace. + let warmup = PgLSPath::new("warmup.sql"); + let content = migration_file(usize::MAX, 1); + workspace + .open_file(OpenFileParams { + project_key, + path: warmup.clone(), + content: content.clone(), + version: 1, + }) + .unwrap(); + use_document(&*workspace, project_key, &warmup, &content); + workspace + .close_file(CloseFileParams { + project_key, + path: warmup, + }) + .unwrap(); + + let contents: Vec<_> = (0..documents) + .map(|document| migration_file(document, groups)) + .collect(); + let text: usize = contents.iter().map(String::len).sum(); + let baseline = reset_peak(); + + // An editor with every document open. + let mut first = (0, 0); + for (document, content) in contents.iter().enumerate() { + let path = PgLSPath::new(format!("{document}.sql")); + workspace + .open_file(OpenFileParams { + project_key, + path: path.clone(), + content: content.clone(), + version: 1, + }) + .unwrap(); + use_document(&*workspace, project_key, &path, content); + if document == 0 { + first = (peak() - baseline, current() - baseline); + } + } + let open_peak = peak() - baseline; + let open_retained = current() - baseline; + for document in 0..documents { + workspace + .close_file(CloseFileParams { + project_key, + path: PgLSPath::new(format!("{document}.sql")), + }) + .unwrap(); + } + let closed_retained = current().saturating_sub(baseline); + + // The CLI: each file is opened, checked and closed. + let reset = reset_peak(); + let cli_started = std::time::Instant::now(); + for (document, content) in contents.iter().enumerate() { + let path = PgLSPath::new(format!("cli_{document}.sql")); + workspace + .open_file(OpenFileParams { + project_key, + path: path.clone(), + content: content.clone(), + version: 1, + }) + .unwrap(); + workspace + .pull_file_diagnostics(PullFileDiagnosticsParams { + project_key, + path: path.clone(), + categories: RuleCategories::all(), + max_diagnostics: u32::MAX, + only: vec![], + skip: vec![], + }) + .unwrap(); + workspace + .close_file(CloseFileParams { project_key, path }) + .unwrap(); + } + let cli_duration = cli_started.elapsed(); + let cli_peak = peak().saturating_sub(reset); + let cli_retained = current().saturating_sub(baseline); + + println!( + "{documents} documents, {:.1} MiB of SQL: one document peak {:.1} MiB, retained {:.1} MiB; \ + all open: peak {:.1} MiB, retained {:.1} MiB; retained after closing {:.1} MiB; \ + CLI peak {:.1} MiB, retained {:.1} MiB, {cli_duration:.1?}", + mib(text), + mib(first.0), + mib(first.1), + mib(open_peak), + mib(open_retained), + mib(closed_retained), + mib(cli_peak), + mib(cli_retained), + ); + + // Every statement has a parse result and a tree, and so does the body of each function. The + // requests at the end of the file resolve a few statements and record a few checkpoints. + let entries_per_document = 2 * 5 * groups; + let cached_documents = DOCUMENT_CACHE_ENTRIES_MAX / entries_per_document + 1; + // A document whose caches were dropped still has its text, statement ranges and + // suppressions. + let uncached_document = 4 * text / documents; + let budget = + (cached_documents.min(documents) * first.1 + documents * uncached_document) * 5 / 4; + assert!( + open_retained <= budget, + "{:.1} MiB retained, budget {:.1} MiB", + mib(open_retained), + mib(budget) + ); + assert!( + open_peak <= budget + first.0, + "{:.1} MiB peak, budget {:.1} MiB", + mib(open_peak), + mib(budget + first.0) + ); + // Closing the documents drops everything derived from them. + assert!( + closed_retained <= 4 * MIB, + "{:.1} MiB retained after closing every document", + mib(closed_retained) + ); + // The CLI only ever has one file open. + assert!( + cli_peak <= 2 * first.0, + "{:.1} MiB peak for the CLI, {:.1} MiB for one document", + mib(cli_peak), + mib(first.0) + ); + assert!( + cli_retained <= 4 * MIB, + "{:.1} MiB retained after the CLI run", + mib(cli_retained) + ); +} From 28fa5fe1686367ee5747bd84936d255ee23e00b1 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Sat, 31 Oct 2026 09:40:44 +0000 Subject: [PATCH 2/4] fix(workspace): charge the statements that executing a statement parses Finding the statement to execute parses the statements before it, and the early returns, like without a database connection, skipped the refresh that would charge them to the cache budget. --- crates/pgls_workspace/src/workspace/server.rs | 2 + .../src/workspace/server.tests.rs | 63 +++++++++++++++++++ 2 files changed, 65 insertions(+) diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index d95035027..8b8bba8c5 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -921,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, ¶ms.path); if stmt.is_none() { return Ok(ExecuteStatementResult { diff --git a/crates/pgls_workspace/src/workspace/server.tests.rs b/crates/pgls_workspace/src/workspace/server.tests.rs index 5217e83d0..f3f89c369 100644 --- a/crates/pgls_workspace/src/workspace/server.tests.rs +++ b/crates/pgls_workspace/src/workspace/server.tests.rs @@ -2657,6 +2657,69 @@ fn document_caches_stay_within_the_budget() { 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] From 9899b9ab099112df67c5ce546de1cd1d941cfea8 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Sat, 31 Oct 2026 09:40:44 +0000 Subject: [PATCH 3/4] fix(workspace): release the memory of the statement caches when they are dropped Clearing a map keeps its capacity, so an evicted document kept the tables of its parse results and trees while the budget counted it as empty. Eviction now replaces the maps, and a map that is less than a quarter full after an edit is shrunk. --- .../workspace/server/document_cache/parse.rs | 59 ++++++++++++++++--- .../workspace/server/document_cache/tree.rs | 39 +++++++++++- 2 files changed, 88 insertions(+), 10 deletions(-) diff --git a/crates/pgls_workspace/src/workspace/server/document_cache/parse.rs b/crates/pgls_workspace/src/workspace/server/document_cache/parse.rs index 13cc6b927..898e29684 100644 --- a/crates/pgls_workspace/src/workspace/server/document_cache/parse.rs +++ b/crates/pgls_workspace/src/workspace/server/document_cache/parse.rs @@ -46,18 +46,26 @@ impl PgQueryStore { self.ast_db.lock().unwrap().len() } - /// Drops every cached result. + /// Drops every cached result, and the memory of the maps. pub fn clear(&self) { - self.ast_db.lock().unwrap().clear(); - self.plpgsql_db.lock().unwrap().clear(); - self.embedded_db.lock().unwrap().clear(); + drop(std::mem::take(&mut *self.ast_db.lock().unwrap())); + drop(std::mem::take(&mut *self.plpgsql_db.lock().unwrap())); + drop(std::mem::take(&mut *self.embedded_db.lock().unwrap())); } /// Drops the results of the statements for which `keep` returns `false`. pub fn retain(&self, keep: impl Fn(&StatementId) -> bool) { - self.ast_db.lock().unwrap().retain(|id, _| keep(id)); - self.plpgsql_db.lock().unwrap().retain(|id, _| keep(id)); - self.embedded_db.lock().unwrap().retain(|id, _| keep(id)); + retain(&mut self.ast_db.lock().unwrap(), &keep); + retain(&mut self.plpgsql_db.lock().unwrap(), &keep); + retain(&mut self.embedded_db.lock().unwrap(), &keep); + } + + /// The number of entries that the maps have room for. + #[cfg(test)] + pub fn capacity(&self) -> usize { + self.ast_db.lock().unwrap().capacity() + + self.plpgsql_db.lock().unwrap().capacity() + + self.embedded_db.lock().unwrap().capacity() } /// The statements embedded in the top-level statement `statement`, in order: the statements @@ -143,6 +151,16 @@ impl PgQueryStore { static COMPOSITE_TYPE_ERROR_RE: LazyLock = LazyLock::new(|| Regex::new(r#"\\?"([^"\\]+\.[^"\\]+)\\?" is not a known variable"#).unwrap()); +/// Keeps the entries of `map` for which `keep` returns `true`. A map that is less than a quarter +/// full is shrunk, so that a document that got much shorter doesn't keep the memory of the +/// statements it had. +pub(super) fn retain(map: &mut FxHashMap, keep: impl Fn(&StatementId) -> bool) { + map.retain(|id, _| keep(id)); + if map.capacity() > 4 * map.len().max(16) { + map.shrink_to_fit(); + } +} + fn is_composite_type_error(err: &str) -> bool { COMPOSITE_TYPE_ERROR_RE.is_match(err) } @@ -151,6 +169,33 @@ fn is_composite_type_error(err: &str) -> bool { mod tests { use super::*; + #[test] + fn dropping_results_releases_the_maps() { + let store = PgQueryStore::new(); + let statements: Vec = (0..200) + .map(|i| { + StatementId::new(&format!( + "create function f{i}() returns int language sql as 'select {i}';" + )) + }) + .collect(); + let fill = || { + for id in &statements { + store.get_or_cache_embedded(id); + store.get_or_cache_plpgsql_parse(id); + } + }; + + fill(); + store.retain(|id| *id == statements[0]); + assert_eq!(store.len(), 1); + assert!(store.capacity() < 3 * 64, "{}", store.capacity()); + + fill(); + store.clear(); + assert_eq!(store.capacity(), 0); + } + #[test] fn test_positional_params_in_grant() { let input = "grant usage on schema public, app_public, app_hidden to :DB_ROLE;"; diff --git a/crates/pgls_workspace/src/workspace/server/document_cache/tree.rs b/crates/pgls_workspace/src/workspace/server/document_cache/tree.rs index eb8b4dafe..20ac3f44a 100644 --- a/crates/pgls_workspace/src/workspace/server/document_cache/tree.rs +++ b/crates/pgls_workspace/src/workspace/server/document_cache/tree.rs @@ -29,14 +29,20 @@ impl TreeSitterStore { self.db.lock().unwrap().len() } - /// Drops every cached tree. + /// Drops every cached tree, and the memory of the map. pub fn clear(&self) { - self.db.lock().unwrap().clear(); + drop(std::mem::take(&mut *self.db.lock().unwrap())); } /// Drops the trees of the statements for which `keep` returns `false`. pub fn retain(&self, keep: impl Fn(&StatementId) -> bool) { - self.db.lock().unwrap().retain(|id, _| keep(id)); + super::parse::retain(&mut self.db.lock().unwrap(), keep); + } + + /// The number of trees that the map has room for. + #[cfg(test)] + pub fn capacity(&self) -> usize { + self.db.lock().unwrap().capacity() } pub fn get_or_cache_tree(&self, statement: &StatementId) -> Arc { @@ -66,3 +72,30 @@ impl TreeSitterStore { tree } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn dropping_trees_releases_the_map() { + let store = TreeSitterStore::new(); + let statements: Vec = (0..200) + .map(|i| StatementId::new(&format!("select {i};"))) + .collect(); + let fill = || { + for id in &statements { + store.get_or_cache_tree(id); + } + }; + + fill(); + store.retain(|id| *id == statements[0]); + assert_eq!(store.len(), 1); + assert!(store.capacity() < 64, "{}", store.capacity()); + + fill(); + store.clear(); + assert_eq!(store.capacity(), 0); + } +} From 669b6e9e5658b6e322a7ca444034f9f8bf6cb02f Mon Sep 17 00:00:00 2001 From: psteinroe Date: Sat, 31 Oct 2026 09:40:44 +0000 Subject: [PATCH 4/4] test(workspace): count tree-sitter's allocations in the memory tests Tree-sitter allocates its parsers and trees in C, which the counting allocator didn't see. The memory tests now make tree-sitter allocate through it. --- crates/pgls_workspace/tests/common/mod.rs | 94 +++++++++++++++++++ crates/pgls_workspace/tests/edit_churn.rs | 8 +- crates/pgls_workspace/tests/memory_bounds.rs | 13 +-- .../tests/native_allocations.rs | 25 +++++ 4 files changed, 131 insertions(+), 9 deletions(-) create mode 100644 crates/pgls_workspace/tests/native_allocations.rs diff --git a/crates/pgls_workspace/tests/common/mod.rs b/crates/pgls_workspace/tests/common/mod.rs index 61ea080f1..78f9cf519 100644 --- a/crates/pgls_workspace/tests/common/mod.rs +++ b/crates/pgls_workspace/tests/common/mod.rs @@ -1,11 +1,18 @@ //! A global allocator that counts the bytes allocated by the test binary, for the memory tests. //! Each test binary that uses it runs a single test, so that nothing else allocates at the same //! time. +//! +//! Tree-sitter allocates its parsers and trees in C. [`count_native_allocations`] makes it +//! allocate through the counting allocator, so that the cached trees are counted too. Postgres' +//! parser also allocates in C, but only while it parses: the parse results that are cached are +//! Rust values. // Each test binary uses a different part of it. #![allow(dead_code)] use std::alloc::{GlobalAlloc, Layout, System}; +use std::ffi::c_void; +use std::sync::Once; use std::sync::atomic::{AtomicUsize, Ordering}; struct Counting; @@ -68,3 +75,90 @@ pub fn current() -> usize { pub fn peak() -> usize { PEAK.load(Ordering::Relaxed) } + +/// Makes tree-sitter allocate through the counting allocator. Call it before anything uses +/// tree-sitter: memory that it allocated before can't be freed with the new functions. +pub fn count_native_allocations() { + static SET: Once = Once::new(); + // SAFETY: called before any tree-sitter allocation of the test, and the functions are a + // complete malloc/calloc/realloc/free family. + SET.call_once(|| unsafe { + tree_sitter::set_allocator( + Some(native_malloc), + Some(native_calloc), + Some(native_realloc), + Some(native_free), + ); + }); +} + +/// C's `free` gets no size, so each allocation starts with a header that holds its size, and +/// is aligned like `malloc`'s. +const HEADER: usize = 16; + +fn native_layout(size: usize) -> Option { + Layout::from_size_align(size.checked_add(HEADER)?, HEADER).ok() +} + +/// Writes the size into the header at `base`, and returns the memory after it. +unsafe fn with_header(base: *mut u8, size: usize) -> *mut c_void { + if base.is_null() { + return std::ptr::null_mut(); + } + unsafe { + base.cast::().write(size); + base.add(HEADER).cast() + } +} + +/// The start of the allocation of `ptr`, and its size. +unsafe fn header(ptr: *mut c_void) -> (*mut u8, usize) { + unsafe { + let base = ptr.cast::().sub(HEADER); + (base, base.cast::().read()) + } +} + +unsafe extern "C" fn native_malloc(size: usize) -> *mut c_void { + let Some(layout) = native_layout(size) else { + return std::ptr::null_mut(); + }; + unsafe { with_header(std::alloc::alloc(layout), size) } +} + +unsafe extern "C" fn native_calloc(count: usize, size: usize) -> *mut c_void { + let Some(layout) = count.checked_mul(size).and_then(native_layout) else { + return std::ptr::null_mut(); + }; + unsafe { with_header(std::alloc::alloc_zeroed(layout), layout.size() - HEADER) } +} + +unsafe extern "C" fn native_realloc(ptr: *mut c_void, size: usize) -> *mut c_void { + if ptr.is_null() { + return unsafe { native_malloc(size) }; + } + let (Some(new_layout), (base, old_size)) = (native_layout(size), unsafe { header(ptr) }) else { + return std::ptr::null_mut(); + }; + let Some(old_layout) = native_layout(old_size) else { + return std::ptr::null_mut(); + }; + unsafe { + with_header( + std::alloc::realloc(base, old_layout, new_layout.size()), + size, + ) + } +} + +unsafe extern "C" fn native_free(ptr: *mut c_void) { + if ptr.is_null() { + return; + } + unsafe { + let (base, size) = header(ptr); + if let Some(layout) = native_layout(size) { + std::alloc::dealloc(base, layout); + } + } +} diff --git a/crates/pgls_workspace/tests/edit_churn.rs b/crates/pgls_workspace/tests/edit_churn.rs index bea6ada0d..f08645b50 100644 --- a/crates/pgls_workspace/tests/edit_churn.rs +++ b/crates/pgls_workspace/tests/edit_churn.rs @@ -2,12 +2,13 @@ //! what the walk resolves follow the statements that the document has now, and don't grow with //! the number of edits. //! -//! Counts every allocation of the test binary with the allocator in `common`, so the test runs -//! on its own in this file. `CHURN_STATEMENTS` and `CHURN_EDITS` change the size of the workload. +//! Counts every allocation of the test binary with the allocator in `common`, including +//! tree-sitter's, so the test runs on its own in this file. `CHURN_STATEMENTS` and `CHURN_EDITS` +//! change the size of the workload. mod common; -use common::{current, mib}; +use common::{count_native_allocations, current, mib}; use pgls_analyse::RuleCategories; use pgls_configuration::{Merge, PartialConfiguration, database::PartialDatabaseConfiguration}; use pgls_fs::PgLSPath; @@ -252,6 +253,7 @@ fn all_tokens((workspace, project_key): &(WorkspaceServer, ProjectKey), path: &P #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] async fn caches_follow_the_document_through_many_edits(test_db: PgPool) { + count_native_allocations(); let len = env("CHURN_STATEMENTS", 1_000); let edits = env("CHURN_EDITS", 2_000); let workspace = workspace(&test_db); diff --git a/crates/pgls_workspace/tests/memory_bounds.rs b/crates/pgls_workspace/tests/memory_bounds.rs index 4c067d67e..73e88d1a4 100644 --- a/crates/pgls_workspace/tests/memory_bounds.rs +++ b/crates/pgls_workspace/tests/memory_bounds.rs @@ -1,14 +1,14 @@ //! The memory that the workspace keeps for open documents. //! -//! Counts every allocation of the test binary with the allocator in `common`, so the test runs -//! on its own in this file. `MEMORY_DOCUMENTS` and `MEMORY_GROUPS` (four statements and one SQL function -//! body each) change the size of the workload, e.g. `MEMORY_DOCUMENTS=20 MEMORY_GROUPS=1250` for -//! 20 documents with 5,000 statements each. Run it in a release build for numbers that match -//! what users see. +//! Counts every allocation of the test binary with the allocator in `common`, including +//! tree-sitter's, so the test runs on its own in this file. `MEMORY_DOCUMENTS` and +//! `MEMORY_GROUPS` (four statements and one SQL function body each) change the size of the +//! workload, e.g. `MEMORY_DOCUMENTS=20 MEMORY_GROUPS=1250` for 20 documents with 5,000 +//! statements each. Run it in a release build for numbers that match what users see. mod common; -use common::{MIB, current, mib, peak, reset_peak}; +use common::{MIB, count_native_allocations, current, mib, peak, reset_peak}; use pgls_analyse::RuleCategories; use pgls_configuration::{Merge, PartialConfiguration, database::PartialDatabaseConfiguration}; use pgls_fs::PgLSPath; @@ -124,6 +124,7 @@ fn use_document( #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] async fn open_documents_stay_within_the_budget(test_db: PgPool) { + count_native_allocations(); let documents = env("MEMORY_DOCUMENTS", 20); let groups = env("MEMORY_GROUPS", 200); let (workspace, project_key) = workspace(&test_db); diff --git a/crates/pgls_workspace/tests/native_allocations.rs b/crates/pgls_workspace/tests/native_allocations.rs new file mode 100644 index 000000000..3ce1bf375 --- /dev/null +++ b/crates/pgls_workspace/tests/native_allocations.rs @@ -0,0 +1,25 @@ +//! The memory tests count what tree-sitter allocates in C. + +mod common; + +use common::{count_native_allocations, current}; + +#[test] +fn tree_sitter_trees_are_counted() { + count_native_allocations(); + let sql = "select id, name from users where id = 1;\n".repeat(2_000); + let mut parser = tree_sitter::Parser::new(); + parser + .set_language(&pgls_treesitter_grammar::LANGUAGE.into()) + .unwrap(); + + let before = current(); + let tree = parser.parse(&sql, None).unwrap(); + let with_tree = current(); + // A tree has a node per token, which takes more than a byte of the text. + assert!(with_tree - before > sql.len(), "{before} {with_tree}"); + + drop(tree); + drop(parser); + assert!(current() < with_tree); +}