xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(duckdb): close worker process reliably

This commit is contained in:
miracle 2026-07-08 20:31:02 +08:00 committed by GitHub
parent da88d03ad4
commit c32d0b7d6e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 600 additions and 47 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
crates/dbx-core/Cargo.toml
View File

@ -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"]

195
crates/dbx-core/src/db/duckdb_worker_process.rs
View File

@ -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<AttachedDatabaseConfig>,
process_limit: usize,
) -> Result<Self, String> {
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<AttachedDatabaseConfig>,
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<db::QueryResult, String> {
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::<db::QueryResult>(
client.ensure_connected().await?;
match client
.send_request_structured::<db::QueryResult>(
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<Vec<db::DatabaseInfo>, 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<Vec<String>, 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<Vec<db::TableInfo>, 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<Vec<db::ColumnInfo>, 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::<serde_json::Value>(
DuckDbWorkerMethod::AttachDatabase,
attached,
Some(self.inner.request_timeout),
)
.await?;
self.metadata_request::<serde_json::Value>(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<T>(&self, method: DuckDbWorkerMethod, params: impl serde::Serialize) -> Result<T, String>
where
T: serde::de::DeserializeOwned,
{
let timeout = self.inner.request_timeout;
match tokio::time::timeout(timeout, self.request::<T>(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::<serde_json::Value>(
DuckDbWorkerMethod::Connect,
self.inner.connect_params.clone(),
Some(self.inner.request_timeout),
)
.await?;
let mut attempts = 0;
loop {
match self
.send_request_structured::<serde_json::Value>(
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<Duration>,
) -> Result<T, String>
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<T>(
&self,
method: DuckDbWorkerMethod,
params: impl serde::Serialize,
timeout: Option<Duration>,
) -> Result<T, DuckDbWorkerError>
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<Duration> {
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,

61
crates/dbx-core/src/db/duckdb_worker_runtime.rs
View File

@ -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::<DuckDbWorkerExecuteParams>();
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;

345
crates/dbx-core/tests/duckdb_worker_process.rs
View File

@ -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<BufReader<tokio::process::ChildStdout>>,
) -> 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<u32> {
for _ in 0..20 {
if let Ok(contents) = std::fs::read_to_string(path) {
if let Ok(pid) = contents.trim().parse::<u32>() {
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()
}

20
crates/dbx-core/tests/support/duckdb_worker_hanging_connect_test_host.rs Normal file
View File

@ -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));
}
}

13
crates/dbx-core/tests/support/duckdb_worker_pid_test_host.rs Normal file
View File

@ -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);
}
}