diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 51000799a..b4155c9d4 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -190,6 +190,7 @@ pub struct AppState { } #[derive(Clone, Copy)] +#[cfg_attr(not(test), allow(dead_code))] struct PoolActivity { last_used_at: Instant, } @@ -832,34 +833,25 @@ impl AppState { async fn start_keepalive_task(&self, pool_key: &str, pool: &PoolKind, config: &ConnectionConfig) { let interval_secs = config.keepalive_interval_secs; - let idle_timeout_secs = config.idle_timeout_secs; - let idle_cleanup_enabled = is_session_scoped_pool_key(pool_key) && idle_timeout_secs > 0; let mut target = keepalive_target_from_pool(pool, config); - if interval_secs == 0 && !idle_cleanup_enabled { + if interval_secs == 0 { return; } if interval_secs > 0 && target.is_none() { log::debug!( "Connection keepalive requested for '{pool_key}', but this database driver does not keep a pingable client handle." ); - if !idle_cleanup_enabled { - return; - } + return; }; let key = pool_key.to_string(); - let interval = if interval_secs > 0 { - Duration::from_secs(interval_secs.max(1)) - } else { - Duration::from_secs(idle_timeout_secs.min(60).max(1)) - }; + let interval = Duration::from_secs(interval_secs.max(1)); let timeout = Duration::from_secs(config.effective_connect_timeout_secs().max(1)); let connections = self.connections.clone(); let keepalive_tasks = self.keepalive_tasks.clone(); let pool_activity = self.pool_activity.clone(); let cancel_contexts = self.postgres_cancel_contexts.clone(); let running_queries = self.running_queries.clone(); - let idle_timeout = Duration::from_secs(idle_timeout_secs.max(1)); let handle = tokio::spawn(async move { loop { tokio::time::sleep(interval).await; @@ -868,27 +860,6 @@ impl AppState { continue; } - if idle_cleanup_enabled { - let idle_for = { - let activity = pool_activity.read().await; - activity.get(&key).map(|activity| activity.last_used_at.elapsed()) - }; - if idle_for.is_some_and(|elapsed| elapsed >= idle_timeout) { - log::info!( - "Closing idle session-scoped connection pool '{key}' after {}s", - idle_timeout.as_secs() - ); - keepalive_tasks.write().await.remove(&key); - pool_activity.write().await.remove(&key); - cancel_contexts.write().await.remove(&key); - let removed = connections.write().await.remove(&key); - if let Some(pool) = removed { - close_pool_kind_with_timeout(key, pool).await; - } - break; - } - } - if let Some(target) = target.as_mut() { let result = tokio::time::timeout(timeout, ping_keepalive_target(target, timeout)).await; match result { @@ -2721,6 +2692,7 @@ fn session_scoped_pool_key(base_pool_key: String, client_session_id: Option<&str .unwrap_or(base_pool_key) } +#[cfg(test)] fn is_session_scoped_pool_key(pool_key: &str) -> bool { pool_key.contains(":session:") } @@ -4186,6 +4158,51 @@ mod tests { let _ = std::fs::remove_dir_all(dir); } + #[tokio::test] + async fn session_scoped_pool_is_not_closed_by_idle_timeout() { + let (state, dir) = test_app_state().await; + let pool_key = "conn:session:tab-1"; + let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap(); + let mut config = mysql_config(None); + config.idle_timeout_secs = 1; + config.keepalive_interval_secs = 0; + + state.connections.write().await.insert(pool_key.to_string(), PoolKind::Sqlite(pool)); + state.pool_activity.write().await.insert( + pool_key.to_string(), + super::PoolActivity { last_used_at: std::time::Instant::now() - std::time::Duration::from_secs(10) }, + ); + let pool = super::clone_pool_kind(state.connections.read().await.get(pool_key).unwrap()); + state.start_keepalive_task(pool_key, &pool, &config).await; + + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + + assert!(state.connections.read().await.contains_key(pool_key)); + assert!(!state.keepalive_tasks.read().await.contains_key(pool_key)); + + let _ = std::fs::remove_dir_all(dir); + } + + #[tokio::test] + async fn close_client_session_pool_releases_session_scoped_pool() { + let (state, dir) = test_app_state().await; + let mut config = mysql_config(None); + config.id = "conn".to_string(); + config.db_type = DatabaseType::Sqlite; + state.configs.write().await.insert(config.id.clone(), config); + + let pool_key = "conn:session:tab-1"; + let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap(); + state.connections.write().await.insert(pool_key.to_string(), PoolKind::Sqlite(pool)); + state.pool_activity.write().await.insert(pool_key.to_string(), super::PoolActivity::now()); + + assert!(state.close_client_session_pool("conn", None, "tab-1").await.unwrap()); + assert!(!state.connections.read().await.contains_key(pool_key)); + assert!(!state.pool_activity.read().await.contains_key(pool_key)); + + let _ = std::fs::remove_dir_all(dir); + } + #[test] fn agent_validate_connection_unknown_method_is_not_stale() { assert!(super::is_agent_validate_connection_unsupported(