From 884ef6ba176ba29479f09ba3da35312d0785ab1a Mon Sep 17 00:00:00 2001 From: Jinpy Date: Mon, 19 Oct 2026 21:56:22 +0800 Subject: [PATCH] fix(plugin): preserve reconnected sidecar sessions Preserve plugin sidecar registrations across reconnects by serializing connection lifecycle calls and isolating unconfirmed disconnects. --- crates/dbx-plugin-runtime/src/plugins/host.rs | 162 +++++++++++- .../dbx-plugin-runtime/src/plugins/runtime.rs | 8 + crates/dbx-web/src/routes/connection.rs | 233 +++++++++++++++++- 3 files changed, 392 insertions(+), 11 deletions(-) diff --git a/crates/dbx-plugin-runtime/src/plugins/host.rs b/crates/dbx-plugin-runtime/src/plugins/host.rs index 22d4878451..c4d9b4ba87 100644 --- a/crates/dbx-plugin-runtime/src/plugins/host.rs +++ b/crates/dbx-plugin-runtime/src/plugins/host.rs @@ -1,5 +1,5 @@ use std::collections::{BTreeMap, HashMap, HashSet}; -use std::sync::Arc; +use std::sync::{Arc, Mutex as StdMutex, Weak}; use std::time::Duration; use crate::models::connection::{ConnectionConfig, ConnectionTestResult, DatabaseType}; @@ -45,6 +45,8 @@ pub struct PluginConnectionHandle { session: Option>, disconnect: bool, params: serde_json::Value, + lifecycle: Arc>, + generation: u64, _activity: super::lifecycle::PluginUsageGuard, } @@ -63,28 +65,66 @@ impl PluginConnectionHandle { } pub async fn disconnect(&self) -> Result<(), String> { + // Fast-path filter: an already-retired handle never reaches the plugin. + // The detached worker re-checks under the lifecycle lock. + if *self.lifecycle.lock().await != self.generation { + return Ok(()); + } let Some(session) = &self.session else { return Ok(()); }; if !self.disconnect || session.status().state != PluginSessionState::Running { return Ok(()); } - let result: serde_json::Value = session - .invoke_with_timeout( - PLUGIN_CONNECTION_DISCONNECT_METHOD, - self.params.clone(), - None, - Some(super::PLUGIN_REQUEST_TIMEOUT), - ) - .await?; - ensure_plugin_operation_succeeded(result) + // The RPC wait and the generation update run in a detached task that + // holds the lifecycle lock: a caller abandoned by a pool-close timeout + // must still retire this generation, and a later connect must wait for + // the in-flight disconnect to settle before re-registering. + let worker = tokio::spawn(disconnect_under_lifecycle( + self.lifecycle.clone(), + self.generation, + session.clone(), + self.params.clone(), + super::PLUGIN_REQUEST_TIMEOUT, + )); + worker.await.map_err(|error| format!("Plugin connection/disconnect worker failed: {error}"))? } } +async fn disconnect_under_lifecycle( + lifecycle: Arc>, + generation: u64, + session: Arc, + params: serde_json::Value, + deadline: Duration, +) -> Result<(), String> { + let mut current = lifecycle.lock().await; + if *current != generation { + return Ok(()); + } + let result: Result = + session.invoke_with_timeout(PLUGIN_CONNECTION_DISCONNECT_METHOD, params, None, Some(deadline)).await; + // A request that never got its answer may still be executing inside the + // plugin. Stop the sidecar so a reconnect starts from a fresh process + // instead of one whose registration the stale disconnect can clobber. + if result.as_ref().is_err_and(|error| super::runtime::request_may_still_execute(error)) { + log::warn!("connection/disconnect ended without confirmation; stopping the sidecar"); + let _ = session.shutdown().await; + } + // A handle may have clones in the pool registry. Only one may send + // connection/disconnect, including when the caller is cancelled. + *current = current.wrapping_add(1); + drop(current); + result.and_then(ensure_plugin_operation_succeeded) +} + +type ConnectionLifecycleTable = HashMap<(String, String), Weak>>; + struct PluginHostInner { registry: PluginRegistry, sessions: RwLock>>, activation_lock: Mutex<()>, + connection_lifecycles: StdMutex, events: broadcast::Sender, binary_messages: broadcast::Sender, } @@ -98,12 +138,25 @@ impl PluginHost { registry, sessions: RwLock::new(HashMap::new()), activation_lock: Mutex::new(()), + connection_lifecycles: StdMutex::new(HashMap::new()), events, binary_messages, }), } } + fn connection_lifecycle(&self, plugin_id: &str, connection_id: &str) -> Arc> { + let mut lifecycles = self.inner.connection_lifecycles.lock().unwrap_or_else(|poisoned| poisoned.into_inner()); + lifecycles.retain(|_, lifecycle| lifecycle.strong_count() > 0); + let key = (plugin_id.to_string(), connection_id.to_string()); + if let Some(lifecycle) = lifecycles.get(&key).and_then(Weak::upgrade) { + return lifecycle; + } + let lifecycle = Arc::new(Mutex::new(0)); + lifecycles.insert(key, Arc::downgrade(&lifecycle)); + lifecycle + } + pub fn subscribe_events(&self) -> broadcast::Receiver { self.inner.events.subscribe() } @@ -353,6 +406,10 @@ impl PluginHost { .begin_connection(config.plugin_id.as_deref().unwrap_or_default(), &config.name)?; let (_, provider) = self.resolve_connection_provider(config)?; validate_plugin_connection_values(config, &provider)?; + let lifecycle = self.connection_lifecycle(config.plugin_id.as_deref().unwrap_or_default(), &config.id); + // Keep connect and disconnect ordered even when a detached pool close + // races a reconnect. A new generation makes older handles inert. + let mut generation = lifecycle.lock().await; let params = plugin_connection_params(config, &provider, runtime_host, runtime_port, runtime_proxy.as_ref())?; let needs_session = provider.has_capability(PluginConnectionCapability::Connect) || provider.has_capability(PluginConnectionCapability::Disconnect); @@ -375,6 +432,9 @@ impl PluginHost { } } let disconnect = provider.has_capability(PluginConnectionCapability::Disconnect); + *generation = generation.wrapping_add(1); + let handle_generation = *generation; + drop(generation); Ok(PluginConnectionHandle { plugin_id: config.plugin_id.clone().unwrap_or_default(), provider_id: provider.id, @@ -382,6 +442,8 @@ impl PluginHost { session, disconnect, params, + lifecycle, + generation: handle_generation, _activity: activity, }) } @@ -1757,4 +1819,84 @@ mod tests { ) .is_ok()); } + + #[cfg(unix)] + #[tokio::test] + async fn unconfirmed_disconnect_stops_the_sidecar_before_releasing_the_lifecycle() { + use std::os::unix::fs::PermissionsExt; + use std::sync::Arc; + + use super::{disconnect_under_lifecycle, PluginHost, PluginSessionState}; + use crate::plugins::PluginRegistry; + + let temp = tempfile::tempdir().unwrap(); + let plugin_dir = temp.path().join("plugins/sample.quarantine"); + std::fs::create_dir_all(&plugin_dir).unwrap(); + std::fs::write( + plugin_dir.join("manifest.json"), + serde_json::json!({ + "manifest_version": 1, + "id": "sample.quarantine", + "name": "Quarantine test", + "version": "1.0.0", + "publisher": "sample", + "engines": { "dbx": ">=0.1.0", "host_api": "^1.0" }, + "entrypoints": { "backend": { + "protocol_versions": [1], + "transport": "stdio-jsonl", + "executable": "plugin.sh" + } }, + "contributions": [{ + "type": "connection-provider", + "id": "sample.connection", + "database_type": "sample", + "capabilities": ["connect", "disconnect"] + }] + }) + .to_string(), + ) + .unwrap(); + let script = plugin_dir.join("plugin.sh"); + std::fs::write( + &script, + r#"#!/bin/sh +IFS= read -r initialize +initialize_id=$(printf '%s' "$initialize" | sed -E 's/.*"id":([0-9]+).*/\1/') +printf '{"jsonrpc":"2.0","id":%s,"result":{"protocolVersion":1,"capabilities":["connections"],"plugin":{"id":"sample.quarantine","version":"1.0.0"}}}\n' "$initialize_id" +# Never answer lifecycle requests: every disconnect stays unconfirmed. +while IFS= read -r _; do :; done +"#, + ) + .unwrap(); + let mut permissions = std::fs::metadata(&script).unwrap().permissions(); + permissions.set_mode(0o755); + std::fs::set_permissions(&script, permissions).unwrap(); + + let registry = PluginRegistry::new_with_app_version(temp.path().join("plugins"), "0.5.67"); + let host = PluginHost::new(registry); + let session = host.activate("sample.quarantine").await.unwrap(); + + // The timeout fires while the plugin is still executing the request; + // the sidecar must be quarantined before the lifecycle lock reopens, + // or a reconnect would reuse a session the stale disconnect clobbers. + let lifecycle = Arc::new(tokio::sync::Mutex::new(0u64)); + let error = disconnect_under_lifecycle( + lifecycle.clone(), + 0, + session.clone(), + serde_json::json!({}), + std::time::Duration::from_millis(500), + ) + .await + .unwrap_err(); + assert!(error.contains("timed out after"), "{error}"); + assert_ne!( + session.status().state, + PluginSessionState::Running, + "an unconfirmed disconnect must stop the sidecar" + ); + // A follow-up activate starts a fresh process with a clean state. + host.activate("sample.quarantine").await.unwrap(); + host.stop("sample.quarantine").await.unwrap(); + } } diff --git a/crates/dbx-plugin-runtime/src/plugins/runtime.rs b/crates/dbx-plugin-runtime/src/plugins/runtime.rs index 6fd4200cef..f1f2ca907a 100644 --- a/crates/dbx-plugin-runtime/src/plugins/runtime.rs +++ b/crates/dbx-plugin-runtime/src/plugins/runtime.rs @@ -19,6 +19,14 @@ use super::{ }; use dbx_platform::ssh_prompt; +/// A request failed this way without the host observing the end of the +/// handler, so the plugin may still be executing it. Lifecycle callers treat +/// the outcome as unconfirmed instead of assuming the request died with the +/// abandoned caller. +pub(super) fn request_may_still_execute(error: &str) -> bool { + error.contains("timed out after") || error.contains("was waiting for user input") +} + pub const PLUGIN_REQUEST_TIMEOUT: Duration = Duration::from_secs(30); /// Host API method a plugin backend calls to ask the user for input /// (`user_input.rs` semantics). Plugin-initiated requests use string ids, which diff --git a/crates/dbx-web/src/routes/connection.rs b/crates/dbx-web/src/routes/connection.rs index 96f9125b4c..3ff3549ef4 100644 --- a/crates/dbx-web/src/routes/connection.rs +++ b/crates/dbx-web/src/routes/connection.rs @@ -503,7 +503,14 @@ pub async fn connect_db( config.db_type == DatabaseType::Nacos, ); - app.remove_connection_pools_detached(&connection_id).await; + // Plugin connection/connect refreshes the sidecar's registration in place. + // Closing the old pool after this request would send a late + // connection/disconnect and remove the refreshed registration. + if config.db_type == DatabaseType::Plugin { + app.drop_connection_pools_without_close(&connection_id).await; + } else { + app.remove_connection_pools_detached(&connection_id).await; + } app.nacos_registry.drop_connection(&connection_id).await; app.reset_connection_transport_for_config(&connection_id, &runtime_config).await; app.configs.write().await.insert(connection_id.clone(), runtime_config); @@ -1524,6 +1531,230 @@ mod tests { let _ = std::fs::remove_dir_all(dir); } + #[cfg(unix)] + #[tokio::test] + async fn plugin_reconnect_ignores_late_disconnect_from_old_pool() { + use std::os::unix::fs::PermissionsExt; + + let (state, dir) = test_web_state().await; + let plugin_dir = dir.join("plugins/sample.reconnect"); + std::fs::create_dir_all(&plugin_dir).unwrap(); + std::fs::write( + plugin_dir.join("manifest.json"), + serde_json::json!({ + "manifest_version": 1, + "id": "sample.reconnect", + "name": "Reconnect test", + "version": "1.0.0", + "publisher": "sample", + "engines": { "dbx": ">=0.1.0", "host_api": "^1.0" }, + "entrypoints": { "backend": { + "protocol_versions": [1], + "transport": "stdio-jsonl", + "executable": "plugin.sh" + } }, + "contributions": [{ + "type": "connection-provider", + "id": "sample.connection", + "database_type": "sample", + "capabilities": ["connect", "disconnect"] + }] + }) + .to_string(), + ) + .unwrap(); + let script = plugin_dir.join("plugin.sh"); + std::fs::write( + &script, + r#"#!/bin/sh +IFS= read -r initialize +initialize_id=$(printf '%s' "$initialize" | sed -E 's/.*"id":([0-9]+).*/\1/') +printf '{"jsonrpc":"2.0","id":%s,"result":{"protocolVersion":1,"capabilities":["connections"],"plugin":{"id":"sample.reconnect","version":"1.0.0"}}}\n' "$initialize_id" +while IFS= read -r request; do + request_id=$(printf '%s' "$request" | sed -E 's/.*"id":([0-9]+).*/\1/') + method=$(printf '%s' "$request" | sed -E 's/.*"method":"([^"]+)".*/\1/') + printf '%s\n' "$method" >> "$(dirname "$0")/events.log" + printf '{"jsonrpc":"2.0","id":%s,"result":{"success":true}}\n' "$request_id" +done +"#, + ) + .unwrap(); + let mut permissions = std::fs::metadata(&script).unwrap().permissions(); + permissions.set_mode(0o755); + std::fs::set_permissions(&script, permissions).unwrap(); + + let mut config = plugin_config("plugin-reconnect"); + config.host = "127.0.0.1".to_string(); + config.plugin_id = Some("sample.reconnect".to_string()); + config.plugin_connection_provider = Some("sample.connection".to_string()); + config.plugin_connection_type = Some("sample".to_string()); + + let request = || ConnectRequest { config: config.clone(), client_attempt: None }; + let _ = connect_db(State(state.clone()), HeaderMap::new(), Json(request())).await.unwrap(); + let old_handle = match state.app.pool_handle(&config.id).await.unwrap() { + PoolKind::PluginConnection(handle) => handle, + _ => panic!("expected a plugin pool"), + }; + let _ = connect_db(State(state.clone()), HeaderMap::new(), Json(request())).await.unwrap(); + old_handle.disconnect().await.unwrap(); + let events = std::fs::read_to_string(plugin_dir.join("events.log")).unwrap(); + assert_eq!(events.lines().collect::>(), ["connection/connect", "connection/connect"]); + + let _ = disconnect_db( + State(state.clone()), + Json(DisconnectRequest { connection_id: config.id.clone(), client_attempt: None }), + ) + .await + .unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + let events = std::fs::read_to_string(plugin_dir.join("events.log")).unwrap(); + if events.lines().count() == 3 { + assert_eq!(events.lines().last(), Some("connection/disconnect")); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await + .expect("explicit disconnect should reach the sidecar"); + drop(old_handle); + state.app.plugin_host.stop("sample.reconnect").await.unwrap(); + let _ = std::fs::remove_dir_all(dir); + } + + #[cfg(unix)] + #[tokio::test] + async fn cancelled_disconnect_waits_for_the_in_flight_request_before_reconnect() { + use std::os::unix::fs::PermissionsExt; + + let (state, dir) = test_web_state().await; + let plugin_dir = dir.join("plugins/sample.reconnect.race"); + std::fs::create_dir_all(&plugin_dir).unwrap(); + std::fs::write( + plugin_dir.join("manifest.json"), + serde_json::json!({ + "manifest_version": 1, + "id": "sample.reconnect.race", + "name": "Reconnect race test", + "version": "1.0.0", + "publisher": "sample", + "engines": { "dbx": ">=0.1.0", "host_api": "^1.0" }, + "entrypoints": { "backend": { + "protocol_versions": [1], + "transport": "stdio-jsonl", + "executable": "plugin.sh" + } }, + "contributions": [{ + "type": "connection-provider", + "id": "sample.connection", + "database_type": "sample", + "capabilities": ["connect", "disconnect"] + }] + }) + .to_string(), + ) + .unwrap(); + // The disconnect handler runs concurrently and completes only after a + // delay, removing the registration marker as its side effect: exactly + // the late write a reconnect must not be scheduled under. + let script = plugin_dir.join("plugin.sh"); + std::fs::write( + &script, + r#"#!/bin/sh +IFS= read -r initialize +initialize_id=$(printf '%s' "$initialize" | sed -E 's/.*"id":([0-9]+).*/\1/') +printf '{"jsonrpc":"2.0","id":%s,"result":{"protocolVersion":1,"capabilities":["connections"],"plugin":{"id":"sample.reconnect.race","version":"1.0.0"}}}\n' "$initialize_id" +while IFS= read -r request; do + request_id=$(printf '%s' "$request" | sed -E 's/.*"id":([0-9]+).*/\1/') + method=$(printf '%s' "$request" | sed -E 's/.*"method":"([^"]+)".*/\1/') + printf '%s\n' "$method" >> "$(dirname "$0")/events.log" + if [ "$method" = "connection/disconnect" ]; then + ( + sleep 1 + rm -f "$(dirname "$0")/registration" + touch "$(dirname "$0")/disconnect-done" + printf '{"jsonrpc":"2.0","id":%s,"result":{"success":true}}\n' "$request_id" + ) & + continue + fi + printf '{"jsonrpc":"2.0","id":%s,"result":{"success":true}}\n' "$request_id" + if [ "$method" = "connection/connect" ]; then + touch "$(dirname "$0")/registration" + fi +done +"#, + ) + .unwrap(); + let mut permissions = std::fs::metadata(&script).unwrap().permissions(); + permissions.set_mode(0o755); + std::fs::set_permissions(&script, permissions).unwrap(); + + let mut config = plugin_config("plugin-reconnect-race"); + config.host = "127.0.0.1".to_string(); + config.plugin_id = Some("sample.reconnect.race".to_string()); + config.plugin_connection_provider = Some("sample.connection".to_string()); + config.plugin_connection_type = Some("sample".to_string()); + + let request = || ConnectRequest { config: config.clone(), client_attempt: None }; + let _ = connect_db(State(state.clone()), HeaderMap::new(), Json(request())).await.unwrap(); + let old_handle = match state.app.pool_handle(&config.id).await.unwrap() { + PoolKind::PluginConnection(handle) => handle, + _ => panic!("expected a plugin pool"), + }; + + // The pool-close path abandons a disconnect after 3 seconds while the + // plugin keeps executing it; a reconnect must wait for that request to + // settle instead of re-registering underneath it. + let abandoned = tokio::time::timeout(std::time::Duration::from_millis(300), old_handle.disconnect()).await; + assert!(abandoned.is_err(), "the disconnect must still be in flight when the caller gives up"); + + let _ = connect_db(State(state.clone()), HeaderMap::new(), Json(request())).await.unwrap(); + // Wait until the stale handler has actually finished before judging: + // `disconnect-done` proves the late disconnect already ran. + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + if plugin_dir.join("disconnect-done").is_file() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + }) + .await + .expect("the abandoned disconnect should eventually complete in the plugin"); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + let events = std::fs::read_to_string(plugin_dir.join("events.log")).unwrap_or_default(); + if events.lines().count() == 3 && plugin_dir.join("registration").is_file() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + }) + .await + .expect("the reconnect should re-register after the abandoned disconnect settled"); + + let events = std::fs::read_to_string(plugin_dir.join("events.log")).unwrap(); + assert_eq!( + events.lines().collect::>(), + ["connection/connect", "connection/disconnect", "connection/connect"], + "the reconnect must reach the plugin only after the abandoned disconnect settled" + ); + assert!( + plugin_dir.join("registration").is_file(), + "the late disconnect must not remove the refreshed registration" + ); + + // The abandoned handle is spent: a retry must not reach the plugin again. + old_handle.disconnect().await.unwrap(); + let events = std::fs::read_to_string(plugin_dir.join("events.log")).unwrap(); + assert_eq!(events.lines().count(), 3); + + drop(old_handle); + state.app.plugin_host.stop("sample.reconnect.race").await.unwrap(); + let _ = std::fs::remove_dir_all(dir); + } + #[tokio::test] async fn failed_consul_agent_target_validation_cleans_up_temporary_state() { let (state, dir) = test_web_state().await;