From a7f8618254cd87dc64718a2a15b1dc19d40d1f25 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Wed, 16 Jul 2025 21:54:48 +0200 Subject: [PATCH 1/2] fix: deadlock --- .../workspace/server/connection_manager.rs | 15 ++++++---- .../src/workspace/server/tree_sitter.rs | 30 ++++++++++--------- 2 files changed, 25 insertions(+), 20 deletions(-) diff --git a/crates/pgt_workspace/src/workspace/server/connection_manager.rs b/crates/pgt_workspace/src/workspace/server/connection_manager.rs index 8955b378..7d4275a7 100644 --- a/crates/pgt_workspace/src/workspace/server/connection_manager.rs +++ b/crates/pgt_workspace/src/workspace/server/connection_manager.rs @@ -34,9 +34,6 @@ 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; @@ -59,6 +56,9 @@ impl ConnectionManager { return Some(cached_pool.pool.clone()); } + // Clean up idle connections before creating new ones to avoid unbounded growth + self.cleanup_idle_pools_with_lock(&mut pools, &key); + // Create a new pool let config = PgConnectOptions::new() .host(&settings.host) @@ -87,11 +87,14 @@ impl ConnectionManager { } /// Remove pools that haven't been accessed for longer than the idle timeout - fn cleanup_idle_pools(&self, ignore_key: &ConnectionKey) { + /// This version works with an existing write lock to avoid deadlocks + fn cleanup_idle_pools_with_lock( + &self, + pools: &mut HashMap, + 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); 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 } } From 28714a53cfc49d05db0453ef19dbf549597059ce Mon Sep 17 00:00:00 2001 From: psteinroe Date: Wed, 16 Jul 2025 22:00:14 +0200 Subject: [PATCH 2/2] progress --- .../workspace/server/connection_manager.rs | 49 +++++++------------ 1 file changed, 18 insertions(+), 31 deletions(-) diff --git a/crates/pgt_workspace/src/workspace/server/connection_manager.rs b/crates/pgt_workspace/src/workspace/server/connection_manager.rs index 7d4275a7..145b6fa0 100644 --- a/crates/pgt_workspace/src/workspace/server/connection_manager.rs +++ b/crates/pgt_workspace/src/workspace/server/connection_manager.rs @@ -39,15 +39,14 @@ impl ConnectionManager { 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 @@ -57,7 +56,19 @@ impl ConnectionManager { } // Clean up idle connections before creating new ones to avoid unbounded growth - self.cleanup_idle_pools_with_lock(&mut pools, &key); + 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() @@ -85,28 +96,4 @@ impl ConnectionManager { Some(pool) } - - /// Remove pools that haven't been accessed for longer than the idle timeout - /// This version works with an existing write lock to avoid deadlocks - fn cleanup_idle_pools_with_lock( - &self, - pools: &mut HashMap, - ignore_key: &ConnectionKey, - ) { - let now = Instant::now(); - - // 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 - } - }); - } }