From 8475c84d4892da6ec6a08782a5c368cc0cc1ec77 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Wed, 26 Aug 2026 15:23:29 +0800 Subject: [PATCH] fix(duckdb): checkpoint wal on worker shutdown --- .gitignore | 1 + agents/drivers/duckdb/src/connection.rs | 4 +- agents/drivers/duckdb/src/runtime.rs | 62 +++++++++++++++++-- .../duckdb/tests/duckdb_worker_process.rs | 50 +++++++++++++++ .../dbx-core/src/db/duckdb_worker_process.rs | 60 +++++++++++++++++- 5 files changed, 168 insertions(+), 9 deletions(-) diff --git a/.gitignore b/.gitignore index 1a715f8baf..bdc5c2b75f 100644 --- a/.gitignore +++ b/.gitignore @@ -22,6 +22,7 @@ src-tauri/webview2-fixed-runtime/ plugins/jdbc/target/ plugins/jdbc/dependency-reduced-pom.xml plugins/jdbc/lib/*.jar +agents/drivers/duckdb/target/ plugins/jdbc/.gradle/ plugins/jdbc/build/ plugins/jdbc/dist/ diff --git a/agents/drivers/duckdb/src/connection.rs b/agents/drivers/duckdb/src/connection.rs index 9981b69838..9f042f70dc 100644 --- a/agents/drivers/duckdb/src/connection.rs +++ b/agents/drivers/duckdb/src/connection.rs @@ -23,7 +23,6 @@ impl DuckDbConnection { } // Preserve PoisonError ownership so close_connection can still release the contained DuckDB handle. - #[cfg(test)] #[allow(clippy::result_large_err)] fn into_inner(self) -> std::sync::LockResult { self.connection.into_inner() @@ -116,8 +115,7 @@ pub fn is_memory_database_path(path: &str) -> bool { /// Unlike relying on Drop, this calls `duckdb_disconnect` synchronously /// so the file handle is released before this function returns. /// On Windows this prevents "file already in use" errors when reconnecting. -#[cfg(test)] -fn close_connection(con: Arc) { +pub(crate) fn close_connection(con: Arc) { match Arc::try_unwrap(con) { Ok(handle) => match handle.into_inner() { Ok(conn) => { diff --git a/agents/drivers/duckdb/src/runtime.rs b/agents/drivers/duckdb/src/runtime.rs index e6d426ab52..9d70df898a 100644 --- a/agents/drivers/duckdb/src/runtime.rs +++ b/agents/drivers/duckdb/src/runtime.rs @@ -148,6 +148,27 @@ impl DuckDbWorkerSession { Ok(connection.interrupt_handle()) } + /// Takes and closes the connection so DuckDB's shutdown checkpoint runs and + /// the database WAL is removed before this worker exits. Returns false — + /// leaving the connection untouched — when the connection is poisoned + /// (dropping it aborts the process) so the caller falls back to an + /// immediate exit and the WAL replays on the next open. + fn close_connection_for_shutdown(&mut self) -> bool { + if self.connection.is_none() { + return true; + } + if !self.is_connection_healthy() { + return false; + } + match self.connection.take() { + Some(connection) => { + crate::connection::close_connection(connection); + true + } + None => true, + } + } + /// Probes whether the connection is still usable after an execute error. /// /// duckdb-rs 1.10503.1 has a bug where `Connection::prepare()` failing at the @@ -205,6 +226,21 @@ struct WorkerHandleResult { } impl DuckDbWorkerRuntime { + /// Closes the session connection when the worker is idle and the connection + /// is healthy, so the DuckDB shutdown checkpoint removes the WAL. Returns + /// false when a query is running (the session is locked or an interrupt is + /// registered) or the connection is poisoned; callers then keep the legacy + /// immediate-exit behavior and rely on WAL replay at the next open. + fn close_session_for_shutdown(&self) -> bool { + if self.active_interrupt.lock().unwrap_or_else(|e| e.into_inner()).is_some() { + return false; + } + match self.session.try_lock() { + Ok(mut session) => session.close_connection_for_shutdown(), + Err(_) => false, + } + } + async fn handle_request( &self, request: DuckDbWorkerRequest, @@ -327,10 +363,18 @@ impl DuckDbWorkerRuntime { shutdown: false, } } - DuckDbWorkerMethod::Shutdown => WorkerHandleResult { - response: Some(DuckDbWorkerResponse::ok(request.id, serde_json::json!({ "shutdown": true }))), - shutdown: true, - }, + DuckDbWorkerMethod::Shutdown => { + // Close the connection before acking: the checkpoint that removes + // the WAL must finish before the parent observes this response + // (or the process exit it waits for). Busy or poisoned sessions + // skip the close; the parent's kill fallback and WAL replay on + // the next open cover them. + self.close_session_for_shutdown(); + WorkerHandleResult { + response: Some(DuckDbWorkerResponse::ok(request.id, serde_json::json!({ "shutdown": true }))), + shutdown: true, + } + } } } @@ -377,7 +421,15 @@ pub async fn run_stdio_worker() -> Result<(), String> { loop { let line = lines.next_line().await.map_err(|e| e.to_string())?; let Some(line) = line else { - std::process::exit(0); + // Parent closed stdin (exit or crash). Idle workers checkpoint and + // remove their WAL on the normal return path; a worker still running + // a query must not wait for it (nor drop a poisoned connection), + // so it keeps the legacy immediate exit and the WAL replays on the + // next open. + if !runtime.close_session_for_shutdown() { + std::process::exit(0); + } + break; }; if line.trim().is_empty() { continue; diff --git a/agents/drivers/duckdb/tests/duckdb_worker_process.rs b/agents/drivers/duckdb/tests/duckdb_worker_process.rs index a358f53b87..39d2c026fd 100644 --- a/agents/drivers/duckdb/tests/duckdb_worker_process.rs +++ b/agents/drivers/duckdb/tests/duckdb_worker_process.rs @@ -499,6 +499,52 @@ async fn worker_process_retries_connect_after_transient_file_lock() { assert_eq!(result.rows, vec![vec![serde_json::json!(42)]]); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn worker_shutdown_checkpoints_and_removes_wal() { + let _guard = duckdb_worker_process_test_guard().await; + let executable = PathBuf::from(env!("CARGO_BIN_EXE_dbx-duckdb-driver")); + let db_path = temp_duckdb_path(); + let wal_path = wal_path_for(&db_path); + let _ = std::fs::remove_file(&db_path); + let _ = std::fs::remove_file(&wal_path); + + let client = + DuckDbWorkerClient::open_with_executable(executable.clone(), db_path.to_string_lossy().to_string(), Vec::new(), None) + .await + .expect("worker process connects"); + client + .execute( + None, + "CREATE TABLE events (id INTEGER); INSERT INTO events VALUES (7);".to_string(), + Some(10), + None, + Some(Duration::from_secs(5)), + ) + .await + .expect("write events"); + + // Committed writes sit in the WAL while the connection is open. + assert!(wal_path.exists(), "committed writes should leave a WAL file next to the database"); + + client.shutdown().await; + + assert!(!wal_path.exists(), "shutdown should checkpoint and remove the WAL file"); + + // The checkpointed data survives a fresh open of the database file. + let client = + DuckDbWorkerClient::open_with_executable(executable, db_path.to_string_lossy().to_string(), Vec::new(), None) + .await + .expect("reopen after shutdown"); + let result = client + .execute(None, "SELECT id FROM events".to_string(), Some(10), None, Some(Duration::from_secs(5))) + .await + .expect("read events back"); + assert_eq!(result.rows, vec![vec![serde_json::json!(7)]]); + client.shutdown().await; + assert!(!wal_path.exists(), "second shutdown should also remove the WAL file"); + let _ = std::fs::remove_file(&db_path); +} + async fn write_worker_request(stdin: &mut tokio::process::ChildStdin, request: DuckDbWorkerRequest) { let line = serde_json::to_string(&request).expect("request json"); stdin.write_all(line.as_bytes()).await.expect("write request"); @@ -517,6 +563,10 @@ async fn read_worker_response( serde_json::from_str(&line).expect("parse worker response") } +fn wal_path_for(db_path: &PathBuf) -> PathBuf { + PathBuf::from(format!("{}.wal", db_path.display())) +} + fn temp_duckdb_path() -> PathBuf { let suffix = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_nanos(); let pid = std::process::id(); diff --git a/crates/dbx-core/src/db/duckdb_worker_process.rs b/crates/dbx-core/src/db/duckdb_worker_process.rs index 4ac90c3b88..d3b8775797 100644 --- a/crates/dbx-core/src/db/duckdb_worker_process.rs +++ b/crates/dbx-core/src/db/duckdb_worker_process.rs @@ -28,6 +28,10 @@ const DUCKDB_WORKER_POISONED_CODE: &str = "duckdb_worker_poisoned"; const DUCKDB_WORKER_REQUEST_TIMEOUT_CODE: &str = "duckdb_worker_request_timeout"; const DEFAULT_WORKER_REQUEST_TIMEOUT: Duration = Duration::from_secs(60); const DEFAULT_WORKER_KILL_WAIT: Duration = Duration::from_secs(3); +/// Grace period for a shutdown-acked worker to exit by itself (session drop +/// plus DuckDB shutdown checkpoint). Kept below the 3s pool close timeout in +/// `connection.rs` so pool close cancellation cannot cut the checkpoint short. +const WORKER_SHUTDOWN_EXIT_WAIT: Duration = Duration::from_millis(2500); const DEFAULT_WORKER_START_WAIT: Duration = Duration::from_secs(5); pub const DUCKDB_DRIVER_PATH_ENV: &str = "DBX_DUCKDB_DRIVER_PATH"; type PendingRequests = Arc>>; @@ -362,7 +366,61 @@ impl DuckDbWorkerClient { Some(self.inner.request_timeout), ) .await; - self.kill().await; + self.wait_for_exit_or_kill().await; + } + + /// Lets a shutdown-acked worker exit on its own so its session drop runs + /// the DuckDB shutdown checkpoint that removes the database WAL; SIGKILL + /// stays the fallback for stuck workers. Killing immediately after the + /// ack would truncate that checkpoint and leave the WAL file behind. + /// The wait must fit inside the caller's pool close timeout so the + /// detached close task is not cancelled mid-checkpoint. + async fn wait_for_exit_or_kill(&self) { + let (child, generation) = { + let mut state = self.inner.state.lock().await; + let child = state.child.take(); + let generation = state.generation; + state.stdin = None; + state.connected = false; + (child, generation) + }; + + let Some(mut child) = child else { return }; + match tokio::time::timeout(WORKER_SHUTDOWN_EXIT_WAIT, child.wait()).await { + Ok(Ok(status)) => { + log::info!("[duckdb-worker:shutdown:exit] status={status}"); + self.fail_pending_for_generation( + generation, + "duckdb_worker_killed", + "DuckDB worker process was killed", + ) + .await; + return; + } + Ok(Err(err)) => { + log::warn!("[duckdb-worker:shutdown:wait-failed] error={err}"); + self.fail_pending_for_generation( + generation, + "duckdb_worker_killed", + "DuckDB worker process was killed", + ) + .await; + return; + } + Err(_) => { + log::warn!("[duckdb-worker:shutdown:wait-timeout] wait_ms={}", WORKER_SHUTDOWN_EXIT_WAIT.as_millis()); + } + } + + let _ = child.start_kill(); + match tokio::time::timeout(DEFAULT_WORKER_KILL_WAIT, child.wait()).await { + Ok(Ok(status)) => log::info!("[duckdb-worker:kill:exit] status={status}"), + Ok(Err(err)) => log::warn!("[duckdb-worker:kill:wait-failed] error={err}"), + Err(_) => { + log::warn!("[duckdb-worker:kill:wait-timeout] wait_ms={}", DEFAULT_WORKER_KILL_WAIT.as_millis()) + } + } + self.fail_pending_for_generation(generation, "duckdb_worker_killed", "DuckDB worker process was killed").await; } pub async fn kill(&self) {