From 47915c243d133d6401cf2a8363a611d66159cf53 Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Tue, 16 Jun 2026 10:45:14 +0800
Subject: [PATCH] fix(connection): avoid blocking disconnect cleanup

---
 crates/dbx-core/src/connection.rs    | 42 ++++++++++++++++++++++++----
 src-tauri/src/commands/connection.rs |  2 +-
 2 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs
index 3a4deeb8c5..ec64f8933f 100644
--- a/crates/dbx-core/src/connection.rs
+++ b/crates/dbx-core/src/connection.rs
@@ -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()
diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs
index 135594fb7e..6bd5b6efca 100644
--- a/src-tauri/src/commands/connection.rs
+++ b/src-tauri/src/commands/connection.rs
@@ -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);