From 996f4c2d385de67161f515bf872812004a4ed5c6 Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Sat, 11 Jul 2026 00:02:33 +0800
Subject: [PATCH] fix(connection): preserve saturated connection pools

---
 crates/dbx-core/src/connection.rs | 57 +++++++++++++++++++++++++++----
 1 file changed, 50 insertions(+), 7 deletions(-)

diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs
index 6a5a59196..51000799a 100644
--- a/crates/dbx-core/src/connection.rs
+++ b/crates/dbx-core/src/connection.rs
@@ -37,6 +37,7 @@ const SQLSERVER_LEGACY_DRIVER_INSTALL_HINT: &str =
 const DEFAULT_AGENT_CONNECT_TIMEOUT_SECS: u64 = 30;
 const ACCESS_AGENT_CONNECT_TIMEOUT_SECS: u64 = 30;
 const POOL_CLOSE_TIMEOUT_SECS: u64 = 3;
+const HEALTH_CHECK_POOL_ACQUIRE_TIMEOUT: Duration = Duration::from_millis(500);
 
 #[cfg(feature = "duckdb-bundled")]
 mod duckdb_types {
@@ -1664,19 +1665,38 @@ impl AppState {
                 PoolKind::Mysql(pool, _) => {
                     let pool = pool.clone();
                     drop(connections);
-                    match db::mysql::get_conn_with_health_check(&pool).await {
-                        Ok(_) => false,
-                        Err(err) => {
+                    match tokio::time::timeout(HEALTH_CHECK_POOL_ACQUIRE_TIMEOUT, pool.get_conn()).await {
+                        // Pool saturation means active work, not a dead connection. Removing this pool would
+                        // start a competing reconnect while foreground queries and metadata are still running.
+                        Err(_) => {
+                            log::debug!("MySQL connection pool '{pool_key}' is busy; skipping health probe");
+                            false
+                        }
+                        Ok(Err(err)) => {
                             log::warn!("MySQL connection pool '{pool_key}' is stale: {err}");
                             true
                         }
+                        Ok(Ok(mut conn)) => {
+                            let timeout = crate::db::connection_timeout();
+                            match tokio::time::timeout(timeout, conn.ping()).await {
+                                Ok(Ok(())) => false,
+                                Ok(Err(err)) => {
+                                    log::warn!("MySQL connection pool '{pool_key}' is stale: {err}");
+                                    true
+                                }
+                                Err(_) => {
+                                    log::warn!("MySQL connection pool '{pool_key}' is stale: health check timed out");
+                                    true
+                                }
+                            }
+                        }
                     }
                 }
                 PoolKind::Postgres(pool) => {
                     let pool = pool.clone();
                     drop(connections);
                     let timeout = crate::db::connection_timeout();
-                    match tokio::time::timeout(timeout, pool.get()).await {
+                    match tokio::time::timeout(HEALTH_CHECK_POOL_ACQUIRE_TIMEOUT, pool.get()).await {
                         Ok(Ok(client)) => match tokio::time::timeout(timeout, client.simple_query("SELECT 1")).await {
                             Ok(Ok(_)) => false,
                             Ok(Err(err)) => {
@@ -1693,8 +1713,8 @@ impl AppState {
                             true
                         }
                         Err(_) => {
-                            log::warn!("PostgreSQL connection pool '{pool_key}' is stale: get connection timed out");
-                            true
+                            log::debug!("PostgreSQL connection pool '{pool_key}' is busy; skipping health probe");
+                            false
                         }
                     }
                 }
@@ -3087,7 +3107,7 @@ mod tests {
         oceanbase_mysql_setup_queries, prestosql_jdbc_config_for_endpoint, redacted_connection_url_for_endpoint,
         redis_sentinel_transport_id, redis_sentinel_transport_prefix, sqlserver_legacy_agent_config,
         sqlserver_legacy_agent_error, uses_bare_mysql_pool, uses_tcp_probe, validate_h2_database_path, AppState,
-        PoolKind, PRESTOSQL_JDBC_DRIVER_CLASS,
+        MysqlMode, PoolKind, PRESTOSQL_JDBC_DRIVER_CLASS,
     };
     use crate::agent_connection::{
         agent_connect_params, mongo_legacy_error_with_auth_hint, mongo_uses_legacy_driver,
@@ -3103,6 +3123,7 @@ mod tests {
     use crate::query;
     use crate::schema;
     use crate::storage::Storage;
+    use std::time::{Duration, Instant};
 
     fn mysql_config(database: Option<&str>) -> ConnectionConfig {
         ConnectionConfig {
@@ -3592,6 +3613,28 @@ mod tests {
         let _ = std::fs::remove_dir_all(dir);
     }
 
+    #[tokio::test]
+    #[ignore = "requires DBX_TEST_MYSQL_URL"]
+    async fn live_mysql_health_check_keeps_saturated_pool() {
+        let url = std::env::var("DBX_TEST_MYSQL_URL").expect("DBX_TEST_MYSQL_URL is required");
+        let (state, dir) = test_app_state().await;
+        let config = mysql_config(Some("testdb"));
+        let pool = db::mysql::connect_bare_with_pool_limit(&url, Duration::from_secs(5), 1).await.unwrap();
+        state
+            .insert_connection_pool("conn".to_string(), PoolKind::Mysql(pool.clone(), MysqlMode::Normal), &config)
+            .await;
+        let held_connection = pool.get_conn().await.unwrap();
+
+        let started = Instant::now();
+        state.check_connection_health("conn").await.unwrap();
+
+        assert!(started.elapsed() < Duration::from_secs(2));
+        assert!(state.connections.read().await.contains_key("conn"));
+        drop(held_connection);
+        state.remove_connection_pools_detached("conn").await;
+        let _ = std::fs::remove_dir_all(dir);
+    }
+
     #[tokio::test]
     async fn jdbc_plugin_env_uses_managed_jre_when_installed() {
         let dir = std::env::temp_dir().join(format!("dbx-core-jdbc-managed-jre-{}", uuid::Uuid::new_v4()));