diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index dd57ba3eb..8b8bba8c5 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(()) } @@ -784,6 +819,7 @@ impl Workspace for WorkspaceServer { disabled_reason: invalidate_disabled_reason, }); + self.charge_document_caches(&documents, ¶ms.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, ¶ms.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, ¶ms.path); info!("Pulled {:?} diagnostic(s)", diagnostics.len()); Ok(PullDiagnosticsResult { diagnostics, @@ -1209,12 +1248,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")] @@ -1317,6 +1358,7 @@ impl Workspace for WorkspaceServer { text: stmt.id.content().to_string(), }); + self.charge_document_caches(&documents, ¶ms.path); Ok(CompletionsResult::in_document(items, &stmt)) } } @@ -1371,6 +1413,7 @@ impl Workspace for WorkspaceServer { }) .unwrap_or_default(); + self.charge_document_caches(&documents, ¶ms.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, ¶ms.path); Ok(PullInlayHintsResult { hints }) } @@ -1516,6 +1560,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 dc250d082..f3f89c369 100644 --- a/crates/pgls_workspace/src/workspace/server.tests.rs +++ b/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 diff --git a/crates/pgls_workspace/src/workspace/server/document.rs b/crates/pgls_workspace/src/workspace/server/document.rs index 2712286ce..b5a8fbf27 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, 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 = 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 f64cfb643..6847982c4 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, MigrationContext, 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 an edit invalidated in a document's cache. On a long file, the analysis, the copies of /// the catalog in the walk checkpoints and the snapshot hold every object the file defines, and /// dropping them takes tens of milliseconds, so [`StaleEntries::drop_in_background`] drops them @@ -142,18 +181,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..898e29684 100644 --- a/crates/pgls_workspace/src/workspace/server/document_cache/parse.rs +++ b/crates/pgls_workspace/src/workspace/server/document_cache/parse.rs @@ -32,11 +32,40 @@ 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, and the memory of the maps. + pub fn clear(&self) { + 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 @@ -122,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) } @@ -130,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 c699d2f67..20ac3f44a 100644 --- a/crates/pgls_workspace/src/workspace/server/document_cache/tree.rs +++ b/crates/pgls_workspace/src/workspace/server/document_cache/tree.rs @@ -24,9 +24,25 @@ impl TreeSitterStore { } } + /// The number of statements with a cached tree. + pub fn len(&self) -> usize { + self.db.lock().unwrap().len() + } + + /// Drops every cached tree, and the memory of the map. + pub fn clear(&self) { + 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 { @@ -56,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); + } +} 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 7035ec2c9..1fbab0296 100644 --- a/crates/pgls_workspace/src/workspace/server/document_cache/walk.rs +++ b/crates/pgls_workspace/src/workspace/server/document_cache/walk.rs @@ -242,7 +242,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..78f9cf519 --- /dev/null +++ b/crates/pgls_workspace/tests/common/mod.rs @@ -0,0 +1,164 @@ +//! 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; + +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) +} + +/// 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 new file mode 100644 index 000000000..f08645b50 --- /dev/null +++ b/crates/pgls_workspace/tests/edit_churn.rs @@ -0,0 +1,354 @@ +//! 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`, 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::{count_native_allocations, 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) { + count_native_allocations(); + 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..73e88d1a4 --- /dev/null +++ b/crates/pgls_workspace/tests/memory_bounds.rs @@ -0,0 +1,270 @@ +//! The memory that the workspace keeps for open documents. +//! +//! 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, count_native_allocations, 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) { + count_native_allocations(); + 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) + ); +} 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); +}