github/dbx
4
0 Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-11-15 14:52:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(agents): concurrent all-agent-driver upgrade

Browse Source
This commit is contained in:
Abeautifulsnow authored and GitHub committed 2026-07-24 23:50:50 +08:00
1 parent 2a99f91099
commit 95db01e359
15 files changed
+1083 -282
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/dbx-core/src/agent_manager.rs
+46 -15
View File
@@ -2,6 +2,7 @@ use std::ffi::OsStr;
use std::fs::File;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex as StdMutex};
use serde::{Deserialize, Serialize};
use tokio::sync::Mutex;
@@ -435,6 +436,18 @@ pub struct AgentManager {
pub(crate) connection_runtimes: Mutex<
std::collections::HashMap<String, std::sync::Arc<tokio::sync::OnceCell<std::sync::Arc<AgentRuntimeClient>>>>,
>,
/// Serializes `load_state` → modify → `save_state` to prevent lost updates
/// when multiple driver installs run concurrently.
pub(crate) state_lock: StdMutex<()>,
/// Per-JRE-key install locks so that concurrent driver installs sharing the
/// same JRE download it only once (DCL pattern: lock → re-check installed → download).
pub(crate) jre_install_locks: Mutex<std::collections::HashMap<String, Arc<Mutex<()>>>>,
/// Per-driver locks serialize install, import, and uninstall operations
/// targeting the same on-disk agent files.
pub(crate) driver_operation_locks: Mutex<std::collections::HashMap<String, Arc<Mutex<()>>>>,
/// Driver operations may run concurrently, but JRE replacement/removal
/// must exclude them until their dependent driver state is persisted.
pub(crate) installation_operation_lock: tokio::sync::RwLock<()>,
}
impl Default for AgentManager {
@@ -460,6 +473,10 @@ impl AgentManager {
app_version: app_version.into(),
daemons: Mutex::new(std::collections::HashMap::new()),
connection_runtimes: Mutex::new(std::collections::HashMap::new()),
state_lock: StdMutex::new(()),
jre_install_locks: Mutex::new(std::collections::HashMap::new()),
driver_operation_locks: Mutex::new(std::collections::HashMap::new()),
installation_operation_lock: tokio::sync::RwLock::new(()),
};
mgr.migrate_legacy_jre();
mgr.cleanup_pending_jre_dirs();
@@ -467,6 +484,19 @@ impl AgentManager {
mgr
}
/// Atomically load, modify, and persist the installation state.
///
/// Keep the closure free of async work: this lock exists specifically to
/// prevent one installation operation from saving an older snapshot over
/// another operation's successful update.
pub fn mutate_state<T>(&self, mutate: impl FnOnce(&mut AgentState) -> T) -> Result<T, String> {
let _guard = self.state_lock.lock().map_err(|_| "Agent installation state lock was poisoned".to_string())?;
let mut state = self.load_state();
let result = mutate(&mut state);
self.save_state(&state)?;
Ok(result)
}
fn migrate_legacy_jre(&self) {
let legacy = self.base_dir.join("jre");
let versioned = self.jre_dir(DEFAULT_JRE_KEY);
@@ -480,25 +510,26 @@ impl AgentManager {
/// removals are pruned from the persisted state. Failures are kept for the
/// next launch and never block startup. (Issue #1100, D6.)
fn cleanup_pending_jre_dirs(&self) {
let mut state = self.load_state();
if state.pending_jre_cleanup.is_empty() {
if self.load_state().pending_jre_cleanup.is_empty() {
return;
}
let mut remaining = Vec::new();
for path in std::mem::take(&mut state.pending_jre_cleanup) {
if !path.exists() {
continue;
}
match std::fs::remove_dir_all(&path) {
Ok(()) => log::info!("Cleaned up pending JRE stash: {}", path.display()),
Err(err) => {
log::warn!("Pending JRE cleanup failed for {}: {err}", path.display());
remaining.push(path);
if let Err(err) = self.mutate_state(|state| {
let mut remaining = Vec::new();
for path in std::mem::take(&mut state.pending_jre_cleanup) {
if !path.exists() {
continue;
}
match std::fs::remove_dir_all(&path) {
Ok(()) => log::info!("Cleaned up pending JRE stash: {}", path.display()),
Err(err) => {
log::warn!("Pending JRE cleanup failed for {}: {err}", path.display());
remaining.push(path);
}
}
}
}
state.pending_jre_cleanup = remaining;
if let Err(err) = self.save_state(&state) {
state.pending_jre_cleanup = remaining;
}) {
log::warn!("Failed to persist post-cleanup AgentState: {err}");
}
}
crates/dbx-core/src/agent_service.rs
+619 -82
View File
@@ -1,8 +1,11 @@
use std::hash::{Hash, Hasher};
use std::io::{Read, Seek, SeekFrom};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;
use futures::{stream, StreamExt};
use crate::agent_catalog;
use crate::agent_manager::{
AgentDriverInfo, AgentManager, AgentRegistry, InstalledDriver, JavaRuntimeMode, DEFAULT_JRE_KEY,
@@ -18,6 +21,10 @@ const JRE_REMOVE_ATTEMPTS: usize = if cfg!(windows) { 6 } else { 1 };
/// Exponential-ish backoff between retries. Total wait ≈ 1.55s on Windows.
const JRE_REMOVE_BACKOFF_MS: &[u64] = &[50, 100, 200, 400, 400, 400];
/// Keep batch updates concurrent without allowing a large registry to exhaust
/// the download server, local disk, or the application's file descriptors.
const MAX_CONCURRENT_AGENT_UPDATES: usize = 4;
/// Delete an old JRE directory, retrying on Windows to cover the daemon-exit
/// and AV-scan release window. Returns the original `std::io::Error` when all
/// retries fail so callers can decide whether to fall back to rename-stash.
@@ -93,7 +100,7 @@ fn stash_old_jre_dir(path: &Path) -> std::io::Result<PathBuf> {
/// path (Some) if the rename fallback was used so the caller can persist it
/// for startup cleanup, or None if the directory was deleted outright (or
/// did not exist).
fn replace_old_jre_dir(am: &AgentManager, path: &Path) -> Result<Option<PathBuf>, String> {
fn replace_old_jre_dir(path: &Path) -> Result<Option<PathBuf>, String> {
match remove_jre_dir_with_retry(path) {
Ok(()) => Ok(None),
Err(remove_err) => {
@@ -102,13 +109,8 @@ fn replace_old_jre_dir(am: &AgentManager, path: &Path) -> Result<Option<PathBuf>
match stash_old_jre_dir(path) {
Ok(stash) => {
log::warn!("remove_dir_all failed, stashed old JRE at {} ({remove_err})", stash.display());
// Persist immediately so a crash before extraction
// still leaves the stash recorded for cleanup.
let mut state = am.load_state();
state.pending_jre_cleanup.push(stash.clone());
if let Err(save_err) = am.save_state(&state) {
log::warn!("Failed to persist pending_jre_cleanup: {save_err}");
}
// The caller will persist this stash under
// state_lock after extraction succeeds.
Ok(Some(stash))
}
Err(rename_err) => {
@@ -122,7 +124,6 @@ fn replace_old_jre_dir(am: &AgentManager, path: &Path) -> Result<Option<PathBuf>
}
#[cfg(not(windows))]
{
let _ = am; // silence unused warning on POSIX
Err(format_jre_dir_remove_error(path, &remove_err))
}
}
@@ -138,6 +139,8 @@ static REGISTRY_CACHE: std::sync::LazyLock<
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
pub struct AgentProgressEvent {
#[serde(skip_serializing_if = "Option::is_none")]
pub operation_id: Option<String>,
pub step: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub downloaded: Option<u64>,
@@ -165,7 +168,15 @@ pub struct UpgradeAllAgentDriversResult {
impl AgentProgressEvent {
pub fn step(step: impl Into<String>) -> Self {
Self { step: step.into(), downloaded: None, total: None, db_type: None, current: None, total_drivers: None }
Self {
operation_id: None,
step: step.into(),
downloaded: None,
total: None,
db_type: None,
current: None,
total_drivers: None,
}
}
pub fn transfer(step: impl Into<String>, downloaded: u64, total: u64) -> Self {
@@ -178,6 +189,11 @@ impl AgentProgressEvent {
self.total_drivers = total_drivers;
self
}
pub fn with_operation_id(mut self, operation_id: &str) -> Self {
self.operation_id = Some(operation_id.to_string());
self
}
}
pub fn build_agent_list(am: &AgentManager, registry: Option<&AgentRegistry>) -> Vec<AgentDriverInfo> {
@@ -303,27 +319,32 @@ pub fn find_local_agent_jar(db_type: &str) -> Option<PathBuf> {
}
pub fn install_local_agent(am: &AgentManager, db_type: &str, source: PathBuf) -> Result<(), String> {
install_local_agent_file(am, db_type, &source)?;
am.mutate_state(|state| record_local_agent_install(state, db_type, DEFAULT_JRE_KEY))
}
fn install_local_agent_file(am: &AgentManager, db_type: &str, source: &Path) -> Result<(), String> {
let jar_path = am.driver_jar_path(db_type);
let parent = jar_path.parent().ok_or_else(|| format!("Invalid driver path: {}", jar_path.display()))?;
std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
let staging_path = parent.join(format!(".agent-jar-import-{}", uuid::Uuid::new_v4()));
std::fs::copy(&source, &staging_path).map_err(|e| format!("Failed to copy local agent jar: {e}"))?;
std::fs::copy(source, &staging_path).map_err(|e| format!("Failed to copy local agent jar: {e}"))?;
if !is_valid_agent_jar(&staging_path) {
std::fs::remove_file(&staging_path).ok();
return Err(format!("Local agent jar is invalid or corrupt: {}", source.display()));
}
replace_imported_agent_file(&staging_path, &jar_path)?;
replace_imported_agent_file(&staging_path, &jar_path)
}
let mut local_state = am.load_state();
local_state.installed_drivers.insert(
fn record_local_agent_install(state: &mut crate::agent_manager::AgentState, db_type: &str, jre_key: &str) {
state.installed_drivers.insert(
db_type.to_string(),
InstalledDriver {
version: "0.1.0-local".to_string(),
installed_at: chrono::Utc::now().to_rfc3339(),
jre: DEFAULT_JRE_KEY.to_string(),
jre: jre_key.to_string(),
},
);
am.save_state(&local_state)
}
fn is_valid_agent_jar(path: &Path) -> bool {
@@ -423,27 +444,45 @@ pub async fn upgrade_all_agent_drivers_from(
progress: impl Fn(AgentProgressEvent),
) -> Result<UpgradeAllAgentDriversResult, String> {
let registry = fetch_registry_from(source).await?;
let agents = build_agent_list(am, Some(&registry));
let updatable: Vec<&AgentDriverInfo> = agents.iter().filter(|agent| agent.update_available).collect();
let total = updatable.len() as u32;
let mut result = UpgradeAllAgentDriversResult::default();
upgrade_all_agent_drivers_with_registry(am, &registry, source, &progress).await
}
for (index, agent) in updatable.iter().enumerate() {
match install_agent_driver_from_registry(
async fn upgrade_all_agent_drivers_with_registry(
am: &AgentManager,
registry: &AgentRegistry,
source: DownloadSource,
progress: &impl Fn(AgentProgressEvent),
) -> Result<UpgradeAllAgentDriversResult, String> {
let agents = build_agent_list(am, Some(registry));
let updatable: Vec<String> =
agents.iter().filter(|agent| agent.update_available).map(|agent| agent.db_type.clone()).collect();
let total = updatable.len() as u32;
// Run independent driver installs concurrently, with a fixed upper bound
// so a large registry cannot saturate download and file-system resources.
let installs = updatable.into_iter().enumerate().map(|(index, db_type)| async move {
let result = install_agent_driver_from_registry_locked(
am,
&registry,
registry,
source,
&agent.db_type,
&progress,
&db_type,
progress,
Some((index + 1) as u32),
Some(total),
)
.await
{
.await;
(db_type, result)
});
let outcomes = stream::iter(installs).buffer_unordered(MAX_CONCURRENT_AGENT_UPDATES).collect::<Vec<_>>().await;
let mut result = UpgradeAllAgentDriversResult::default();
for (db_type, outcome) in outcomes {
match outcome {
Ok(()) => result.upgraded += 1,
Err(error) => {
log::warn!("Failed to update {} agent driver: {}", agent.db_type, error);
result.failed.push(AgentDriverUpdateIssue { db_type: agent.db_type.clone(), error });
log::warn!("Failed to update {} agent driver: {}", db_type, error);
result.failed.push(AgentDriverUpdateIssue { db_type, error });
}
}
}
@@ -452,7 +491,20 @@ pub async fn upgrade_all_agent_drivers_from(
Ok(result)
}
async fn driver_operation_lock(am: &AgentManager, db_type: &str) -> Arc<tokio::sync::Mutex<()>> {
let mut locks = am.driver_operation_locks.lock().await;
locks.entry(db_type.to_string()).or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))).clone()
}
async fn jre_operation_lock(am: &AgentManager, jre_key: &str) -> Arc<tokio::sync::Mutex<()>> {
let mut locks = am.jre_install_locks.lock().await;
locks.entry(jre_key.to_string()).or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))).clone()
}
pub async fn uninstall_agent_driver(am: &AgentManager, db_type: &str) -> Result<(), String> {
let _installation_guard = am.installation_operation_lock.read().await;
let driver_lock = driver_operation_lock(am, db_type).await;
let _driver_guard = driver_lock.lock().await;
prune_driver_download_cache(am, db_type)?;
let jar_path = am.driver_jar_path(db_type);
if jar_path.exists() {
@@ -463,9 +515,7 @@ pub async fn uninstall_agent_driver(am: &AgentManager, db_type: &str) -> Result<
std::fs::remove_dir_all(driver_dir).map_err(|err| err.to_string())?;
}
}
let mut local_state = am.load_state();
local_state.installed_drivers.remove(db_type);
am.save_state(&local_state)?;
am.mutate_state(|state| state.installed_drivers.remove(db_type))?;
am.stop_daemon_by_key(db_type).await;
Ok(())
}
@@ -475,6 +525,11 @@ pub fn clear_agent_download_cache(am: &AgentManager) -> Result<(), String> {
}
pub async fn uninstall_agent_jre(am: &AgentManager, jre_key: &str) -> Result<(), String> {
// Keep the dependency check and removal atomic with respect to driver
// installs/uninstalls that may add or remove a dependency on this JRE.
let _installation_guard = am.installation_operation_lock.write().await;
let jre_lock = jre_operation_lock(am, jre_key).await;
let _jre_guard = jre_lock.lock().await;
let local_state = am.load_state();
let dependents: Vec<&str> = local_state
.installed_drivers
@@ -492,9 +547,7 @@ pub async fn uninstall_agent_jre(am: &AgentManager, jre_key: &str) -> Result<(),
if let Err(err) = remove_jre_dir_with_retry(&jre_dir) {
return Err(format_jre_dir_remove_error(&jre_dir, &err));
}
let mut local_state = am.load_state();
local_state.jre_versions.remove(jre_key);
am.save_state(&local_state)?;
am.mutate_state(|state| state.jre_versions.remove(jre_key))?;
Ok(())
}
@@ -512,6 +565,11 @@ pub async fn reinstall_agent_jre_from(
source: DownloadSource,
progress: impl Fn(AgentProgressEvent),
) -> Result<(), String> {
// Replacing a JRE must not race a driver operation that is using or about
// to persist a dependency on the same runtime.
let _installation_guard = am.installation_operation_lock.write().await;
let jre_lock = jre_operation_lock(am, jre_key).await;
let _jre_guard = jre_lock.lock().await;
let registry = fetch_registry_from(source).await?;
let jre_info = registry.resolve_jre(jre_key).ok_or_else(|| format!("No JRE definition for version: {jre_key}"))?;
let platform = AgentManager::current_platform();
@@ -540,32 +598,33 @@ pub async fn reinstall_agent_jre_from(
// handles on Windows (Issue #1100). Falls back to a rename-stash if the
// directory still cannot be removed.
am.stop_daemons().await;
replace_old_jre_dir(am, &jre_dir)?;
let stash = replace_old_jre_dir(&jre_dir)?;
persist_pending_jre_cleanup(am, stash.as_ref()).await?;
extract_tar_gz(&jre_archive, &jre_dir)?;
std::fs::remove_file(&jre_archive).ok();
let mut local_state = am.load_state();
local_state.jre_versions.insert(jre_key.to_string(), jre_info.version.clone());
am.save_state(&local_state)?;
am.mutate_state(|state| state.jre_versions.insert(jre_key.to_string(), jre_info.version.clone()))?;
cleanup_jre_download_cache_after_success(am, jre_key);
progress(AgentProgressEvent::step("done"));
Ok(())
}
pub fn import_agents_from_zip(
pub async fn import_agents_from_zip(
am: &AgentManager,
zip_path: &Path,
progress: impl Fn(AgentProgressEvent),
) -> Result<OfflineImportResult, String> {
import_offline_zip(am, zip_path, |p| {
progress(AgentProgressEvent {
operation_id: None,
step: p.step,
downloaded: Some(p.current as u64),
total: Some(p.total as u64),
db_type: Some(p.label),
db_type: p.db_type,
current: Some(p.current),
total_drivers: Some(p.total),
});
})
.await
}
async fn install_agent_driver_with_batch(
@@ -575,6 +634,35 @@ async fn install_agent_driver_with_batch(
progress: &impl Fn(AgentProgressEvent),
current: Option<u32>,
total_drivers: Option<u32>,
) -> Result<(), String> {
let _installation_guard = am.installation_operation_lock.read().await;
let driver_lock = driver_operation_lock(am, db_type).await;
let _driver_guard = driver_lock.lock().await;
install_agent_driver_with_batch_unlocked(am, db_type, source, progress, current, total_drivers).await
}
async fn install_agent_driver_from_registry_locked(
am: &AgentManager,
registry: &AgentRegistry,
source: DownloadSource,
db_type: &str,
progress: &impl Fn(AgentProgressEvent),
current: Option<u32>,
total_drivers: Option<u32>,
) -> Result<(), String> {
let _installation_guard = am.installation_operation_lock.read().await;
let driver_lock = driver_operation_lock(am, db_type).await;
let _driver_guard = driver_lock.lock().await;
install_agent_driver_from_registry(am, registry, source, db_type, progress, current, total_drivers).await
}
async fn install_agent_driver_with_batch_unlocked(
am: &AgentManager,
db_type: &str,
source: DownloadSource,
progress: &impl Fn(AgentProgressEvent),
current: Option<u32>,
total_drivers: Option<u32>,
) -> Result<(), String> {
match fetch_registry_from(source).await {
Ok(registry) => {
@@ -623,13 +711,30 @@ async fn ensure_jre_from_registry(
current: Option<u32>,
total_drivers: Option<u32>,
) -> Result<(), String> {
// Fast path: already installed — return immediately without acquiring the
// per-JRE lock. The lock is only needed when a download + extract may be
// required.
if !jre_needs_install(am, registry, jre_key) {
return Ok(());
}
// Acquire (or create) the per-JRE-key mutex so that concurrent driver
// installs sharing the same JRE download it exactly once.
let lock = jre_operation_lock(am, jre_key).await;
let _jre_guard = lock.lock().await;
// Double-check: the previous lock holder may have already installed.
if !jre_needs_install(am, registry, jre_key) {
return Ok(());
}
let jre_info = registry.resolve_jre(jre_key).ok_or_else(|| format!("No JRE definition for version: {jre_key}"))?;
let platform = AgentManager::current_platform();
let platform_jre = jre_info
.platforms
.get(platform)
.ok_or_else(|| format!("No JRE {jre_key} available for platform: {platform}"))?;
let jre_archive = am.base_dir().join("jre-download.tar.gz");
let jre_archive = am.base_dir().join(format!("jre-{}-download.tar.gz", jre_key));
progress(AgentProgressEvent::transfer("jre", 0, platform_jre.size).with_batch(
Some(db_type),
current,
@@ -652,15 +757,67 @@ async fn ensure_jre_from_registry(
.await?;
progress(AgentProgressEvent::transfer("jre-extract", 0, 0).with_batch(Some(db_type), current, total_drivers));
let jre_dir = am.jre_dir(jre_key);
// Stop daemons first (Windows ERROR_ACCESS_DENIED, Issue #1100).
am.stop_daemons().await;
replace_old_jre_dir(am, &jre_dir)?;
// Stop only daemons that use this JRE before replacing its directory
// (Windows ERROR_ACCESS_DENIED, Issue #1100). In a concurrent
// upgrade-all this avoids killing unrelated daemons mid-install.
stop_daemons_using_jre(am, jre_key).await;
let stash = replace_old_jre_dir(&jre_dir)?;
// Persist the stash path *before* extraction so that a crash during
// extract_tar_gz (or a process kill) doesn't leave the renamed-stash
// directory as an orphan that never gets cleaned up.
persist_pending_jre_cleanup(am, stash.as_ref()).await?;
extract_tar_gz(&jre_archive, &jre_dir)?;
std::fs::remove_file(&jre_archive).ok();
cleanup_jre_download_cache_after_success(am, jre_key);
// Persist the JRE version after extraction succeeds, while still holding
// the per-JRE lock. This guarantees the DCL in another task's
// jre_needs_install() sees the installed version and skips download.
am.mutate_state(|state| state.jre_versions.insert(jre_key.to_string(), jre_info.version.clone()))?;
Ok(())
}
/// Stop daemons whose installed driver lists `jre_key` as its runtime.
async fn stop_daemons_using_jre(am: &AgentManager, jre_key: &str) {
let state = am.load_state();
for (db_type, driver) in &state.installed_drivers {
if driver.jre == jre_key {
am.stop_daemon_by_key(db_type).await;
}
}
}
/// Record a rename-stashed JRE before extraction so startup can clean it up
/// even if the process exits mid-install.
async fn persist_pending_jre_cleanup(am: &AgentManager, stash: Option<&PathBuf>) -> Result<(), String> {
let Some(stash_path) = stash else {
return Ok(());
};
am.mutate_state(|state| {
if !state.pending_jre_cleanup.contains(stash_path) {
state.pending_jre_cleanup.push(stash_path.clone());
}
})
}
async fn persist_local_agent_install_state(
am: &AgentManager,
db_type: &str,
jre_key: &str,
jre_version: Option<&str>,
) -> Result<(), String> {
am.mutate_state(|state| {
if let Some(version) = jre_version {
state.jre_versions.insert(jre_key.to_string(), version.to_string());
}
record_local_agent_install(state, db_type, jre_key);
})
}
async fn install_local_agent_with_registry_jre(
am: &AgentManager,
registry: &AgentRegistry,
@@ -675,12 +832,16 @@ async fn install_local_agent_with_registry_jre(
if jre_needs_install(am, registry, jre_key) {
ensure_jre_from_registry(am, registry, source, jre_key, db_type, progress, current, total_drivers).await?;
}
install_local_agent(am, db_type, local_jar)?;
if let Some(jre_info) = registry.resolve_jre(jre_key) {
let mut local_state = am.load_state();
local_state.jre_versions.insert(jre_key.to_string(), jre_info.version.clone());
am.save_state(&local_state)?;
}
install_local_agent_file(am, db_type, &local_jar)?;
// This fallback can run for several drivers during upgrade-all. Keep its
// driver and JRE updates in one state_lock-protected transaction.
persist_local_agent_install_state(
am,
db_type,
jre_key,
registry.resolve_jre(jre_key).map(|jre| jre.version.as_str()),
)
.await?;
am.stop_daemon_by_key(db_type).await;
progress(AgentProgressEvent::step("done").with_batch(Some(db_type), current, total_drivers));
Ok(())
@@ -765,21 +926,21 @@ async fn install_agent_driver_from_registry(
std::fs::remove_file(am.driver_native_path(db_type)).ok();
}
let mut local_state = am.load_state();
if requires_java_runtime {
if let Some(jre_info) = registry.resolve_jre(jre_key) {
local_state.jre_versions.insert(jre_key.clone(), jre_info.version.clone());
am.mutate_state(|state| {
if requires_java_runtime {
if let Some(jre_info) = registry.resolve_jre(jre_key) {
state.jre_versions.insert(jre_key.clone(), jre_info.version.clone());
}
}
}
local_state.installed_drivers.insert(
db_type.to_string(),
InstalledDriver {
version: driver.version.clone(),
installed_at: chrono::Utc::now().to_rfc3339(),
jre: jre_key.clone(),
},
);
am.save_state(&local_state)?;
state.installed_drivers.insert(
db_type.to_string(),
InstalledDriver {
version: driver.version.clone(),
installed_at: chrono::Utc::now().to_rfc3339(),
jre: jre_key.clone(),
},
);
})?;
am.stop_daemon_by_key(db_type).await;
cleanup_driver_download_cache_after_success(am, db_type);
progress(AgentProgressEvent::step("done").with_batch(Some(db_type), current, total_drivers));
@@ -1234,7 +1395,12 @@ pub struct OfflineImportProgress {
pub step: String,
pub current: u32,
pub total: u32,
/// Display label for the current item (e.g. "MySQL", "JRE 21.0.12").
pub label: String,
/// The real database-type key (e.g. "mysql"), used by the frontend for
/// per-driver progress routing. `None` for JRE-only steps.
#[serde(skip_serializing_if = "Option::is_none")]
pub db_type: Option<String>,
}
#[derive(Debug, Clone)]
@@ -1265,11 +1431,16 @@ pub fn inspect_offline_zip(zip_path: &Path) -> Result<OfflineImportPlan, String>
})
}
pub fn import_offline_zip(
pub async fn import_offline_zip(
am: &AgentManager,
zip_path: &Path,
progress: impl Fn(OfflineImportProgress),
) -> Result<OfflineImportResult, String> {
// Offline import can touch both JRE and driver directories — hold an
// exclusive installation-operation lock so that concurrent driver installs,
// JRE installs, Upgrade All, and uninstall operations are serialised.
let _installation_guard = am.installation_operation_lock.write().await;
let file = std::fs::File::open(zip_path).map_err(|e| format!("Failed to open ZIP file: {e}"))?;
let mut archive = zip::ZipArchive::new(file).map_err(|e| format!("Invalid ZIP file: {e}"))?;
@@ -1298,7 +1469,13 @@ pub fn import_offline_zip(
continue;
}
progress(OfflineImportProgress { step: "jre-extract".into(), current, total, label: format!("JRE {jre_key}") });
progress(OfflineImportProgress {
step: "jre-extract".into(),
current,
total,
label: format!("JRE {jre_key}"),
db_type: None,
});
let mut entry = archive.by_name(entry_name).map_err(|e| format!("Failed to read {entry_name}: {e}"))?;
let tmp_archive = am.base_dir().join(format!("jre-offline-{jre_key}.tar.gz"));
@@ -1352,6 +1529,7 @@ pub fn import_offline_zip(
current,
total,
label: agent_catalog::label_for_key(db_type).unwrap_or(db_type).to_string(),
db_type: Some(db_type.clone()),
});
let driver_path = if *is_native { am.driver_native_path(db_type) } else { am.driver_jar_path(db_type) };
@@ -1396,7 +1574,23 @@ pub fn import_offline_zip(
result.drivers_installed.push(db_type.clone());
}
am.save_state(&local_state)?;
am.mutate_state(|state| {
for jre_key in &result.jre_installed {
if let Some(version) = local_state.jre_versions.get(jre_key) {
state.jre_versions.insert(jre_key.clone(), version.clone());
}
}
for path in &local_state.pending_jre_cleanup {
if !state.pending_jre_cleanup.contains(path) {
state.pending_jre_cleanup.push(path.clone());
}
}
for db_type in &result.drivers_installed {
if let Some(driver) = local_state.installed_drivers.get(db_type) {
state.installed_drivers.insert(db_type.clone(), driver.clone());
}
}
})?;
Ok(result)
}
@@ -1541,7 +1735,14 @@ fn extract_tar_gz(archive: &Path, dest: &Path) -> Result<(), String> {
Ok(())
}
pub fn import_agent_driver(am: &AgentManager, db_type: &str, source_path: &Path) -> Result<(), String> {
pub async fn import_agent_driver(am: &AgentManager, db_type: &str, source_path: &Path) -> Result<(), String> {
// Manual imports replace the same artifact paths as downloads. Reuse the
// install operation and per-driver locks so an import cannot race an
// install, Upgrade All, or uninstall for this driver.
let _installation_guard = am.installation_operation_lock.read().await;
let driver_lock = driver_operation_lock(am, db_type).await;
let _driver_guard = driver_lock.lock().await;
if !source_path.is_file() {
return Err(format!("File not found: {}", source_path.display()));
}
@@ -1562,20 +1763,20 @@ pub fn import_agent_driver(am: &AgentManager, db_type: &str, source_path: &Path)
replace_imported_agent_file(&staging_path, &native_path)?;
std::fs::remove_file(am.driver_jar_path(db_type)).ok();
let mut local_state = am.load_state();
local_state.installed_drivers.insert(
db_type.to_string(),
InstalledDriver {
version: "0.1.0-local".to_string(),
installed_at: chrono::Utc::now().to_rfc3339(),
jre: DEFAULT_JRE_KEY.to_string(),
},
);
am.save_state(&local_state)
am.mutate_state(|state| {
state.installed_drivers.insert(
db_type.to_string(),
InstalledDriver {
version: "0.1.0-local".to_string(),
installed_at: chrono::Utc::now().to_rfc3339(),
jre: DEFAULT_JRE_KEY.to_string(),
},
);
})
}
pub fn import_agent_jar(am: &AgentManager, db_type: &str, jar_path: &Path) -> Result<(), String> {
import_agent_driver(am, db_type, jar_path)
pub async fn import_agent_jar(am: &AgentManager, db_type: &str, jar_path: &Path) -> Result<(), String> {
import_agent_driver(am, db_type, jar_path).await
}
fn replace_imported_agent_file(staging_path: &Path, target_path: &Path) -> Result<(), String> {
@@ -1804,13 +2005,23 @@ mod agent_download_url_tests {
#[cfg(test)]
mod agent_registry_install_tests {
use super::*;
use crate::agent_manager::{ArtifactInfo, DriverInfo, JavaRuntimeConfig};
use crate::agent_manager::{ArtifactInfo, DriverInfo, InstalledDriver, JavaRuntimeConfig, JreInfo};
fn test_manager(name: &str) -> AgentManager {
let dir = std::env::temp_dir().join(format!("dbx-agent-registry-install-{name}-{}", uuid::Uuid::new_v4()));
AgentManager::new_with_base_dir(dir)
}
fn write_test_agent_jar(path: &Path) {
use std::io::Write;
let file = std::fs::File::create(path).unwrap();
let mut archive = zip::ZipWriter::new(file);
archive.start_file("META-INF/MANIFEST.MF", zip::write::SimpleFileOptions::default()).unwrap();
archive.write_all(b"Manifest-Version: 1.0\nMain-Class: com.dbx.Agent\n").unwrap();
archive.finish().unwrap();
}
fn registry_with_native_and_legacy_jar(
db_type: &str,
version: &str,
@@ -1871,6 +2082,57 @@ mod agent_registry_install_tests {
cache_path
}
fn registry_with_jre(jre_key: &str, version: &str, url: &str, size: u64) -> AgentRegistry {
AgentRegistry {
jre: None,
jres: [(
jre_key.to_string(),
JreInfo {
version: version.to_string(),
platforms: [(
AgentManager::current_platform().to_string(),
ArtifactInfo { url: url.to_string(), size },
)]
.into_iter()
.collect(),
},
)]
.into_iter()
.collect(),
drivers: std::collections::HashMap::new(),
}
}
fn build_jre_archive(am: &AgentManager, jre_key: &str) -> Vec<u8> {
let archive_root = am.base_dir().join("jre-test-archive");
let payload = archive_root.join("payload");
let java_path = am.jre_java_path(jre_key);
let relative_java_path = java_path.strip_prefix(am.jre_dir(jre_key)).unwrap();
let java_path = payload.join(relative_java_path);
std::fs::create_dir_all(java_path.parent().unwrap()).unwrap();
std::fs::write(java_path, b"java").unwrap();
let archive = archive_root.join("runtime.tar.gz");
let status = crate::process::new_std_command("tar")
.args(["czf", &archive.to_string_lossy(), "-C", &archive_root.to_string_lossy(), "payload"])
.status()
.unwrap();
assert!(status.success());
std::fs::read(archive).unwrap()
}
fn write_cached_jre_download(am: &AgentManager, jre_key: &str, version: &str, url: &str, archive: &[u8]) {
let dest = am.base_dir().join(format!("jre-{jre_key}-download.tar.gz"));
let cache_path = cached_download_path(
am,
url,
archive.len() as u64,
Some(CacheIdentity::Jre { key: jre_key, version }),
&dest,
);
std::fs::create_dir_all(cache_path.parent().unwrap()).unwrap();
std::fs::write(cache_path, archive).unwrap();
}
#[tokio::test]
async fn registry_install_accepts_native_driver_with_legacy_jar_fallback() {
let manager = test_manager("native-with-jar-fallback");
@@ -1908,6 +2170,263 @@ mod agent_registry_install_tests {
.any(|event| event.step == "done" && event.db_type.as_deref() == Some(db_type)));
}
#[tokio::test]
async fn batch_upgrade_preserves_successful_state_and_reports_independent_failure() {
let manager = test_manager("batch-upgrade");
let oracle_url = "https://example.com/dbx-agent-oracle";
let dameng_url = "https://example.com/dbx-agent-dameng";
let kingbase_url = "https://example.com/dbx-agent-kingbase.jar";
let oracle_bytes = b"oracle-native-agent";
let dameng_bytes = b"dameng-native-agent";
let corrupt_jar = b"not-a-jar";
let mut registry =
registry_with_native_and_legacy_jar("oracle", "2.0.0", oracle_url, oracle_bytes.len() as u64);
registry.drivers.extend(
registry_with_native_and_legacy_jar("dameng", "2.0.0", dameng_url, dameng_bytes.len() as u64).drivers,
);
registry.drivers.extend(registry_with_jar("kingbase", "2.0.0", kingbase_url, corrupt_jar.len() as u64).drivers);
let mut state = manager.load_state();
state.java_runtime = JavaRuntimeConfig { mode: JavaRuntimeMode::System, custom_java_path: None };
for (db_type, version) in [("oracle", "1.0.0"), ("dameng", "1.0.0"), ("kingbase", "1.0.0")] {
state.installed_drivers.insert(
db_type.to_string(),
InstalledDriver {
version: version.to_string(),
installed_at: "2026-01-01T00:00:00Z".to_string(),
jre: DEFAULT_JRE_KEY.to_string(),
},
);
}
manager.save_state(&state).unwrap();
write_cached_driver_download(
&manager,
"oracle",
"2.0.0",
oracle_url,
&manager.driver_native_path("oracle"),
oracle_bytes,
);
write_cached_driver_download(
&manager,
"dameng",
"2.0.0",
dameng_url,
&manager.driver_native_path("dameng"),
dameng_bytes,
);
write_cached_driver_download(
&manager,
"kingbase",
"2.0.0",
kingbase_url,
&manager.driver_jar_path("kingbase"),
corrupt_jar,
);
let events = std::sync::Mutex::new(Vec::new());
let progress = |event| events.lock().unwrap().push(event);
let result = upgrade_all_agent_drivers_with_registry(&manager, &registry, DownloadSource::Official, &progress)
.await
.unwrap();
assert_eq!(result.upgraded, 2);
assert_eq!(result.failed.len(), 1);
assert_eq!(result.failed[0].db_type, "kingbase");
let state = manager.load_state();
assert_eq!(state.installed_drivers["oracle"].version, "2.0.0");
assert_eq!(state.installed_drivers["dameng"].version, "2.0.0");
assert_eq!(state.installed_drivers["kingbase"].version, "1.0.0");
assert_eq!(events.lock().unwrap().iter().filter(|event| event.step == "done").count(), 2);
}
#[tokio::test]
async fn shared_jre_is_not_downloaded_again_after_the_first_install_persists_its_version() {
let manager = test_manager("shared-jre-deduplication");
let jre_key = DEFAULT_JRE_KEY;
let version = "21.0.12";
let url = "https://example.com/dbx-jre.tar.gz";
let archive = build_jre_archive(&manager, jre_key);
let registry = registry_with_jre(jre_key, version, url, archive.len() as u64);
let events = std::sync::Mutex::new(Vec::new());
let progress = |event| events.lock().unwrap().push(event);
write_cached_jre_download(&manager, jre_key, version, url, &archive);
ensure_jre_from_registry(
&manager,
&registry,
DownloadSource::Official,
jre_key,
"oracle",
&progress,
Some(1),
Some(2),
)
.await
.unwrap();
// The successful install cleans its cache. Re-add it so the old
// implementation fails deterministically by consuming it again.
write_cached_jre_download(&manager, jre_key, version, url, &archive);
ensure_jre_from_registry(
&manager,
&registry,
DownloadSource::Official,
jre_key,
"dameng",
&progress,
Some(2),
Some(2),
)
.await
.unwrap();
assert_eq!(manager.load_state().jre_versions[jre_key], version);
assert!(events
.lock()
.unwrap()
.iter()
.any(|event| event.step == "jre" && event.db_type.as_deref() == Some("oracle")));
assert!(!events
.lock()
.unwrap()
.iter()
.any(|event| event.step == "jre" && event.db_type.as_deref() == Some("dameng")));
}
#[tokio::test]
async fn stash_is_recorded_before_jre_extraction() {
let manager = test_manager("stash-before-extract");
let stash = manager.base_dir().join("jre-21.old-test");
persist_pending_jre_cleanup(&manager, Some(&stash)).await.unwrap();
assert_eq!(manager.load_state().pending_jre_cleanup, vec![stash]);
}
#[test]
fn concurrent_local_agent_state_updates_preserve_each_driver() {
let manager = test_manager("concurrent-local-agent-state");
let start = std::sync::Arc::new(std::sync::Barrier::new(3));
std::thread::scope(|scope| {
for db_type in ["oracle", "dameng"] {
let start = start.clone();
let manager = &manager;
scope.spawn(move || {
start.wait();
manager
.mutate_state(|state| {
state.jre_versions.insert(DEFAULT_JRE_KEY.to_string(), "21.0.12".to_string());
record_local_agent_install(state, db_type, DEFAULT_JRE_KEY);
})
.unwrap();
});
}
start.wait();
});
let state = manager.load_state();
assert!(state.installed_drivers.contains_key("oracle"));
assert!(state.installed_drivers.contains_key("dameng"));
assert_eq!(state.jre_versions[DEFAULT_JRE_KEY], "21.0.12");
}
#[tokio::test]
async fn batch_registry_install_waits_for_an_existing_driver_operation() {
let manager = test_manager("batch-driver-operation-lock");
let db_type = "oracle";
let version = "0.1.31";
let native_url = "https://example.com/dbx-agent-oracle";
let native_bytes = b"native-agent";
let registry = registry_with_native_and_legacy_jar(db_type, version, native_url, native_bytes.len() as u64);
write_cached_driver_download(
&manager,
db_type,
version,
native_url,
&manager.driver_native_path(db_type),
native_bytes,
);
let first_lock = driver_operation_lock(&manager, "oracle").await;
let first_guard = first_lock.lock().await;
let progress = |_| {};
let blocked = tokio::time::timeout(std::time::Duration::from_millis(50), async {
install_agent_driver_from_registry_locked(
&manager,
&registry,
DownloadSource::Official,
db_type,
&progress,
Some(1),
Some(1),
)
.await
})
.await;
assert!(blocked.is_err(), "batch install entered while another operation owned the driver files");
drop(first_guard);
tokio::time::timeout(
std::time::Duration::from_secs(1),
install_agent_driver_from_registry_locked(
&manager,
&registry,
DownloadSource::Official,
db_type,
&progress,
Some(1),
Some(1),
),
)
.await
.expect("batch install did not resume after the driver lock was released")
.unwrap();
assert_eq!(std::fs::read(manager.driver_native_path(db_type)).unwrap(), native_bytes);
}
#[tokio::test]
async fn manual_import_waits_for_an_existing_driver_operation() {
let manager = test_manager("manual-import-driver-operation-lock");
let db_type = "h2";
let source = manager.base_dir().join("dbx-agent-h2.jar");
std::fs::create_dir_all(source.parent().unwrap()).unwrap();
write_test_agent_jar(&source);
let lock = driver_operation_lock(&manager, db_type).await;
let first_guard = lock.lock().await;
let blocked =
tokio::time::timeout(std::time::Duration::from_millis(50), import_agent_driver(&manager, db_type, &source))
.await;
assert!(blocked.is_err(), "manual import entered while another operation owned the driver files");
drop(first_guard);
tokio::time::timeout(std::time::Duration::from_secs(1), import_agent_driver(&manager, db_type, &source))
.await
.expect("manual import did not resume after the driver lock was released")
.unwrap();
assert_eq!(std::fs::read(manager.driver_jar_path(db_type)).unwrap(), std::fs::read(source).unwrap());
}
#[tokio::test]
async fn jre_exclusive_operation_waits_for_in_flight_driver_operation() {
let manager = test_manager("jre-exclusive-operation-lock");
let driver_guard = manager.installation_operation_lock.read().await;
let blocked =
tokio::time::timeout(std::time::Duration::from_millis(50), manager.installation_operation_lock.write())
.await;
assert!(blocked.is_err(), "JRE replacement entered before an in-flight driver operation completed");
drop(driver_guard);
let _jre_guard =
tokio::time::timeout(std::time::Duration::from_secs(1), manager.installation_operation_lock.write())
.await
.expect("JRE replacement did not resume after driver operations completed");
}
#[tokio::test]
async fn registry_install_rejects_corrupt_downloaded_jar() {
let manager = test_manager("corrupt-jar");
@@ -1943,6 +2462,24 @@ mod agent_registry_install_tests {
assert!(!jar_path.exists());
assert!(!manager.load_state().installed_drivers.contains_key(db_type));
}
#[tokio::test]
async fn offline_import_exclusive_lock_waits_for_in_flight_driver_operation() {
let manager = test_manager("offline-import-lock");
// Simulate an in-flight driver operation holding a read lock.
let driver_guard = manager.installation_operation_lock.read().await;
// import_offline_zip acquires the write lock — it must wait.
let blocked =
tokio::time::timeout(std::time::Duration::from_millis(50), manager.installation_operation_lock.write())
.await;
assert!(blocked.is_err(), "offline import entered before an in-flight driver operation completed");
drop(driver_guard);
let _offline_guard =
tokio::time::timeout(std::time::Duration::from_secs(1), manager.installation_operation_lock.write())
.await
.expect("offline import did not resume after driver operations completed");
}
}
#[cfg(test)]
crates/dbx-core/src/db/ssh_tunnel.rs
+4 -4
View File
@@ -163,11 +163,11 @@ async fn connect_and_authenticate(
async fn try_authenticate_with_agent(
session: &mut Handle<SshClient>,
ssh_user: &str,
ssh_agent_sock_path: &str,
_ssh_agent_sock_path: &str,
connect_timeout: &Duration,
) -> Result<(), String> {
#[cfg(unix)]
let mut agent = if ssh_agent_sock_path.is_empty() {
let mut agent = if _ssh_agent_sock_path.is_empty() {
match AgentClient::connect_env().await {
Ok(a) => a,
Err(e) => {
@@ -175,12 +175,12 @@ async fn try_authenticate_with_agent(
}
}
} else {
match AgentClient::connect_uds(ssh_agent_sock_path).await {
match AgentClient::connect_uds(_ssh_agent_sock_path).await {
Ok(a) => a,
Err(e) => {
return Err(format!(
"No SSH password or key provided, and ssh-agent at '{}' is unavailable: {e}",
ssh_agent_sock_path
_ssh_agent_sock_path
));
}
}
crates/dbx-core/tests/agent_service.rs
+88 -52
View File
@@ -364,7 +364,9 @@ fn atomic_replace_moves_download_into_place() {
#[test]
fn agent_progress_event_serializes_backward_compatible_fields() {
let event = AgentProgressEvent::transfer("driver", 512, 1024).with_batch(Some("h2"), Some(1), Some(2));
let event = AgentProgressEvent::transfer("driver", 512, 1024)
.with_batch(Some("h2"), Some(1), Some(2))
.with_operation_id("upgrade-123");
let value = serde_json::to_value(event).unwrap();
@@ -374,16 +376,17 @@ fn agent_progress_event_serializes_backward_compatible_fields() {
assert_eq!(value["db_type"], "h2");
assert_eq!(value["current"], 1);
assert_eq!(value["total_drivers"], 2);
assert_eq!(value["operation_id"], "upgrade-123");
}
#[test]
fn local_jar_import_updates_driver_state() {
#[tokio::test]
async fn local_jar_import_updates_driver_state() {
let manager = test_manager("local-import");
let source = test_path("local-import-source").join("dbx-agent-h2.jar");
std::fs::create_dir_all(source.parent().unwrap()).unwrap();
write_test_agent_jar(&source);
import_agent_jar(&manager, "h2", &source).unwrap();
import_agent_jar(&manager, "h2", &source).await.unwrap();
assert_eq!(std::fs::read(manager.driver_jar_path("h2")).unwrap(), std::fs::read(&source).unwrap());
let state = manager.load_state();
@@ -392,22 +395,22 @@ fn local_jar_import_updates_driver_state() {
assert_eq!(installed.jre, DEFAULT_JRE_KEY);
}
#[test]
fn local_jar_import_rejects_corrupt_jar() {
#[tokio::test]
async fn local_jar_import_rejects_corrupt_jar() {
let manager = test_manager("local-import-corrupt");
let source = test_path("local-import-corrupt-source").join("dbx-agent-h2.jar");
std::fs::create_dir_all(source.parent().unwrap()).unwrap();
std::fs::write(&source, b"jar").unwrap();
let err = import_agent_jar(&manager, "h2", &source).unwrap_err();
let err = import_agent_jar(&manager, "h2", &source).await.unwrap_err();
assert!(err.contains("invalid or corrupt"));
assert!(!manager.driver_jar_path("h2").exists());
assert!(!manager.load_state().installed_drivers.contains_key("h2"));
}
#[test]
fn local_native_import_installs_current_platform_executable() {
#[tokio::test]
async fn local_native_import_installs_current_platform_executable() {
let manager = test_manager("local-native-import");
let source = test_path("local-native-import-source").join(if cfg!(windows) {
"dbx-agent-kingbase-windows.exe"
@@ -417,35 +420,35 @@ fn local_native_import_installs_current_platform_executable() {
std::fs::create_dir_all(source.parent().unwrap()).unwrap();
std::fs::write(&source, current_platform_native_binary()).unwrap();
import_agent_driver(&manager, "kingbase", &source).unwrap();
import_agent_driver(&manager, "kingbase", &source).await.unwrap();
assert_eq!(std::fs::read(manager.driver_native_path("kingbase")).unwrap(), std::fs::read(&source).unwrap());
assert!(!manager.driver_jar_path("kingbase").exists());
assert_eq!(manager.load_state().installed_drivers["kingbase"].version, "0.1.0-local");
}
#[test]
fn local_native_import_rejects_wrong_platform_binary() {
#[tokio::test]
async fn local_native_import_rejects_wrong_platform_binary() {
let manager = test_manager("local-native-import-invalid");
let source = test_path("local-native-import-invalid-source").join("dbx-agent-kingbase");
std::fs::create_dir_all(source.parent().unwrap()).unwrap();
std::fs::write(&source, b"not-an-executable").unwrap();
let err = import_agent_driver(&manager, "kingbase", &source).unwrap_err();
let err = import_agent_driver(&manager, "kingbase", &source).await.unwrap_err();
assert!(err.contains(AgentManager::current_platform()));
assert!(!manager.driver_native_path("kingbase").exists());
assert!(!manager.load_state().installed_drivers.contains_key("kingbase"));
}
#[test]
fn local_native_import_rejects_wrong_arch_binary() {
#[tokio::test]
async fn local_native_import_rejects_wrong_arch_binary() {
let manager = test_manager("local-native-wrong-arch");
let native_path = test_path("local-native-wrong-arch-file").join("agent");
std::fs::create_dir_all(native_path.parent().unwrap()).unwrap();
std::fs::write(&native_path, native_binary_for_arch(!cfg!(target_arch = "aarch64"))).unwrap();
let err = import_agent_driver(&manager, "kingbase", &native_path).unwrap_err();
let err = import_agent_driver(&manager, "kingbase", &native_path).await.unwrap_err();
assert!(err.contains("not a"));
assert!(!manager.driver_native_path("kingbase").exists());
@@ -515,8 +518,8 @@ fn clear_download_cache_removes_only_cache_entries() {
assert!(installed_driver.exists());
}
#[test]
fn offline_zip_import_emits_progress_and_updates_state() {
#[tokio::test]
async fn offline_zip_import_emits_progress_and_updates_state() {
let manager = test_manager("offline-progress");
let zip_path = test_path("offline-progress-zip").join("agents.zip");
std::fs::create_dir_all(zip_path.parent().unwrap()).unwrap();
@@ -526,6 +529,7 @@ fn offline_zip_import_emits_progress_and_updates_state() {
let result = import_agents_from_zip(&manager, &zip_path, |event| {
events.lock().unwrap().push(event);
})
.await
.unwrap();
assert_eq!(result.drivers_installed, vec!["h2"]);
@@ -533,11 +537,12 @@ fn offline_zip_import_emits_progress_and_updates_state() {
assert!(installed_jar.windows(b"Main-Class:".len()).any(|window| window == b"Main-Class:"));
assert_eq!(manager.load_state().installed_drivers.get("h2").unwrap().version, "0.2.0");
let events = events.lock().unwrap();
assert!(events.iter().any(|event| event.step == "driver" && event.db_type.as_deref() == Some("H2")));
// db_type carries the real database key, not the display label.
assert!(events.iter().any(|event| event.step == "driver" && event.db_type.as_deref() == Some("h2")));
}
#[test]
fn offline_zip_import_installs_release_named_jre() {
#[tokio::test]
async fn offline_zip_import_installs_release_named_jre() {
let manager = test_manager("offline-release-jre");
let zip_path = test_path("offline-release-jre-zip").join("agents.zip");
std::fs::create_dir_all(zip_path.parent().unwrap()).unwrap();
@@ -547,6 +552,7 @@ fn offline_zip_import_installs_release_named_jre() {
let result = import_agents_from_zip(&manager, &zip_path, |event| {
events.lock().unwrap().push(event);
})
.await
.unwrap();
assert_eq!(result.jre_installed, vec![DEFAULT_JRE_KEY]);
@@ -556,35 +562,35 @@ fn offline_zip_import_installs_release_named_jre() {
assert!(events.lock().unwrap().iter().any(|event| event.step == "jre-extract"));
}
#[test]
fn offline_zip_import_preserves_existing_jre_when_archive_is_corrupt() {
#[tokio::test]
async fn offline_zip_import_preserves_existing_jre_when_archive_is_corrupt() {
let manager = test_manager("offline-corrupt-jre-preserves-existing");
let root = test_path("offline-corrupt-jre-preserves-existing-zip");
let valid_zip = root.join("valid.zip");
let corrupt_zip = root.join("corrupt.zip");
std::fs::create_dir_all(&root).unwrap();
write_offline_driver_zip_with_jre(&valid_zip, "h2", "0.2.0", "21.0.12");
import_agents_from_zip(&manager, &valid_zip, |_| {}).unwrap();
import_agents_from_zip(&manager, &valid_zip, |_| {}).await.unwrap();
let java_path = manager.jre_java_path(DEFAULT_JRE_KEY);
let original_java = std::fs::read(&java_path).unwrap();
write_offline_driver_zip_with_jre_bytes(&corrupt_zip, "h2", "0.3.0", "21.0.13", b"not-a-tar-gz".to_vec());
let err = import_agents_from_zip(&manager, &corrupt_zip, |_| {}).unwrap_err();
let err = import_agents_from_zip(&manager, &corrupt_zip, |_| {}).await.unwrap_err();
assert!(err.contains("Failed to extract JRE archive"));
assert_eq!(std::fs::read(java_path).unwrap(), original_java);
assert_eq!(manager.load_state().jre_versions.get(DEFAULT_JRE_KEY).map(String::as_str), Some("21.0.12"));
}
#[test]
fn offline_zip_import_preserves_existing_jre_when_driver_is_corrupt() {
#[tokio::test]
async fn offline_zip_import_preserves_existing_jre_when_driver_is_corrupt() {
let manager = test_manager("offline-corrupt-driver-preserves-jre");
let root = test_path("offline-corrupt-driver-preserves-jre-zip");
let valid_zip = root.join("valid.zip");
let corrupt_zip = root.join("corrupt.zip");
std::fs::create_dir_all(&root).unwrap();
write_offline_driver_zip_with_jre(&valid_zip, "h2", "0.2.0", "21.0.12");
import_agents_from_zip(&manager, &valid_zip, |_| {}).unwrap();
import_agents_from_zip(&manager, &valid_zip, |_| {}).await.unwrap();
let java_path = manager.jre_java_path(DEFAULT_JRE_KEY);
let original_java = std::fs::read(&java_path).unwrap();
@@ -596,48 +602,48 @@ fn offline_zip_import_preserves_existing_jre_when_driver_is_corrupt() {
test_jre_archive_bytes(),
b"jar".to_vec(),
);
let err = import_agents_from_zip(&manager, &corrupt_zip, |_| {}).unwrap_err();
let err = import_agents_from_zip(&manager, &corrupt_zip, |_| {}).await.unwrap_err();
assert!(err.contains("invalid or corrupt"));
assert_eq!(std::fs::read(java_path).unwrap(), original_java);
assert_eq!(manager.load_state().jre_versions.get(DEFAULT_JRE_KEY).map(String::as_str), Some("21.0.12"));
}
#[test]
fn offline_zip_import_rejects_corrupt_jar() {
#[tokio::test]
async fn offline_zip_import_rejects_corrupt_jar() {
let manager = test_manager("offline-corrupt-driver");
let zip_path = test_path("offline-corrupt-driver-zip").join("agents.zip");
std::fs::create_dir_all(zip_path.parent().unwrap()).unwrap();
write_offline_driver_zip_with_jar(&zip_path, "h2", "0.2.0", b"jar".to_vec());
let err = import_agents_from_zip(&manager, &zip_path, |_| {}).unwrap_err();
let err = import_agents_from_zip(&manager, &zip_path, |_| {}).await.unwrap_err();
assert!(err.contains("invalid or corrupt"));
assert!(!manager.driver_jar_path("h2").exists());
assert!(!manager.load_state().installed_drivers.contains_key("h2"));
}
#[test]
fn offline_zip_import_preserves_existing_driver_when_jar_is_corrupt() {
#[tokio::test]
async fn offline_zip_import_preserves_existing_driver_when_jar_is_corrupt() {
let manager = test_manager("offline-corrupt-driver-preserves-existing");
let root = test_path("offline-corrupt-driver-preserves-existing-zip");
let valid_zip = root.join("valid.zip");
let corrupt_zip = root.join("corrupt.zip");
std::fs::create_dir_all(&root).unwrap();
write_offline_driver_zip(&valid_zip, "h2", "0.2.0");
import_agents_from_zip(&manager, &valid_zip, |_| {}).unwrap();
import_agents_from_zip(&manager, &valid_zip, |_| {}).await.unwrap();
let original = std::fs::read(manager.driver_jar_path("h2")).unwrap();
write_offline_driver_zip_with_jar(&corrupt_zip, "h2", "0.3.0", b"jar".to_vec());
let err = import_agents_from_zip(&manager, &corrupt_zip, |_| {}).unwrap_err();
let err = import_agents_from_zip(&manager, &corrupt_zip, |_| {}).await.unwrap_err();
assert!(err.contains("invalid or corrupt"));
assert_eq!(std::fs::read(manager.driver_jar_path("h2")).unwrap(), original);
assert_eq!(manager.load_state().installed_drivers["h2"].version, "0.2.0");
}
#[test]
fn offline_zip_import_keeps_legacy_unversioned_jar_compatibility() {
#[tokio::test]
async fn offline_zip_import_keeps_legacy_unversioned_jar_compatibility() {
let manager = test_manager("offline-legacy-jar-zip");
let zip_path = test_path("offline-legacy-jar-zip").join("h2.zip");
std::fs::create_dir_all(zip_path.parent().unwrap()).unwrap();
@@ -662,19 +668,19 @@ fn offline_zip_import_keeps_legacy_unversioned_jar_compatibility() {
std::io::Write::write_all(&mut zip, &jar).unwrap();
zip.finish().unwrap();
import_agents_from_zip(&manager, &zip_path, |_| {}).unwrap();
import_agents_from_zip(&manager, &zip_path, |_| {}).await.unwrap();
assert_eq!(manager.load_state().installed_drivers["h2"].version, "0.1.9");
}
#[test]
fn offline_zip_import_installs_versioned_native_driver_package() {
#[tokio::test]
async fn offline_zip_import_installs_versioned_native_driver_package() {
let manager = test_manager("offline-native-driver-zip");
let zip_path = test_path("offline-native-driver-zip").join("kingbase.zip");
std::fs::create_dir_all(zip_path.parent().unwrap()).unwrap();
write_offline_native_driver_zip(&zip_path, "kingbase", "0.1.34");
let result = import_agents_from_zip(&manager, &zip_path, |_| {}).unwrap();
let result = import_agents_from_zip(&manager, &zip_path, |_| {}).await.unwrap();
assert_eq!(result.drivers_installed, vec!["kingbase"]);
assert_eq!(manager.load_state().installed_drivers["kingbase"].version, "0.1.34");
@@ -682,30 +688,30 @@ fn offline_zip_import_installs_versioned_native_driver_package() {
assert!(!manager.driver_jar_path("kingbase").exists());
}
#[test]
fn offline_zip_import_rejects_native_driver_for_another_platform() {
#[tokio::test]
async fn offline_zip_import_rejects_native_driver_for_another_platform() {
let manager = test_manager("offline-wrong-platform-native-driver-zip");
let zip_path = test_path("offline-wrong-platform-native-driver-zip").join("kingbase.zip");
std::fs::create_dir_all(zip_path.parent().unwrap()).unwrap();
let other_platform = if AgentManager::current_platform() == "windows-x64" { "linux-x64" } else { "windows-x64" };
write_offline_native_driver_zip_for_platform(&zip_path, "kingbase", "0.1.34", other_platform);
let err = import_agents_from_zip(&manager, &zip_path, |_| {}).unwrap_err();
let err = import_agents_from_zip(&manager, &zip_path, |_| {}).await.unwrap_err();
assert!(err.contains("no drivers compatible"));
assert!(err.contains(AgentManager::current_platform()));
assert!(!manager.is_driver_installed("kingbase"));
}
#[test]
fn offline_zip_import_preserves_existing_native_driver_when_binary_is_invalid() {
#[tokio::test]
async fn offline_zip_import_preserves_existing_native_driver_when_binary_is_invalid() {
let manager = test_manager("offline-invalid-native-preserves-existing");
let root = test_path("offline-invalid-native-preserves-existing-zip");
let valid_zip = root.join("valid.zip");
let invalid_zip = root.join("invalid.zip");
std::fs::create_dir_all(&root).unwrap();
write_offline_native_driver_zip(&valid_zip, "kingbase", "0.1.34");
import_agents_from_zip(&manager, &valid_zip, |_| {}).unwrap();
import_agents_from_zip(&manager, &valid_zip, |_| {}).await.unwrap();
let original = std::fs::read(manager.driver_native_path("kingbase")).unwrap();
write_offline_native_driver_zip_with_bytes(
@@ -715,7 +721,7 @@ fn offline_zip_import_preserves_existing_native_driver_when_binary_is_invalid()
AgentManager::current_platform(),
b"not-a-native-agent".to_vec(),
);
let err = import_agents_from_zip(&manager, &invalid_zip, |_| {}).unwrap_err();
let err = import_agents_from_zip(&manager, &invalid_zip, |_| {}).await.unwrap_err();
assert!(err.contains("not a"));
assert_eq!(std::fs::read(manager.driver_native_path("kingbase")).unwrap(), original);
@@ -763,15 +769,15 @@ fn offline_zip_import_rejects_unsafe_entry_path() {
assert!(err.contains("unsafe path"));
}
#[test]
fn offline_zip_import_prefers_native_artifact_over_java_fallback() {
#[tokio::test]
async fn offline_zip_import_prefers_native_artifact_over_java_fallback() {
let manager = test_manager("offline-native-preferred");
let zip_path = test_path("offline-native-preferred-zip").join("kingbase.zip");
std::fs::create_dir_all(zip_path.parent().unwrap()).unwrap();
write_offline_hybrid_driver_zip(&zip_path, "kingbase", "0.1.34");
let plan = inspect_offline_zip(&zip_path).unwrap();
let result = import_agents_from_zip(&manager, &zip_path, |_| {}).unwrap();
let result = import_agents_from_zip(&manager, &zip_path, |_| {}).await.unwrap();
assert_eq!(plan.driver_keys, vec!["kingbase"]);
assert_eq!(result.drivers_installed, vec!["kingbase"]);
@@ -779,6 +785,36 @@ fn offline_zip_import_prefers_native_artifact_over_java_fallback() {
assert!(!manager.driver_jar_path("kingbase").exists());
}
#[tokio::test]
async fn offline_zip_import_progress_carries_real_db_type_not_label() {
// Verify that the AgentProgressEvent.db_type field carries the real
// database key (e.g. "kingbase") rather than the display label
// (e.g. "KingbaseES"), so the frontend per-driver progress routing
// matches.
let manager = test_manager("offline-progress-db-type");
let zip_path = test_path("offline-progress-db-type-zip").join("kingbase.zip");
std::fs::create_dir_all(zip_path.parent().unwrap()).unwrap();
write_offline_native_driver_zip(&zip_path, "kingbase", "0.1.34");
let events = std::sync::Mutex::new(Vec::new());
import_agents_from_zip(&manager, &zip_path, |event| {
events.lock().unwrap().push(event);
})
.await
.unwrap();
let events = events.lock().unwrap();
let driver_events: Vec<_> = events.iter().filter(|e| e.step == "driver").collect();
assert!(!driver_events.is_empty(), "expected at least one driver progress event");
for event in &driver_events {
assert_eq!(
event.db_type.as_deref(),
Some("kingbase"),
"db_type must be the real key 'kingbase', not the display label"
);
}
}
fn write_offline_driver_zip(path: &std::path::Path, db_type: &str, version: &str) {
write_offline_driver_zip_with_jar(path, db_type, version, test_agent_jar_bytes());
}