test(workspace): the document caches follow a long editing session

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<DocumentCacheStats> {
+        let documents = self.documents.read().unwrap();
+        documents.get(path).map(Document::cache_stats)
+    }
+
     /// Records what the caches of the document at `path` keep after a request, and drops the
     /// caches of the least recently used other documents to stay within
     /// [`DOCUMENT_CACHE_ENTRIES_MAX`].

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<StatementId> =
+            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) {

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)
+}