fix(mysql): prevent stale export pool cleanup

2026-11-15 14:52:36 +00:00 · 2026-08-25 22:50:40 +08:00
1 parent aa91a0aaa9
commit 0658d18922
@@ -3325,6 +3325,7 @@ impl AppState {
            return false;
        }

+        let mut checked_mysql_pool = None;
        let stale = {
            let connections = self.connections.read().await;
            let Some(pool) = connections.get(pool_key) else {
@@ -3333,6 +3334,7 @@ impl AppState {
            match pool {
                PoolKind::Mysql(pool, _) => {
                    let pool = pool.clone();
+                    checked_mysql_pool = Some(pool.clone());
                    drop(connections);
                    match db::mysql::checkout_mysql_conn(&pool, HEALTH_CHECK_POOL_ACQUIRE_TIMEOUT).await {
                        // Pool saturation means active work, not a dead connection. Removing this pool would
@@ -3624,16 +3626,75 @@ impl AppState {
            return false;
        }

+        if let Some(checked_pool) = checked_mysql_pool {
+            return self.remove_stale_mysql_pool_if_current(pool_key, &checked_pool).await;
+        }
+
+        let removed = self.connections.write().await.remove(pool_key);
+
+        let Some(pool) = removed else {
+            return false;
+        };
+
        self.stop_keepalive_task(pool_key).await;
        self.pool_activity.write().await.remove(pool_key);
        self.postgres_cancel_contexts.write().await.remove(pool_key);
-        let removed = self.connections.write().await.remove(pool_key);
-        if let Some(pool) = removed {
-            self.pool_routing_control().close_pool_with_timeout(pool_key.to_string(), pool).await;
-            true
-        } else {
-            false
-        }
+        self.pool_routing_control().close_pool_with_timeout(pool_key.to_string(), pool).await;
+        true
+    }
+
+    async fn remove_stale_mysql_pool_if_current(&self, pool_key: &str, checked_pool: &db::mysql::MySqlPool) -> bool {
+        self.remove_stale_mysql_pool_if_current_inner(
+            pool_key,
+            checked_pool,
+            #[cfg(test)]
+            None,
+        )
+        .await
+    }
+
+    async fn remove_stale_mysql_pool_if_current_inner(
+        &self,
+        pool_key: &str,
+        checked_pool: &db::mysql::MySqlPool,
+        #[cfg(test)] cleanup_barriers: Option<(
+            std::sync::Arc<tokio::sync::Barrier>,
+            std::sync::Arc<tokio::sync::Barrier>,
+        )>,
+    ) -> bool {
+        let routing = self.pool_routing_control();
+        let removed = loop {
+            let mut connections = self.connections.write().await;
+            let is_current = matches!(
+                connections.get(pool_key),
+                Some(PoolKind::Mysql(current, _)) if checked_pool.is_same_pool(current)
+            );
+            if !is_current {
+                log::debug!(
+                    "MySQL connection pool '{pool_key}' was replaced while its health check was running; keeping the current route"
+                );
+                return false;
+            }
+            let Ok(mut activity) = self.pool_activity.try_write() else {
+                drop(connections);
+                tokio::task::yield_now().await;
+                continue;
+            };
+            let removed = remove_mysql_pool_if_current(&mut connections, pool_key, checked_pool)
+                .expect("checked MySQL pool must remain current while routing is locked");
+            #[cfg(test)]
+            if let Some((route_removed, continue_cleanup)) = cleanup_barriers.as_ref() {
+                route_removed.wait().await;
+                continue_cleanup.wait().await;
+            }
+            routing.stop_keepalive(pool_key);
+            activity.remove(pool_key);
+            break removed;
+        };
+
+        self.postgres_cancel_contexts.write().await.remove(pool_key);
+        routing.close_pool_with_timeout(pool_key.to_string(), removed).await;
+        true
    }

    pub async fn reconnect_pool(&self, connection_id: &str, database: Option<&str>) -> Result<String, String> {
@@ -5304,6 +5365,18 @@ fn clone_pool_kind(pool: &PoolKind) -> PoolKind {
    }
 }

+fn remove_mysql_pool_if_current(
+    connections: &mut HashMap<String, PoolKind>,
+    pool_key: &str,
+    expected: &db::mysql::MySqlPool,
+) -> Option<PoolKind> {
+    let is_current = matches!(
+        connections.get(pool_key),
+        Some(PoolKind::Mysql(current, _)) if expected.is_same_pool(current)
+    );
+    is_current.then(|| connections.remove(pool_key)).flatten()
+}
+
 async fn close_pool_kind(pool: PoolKind) -> Result<(), String> {
    match pool {
        PoolKind::Mysql(p, _) => {
@@ -7334,6 +7407,102 @@ mod tests {
        assert_eq!(super::mysql_pool_max_connections_for_session(Some("tab-1")), 1);
    }

+    #[test]
+    fn stale_mysql_pool_observation_does_not_remove_replacement_generation() {
+        let checked = crate::db::mysql::MySqlPool::new("mysql://root@127.0.0.1:3306/app", 10);
+        let checked_clone = checked.clone();
+        let replacement = crate::db::mysql::MySqlPool::new("mysql://root@127.0.0.1:3306/app", 10);
+        assert!(checked.is_same_pool(&checked_clone));
+        assert!(!checked.is_same_pool(&replacement));
+
+        let mut connections = std::collections::HashMap::from([(
+            "conn:app".to_string(),
+            PoolKind::Mysql(replacement.clone(), MysqlMode::Normal),
+        )]);
+
+        assert!(super::remove_mysql_pool_if_current(&mut connections, "conn:app", &checked).is_none());
+        assert!(
+            matches!(connections.get("conn:app"), Some(PoolKind::Mysql(current, _)) if replacement.is_same_pool(current))
+        );
+
+        assert!(super::remove_mysql_pool_if_current(&mut connections, "conn:app", &replacement).is_some());
+        assert!(!connections.contains_key("conn:app"));
+    }
+
+    #[tokio::test]
+    async fn stale_mysql_cleanup_preserves_replacement_support_state() {
+        let (state, dir) = test_app_state().await;
+        let state = std::sync::Arc::new(state);
+        let pool_key = "conn:app";
+        let checked = crate::db::mysql::MySqlPool::new("mysql://root@127.0.0.1:3306/app", 10);
+        let replacement = crate::db::mysql::MySqlPool::new("mysql://root@127.0.0.1:3306/app", 10);
+        let mut config = mysql_config(Some("app"));
+        config.keepalive_interval_secs = 60;
+        state
+            .connections
+            .write()
+            .await
+            .insert(pool_key.to_string(), PoolKind::Mysql(checked.clone(), MysqlMode::Normal));
+        state.pool_activity.write().await.insert(pool_key.to_string(), super::PoolActivity::now());
+        state.start_keepalive_task(
+            pool_key,
+            &PoolKind::Mysql(checked.clone(), MysqlMode::Normal),
+            &config,
+            #[cfg(feature = "mq-admin")]
+            None,
+        );
+        assert_eq!(state.supervised_task_count(), 1);
+
+        let route_removed = std::sync::Arc::new(tokio::sync::Barrier::new(2));
+        let continue_cleanup = std::sync::Arc::new(tokio::sync::Barrier::new(2));
+        let cleanup_state = state.clone();
+        let cleanup_checked = checked.clone();
+        let cleanup_route_removed = route_removed.clone();
+        let cleanup_continue = continue_cleanup.clone();
+        let cleanup = tokio::spawn(async move {
+            cleanup_state
+                .remove_stale_mysql_pool_if_current_inner(
+                    pool_key,
+                    &cleanup_checked,
+                    Some((cleanup_route_removed, cleanup_continue)),
+                )
+                .await
+        });
+        route_removed.wait().await;
+
+        let publish_state = state.clone();
+        let publish_config = config.clone();
+        let publish_replacement = replacement.clone();
+        let mut publish = tokio::spawn(async move {
+            publish_state
+                .insert_connection_pool(
+                    pool_key.to_string(),
+                    PoolKind::Mysql(publish_replacement, MysqlMode::Normal),
+                    &publish_config,
+                )
+                .await
+        });
+
+        assert!(
+            tokio::time::timeout(Duration::from_millis(50), &mut publish).await.is_err(),
+            "replacement publication must wait while stale support state is cleaned"
+        );
+        continue_cleanup.wait().await;
+        assert!(cleanup.await.unwrap());
+        publish.await.unwrap().unwrap();
+
+        let connections = state.connections.read().await;
+        assert!(
+            matches!(connections.get(pool_key), Some(PoolKind::Mysql(current, _)) if replacement.is_same_pool(current))
+        );
+        drop(connections);
+        assert!(state.pool_activity.read().await.contains_key(pool_key));
+        assert_eq!(state.supervised_task_count(), 1);
+
+        state.shutdown(Duration::from_secs(1)).await;
+        let _ = std::fs::remove_dir_all(dir);
+    }
+
    #[test]
    fn metadata_concurrency_reserves_query_capacity() {
        assert_eq!(super::metadata_concurrency_limit(DatabaseType::Mysql, 10), 6);
@@ -3,7 +3,8 @@ use serde_json::Value;
 use std::collections::{HashMap, HashSet};
 use std::future::Future;
 use std::io::{BufWriter, Write};
-use std::sync::RwLock;
+use std::pin::Pin;
+use std::sync::{Arc, RwLock};
 use std::time::Duration;
 use tokio_util::sync::CancellationToken;

@@ -1586,10 +1587,10 @@ pub async fn clear_export_cancelled(export_id: &str) {
 /// not accept a cancellation token; polling here keeps the export task
 /// responsive and dropping the pending future follows the same bounded
 /// prefetch cancellation behavior used below.
-async fn await_export_operation<T, F>(export_id: &str, operation: F) -> Result<T, String>
-where
-    F: Future<Output = Result<T, String>>,
-{
+async fn await_export_operation<T>(
+    export_id: &str,
+    operation: Pin<Box<dyn Future<Output = Result<T, String>> + Send + '_>>,
+) -> Result<T, String> {
    tokio::pin!(operation);
    loop {
        if is_export_cancelled_now(export_id) {
@@ -1609,6 +1610,43 @@ where
    }
 }

+struct AbortExportTaskOnDrop(tokio::task::AbortHandle);
+
+impl Drop for AbortExportTaskOnDrop {
+    fn drop(&mut self) {
+        self.0.abort();
+    }
+}
+
+async fn get_export_table_ddl_isolated(
+    state: Arc<crate::connection::AppState>,
+    connection_id: String,
+    database: String,
+    schema: String,
+    table: String,
+) -> Result<String, String> {
+    let task = tokio::spawn(async move {
+        crate::schema::get_table_relation_export_ddl_core(&state, &connection_id, &database, &schema, &table, None)
+            .await
+    });
+    let _abort_on_drop = AbortExportTaskOnDrop(task.abort_handle());
+    task.await.map_err(|error| format!("Database export metadata task failed: {error}"))?
+}
+
+async fn get_export_table_columns_isolated(
+    state: Arc<crate::connection::AppState>,
+    connection_id: String,
+    database: String,
+    schema: String,
+    table: String,
+) -> Result<Vec<crate::db::ColumnInfo>, String> {
+    let task = tokio::spawn(async move {
+        crate::schema::get_columns_core(&state, &connection_id, &database, &schema, &table).await
+    });
+    let _abort_on_drop = AbortExportTaskOnDrop(task.abort_handle());
+    task.await.map_err(|error| format!("Database export metadata task failed: {error}"))?
+}
+
 fn snapshot_batch_cancelled(db_type: &DatabaseType, export_id: &str) -> bool {
    *db_type != DatabaseType::Postgres && is_export_cancelled_now(export_id)
 }
@@ -2106,7 +2144,7 @@ mod windows_export_destination {
 }

 pub async fn export_database_sql_core(
-    state: &crate::connection::AppState,
+    state: &Arc<crate::connection::AppState>,
    request: &DatabaseExportRequest,
    on_progress: impl Fn(ExportProgress) + Sync,
 ) -> Result<(), String> {
@@ -2126,7 +2164,7 @@ pub async fn export_database_sql_core(
 }

 async fn export_database_sql_core_inner(
-    state: &crate::connection::AppState,
+    state: &Arc<crate::connection::AppState>,
    request: &DatabaseExportRequest,
    on_progress: impl Fn(ExportProgress) + Sync,
 ) -> Result<(), String> {
@@ -2275,7 +2313,7 @@ async fn export_database_sql_core_inner(
    let postgres_sequences = if request.include_structure && matches!(db_type, DatabaseType::Postgres) {
        match await_export_operation(
            &request.export_id,
-            list_postgres_export_sequences(
+            Box::pin(list_postgres_export_sequences(
                state,
                &pool_key,
                &request.schema,
@@ -2283,7 +2321,7 @@ async fn export_database_sql_core_inner(
                &request.excluded_tables,
                request.include_objects,
                request.fail_on_error,
-            ),
+            )),
        )
        .await
        {
@@ -2416,8 +2454,8 @@ async fn export_database_sql_core_inner(
        use futures::StreamExt;
        let prefetch_targets: Vec<(usize, String)> =
            tables.iter().enumerate().map(|(index, table_info)| (index, table_info.name.clone())).collect();
-        let mut prefetch_stream =
-            futures::stream::iter(prefetch_targets.into_iter().map(|(index, table_name)| async move {
+        let mut prefetch_stream = futures::stream::iter(prefetch_targets.into_iter().map(|(index, table_name)| {
+            Box::pin(async move {
                if is_export_cancelled_now(&request.export_id) {
                    return (index, PrefetchedTableMetadata { ddl: None, columns: None });
                }
@@ -2425,14 +2463,13 @@ async fn export_database_sql_core_inner(
                    Some(
                        await_export_operation(
                            &request.export_id,
-                            crate::schema::get_table_relation_export_ddl_core(
-                                state,
-                                &request.connection_id,
-                                &request.database,
-                                &request.schema,
-                                &table_name,
-                                None,
-                            ),
+                            Box::pin(get_export_table_ddl_isolated(
+                                state.clone(),
+                                request.connection_id.clone(),
+                                request.database.clone(),
+                                request.schema.clone(),
+                                table_name.clone(),
+                            )),
                        )
                        .await,
                    )
@@ -2446,13 +2483,13 @@ async fn export_database_sql_core_inner(
                    Some(
                        await_export_operation(
                            &request.export_id,
-                            crate::schema::get_columns_core(
-                                state,
-                                &request.connection_id,
-                                &request.database,
-                                &request.schema,
-                                &table_name,
-                            ),
+                            Box::pin(get_export_table_columns_isolated(
+                                state.clone(),
+                                request.connection_id.clone(),
+                                request.database.clone(),
+                                request.schema.clone(),
+                                table_name.clone(),
+                            )),
                        )
                        .await,
                    )
@@ -2460,8 +2497,9 @@ async fn export_database_sql_core_inner(
                    None
                };
                (index, PrefetchedTableMetadata { ddl, columns })
-            }))
-            .buffer_unordered(database_export_metadata_prefetch_concurrency(db_type));
+            })
+        }))
+        .buffer_unordered(database_export_metadata_prefetch_concurrency(db_type));
        while let Some((index, metadata)) = prefetch_stream.next().await {
            if metadata
                .ddl
@@ -2597,14 +2635,13 @@ async fn export_database_sql_core_inner(
                None => {
                    await_export_operation(
                        &request.export_id,
-                        crate::schema::get_table_relation_export_ddl_core(
-                            state,
-                            &request.connection_id,
-                            &request.database,
-                            &request.schema,
-                            table_name,
-                            None,
-                        ),
+                        Box::pin(get_export_table_ddl_isolated(
+                            state.clone(),
+                            request.connection_id.clone(),
+                            request.database.clone(),
+                            request.schema.clone(),
+                            table_name.clone(),
+                        )),
                    )
                    .await
                }
@@ -2645,13 +2682,13 @@ async fn export_database_sql_core_inner(
                None => {
                    await_export_operation(
                        &request.export_id,
-                        crate::schema::get_columns_core(
-                            state,
-                            &request.connection_id,
-                            &request.database,
-                            &request.schema,
-                            table_name,
-                        ),
+                        Box::pin(get_export_table_columns_isolated(
+                            state.clone(),
+                            request.connection_id.clone(),
+                            request.database.clone(),
+                            request.schema.clone(),
+                            table_name.clone(),
+                        )),
                    )
                    .await
                }
@@ -3124,10 +3161,13 @@ mod tests {
        clear_export_cancelled(&export_id).await;
        let task_export_id = export_id.clone();
        let task = tokio::spawn(async move {
-            await_export_operation(&task_export_id, async {
-                tokio::time::sleep(std::time::Duration::from_secs(30)).await;
-                Ok::<_, String>(())
-            })
+            await_export_operation(
+                &task_export_id,
+                Box::pin(async {
+                    tokio::time::sleep(std::time::Duration::from_secs(30)).await;
+                    Ok::<_, String>(())
+                }),
+            )
            .await
        });

@@ -3145,7 +3185,7 @@ mod tests {
    async fn await_export_operation_preserves_completed_metadata() {
        let export_id = format!("complete-metadata-{}", uuid::Uuid::new_v4());
        clear_export_cancelled(&export_id).await;
-        let result = await_export_operation(&export_id, async { Ok::<_, String>(42_u32) }).await;
+        let result = await_export_operation(&export_id, Box::pin(async { Ok::<_, String>(42_u32) })).await;
        assert_eq!(result, Ok(42));
    }

@@ -55,6 +55,14 @@ impl MySqlPool {
    pub async fn disconnect(self) -> Result<(), mysql_async::Error> {
        self.inner.disconnect().await
    }
+
+    /// Returns whether both handles refer to the same underlying pool
+    /// generation. Pool options are not an identity: a reconnect creates a new
+    /// pool with identical options, and a late health probe for the old pool
+    /// must not be allowed to remove that replacement from routing.
+    pub(crate) fn is_same_pool(&self, other: &Self) -> bool {
+        std::sync::Arc::ptr_eq(&self.inner.metrics(), &other.inner.metrics())
+    }
 }

 impl Deref for MySqlPool {