From 3f1a10586c13ee92b4587a16dbb01f560f76985e Mon Sep 17 00:00:00 2001 From: psteinroe Date: Sat, 24 Oct 2026 18:56:14 +0000 Subject: [PATCH] test(workspace): the document caches follow a long editing session Opens a document of 1,000 statements and applies 2,000 statement edits (rewrites, inserts, deletes, and changes of SQL function bodies), with diagnostics, hover, inlay hints and code actions in between. Afterwards the caches hold a parse result and a tree for exactly the statements of the document, and the memory retained is what it was after the first pass. The counting allocator moves to tests/common for both memory tests. --- crates/pgls_workspace/src/workspace/server.rs | 19 ++ .../src/workspace/server/document.rs | 12 + crates/pgls_workspace/tests/common/mod.rs | 70 ++++ crates/pgls_workspace/tests/edit_churn.rs | 311 ++++++++++++++++++ crates/pgls_workspace/tests/memory_bounds.rs | 74 +---- 5 files changed, 422 insertions(+), 64 deletions(-) create mode 100644 crates/pgls_workspace/tests/common/mod.rs create mode 100644 crates/pgls_workspace/tests/edit_churn.rs diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index c4383bb46..6879b89d1 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -140,6 +140,18 @@ pub struct WorkspaceServer { /// limit, 20 open files of 5,000 statements each kept 2.7 GB. pub const DOCUMENT_CACHE_ENTRIES_MAX: usize = 20_000; +/// 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 + /// caches are keyed by the text of the statement. + pub statements: usize, + pub parse_results: usize, + pub trees: usize, + pub analysis: bool, +} + /// A database snapshot with the changes of the statements before a statement. struct SnapshotBefore { /// The database snapshot it is built from. @@ -224,6 +236,13 @@ 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`]. diff --git a/crates/pgls_workspace/src/workspace/server/document.rs b/crates/pgls_workspace/src/workspace/server/document.rs index 3857957da..d278b1b3e 100644 --- a/crates/pgls_workspace/src/workspace/server/document.rs +++ b/crates/pgls_workspace/src/workspace/server/document.rs @@ -82,6 +82,18 @@ impl Document { self.ast_db.len() + self.cst_db.len() } + /// What the caches of the document hold. + pub fn cache_stats(&self) -> super::DocumentCacheStats { + let statements: FxHashSet = + self.iter(DefaultMapper).map(|(id, _, _)| id).collect(); + super::DocumentCacheStats { + statements: statements.len(), + parse_results: self.ast_db.len(), + trees: self.cst_db.len(), + analysis: self.analysis.lock().unwrap().is_some(), + } + } + /// Drops everything that is cached for the document. It is recomputed when it is needed /// again. pub fn clear_caches(&self) { 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..a95698545 --- /dev/null +++ b/crates/pgls_workspace/tests/edit_churn.rs @@ -0,0 +1,311 @@ +//! The caches of a document during a long editing session: they 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, + }, + workspace::{ + ChangeFileParams, CloseFileParams, OpenFileParams, 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 { + 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(); + workspace + .register_project_folder(RegisterProjectFolderParams { + path: None, + set_as_current_workspace: true, + }) + .unwrap(); + workspace + .update_settings(UpdateSettingsParams { + configuration: conf, + gitignore_matches: vec![], + vcs_base_path: None, + workspace_directory: None, + }) + .unwrap(); + workspace +} + +fn diagnostics_and_code_actions(workspace: &WorkspaceServer, path: &PgLSPath, at: TextSize) { + workspace + .pull_file_diagnostics(PullFileDiagnosticsParams { + path: path.clone(), + categories: RuleCategories::all(), + max_diagnostics: u32::MAX, + only: vec![], + skip: vec![], + }) + .unwrap(); + workspace + .pull_code_actions(CodeActionsParams { + 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: &WorkspaceServer, path: &PgLSPath, sql: &str, offset: usize) { + let at = TextSize::try_from(offset + 9).unwrap(); + workspace + .on_hover(OnHoverParams { + path: path.clone(), + position: at, + }) + .unwrap(); + let end = TextSize::try_from((offset + 2_000).min(sql.len())).unwrap(); + workspace + .pull_inlay_hints(PullInlayHintsParams { + path: path.clone(), + range: TextRange::new(TextSize::try_from(offset).unwrap(), end), + }) + .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 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 + .open_file(OpenFileParams { + 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); + 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 + .change_file(ChangeFileParams { + 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(), + ); + } + } + + // 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); + let after_edits = current() - baseline; + + let stats = workspace.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 and a + // tree, and nothing else is cached. + assert_eq!(stats.parse_results, stats.statements, "{stats:?}"); + assert_eq!(stats.trees, 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.close_file(CloseFileParams { path }).unwrap(); +} diff --git a/crates/pgls_workspace/tests/memory_bounds.rs b/crates/pgls_workspace/tests/memory_bounds.rs index 47f386ac7..596657b33 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 a global allocator, so the test runs on its -//! own in this file. `MEMORY_DOCUMENTS` and `MEMORY_GROUPS` (four statements and one SQL function +//! 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. -use std::alloc::{GlobalAlloc, Layout, System}; -use std::sync::atomic::{AtomicUsize, Ordering}; +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; @@ -26,57 +26,6 @@ use pgls_workspace::{ }; use sqlx::PgPool; -struct Counting; - -static CURRENT: AtomicUsize = AtomicUsize::new(0); -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; - -const MIB: usize = 1024 * 1024; - -fn mib(bytes: usize) -> f64 { - bytes as f64 / MIB as f64 -} - -/// Resets the peak to the current usage and returns it. -fn reset_peak() -> usize { - let current = CURRENT.load(Ordering::Relaxed); - PEAK.store(current, Ordering::Relaxed); - current -} - fn env(name: &str, default: usize) -> usize { std::env::var(name).map_or(default, |value| value.parse().unwrap()) } @@ -205,14 +154,11 @@ async fn open_documents_stay_within_the_budget(test_db: PgPool) { .unwrap(); use_document(&*workspace, &path, content); if document == 0 { - first = ( - PEAK.load(Ordering::Relaxed) - baseline, - CURRENT.load(Ordering::Relaxed) - baseline, - ); + first = (peak() - baseline, current() - baseline); } } - let open_peak = PEAK.load(Ordering::Relaxed) - baseline; - let open_retained = CURRENT.load(Ordering::Relaxed) - baseline; + let open_peak = peak() - baseline; + let open_retained = current() - baseline; for document in 0..documents { workspace .close_file(CloseFileParams { @@ -220,7 +166,7 @@ async fn open_documents_stay_within_the_budget(test_db: PgPool) { }) .unwrap(); } - let closed_retained = CURRENT.load(Ordering::Relaxed).saturating_sub(baseline); + let closed_retained = current().saturating_sub(baseline); // The CLI: each file is opened, checked and closed. let reset = reset_peak(); @@ -244,8 +190,8 @@ async fn open_documents_stay_within_the_budget(test_db: PgPool) { .unwrap(); workspace.close_file(CloseFileParams { path }).unwrap(); } - let cli_peak = PEAK.load(Ordering::Relaxed).saturating_sub(reset); - let cli_retained = CURRENT.load(Ordering::Relaxed).saturating_sub(baseline); + 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; \