xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(connection): recover idle connection pools

This commit is contained in:
t8y2 2026-06-24 20:12:24 +08:00
parent a87e4be029
commit b356aa8df5
4 changed files with 153 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
crates/dbx-core/src/connection.rs
View File

@ -538,7 +538,19 @@ impl AppState {
}
break;
}
Err(_) => log::warn!("Connection keepalive timed out for '{key}' after {}s", timeout.as_secs()),
Err(_) => {
log::warn!(
"Connection keepalive timed out for '{key}' after {}s; invalidating pool",
timeout.as_secs()
);
keepalive_tasks.write().await.remove(&key);
pool_activity.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;
}
}
}
}

27
crates/dbx-core/src/db/mysql.rs
View File

@ -1831,16 +1831,33 @@ fn skip_mysql_quoted(sql: &str, start: usize, quote: u8) -> usize {
/// Get a connection from the pool with a health check. If the connection is dead
/// (e.g. after app was backgrounded), it tries again with a fresh connection.
pub async fn get_conn_with_health_check(pool: &MySqlPool) -> Result<mysql_async::Conn, String> {
let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?;
match tokio::time::timeout(crate::db::connection_timeout(), conn.ping()).await {
Ok(Ok(())) => Ok(conn),
let timeout = crate::db::connection_timeout();
let mut conn = get_conn_with_timeout(pool, timeout).await?;
match ping_conn_with_timeout(&mut conn, timeout).await {
Ok(()) => Ok(conn),
_ => {
let _ = conn.disconnect().await;
pool.get_conn().await.map_err(|e| e.to_string())
let _ = tokio::time::timeout(timeout, conn.disconnect()).await;
let mut conn = get_conn_with_timeout(pool, timeout).await?;
ping_conn_with_timeout(&mut conn, timeout).await?;
Ok(conn)
}
}
}
async fn get_conn_with_timeout(pool: &MySqlPool, timeout: Duration) -> Result<mysql_async::Conn, String> {
tokio::time::timeout(timeout, pool.get_conn())
.await
.map_err(|_| "MySQL get connection timed out".to_string())?
.map_err(|e| e.to_string())
}
async fn ping_conn_with_timeout(conn: &mut mysql_async::Conn, timeout: Duration) -> Result<(), String> {
tokio::time::timeout(timeout, conn.ping())
.await
.map_err(|_| "MySQL ping timed out".to_string())?
.map_err(|e| e.to_string())
}
async fn execute_result_set_with_text_protocol_on_conn(
conn: &mut mysql_async::Conn,
sql: &str,

92
crates/dbx-core/src/query.rs
View File

@ -694,7 +694,7 @@ fn should_discard_agent_pool_after_error(err: &str) -> bool {
pub fn pool_error_action(db_type: Option<DatabaseType>, err: &str) -> PoolErrorAction {
let lower = err.to_lowercase();
if db::sqlserver::is_driver_panic_error(err)
|| (db_type == Some(DatabaseType::SqlServer) && is_dbx_query_timeout_error(&lower))
|| (is_dbx_query_timeout_error(&lower) && should_discard_pool_after_query_timeout(db_type))
|| (db_type.is_some_and(|db_type| database_capabilities::is_agent_type(&db_type))
&& should_discard_agent_pool_after_error(err)
&& !is_connection_error(err))
@ -709,6 +709,34 @@ pub fn pool_error_action(db_type: Option<DatabaseType>, err: &str) -> PoolErrorA
}
}
fn should_discard_pool_after_query_timeout(db_type: Option<DatabaseType>) -> bool {
let Some(db_type) = db_type else {
return false;
};
database_capabilities::is_agent_type(&db_type)
|| matches!(
db_type,
DatabaseType::Mysql
| DatabaseType::Postgres
| DatabaseType::Redshift
| DatabaseType::Gaussdb
| DatabaseType::Kwdb
| DatabaseType::OpenGauss
| DatabaseType::Questdb
| DatabaseType::Doris
| DatabaseType::StarRocks
| DatabaseType::ManticoreSearch
| DatabaseType::ClickHouse
| DatabaseType::SqlServer
| DatabaseType::Rqlite
| DatabaseType::Turso
| DatabaseType::Elasticsearch
| DatabaseType::Qdrant
| DatabaseType::Milvus
| DatabaseType::InfluxDb
)
}
pub fn should_discard_pool_after_error(db_type: Option<DatabaseType>, err: &str) -> bool {
matches!(pool_error_action(db_type, err), PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry)
}
@ -1238,15 +1266,18 @@ pub async fn execute_sql_statement_with_options(
do_execute(state, &pool_key, mysql_dialect, Some(database), sql, schema, cancel_token.clone(), options.clone())
.await;
match &result {
Err(e)
if pool_error_action(db_type, e) == PoolErrorAction::ReconnectAndRetry && !is_canceled(&cancel_token) =>
{
let action = result.as_ref().err().map(|e| pool_error_action(db_type, e));
match action {
Some(PoolErrorAction::ReconnectAndRetry) if !is_canceled(&cancel_token) => {
let db_opt = if database.is_empty() { None } else { Some(database) };
let new_key =
state.reconnect_pool_for_session(connection_id, db_opt, options.client_session_id.as_deref()).await?;
do_execute(state, &new_key, mysql_dialect, Some(database), sql, schema, cancel_token, options).await
}
Some(PoolErrorAction::Discard) => {
state.remove_pool_by_key(&pool_key).await;
result
}
_ => result,
}
}
@ -1467,7 +1498,18 @@ pub async fn execute_multi_core_with_options(
// Read-only check for MySQL batch path
check_read_only_for_connection_multi(state, &pool_key, &statements).await?;
let mysql_dialect = connection_mysql_query_dialect(state, connection_id).await;
return execute_multi_mysql(&pool, mode, mysql_dialect, &statements, cancel_token, options).await;
return execute_multi_mysql(
state,
&pool_key,
db_type,
&pool,
mode,
mysql_dialect,
&statements,
cancel_token,
options,
)
.await;
}
let mut results = Vec::with_capacity(statements.len());
@ -1498,6 +1540,9 @@ pub async fn execute_multi_core_with_options(
}
async fn execute_multi_mysql(
state: &AppState,
pool_key: &str,
db_type: Option<DatabaseType>,
pool: &db::mysql::MySqlPool,
mode: crate::connection::MysqlMode,
dialect: db::mysql::MySqlQueryDialect,
@ -1510,7 +1555,13 @@ async fn execute_multi_mysql(
let max_rows = options.max_rows;
let mut conn = match db::mysql::get_conn_with_health_check(pool).await {
Ok(conn) => conn,
Err(err) => return Ok(vec![error_query_result(err)]),
Err(err) => {
if matches!(pool_error_action(db_type, &err), PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry)
{
state.remove_pool_by_key(pool_key).await;
}
return Ok(vec![error_query_result(err)]);
}
};
let mut results = Vec::with_capacity(statements.len());
@ -1528,7 +1579,14 @@ async fn execute_multi_mysql(
.await
{
Ok(result) => results.push(result),
Err(err) => results.push(error_query_result(err)),
Err(err) => {
let action = pool_error_action(db_type, &err);
results.push(error_query_result(err));
if matches!(action, PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry) {
state.remove_pool_by_key(pool_key).await;
break;
}
}
}
}
@ -1745,6 +1803,7 @@ pub async fn execute_statements_in_transaction(
check_read_only_for_connection_multi(state, &pool_key, statements).await?;
let start = std::time::Instant::now();
let db_type = connection_database_type(state, connection_id).await;
// Clone the pool handle within the lock, then drop it before any async work.
let path = {
@ -1780,7 +1839,7 @@ pub async fn execute_statements_in_transaction(
})
};
match path {
let result = match path {
Some(TxPath::Pg(pool)) => exec_tx_pg_inner(pool, statements, schema, start).await,
Some(TxPath::Mysql(pool, _bare)) => exec_tx_mysql_inner(pool, statements, start).await,
Some(TxPath::Sqlite(pool)) => exec_tx_sqlite_inner(pool, statements, start).await,
@ -1793,7 +1852,15 @@ pub async fn execute_statements_in_transaction(
exec_tx_none_inner(state, &pool_key, mysql_dialect, Some(database), statements, schema, start).await
}
None => Err("Connection not found for transaction".to_string()),
};
if let Err(err) = result.as_ref() {
if matches!(pool_error_action(db_type, err), PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry) {
state.remove_pool_by_key(&pool_key).await;
}
}
result
}
/// Owned pool variants for safe dispatch across async boundaries.
@ -2201,7 +2268,12 @@ mod tests {
let err = "Query timed out after 30 seconds";
assert_eq!(pool_error_action(Some(DatabaseType::SqlServer), err), PoolErrorAction::Discard);
assert_eq!(pool_error_action(Some(DatabaseType::Mysql), err), PoolErrorAction::Keep);
assert_eq!(pool_error_action(Some(DatabaseType::Mysql), err), PoolErrorAction::Discard);
assert_eq!(pool_error_action(Some(DatabaseType::Postgres), err), PoolErrorAction::Discard);
assert_eq!(pool_error_action(Some(DatabaseType::ClickHouse), err), PoolErrorAction::Discard);
assert_eq!(pool_error_action(Some(DatabaseType::Oracle), err), PoolErrorAction::Discard);
assert_eq!(pool_error_action(Some(DatabaseType::Sqlite), err), PoolErrorAction::Keep);
assert_eq!(pool_error_action(Some(DatabaseType::DuckDb), err), PoolErrorAction::Keep);
}
#[test]

36
crates/dbx-core/tests/live_mysql57.rs
View File

@ -67,3 +67,39 @@ async fn live_mysql_query_cancel_kills_running_sleep() {
let result = result.unwrap();
assert_eq!(result.rows, vec![vec![serde_json::json!("1")]]);
}
#[tokio::test]
#[ignore = "requires a remote MySQL endpoint"]
async fn live_mysql_recovers_after_server_idle_disconnect() {
let url = std::env::var("DBX_LIVE_MYSQL_IDLE_URL").expect("DBX_LIVE_MYSQL_IDLE_URL");
let pool =
dbx_core::db::mysql::connect_with_ca_cert_and_pool_limit(&url, None, std::time::Duration::from_secs(5), 1)
.await
.unwrap();
dbx_core::db::mysql::execute_query_with_max_rows(
&pool,
"SET SESSION wait_timeout = 1",
false,
Some(10),
Default::default(),
)
.await
.unwrap();
tokio::time::sleep(std::time::Duration::from_secs(2)).await;
let result = dbx_core::db::mysql::execute_query_with_max_rows(
&pool,
"SELECT 1 AS recovered",
false,
Some(10),
Default::default(),
)
.await
.unwrap();
assert_eq!(result.columns, vec!["recovered"]);
assert_eq!(result.rows, vec![vec![serde_json::json!("1")]]);
}