diff --git a/crates/pgt_workspace/src/workspace/server/connection_manager.rs b/crates/pgt_workspace/src/workspace/server/connection_manager.rs index 8955b378..145b6fa0 100644 --- a/crates/pgt_workspace/src/workspace/server/connection_manager.rs +++ b/crates/pgt_workspace/src/workspace/server/connection_manager.rs @@ -34,23 +34,19 @@ impl ConnectionManager { pub(crate) fn get_pool(&self, settings: &DatabaseSettings) -> Option { let key = ConnectionKey::from(settings); - // Cleanup idle connections first - self.cleanup_idle_pools(&key); - if !settings.enable_connection { tracing::info!("Database connection disabled."); return None; } - // 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()); + { + if let Ok(pools) = self.pools.read() { + if let Some(cached_pool) = pools.get(&key) { + 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 @@ -59,6 +55,21 @@ impl ConnectionManager { return Some(cached_pool.pool.clone()); } + // Clean up idle connections before creating new ones to avoid unbounded growth + let now = Instant::now(); + pools.retain(|k, cached_pool| { + let idle_duration = now.duration_since(cached_pool.last_accessed); + if idle_duration > cached_pool.idle_timeout && k != &key { + tracing::debug!( + "Removing idle database connection (idle for {:?})", + idle_duration + ); + false + } else { + true + } + }); + // Create a new pool let config = PgConnectOptions::new() .host(&settings.host) @@ -85,25 +96,4 @@ impl ConnectionManager { Some(pool) } - - /// Remove pools that haven't been accessed for longer than the idle timeout - 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 - 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!( - "Removing idle database connection (idle for {:?})", - idle_duration - ); - false - } else { - true - } - }); - } } diff --git a/crates/pgt_workspace/src/workspace/server/tree_sitter.rs b/crates/pgt_workspace/src/workspace/server/tree_sitter.rs index b8f62b63..71411d27 100644 --- a/crates/pgt_workspace/src/workspace/server/tree_sitter.rs +++ b/crates/pgt_workspace/src/workspace/server/tree_sitter.rs @@ -28,27 +28,29 @@ impl TreeSitterStore { } pub fn get_or_cache_tree(&self, statement: &StatementId) -> Arc { - let mut cache = self.db.lock().expect("Failed to lock cache"); - - if let Some(existing) = cache.get(statement) { - return existing.clone(); + // First check cache + { + let mut cache = self.db.lock().unwrap(); + if let Some(existing) = cache.get(statement) { + return existing.clone(); + } } - // Cache miss - drop cache lock, parse, then re-acquire to insert - drop(cache); - - let mut parser = self.parser.lock().expect("Failed to lock parser"); + // Cache miss - parse outside of cache lock to avoid deadlock + let mut parser = self.parser.lock().unwrap(); let tree = Arc::new(parser.parse(statement.content(), None).unwrap()); drop(parser); - let mut cache = self.db.lock().expect("Failed to lock cache"); - - // Double-check after re-acquiring lock - if let Some(existing) = cache.get(statement) { - return existing.clone(); + // Insert into cache + { + let mut cache = self.db.lock().unwrap(); + // Double-check in case another thread inserted while we were parsing + if let Some(existing) = cache.get(statement) { + return existing.clone(); + } + cache.put(statement.clone(), tree.clone()); } - cache.put(statement.clone(), tree.clone()); tree } }