From b356aa8df5764fee20b24152967f6a9cf3672a8b Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Wed, 24 Jun 2026 20:12:24 +0800 Subject: [PATCH] fix(connection): recover idle connection pools --- crates/dbx-core/src/connection.rs | 14 +++- crates/dbx-core/src/db/mysql.rs | 27 ++++++-- crates/dbx-core/src/query.rs | 92 ++++++++++++++++++++++++--- crates/dbx-core/tests/live_mysql57.rs | 36 +++++++++++ 4 files changed, 153 insertions(+), 16 deletions(-) diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 1e2f227a4..069c1fbf2 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -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; + } } } } diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 7b49b18e1..2b0b25b15 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -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 { - 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 { + 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, diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index aba4e7f21..694bfa1c9 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -694,7 +694,7 @@ fn should_discard_agent_pool_after_error(err: &str) -> bool { pub fn pool_error_action(db_type: Option, 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, err: &str) -> PoolErrorA } } +fn should_discard_pool_after_query_timeout(db_type: Option) -> 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, 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, 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] diff --git a/crates/dbx-core/tests/live_mysql57.rs b/crates/dbx-core/tests/live_mysql57.rs index 8977ff871..3768c4834 100644 --- a/crates/dbx-core/tests/live_mysql57.rs +++ b/crates/dbx-core/tests/live_mysql57.rs @@ -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")]]); +}