fix(connection): avoid blocking disconnect cleanup

This commit is contained in:
t8y2
2026-06-16 10:45:14 +08:00
parent 80a390a4c0
commit 47915c243d
2 changed files with 38 additions and 6 deletions
+37 -5
View File
@@ -32,6 +32,7 @@ pub const JDBC_PLUGIN_NOT_INSTALLED: &str =
"JDBC plugin is not installed. Install the optional JDBC plugin to use this connection.";
const DEFAULT_AGENT_CONNECT_TIMEOUT_SECS: u64 = 30;
const ACCESS_AGENT_CONNECT_TIMEOUT_SECS: u64 = 30;
const POOL_CLOSE_TIMEOUT_SECS: u64 = 5;
#[cfg(feature = "duckdb-bundled")]
mod duckdb_types {
@@ -1061,6 +1062,16 @@ impl AppState {
}
pub async fn remove_connection_pools(&self, connection_id: &str) {
let removed = self.drain_connection_pools(connection_id).await;
close_removed_pools(removed).await;
}
pub async fn remove_connection_pools_detached(&self, connection_id: &str) {
let removed = self.drain_connection_pools(connection_id).await;
close_removed_pools_in_background(removed);
}
async fn drain_connection_pools(&self, connection_id: &str) -> Vec<(String, PoolKind)> {
let pool_prefix = format!("{connection_id}:");
let keys_to_remove: Vec<String> = self
.connections
@@ -1075,13 +1086,10 @@ impl AppState {
let mut removed = Vec::with_capacity(keys_to_remove.len());
for key in keys_to_remove {
if let Some(pool) = conns.remove(&key) {
removed.push(pool);
removed.push((key, pool));
}
}
drop(conns);
for pool in removed {
close_pool_kind(pool).await;
}
removed
}
async fn uses_forwarded_transport(&self, connection_id: &str) -> bool {
@@ -1237,6 +1245,30 @@ pub async fn close_pool_kind(pool: PoolKind) {
}
}
async fn close_removed_pools(removed: Vec<(String, PoolKind)>) {
for (pool_key, pool) in removed {
close_pool_kind_with_timeout(pool_key, pool).await;
}
}
fn close_removed_pools_in_background(removed: Vec<(String, PoolKind)>) {
if removed.is_empty() {
return;
}
tokio::spawn(async move {
close_removed_pools(removed).await;
});
}
async fn close_pool_kind_with_timeout(pool_key: String, pool: PoolKind) {
match tokio::time::timeout(Duration::from_secs(POOL_CLOSE_TIMEOUT_SECS), close_pool_kind(pool)).await {
Ok(()) => {}
Err(_) => log::warn!(
"Timed out closing connection pool '{pool_key}' after {POOL_CLOSE_TIMEOUT_SECS}s; cleanup will continue by dropping the pool handle."
),
}
}
fn extract_auth_token_from_params(params: &str) -> Option<String> {
params
.trim()
+1 -1
View File
@@ -734,7 +734,7 @@ pub async fn connection_final_proxy_port(
#[tauri::command]
pub async fn disconnect_db(state: State<'_, Arc<AppState>>, connection_id: String) -> Result<(), String> {
state.remove_connection_pools(&connection_id).await;
state.remove_connection_pools_detached(&connection_id).await;
state.reset_connection_transport(&connection_id).await;
if connection_id.starts_with("__visible_draft_") {
state.configs.write().await.remove(&connection_id);