refactor: simplify caches (#448)

Author: psteinroe

Cargo.lock

@@ -13,9 +13,9 @@ version              = "0.0.0"
 
 [dependencies]
 biome_deserialize = "0.6.0"
-dashmap           = "5.5.3"
 futures           = "0.3.31"
 globset           = "0.4.16"
+lru               = "0.12"
 
 ignore                    = { workspace = true }
 pgt_analyse               = { workspace = true, features = ["serde"] }

crates/pgt_workspace/Cargo.toml

@@ -1,4 +1,5 @@
 use std::{
+    collections::HashMap,
     fs,
     panic::RefUnwindSafe,
     path::{Path, PathBuf},
@@ -8,7 +9,6 @@ use std::{
 use analyser::AnalyserVisitorBuilder;
 use async_helper::run_async;
 use connection_manager::ConnectionManager;
-use dashmap::DashMap;
 use document::{
     AsyncDiagnosticsMapper, CursorPositionFilter, DefaultMapper, Document, ExecuteStatementMapper,
     SyncDiagnosticsMapper,
@@ -67,7 +67,7 @@ pub(super) struct WorkspaceServer {
     /// Stores the schema cache for this workspace
     schema_cache: SchemaCacheManager,
 
-    documents: DashMap<PgTPath, Document>,
+    documents: RwLock<HashMap<PgTPath, Document>>,
 
     connection: ConnectionManager,
 }
@@ -89,7 +89,7 @@ impl WorkspaceServer {
     pub(crate) fn new() -> Self {
         Self {
             settings: RwLock::default(),
-            documents: DashMap::default(),
+            documents: RwLock::new(HashMap::new()),
             schema_cache: SchemaCacheManager::new(),
             connection: ConnectionManager::new(),
         }
@@ -262,7 +262,8 @@ impl Workspace for WorkspaceServer {
     /// Add a new file to the workspace
     #[tracing::instrument(level = "info", skip_all, fields(path = params.path.as_path().as_os_str().to_str()), err)]
     fn open_file(&self, params: OpenFileParams) -> Result<(), WorkspaceError> {
-        self.documents
+        let mut documents = self.documents.write().unwrap();
+        documents
             .entry(params.path.clone())
             .or_insert_with(|| Document::new(params.content, params.version));
 
@@ -275,7 +276,8 @@ impl Workspace for WorkspaceServer {
 
     /// Remove a file from the workspace
     fn close_file(&self, params: super::CloseFileParams) -> Result<(), WorkspaceError> {
-        self.documents
+        let mut documents = self.documents.write().unwrap();
+        documents
             .remove(&params.path)
             .ok_or_else(WorkspaceError::not_found)?;
 
@@ -288,13 +290,15 @@ impl Workspace for WorkspaceServer {
         version = params.version
     ), err)]
     fn change_file(&self, params: super::ChangeFileParams) -> Result<(), WorkspaceError> {
-        match self.documents.entry(params.path.clone()) {
-            dashmap::mapref::entry::Entry::Occupied(mut entry) => {
+        let mut documents = self.documents.write().unwrap();
+
+        match documents.entry(params.path.clone()) {
+            std::collections::hash_map::Entry::Occupied(mut entry) => {
                 entry
                     .get_mut()
                     .update_content(params.content, params.version);
             }
-            dashmap::mapref::entry::Entry::Vacant(entry) => {
+            std::collections::hash_map::Entry::Vacant(entry) => {
                 entry.insert(Document::new(params.content, params.version));
             }
         }
@@ -307,8 +311,8 @@ impl Workspace for WorkspaceServer {
     }
 
     fn get_file_content(&self, params: GetFileContentParams) -> Result<String, WorkspaceError> {
-        let document = self
-            .documents
+        let documents = self.documents.read().unwrap();
+        let document = documents
             .get(&params.path)
             .ok_or(WorkspaceError::not_found())?;
         Ok(document.get_document_content().to_string())
@@ -322,8 +326,8 @@ impl Workspace for WorkspaceServer {
         &self,
         params: code_actions::CodeActionsParams,
     ) -> Result<code_actions::CodeActionsResult, WorkspaceError> {
-        let parser = self
-            .documents
+        let documents = self.documents.read().unwrap();
+        let parser = documents
             .get(&params.path)
             .ok_or(WorkspaceError::not_found())?;
 
@@ -364,8 +368,8 @@ impl Workspace for WorkspaceServer {
         &self,
         params: ExecuteStatementParams,
     ) -> Result<ExecuteStatementResult, WorkspaceError> {
-        let parser = self
-            .documents
+        let documents = self.documents.read().unwrap();
+        let parser = documents
             .get(&params.path)
             .ok_or(WorkspaceError::not_found())?;
 
@@ -422,8 +426,8 @@ impl Workspace for WorkspaceServer {
             }
         };
 
-        let doc = self
-            .documents
+        let documents = self.documents.read().unwrap();
+        let doc = documents
             .get(&params.path)
             .ok_or(WorkspaceError::not_found())?;
 
@@ -607,8 +611,8 @@ impl Workspace for WorkspaceServer {
         &self,
         params: GetCompletionsParams,
     ) -> Result<CompletionsResult, WorkspaceError> {
-        let parsed_doc = self
-            .documents
+        let documents = self.documents.read().unwrap();
+        let parsed_doc = documents
             .get(&params.path)
             .ok_or(WorkspaceError::not_found())?;
 
@@ -621,7 +625,7 @@ impl Workspace for WorkspaceServer {
 
         let schema_cache = self.schema_cache.load(pool)?;
 
-        match get_statement_for_completions(&parsed_doc, params.position) {
+        match get_statement_for_completions(parsed_doc, params.position) {
             None => {
                 tracing::debug!("No statement found.");
                 Ok(CompletionsResult::default())

crates/pgt_workspace/src/workspace/server.rs

@@ -1,17 +1,20 @@
-use std::sync::Arc;
+use std::num::NonZeroUsize;
+use std::sync::{Arc, Mutex};
 
-use dashmap::DashMap;
+use lru::LruCache;
 use pgt_lexer::SyntaxKind;
 
 use super::statement_identifier::StatementId;
 
+const DEFAULT_CACHE_SIZE: usize = 1000;
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct StatementAnnotations {
     ends_with_semicolon: bool,
 }
 
 pub struct AnnotationStore {
-    db: DashMap<StatementId, Arc<StatementAnnotations>>,
+    db: Mutex<LruCache<StatementId, Arc<StatementAnnotations>>>,
 }
 
 const WHITESPACE_TOKENS: [SyntaxKind; 6] = [
@@ -25,7 +28,11 @@ const WHITESPACE_TOKENS: [SyntaxKind; 6] = [
 
 impl AnnotationStore {
     pub fn new() -> AnnotationStore {
-        AnnotationStore { db: DashMap::new() }
+        AnnotationStore {
+            db: Mutex::new(LruCache::new(
+                NonZeroUsize::new(DEFAULT_CACHE_SIZE).unwrap(),
+            )),
+        }
     }
 
     #[allow(unused)]
@@ -34,8 +41,10 @@ impl AnnotationStore {
         statement_id: &StatementId,
         content: &str,
     ) -> Arc<StatementAnnotations> {
-        if let Some(existing) = self.db.get(statement_id).map(|x| x.clone()) {
-            return existing;
+        let mut cache = self.db.lock().unwrap();
+
+        if let Some(existing) = cache.get(statement_id) {
+            return existing.clone();
         }
 
         let lexed = pgt_lexer::lex(content);
@@ -51,7 +60,7 @@ impl AnnotationStore {
             ends_with_semicolon,
         });
 
-        self.db.insert(statement_id.clone(), annotations.clone());
+        cache.put(statement_id.clone(), annotations.clone());
 
         annotations
     }

crates/pgt_workspace/src/workspace/server/annotation.rs

@@ -1,6 +1,7 @@
+use std::collections::HashMap;
+use std::sync::RwLock;
 use std::time::{Duration, Instant};
 
-use dashmap::DashMap;
 use sqlx::{PgPool, Postgres, pool::PoolOptions, postgres::PgConnectOptions};
 
 use crate::settings::DatabaseSettings;
@@ -16,13 +17,13 @@ struct CachedPool {
 
 #[derive(Default)]
 pub struct ConnectionManager {
-    pools: DashMap<ConnectionKey, CachedPool>,
+    pools: RwLock<HashMap<ConnectionKey, CachedPool>>,
 }
 
 impl ConnectionManager {
     pub fn new() -> Self {
         Self {
-            pools: DashMap::new(),
+            pools: RwLock::new(HashMap::new()),
         }
     }
 
@@ -41,8 +42,19 @@ impl ConnectionManager {
             return None;
         }
 
-        // If we have a cached pool, update its last_accessed time and return it
-        if let Some(mut cached_pool) = self.pools.get_mut(&key) {
+        // Try read lock first for cache hit
+        if let Ok(pools) = self.pools.read() {
+            if let Some(cached_pool) = pools.get(&key) {
+                // Can't update last_accessed with read lock, but that's okay for occasional misses
+                return Some(cached_pool.pool.clone());
+            }
+        }
+
+        // Cache miss or need to update timestamp - use write lock
+        let mut pools = self.pools.write().unwrap();
+
+        // Double-check after acquiring write lock
+        if let Some(cached_pool) = pools.get_mut(&key) {
             cached_pool.last_accessed = Instant::now();
             return Some(cached_pool.pool.clone());
         }
@@ -69,7 +81,7 @@ impl ConnectionManager {
             idle_timeout: Duration::from_secs(60 * 5),
         };
 
-        self.pools.insert(key, cached_pool);
+        pools.insert(key, cached_pool);
 
         Some(pool)
     }
@@ -78,8 +90,10 @@ impl ConnectionManager {
     fn cleanup_idle_pools(&self, ignore_key: &ConnectionKey) {
         let now = Instant::now();
 
+        let mut pools = self.pools.write().unwrap();
+
         // Use retain to keep only non-idle connections
-        self.pools.retain(|key, cached_pool| {
+        pools.retain(|key, cached_pool| {
             let idle_duration = now.duration_since(cached_pool.last_accessed);
             if idle_duration > cached_pool.idle_timeout && key != ignore_key {
                 tracing::debug!(

crates/pgt_workspace/src/workspace/server/connection_manager.rs

@@ -1,29 +1,38 @@
-use std::sync::Arc;
+use std::num::NonZeroUsize;
+use std::sync::{Arc, Mutex};
 
-use dashmap::DashMap;
+use lru::LruCache;
 use pgt_query_ext::diagnostics::*;
 
 use super::statement_identifier::StatementId;
 
+const DEFAULT_CACHE_SIZE: usize = 1000;
+
 pub struct PgQueryStore {
-    db: DashMap<StatementId, Arc<Result<pgt_query_ext::NodeEnum, SyntaxDiagnostic>>>,
+    db: Mutex<LruCache<StatementId, Arc<Result<pgt_query_ext::NodeEnum, SyntaxDiagnostic>>>>,
 }
 
 impl PgQueryStore {
     pub fn new() -> PgQueryStore {
-        PgQueryStore { db: DashMap::new() }
+        PgQueryStore {
+            db: Mutex::new(LruCache::new(
+                NonZeroUsize::new(DEFAULT_CACHE_SIZE).unwrap(),
+            )),
+        }
     }
 
     pub fn get_or_cache_ast(
         &self,
         statement: &StatementId,
     ) -> Arc<Result<pgt_query_ext::NodeEnum, SyntaxDiagnostic>> {
-        if let Some(existing) = self.db.get(statement).map(|x| x.clone()) {
-            return existing;
+        let mut cache = self.db.lock().unwrap();
+
+        if let Some(existing) = cache.get(statement) {
+            return existing.clone();
         }
 
         let r = Arc::new(pgt_query_ext::parse(statement.content()).map_err(SyntaxDiagnostic::from));
-        self.db.insert(statement.clone(), r.clone());
+        cache.put(statement.clone(), r.clone());
         r
     }
 }