From c32d0b7d6eefeb58d7e14235fdb523837a460c21 Mon Sep 17 00:00:00 2001 From: miracle Date: Wed, 8 Jul 2026 20:31:02 +0800 Subject: [PATCH] fix(duckdb): close worker process reliably --- crates/dbx-core/Cargo.toml | 13 + .../dbx-core/src/db/duckdb_worker_process.rs | 195 ++++++++-- .../dbx-core/src/db/duckdb_worker_runtime.rs | 61 +++- .../dbx-core/tests/duckdb_worker_process.rs | 345 ++++++++++++++++++ ...duckdb_worker_hanging_connect_test_host.rs | 20 + .../support/duckdb_worker_pid_test_host.rs | 13 + 6 files changed, 600 insertions(+), 47 deletions(-) create mode 100644 crates/dbx-core/tests/support/duckdb_worker_hanging_connect_test_host.rs create mode 100644 crates/dbx-core/tests/support/duckdb_worker_pid_test_host.rs diff --git a/crates/dbx-core/Cargo.toml b/crates/dbx-core/Cargo.toml index b8654ee9a..ef6de41fd 100644 --- a/crates/dbx-core/Cargo.toml +++ b/crates/dbx-core/Cargo.toml @@ -11,6 +11,19 @@ required-features = ["duckdb-bundled"] test = false bench = false +[[bin]] +name = "duckdb-worker-hanging-connect-test-host" +path = "tests/support/duckdb_worker_hanging_connect_test_host.rs" +required-features = ["duckdb-bundled"] +test = false + +[[bin]] +name = "duckdb-worker-pid-test-host" +path = "tests/support/duckdb_worker_pid_test_host.rs" +required-features = ["duckdb-bundled"] +test = false +bench = false + [features] default = ["duckdb-bundled", "mq-admin", "sqlite-sqlcipher"] duckdb-bundled = ["duckdb/bundled"] diff --git a/crates/dbx-core/src/db/duckdb_worker_process.rs b/crates/dbx-core/src/db/duckdb_worker_process.rs index aace96530..796c41b8b 100644 --- a/crates/dbx-core/src/db/duckdb_worker_process.rs +++ b/crates/dbx-core/src/db/duckdb_worker_process.rs @@ -20,6 +20,11 @@ use crate::db::duckdb_worker_protocol::{ use crate::models::connection::AttachedDatabaseConfig; use crate::storage::{normalize_duckdb_worker_max_processes, DUCKDB_WORKER_MAX_PROCESSES_DEFAULT}; +/// Error code the worker reports when a query error left the DuckDB connection poisoned +/// (see the duckdb-rs Parser Error bug). The client kills the worker on this code so the +/// next request starts a fresh one. Must match the code emitted in `duckdb_worker_runtime`. +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); const DEFAULT_WORKER_START_WAIT: Duration = Duration::from_secs(5); @@ -92,8 +97,29 @@ impl DuckDbWorkerClient { attached_databases: Vec, process_limit: usize, ) -> Result { + let client = Self::new_unconnected_with_timeouts( + executable, + path, + attached_databases, + process_limit, + DEFAULT_WORKER_REQUEST_TIMEOUT, + DEFAULT_WORKER_START_WAIT, + ); + client.ensure_connected().await?; + Ok(client) + } + + #[doc(hidden)] + pub fn new_unconnected_with_timeouts( + executable: PathBuf, + path: String, + attached_databases: Vec, + process_limit: usize, + request_timeout: Duration, + worker_start_timeout: Duration, + ) -> Self { let (process_limiter, process_limit) = duckdb_worker_process_limiter(process_limit); - let client = Self { + Self { inner: Arc::new(DuckDbWorkerClientInner { state: Mutex::new(WorkerProcessState::default()), pending: Arc::new(Mutex::new(HashMap::new())), @@ -103,13 +129,11 @@ impl DuckDbWorkerClient { process_limit, executable, connect_params: DuckDbWorkerConnectParams { path, attached_databases }, - request_timeout: DEFAULT_WORKER_REQUEST_TIMEOUT, - worker_start_timeout: DEFAULT_WORKER_START_WAIT, + request_timeout, + worker_start_timeout, next_id: AtomicU64::new(1), }), - }; - client.ensure_connected().await?; - Ok(client) + } } pub async fn execute( @@ -122,14 +146,32 @@ impl DuckDbWorkerClient { ) -> Result { let _query_guard = self.inner.query_lock.lock().await; let client = self.clone(); + // Cancellation and timeout restart the worker via cancel_or_kill below. An ordinary + // query error (parser/binder/catalog/runtime) leaves the worker healthy, so we keep it + // alive to preserve session-local state (:memory: contents, temp tables, SET, raw ATTACH). + // The one exception is a poisoned connection: duckdb-rs can leave the connection unusable + // after a syntax Parser Error, and the worker reports that with `duckdb_worker_poisoned`. + // We kill on that code so the next request restarts a fresh worker. Killing here (rather + // than letting the worker self-exit) avoids racing our own next request against a dying + // worker, and OS-level kill never runs the destructor that would abort the process. let future = async move { - client - .request::( + client.ensure_connected().await?; + match client + .send_request_structured::( DuckDbWorkerMethod::Execute, DuckDbWorkerExecuteParams { sql, database, max_rows }, None, ) .await + { + Ok(result) => Ok(result), + Err(error) => { + if error.code == DUCKDB_WORKER_POISONED_CODE { + client.kill().await; + } + Err(error.message) + } + } }; tokio::pin!(future); @@ -165,23 +207,17 @@ impl DuckDbWorkerClient { } pub async fn list_databases(&self) -> Result, String> { - self.request(DuckDbWorkerMethod::ListDatabases, serde_json::json!({}), Some(self.inner.request_timeout)).await + self.metadata_request(DuckDbWorkerMethod::ListDatabases, serde_json::json!({})).await } pub async fn list_schemas(&self, database: String) -> Result, String> { - self.request( - DuckDbWorkerMethod::ListSchemas, - serde_json::json!({ "database": database }), - Some(self.inner.request_timeout), - ) - .await + self.metadata_request(DuckDbWorkerMethod::ListSchemas, serde_json::json!({ "database": database })).await } pub async fn list_tables(&self, database: String, schema: String) -> Result, String> { - self.request( + self.metadata_request( DuckDbWorkerMethod::ListTables, serde_json::json!({ "database": database, "schema": schema }), - Some(self.inner.request_timeout), ) .await } @@ -192,24 +228,36 @@ impl DuckDbWorkerClient { schema: String, table: String, ) -> Result, String> { - self.request( + self.metadata_request( DuckDbWorkerMethod::ListColumns, serde_json::json!({ "database": database, "schema": schema, "table": table }), - Some(self.inner.request_timeout), ) .await } pub async fn attach_database(&self, attached: AttachedDatabaseConfig) -> Result<(), String> { - self.request::( - DuckDbWorkerMethod::AttachDatabase, - attached, - Some(self.inner.request_timeout), - ) - .await?; + self.metadata_request::(DuckDbWorkerMethod::AttachDatabase, attached).await?; Ok(()) } + /// Runs a metadata request that executes synchronously inside the worker's stdin loop. + /// If it times out the worker is likely stuck inside DuckDB and can no longer read further + /// requests, so we kill it; the next call restarts a fresh worker via `ensure_connected`. + /// Ordinary errors (e.g. a missing table) leave the worker healthy and do not kill it. + async fn metadata_request(&self, method: DuckDbWorkerMethod, params: impl serde::Serialize) -> Result + where + T: serde::de::DeserializeOwned, + { + let timeout = self.inner.request_timeout; + match tokio::time::timeout(timeout, self.request::(method, params, None)).await { + Ok(result) => result, + Err(_) => { + self.kill().await; + Err(format!("DuckDB worker request timed out after {}s", timeout.as_secs())) + } + } + } + pub async fn cancel(&self) -> Result<(), String> { self.kill().await; Ok(()) @@ -272,12 +320,38 @@ impl DuckDbWorkerClient { } } - self.send_request::( - DuckDbWorkerMethod::Connect, - self.inner.connect_params.clone(), - Some(self.inner.request_timeout), - ) - .await?; + let mut attempts = 0; + loop { + match self + .send_request_structured::( + DuckDbWorkerMethod::Connect, + self.inner.connect_params.clone(), + Some(self.inner.request_timeout), + ) + .await + { + Ok(_) => break, + Err(error) => { + let retry_delay = duckdb_connect_file_lock_retry_delay(attempts, &error.message); + // A failed Connect leaves this worker without a valid session. Keep no stale + // child around, especially after OS file-lock errors where the user may retry + // once the competing process exits. + self.kill().await; + if let Some(delay) = retry_delay { + attempts += 1; + log::warn!( + "[duckdb-worker:connect-retry] attempt={} delay_ms={} error={}", + attempts, + delay.as_millis(), + error.message + ); + tokio::time::sleep(delay).await; + continue; + } + return Err(error.message); + } + } + } let mut state = self.inner.state.lock().await; state.connected = true; @@ -340,6 +414,21 @@ impl DuckDbWorkerClient { params: impl serde::Serialize, timeout: Option, ) -> Result + where + T: serde::de::DeserializeOwned, + { + self.send_request_structured(method, params, timeout).await.map_err(|error| error.message) + } + + /// Like `send_request` but preserves the worker error `code` so callers can react to + /// specific conditions (e.g. `duckdb_worker_poisoned`). A protocol/transport failure is + /// surfaced as an error with code `duckdb_worker_error`. + async fn send_request_structured( + &self, + method: DuckDbWorkerMethod, + params: impl serde::Serialize, + timeout: Option, + ) -> Result where T: serde::de::DeserializeOwned, { @@ -365,30 +454,33 @@ impl DuckDbWorkerClient { if let Err(err) = write_result { self.inner.pending.lock().await.remove(&id); - return Err(err); + return Err(err.into()); } let response = match timeout { Some(timeout) => match tokio::time::timeout(timeout, rx).await { Ok(Ok(response)) => response, - Ok(Err(_)) => return Err("DuckDB worker response channel closed".to_string()), + Ok(Err(_)) => return Err("DuckDB worker response channel closed".into()), Err(_) => { self.inner.pending.lock().await.remove(&id); - return Err(format!("DuckDB worker request timed out after {}s", timeout.as_secs())); + return Err(DuckDbWorkerError::new( + DUCKDB_WORKER_REQUEST_TIMEOUT_CODE, + format!("DuckDB worker request timed out after {}s", timeout.as_secs()), + )); } }, - None => rx.await.map_err(|_| "DuckDB worker response channel closed".to_string())?, + None => rx.await.map_err(|_| DuckDbWorkerError::from("DuckDB worker response channel closed"))?, }; if !response.ok { let error = response .error .unwrap_or_else(|| DuckDbWorkerError::new("duckdb_worker_error", "DuckDB worker request failed")); - return Err(error.message); + return Err(error); } let result = response.result.unwrap_or(serde_json::Value::Null); - serde_json::from_value(result).map_err(|e| e.to_string()) + serde_json::from_value(result).map_err(|e| DuckDbWorkerError::from(e.to_string())) } async fn fail_pending_for_generation(&self, generation: u64, code: &'static str, message: &'static str) { @@ -434,6 +526,35 @@ fn duckdb_worker_process_limit_error(process_limit: usize) -> String { ) } +fn duckdb_connect_file_lock_retry_delay(attempts: usize, message: &str) -> Option { + if !is_transient_duckdb_file_lock_error(message) { + return None; + } + match attempts { + 0 => Some(Duration::from_millis(50)), + 1 => Some(Duration::from_millis(100)), + 2 => Some(Duration::from_millis(200)), + 3 => Some(Duration::from_millis(400)), + 4 => Some(Duration::from_millis(800)), + _ => None, + } +} + +fn is_transient_duckdb_file_lock_error(message: &str) -> bool { + let lower = message.to_ascii_lowercase(); + let mentions_file_open = lower.contains("cannot open file") + || lower.contains("could not set lock") + || lower.contains("file is already open"); + let mentions_lock = lower.contains("file is already open") + || lower.contains("being used by another process") + || lower.contains("process cannot access the file") + || lower.contains("sharing violation") + || lower.contains("resource temporarily unavailable") + || message.contains("另一个程序正在使用") + || message.contains("进程无法访问"); + mentions_file_open && mentions_lock +} + fn spawn_stdout_reader( stdout: tokio::process::ChildStdout, pending: PendingRequests, diff --git a/crates/dbx-core/src/db/duckdb_worker_runtime.rs b/crates/dbx-core/src/db/duckdb_worker_runtime.rs index bf28129c2..ac0cdb002 100644 --- a/crates/dbx-core/src/db/duckdb_worker_runtime.rs +++ b/crates/dbx-core/src/db/duckdb_worker_runtime.rs @@ -102,6 +102,27 @@ impl DuckDbWorkerSession { let connection = self.connection.as_ref().ok_or("DuckDB worker is not connected")?.clone(); Ok(connection.interrupt_handle()) } + + /// 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 + /// `duckdb_extract_statements` stage (a syntax `Parser Error`) permanently poisons + /// the connection: every later operation returns `resource deadlock would occur`, + /// and dropping the connection aborts the whole process. Benign errors (binder, + /// catalog, runtime) do not poison the connection. + /// + /// This probe runs a trivial query through `execute_batch`, which uses the + /// non-poisoning `duckdb_query_arrow` path and fails fast on a poisoned connection. + /// It returns `true` when the connection is healthy and safe to reuse. + fn is_connection_healthy(&self) -> bool { + match self.connection.as_ref() { + Some(connection) => match connection.lock() { + Ok(locked) => locked.execute_batch("SELECT 1").is_ok(), + Err(_) => false, + }, + None => false, + } + } } #[derive(Clone, Default)] @@ -175,19 +196,39 @@ impl DuckDbWorkerRuntime { tokio::spawn(async move { let id = request.id.clone(); let params = request.parse_params::(); - let result = match params { - Ok(params) => tokio::task::spawn_blocking(move || { - let mut session = session.lock().unwrap_or_else(|e| e.into_inner()); - session.execute(params) - }) - .await - .map_err(|e| e.to_string()) - .and_then(|result| result), - Err(err) => Err(err.message), + let (result, poisoned) = match params { + Ok(params) => { + let probe_session = session.clone(); + tokio::task::spawn_blocking(move || { + let mut session = probe_session.lock().unwrap_or_else(|e| e.into_inner()); + let result = session.execute(params); + // Only probe connection health when execute failed: a syntax + // Parser Error can poison the connection (duckdb-rs bug), and any + // later use — including drop — would abort the process. + let poisoned = result.is_err() && !session.is_connection_healthy(); + (result, poisoned) + }) + .await + .map_err(|e| e.to_string()) + .unwrap_or_else(|err| (Err(err), true)) + } + Err(err) => (Err(err.message), false), }; *active_interrupt.lock().unwrap_or_else(|e| e.into_inner()) = None; let response = match result { Ok(result) => DuckDbWorkerResponse::ok(id, result), + // A poisoned connection cannot be reused or even safely dropped in-process. + // Report a distinct code so the parent kills this worker (OS-level, skipping + // destructors) and restarts a fresh one on the next request. We must not exit + // here ourselves: process::exit races the parent's next request, which could + // be written to a dying worker. Killing on the parent side has no such window. + Err(err) if poisoned => { + log::warn!("[duckdb-worker:poisoned-connection] reporting to parent for restart"); + DuckDbWorkerResponse::err( + id, + DuckDbWorkerError::from_message("duckdb_worker_poisoned", err), + ) + } Err(err) => { DuckDbWorkerResponse::err(id, DuckDbWorkerError::from_message("duckdb_execute_failed", err)) } @@ -256,7 +297,7 @@ 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 { - break; + std::process::exit(0); }; if line.trim().is_empty() { continue; diff --git a/crates/dbx-core/tests/duckdb_worker_process.rs b/crates/dbx-core/tests/duckdb_worker_process.rs index 741828b81..7995556ef 100644 --- a/crates/dbx-core/tests/duckdb_worker_process.rs +++ b/crates/dbx-core/tests/duckdb_worker_process.rs @@ -1,11 +1,18 @@ #![cfg(feature = "duckdb-bundled")] use std::path::PathBuf; +use std::process::Stdio; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use dbx_core::db::duckdb_worker_process::DuckDbWorkerClient; +use dbx_core::db::duckdb_worker_protocol::{ + DuckDbWorkerConnectParams, DuckDbWorkerExecuteParams, DuckDbWorkerMethod, DuckDbWorkerRequest, DuckDbWorkerResponse, +}; use dbx_core::query_cancel::{RunningQueries, RunningTaskMetadata}; +use sysinfo::{Pid, ProcessRefreshKind, ProcessesToUpdate, RefreshKind, System}; +use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; +use tokio::process::Command; use tokio_util::sync::CancellationToken; static TEMP_DB_COUNTER: AtomicU64 = AtomicU64::new(0); @@ -122,6 +129,303 @@ async fn worker_process_recovers_after_registered_cancel_interrupt() { let _ = std::fs::remove_file(&db_path); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn worker_process_recovers_after_parser_error() { + let executable = PathBuf::from(env!("CARGO_BIN_EXE_duckdb-worker-test-host")); + let db_path = temp_duckdb_path(); + let _ = std::fs::remove_file(&db_path); + + let client = + DuckDbWorkerClient::open_with_executable(executable, db_path.to_string_lossy().to_string(), Vec::new()) + .await + .expect("worker process connects"); + + client + .execute( + None, + "CREATE TABLE events (id INTEGER, name VARCHAR); INSERT INTO events VALUES (1, 'login');".to_string(), + Some(10), + None, + Some(Duration::from_secs(5)), + ) + .await + .expect("create events table"); + + let err = client + .execute(None, "select * from table limit 19;".to_string(), Some(10), None, Some(Duration::from_secs(5))) + .await + .expect_err("reserved word query should fail"); + assert!(err.contains("Parser Error"), "unexpected error: {err}"); + + let probe = tokio::time::timeout( + Duration::from_secs(5), + client.execute( + None, + "SELECT * FROM events LIMIT 100;".to_string(), + Some(100), + None, + Some(Duration::from_secs(5)), + ), + ) + .await + .expect("query after parser error should not hang") + .expect("query after parser error should succeed"); + + assert_eq!(probe.columns, vec!["id".to_string(), "name".to_string()]); + assert_eq!(probe.rows, vec![vec![serde_json::json!(1), serde_json::json!("login")]]); + + client.shutdown().await; + let _ = std::fs::remove_file(&db_path); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn worker_process_keeps_session_state_after_benign_error() { + let executable = PathBuf::from(env!("CARGO_BIN_EXE_duckdb-worker-test-host")); + let db_path = temp_duckdb_path(); + let _ = std::fs::remove_file(&db_path); + + let client = + DuckDbWorkerClient::open_with_executable(executable, db_path.to_string_lossy().to_string(), Vec::new()) + .await + .expect("worker process connects"); + + // Temp tables live only in the worker's session. If a benign error restarted the + // worker, this table would be gone after the error. + client + .execute( + None, + "CREATE TEMP TABLE scratch (id INTEGER); INSERT INTO scratch VALUES (7);".to_string(), + Some(10), + None, + Some(Duration::from_secs(5)), + ) + .await + .expect("create temp table"); + + // Catalog error: prepare succeeds, the query fails at bind/exec time. This does not + // poison the connection, so the worker must stay alive and keep the temp table. + let err = client + .execute(None, "SELECT * FROM does_not_exist;".to_string(), Some(10), None, Some(Duration::from_secs(5))) + .await + .expect_err("missing table query should fail"); + assert!(err.contains("does_not_exist") || err.contains("Catalog Error"), "unexpected error: {err}"); + + let probe = tokio::time::timeout( + Duration::from_secs(5), + client.execute(None, "SELECT id FROM scratch;".to_string(), Some(10), None, Some(Duration::from_secs(5))), + ) + .await + .expect("query after benign error should not hang") + .expect("temp table should still exist after a benign error"); + + assert_eq!(probe.columns, vec!["id".to_string()]); + assert_eq!(probe.rows, vec![vec![serde_json::json!(7)]]); + + client.shutdown().await; + let _ = std::fs::remove_file(&db_path); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn worker_process_exits_when_stdin_closes_during_active_query() { + let executable = PathBuf::from(env!("CARGO_BIN_EXE_duckdb-worker-test-host")); + let db_path = temp_duckdb_path(); + let _ = std::fs::remove_file(&db_path); + let mut child = Command::new(executable) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::inherit()) + .spawn() + .expect("spawn worker"); + let mut stdin = child.stdin.take().expect("worker stdin"); + let stdout = child.stdout.take().expect("worker stdout"); + let mut lines = BufReader::new(stdout).lines(); + + write_worker_request( + &mut stdin, + DuckDbWorkerRequest::new( + "connect", + DuckDbWorkerMethod::Connect, + DuckDbWorkerConnectParams { path: db_path.to_string_lossy().to_string(), attached_databases: Vec::new() }, + ) + .expect("connect request"), + ) + .await; + let connected = read_worker_response(&mut lines).await; + assert!(connected.ok, "connect failed: {connected:?}"); + + write_worker_request( + &mut stdin, + DuckDbWorkerRequest::new( + "long-query", + DuckDbWorkerMethod::Execute, + DuckDbWorkerExecuteParams { + sql: "SELECT sum(sin(i::DOUBLE) * cos(i::DOUBLE / 3.0)) FROM range(100000000000) AS t(i)".to_string(), + database: None, + max_rows: Some(10), + }, + ) + .expect("execute request"), + ) + .await; + + tokio::time::sleep(Duration::from_millis(200)).await; + drop(stdin); + + tokio::time::timeout(Duration::from_secs(5), child.wait()) + .await + .expect("worker should exit after stdin closes") + .expect("wait for worker"); + let _ = std::fs::remove_file(&db_path); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn worker_process_is_killed_after_connect_timeout() { + let executable = PathBuf::from(env!("CARGO_BIN_EXE_duckdb-worker-hanging-connect-test-host")); + let db_path = temp_duckdb_path(); + let pid_file = temp_pid_path(); + let _ = std::fs::remove_file(&pid_file); + let _ = std::fs::remove_file(&db_path); + + std::env::set_var("DBX_DUCKDB_HANGING_CONNECT_PID_FILE", &pid_file); + let client = DuckDbWorkerClient::new_unconnected_with_timeouts( + executable, + db_path.to_string_lossy().to_string(), + Vec::new(), + 4, + Duration::from_millis(200), + Duration::from_secs(5), + ); + + let err = client + .execute(None, "SELECT 1".to_string(), Some(10), None, Some(Duration::from_secs(5))) + .await + .expect_err("connect should time out"); + std::env::remove_var("DBX_DUCKDB_HANGING_CONNECT_PID_FILE"); + assert!(err.contains("timed out"), "unexpected error: {err}"); + + let pid = read_pid_file(&pid_file).expect("hanging worker pid"); + wait_until_process_exits(pid, Duration::from_secs(5)).await; + + let _ = std::fs::remove_file(&pid_file); + let _ = std::fs::remove_file(&db_path); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn worker_process_is_killed_after_connect_error() { + let executable = PathBuf::from(env!("CARGO_BIN_EXE_duckdb-worker-pid-test-host")); + let db_path = temp_duckdb_path(); + let pid_file = temp_pid_path(); + let _ = std::fs::remove_file(&pid_file); + let _ = std::fs::remove_file(&db_path); + + let _file_owner = duckdb::Connection::open(&db_path).expect("open lock owner connection"); + + std::env::set_var("DBX_DUCKDB_PID_TEST_HOST_PID_FILE", &pid_file); + let client = DuckDbWorkerClient::new_unconnected_with_timeouts( + executable, + db_path.to_string_lossy().to_string(), + Vec::new(), + 4, + Duration::from_secs(5), + Duration::from_secs(5), + ); + + let err = client + .execute(None, "SELECT 1".to_string(), Some(10), None, Some(Duration::from_secs(5))) + .await + .expect_err("connect should fail while another process owns the DuckDB file"); + std::env::remove_var("DBX_DUCKDB_PID_TEST_HOST_PID_FILE"); + assert!( + err.contains("Cannot open file") || err.contains("already open") || err.contains("另一个程序正在使用"), + "unexpected error: {err}" + ); + + let pid = read_pid_file(&pid_file).expect("worker pid"); + wait_until_process_exits(pid, Duration::from_secs(5)).await; + + let _ = std::fs::remove_file(&pid_file); + drop(_file_owner); + let _ = std::fs::remove_file(&db_path); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn worker_process_retries_connect_after_transient_file_lock() { + let executable = PathBuf::from(env!("CARGO_BIN_EXE_duckdb-worker-test-host")); + let db_path = temp_duckdb_path(); + let _ = std::fs::remove_file(&db_path); + + let mut owner_child = Command::new(&executable) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::inherit()) + .spawn() + .expect("spawn lock owner worker"); + let mut owner_stdin = owner_child.stdin.take().expect("owner worker stdin"); + let owner_stdout = owner_child.stdout.take().expect("owner worker stdout"); + let mut owner_lines = BufReader::new(owner_stdout).lines(); + write_worker_request( + &mut owner_stdin, + DuckDbWorkerRequest::new( + "connect-owner", + DuckDbWorkerMethod::Connect, + DuckDbWorkerConnectParams { path: db_path.to_string_lossy().to_string(), attached_databases: Vec::new() }, + ) + .expect("owner connect request"), + ) + .await; + let connected = read_worker_response(&mut owner_lines).await; + assert!(connected.ok, "owner connect failed: {connected:?}"); + drop(owner_lines); + + let release_owner = tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(200)).await; + drop(owner_stdin); + let _ = tokio::time::timeout(Duration::from_secs(5), owner_child.wait()).await; + }); + + let client = DuckDbWorkerClient::new_unconnected_with_timeouts( + executable, + db_path.to_string_lossy().to_string(), + Vec::new(), + 4, + Duration::from_secs(5), + Duration::from_secs(5), + ); + let result = tokio::time::timeout( + Duration::from_secs(5), + client.execute(None, "SELECT 42 AS retried".to_string(), Some(10), None, Some(Duration::from_secs(5))), + ) + .await; + + let _ = release_owner.await; + client.shutdown().await; + let _ = std::fs::remove_file(&db_path); + + let result = result + .expect("retrying connect should not hang") + .expect("connect should retry after a transient DuckDB file lock"); + assert_eq!(result.columns, vec!["retried".to_string()]); + assert_eq!(result.rows, vec![vec![serde_json::json!(42)]]); +} + +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"); + stdin.write_all(b"\n").await.expect("write newline"); + stdin.flush().await.expect("flush request"); +} + +async fn read_worker_response( + lines: &mut tokio::io::Lines>, +) -> DuckDbWorkerResponse { + let line = tokio::time::timeout(Duration::from_secs(5), lines.next_line()) + .await + .expect("worker response timed out") + .expect("read worker response") + .expect("worker response line"); + serde_json::from_str(&line).expect("parse worker response") +} + fn temp_duckdb_path() -> PathBuf { let suffix = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_nanos(); let pid = std::process::id(); @@ -130,3 +434,44 @@ fn temp_duckdb_path() -> PathBuf { let counter = TEMP_DB_COUNTER.fetch_add(1, Ordering::Relaxed); std::env::temp_dir().join(format!("dbx-duckdb-worker-process-{pid}-{suffix}-{counter}.duckdb")) } + +fn temp_pid_path() -> PathBuf { + let suffix = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_nanos(); + let pid = std::process::id(); + let counter = TEMP_DB_COUNTER.fetch_add(1, Ordering::Relaxed); + std::env::temp_dir().join(format!("dbx-duckdb-worker-process-{pid}-{suffix}-{counter}.pid")) +} + +fn read_pid_file(path: &PathBuf) -> Option { + for _ in 0..20 { + if let Ok(contents) = std::fs::read_to_string(path) { + if let Ok(pid) = contents.trim().parse::() { + return Some(pid); + } + } + std::thread::sleep(Duration::from_millis(25)); + } + None +} + +async fn wait_until_process_exits(pid: u32, timeout: Duration) { + let deadline = tokio::time::Instant::now() + timeout; + loop { + if !process_exists(pid) { + return; + } + assert!(tokio::time::Instant::now() < deadline, "worker process {pid} should have been killed"); + tokio::time::sleep(Duration::from_millis(50)).await; + } +} + +fn process_exists(pid: u32) -> bool { + let mut system = + System::new_with_specifics(RefreshKind::new().with_processes(ProcessRefreshKind::new().with_memory())); + system.refresh_processes_specifics( + ProcessesToUpdate::Some(&[Pid::from(pid as usize)]), + true, + ProcessRefreshKind::new().with_memory(), + ); + system.process(Pid::from(pid as usize)).is_some() +} diff --git a/crates/dbx-core/tests/support/duckdb_worker_hanging_connect_test_host.rs b/crates/dbx-core/tests/support/duckdb_worker_hanging_connect_test_host.rs new file mode 100644 index 000000000..291923927 --- /dev/null +++ b/crates/dbx-core/tests/support/duckdb_worker_hanging_connect_test_host.rs @@ -0,0 +1,20 @@ +#![cfg(feature = "duckdb-bundled")] + +use std::io::{self, BufRead}; +use std::time::Duration; + +fn main() { + if let Ok(path) = std::env::var("DBX_DUCKDB_HANGING_CONNECT_PID_FILE") { + let _ = std::fs::write(path, std::process::id().to_string()); + } + + let stdin = io::stdin(); + let mut lines = stdin.lock().lines(); + let Some(Ok(_line)) = lines.next() else { + return; + }; + + loop { + std::thread::sleep(Duration::from_secs(60)); + } +} diff --git a/crates/dbx-core/tests/support/duckdb_worker_pid_test_host.rs b/crates/dbx-core/tests/support/duckdb_worker_pid_test_host.rs new file mode 100644 index 000000000..4ce0fc672 --- /dev/null +++ b/crates/dbx-core/tests/support/duckdb_worker_pid_test_host.rs @@ -0,0 +1,13 @@ +#![cfg(feature = "duckdb-bundled")] + +fn main() { + if let Ok(path) = std::env::var("DBX_DUCKDB_PID_TEST_HOST_PID_FILE") { + let _ = std::fs::write(path, std::process::id().to_string()); + } + + let runtime = tokio::runtime::Runtime::new().expect("Failed to create DuckDB worker test runtime"); + if let Err(err) = runtime.block_on(dbx_core::db::duckdb_worker_runtime::run_stdio_worker()) { + eprintln!("DuckDB worker test host failed: {err}"); + std::process::exit(1); + } +}