fix(connection): preserve saturated connection pools

This commit is contained in:
t8y2 2026-07-11 00:02:33 +08:00
parent 375be106f8
commit 996f4c2d38
1 changed files with 50 additions and 7 deletions

View File

@ -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()));