github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-19 16:16:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(plugin): preserve reconnected sidecar sessions

Browse Source 
Preserve plugin sidecar registrations across reconnects by serializing connection lifecycle calls and isolating unconfirmed disconnects.
This commit is contained in:
Jinpy
2026-10-19 21:56:22 +08:00
committed by GitHub
parent 65a3e24aed
commit 884ef6ba17
3 changed files with 392 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/dbx-plugin-runtime/src/plugins/host.rs
+152 -10
View File
@@ -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<Arc<PluginSidecarSession>>,
disconnect: bool,
params: serde_json::Value,
lifecycle: Arc<Mutex<u64>>,
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<Mutex<u64>>,
generation: u64,
session: Arc<PluginSidecarSession>,
params: serde_json::Value,
deadline: Duration,
) -> Result<(), String> {
let mut current = lifecycle.lock().await;
if *current != generation {
return Ok(());
}
let result: Result<serde_json::Value, String> =
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<Mutex<u64>>>;
struct PluginHostInner {
registry: PluginRegistry,
sessions: RwLock<HashMap<String, Arc<PluginSidecarSession>>>,
activation_lock: Mutex<()>,
connection_lifecycles: StdMutex<ConnectionLifecycleTable>,
events: broadcast::Sender<PluginEvent>,
binary_messages: broadcast::Sender<PluginBinaryMessage>,
}
@@ -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<Mutex<u64>> {
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<PluginEvent> {
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();
}
}
crates/dbx-plugin-runtime/src/plugins/runtime.rs
+8
View File
@@ -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
crates/dbx-web/src/routes/connection.rs
+232 -1
View File
@@ -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::<Vec<_>>(), ["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::<Vec<_>>(),
["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;