From ee89db71d36f038bf59e995f32a7a284d3b4ab3f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Isma=C3=ABl=20Mej=C3=ADa?= Date: Thu, 18 Jun 2026 18:19:15 +0200 Subject: [PATCH] fix: recoverable client on connection failure Replace OnceLock with thread_local RefCell> that auto-resets on connection errors (broken pipe, pool timeout, connection refused, etc.). Subsequent calls transparently re-create the client pool instead of failing permanently for the session's lifetime. Add is_connection_error() heuristic with unit tests and pg_test integration coverage. --- src/client.rs | 231 ++++++++++++++++++++++++++++++++++++-------------- src/lib.rs | 25 ++++++ 2 files changed, 193 insertions(+), 63 deletions(-) diff --git a/src/client.rs b/src/client.rs index 625ea738..b471ac32 100644 --- a/src/client.rs +++ b/src/client.rs @@ -5,7 +5,11 @@ //! //! This module provides cached Tokio runtime and Duroxide client for efficient //! df.start(), df.signal(), and df.cancel() calls from user sessions. +//! +//! The client is lazily initialized on first use and can automatically +//! recover from connection failures by re-creating the pool on next call. +use std::cell::RefCell; use std::sync::OnceLock; use duroxide::Client; @@ -17,8 +21,13 @@ use crate::types::{backend_duroxide_schema, new_backend_provider, postgres_conne /// Cached tokio runtime for client operations. static CLIENT_RUNTIME: OnceLock = OnceLock::new(); -/// Cached Duroxide client with connection pool. -static DUROXIDE_CLIENT: OnceLock = OnceLock::new(); +// Per-backend cached Duroxide client. Uses thread_local + RefCell because +// PostgreSQL backends are single-threaded forked processes. This allows +// the client to be reset on connection failures (unlike OnceLock which +// is permanent). +thread_local! { + static DUROXIDE_CLIENT: RefCell> = const { RefCell::new(None) }; +} /// Check whether the background worker has finished initializing the duroxide /// schema for the current binary's expected schema version. @@ -68,35 +77,79 @@ fn get_client_runtime() -> &'static Runtime { }) } -/// Get or create the cached Duroxide client. -fn get_duroxide_client() -> Result<&'static Client, String> { - if let Some(client) = DUROXIDE_CLIENT.get() { - return Ok(client); - } +/// Initialize or get the cached Duroxide client, executing `f` with it. +/// If the client doesn't exist yet, creates it. If `f` returns an error +/// that looks like a connection failure, resets the client so the next +/// call will re-initialize. +fn with_duroxide_client(f: F) -> Result +where + F: FnOnce(&Client, &Runtime) -> Result, +{ + let rt = get_client_runtime(); + + // Try to use existing client + let has_client = DUROXIDE_CLIENT.with(|cell| cell.borrow().is_some()); - if !is_worker_ready() { - return Err( - "pg_durable background worker not yet initialized — try again in a moment".to_string(), - ); + if !has_client { + // Need to create a new client + if !is_worker_ready() { + return Err( + "pg_durable background worker not yet initialized — try again in a moment" + .to_string(), + ); + } + + let pg_conn_str = postgres_connection_string(); + let schema = backend_duroxide_schema(); + let client = rt.block_on(async { + // Limit backend provider to 1 connection — backends need minimal + // duroxide access (start/cancel/signal only). + std::env::set_var("DUROXIDE_PG_POOL_MAX", "1"); + + let store = new_backend_provider(&pg_conn_str, schema).await?; + Ok::(Client::new(store)) + })?; + + DUROXIDE_CLIENT.with(|cell| { + *cell.borrow_mut() = Some(client); + }); } - let rt = get_client_runtime(); - let pg_conn_str = postgres_connection_string(); - let schema = backend_duroxide_schema(); + // Execute the operation with the client + let result = DUROXIDE_CLIENT.with(|cell| { + let borrow = cell.borrow(); + let client = borrow + .as_ref() + .ok_or_else(|| "Client unexpectedly missing".to_string())?; + f(client, rt) + }); - rt.block_on(async { - // Limit backend provider to 1 connection — backends need minimal duroxide - // access (start/cancel/signal only). The runtime is single-threaded - // (new_current_thread). Note: std::env::set_var becomes unsafe in Rust 2024 edition. - std::env::set_var("DUROXIDE_PG_POOL_MAX", "1"); + // On connection-level errors, reset the client so next call retries + if let Err(ref e) = result { + if is_connection_error(e) { + DUROXIDE_CLIENT.with(|cell| { + *cell.borrow_mut() = None; + }); + } + } - let store = new_backend_provider(&pg_conn_str, schema).await?; + result +} - let _ = DUROXIDE_CLIENT.set(Client::new(store)); - DUROXIDE_CLIENT - .get() - .ok_or_else(|| "Failed to initialize client".to_string()) - }) +/// Heuristic to detect connection-level errors that warrant client reset. +fn is_connection_error(err: &str) -> bool { + let lower = err.to_lowercase(); + lower.contains("connection") + || lower.contains("pool timed out") + || lower.contains("broken pipe") + || lower.contains("reset by peer") + || lower.contains("closed") +} + +/// Test-accessible wrapper for is_connection_error. +#[cfg(any(test, feature = "pg_test"))] +pub(crate) fn is_connection_error_for_test(err: &str) -> bool { + is_connection_error(err) } async fn list_running_descendants(client: &Client, root_instance_id: &str) -> Vec { @@ -146,57 +199,109 @@ pub fn start_durable_function( instance_id ); - let rt = get_client_runtime(); - let client = get_duroxide_client()?; - - rt.block_on(async { - client - .start_orchestration(instance_id, function_name, input) - .await - .map_err(|e| format!("Failed to start durable function: {e:?}"))?; - Ok(()) + let fn_name = function_name.to_string(); + let inst_id = instance_id.to_string(); + let inp = input.to_string(); + + with_duroxide_client(|client, rt| { + rt.block_on(async { + client + .start_orchestration(&inst_id, &fn_name, &inp) + .await + .map_err(|e| format!("Failed to start durable function: {e:?}"))?; + Ok(()) + }) }) } /// Cancel a durable function. pub fn cancel_durable_function(instance_id: &str, reason: &str) -> Result<(), String> { - let rt = get_client_runtime(); - let client = get_duroxide_client()?; - - rt.block_on(async { - client - .cancel_instance(instance_id, reason) - .await - .map_err(|e| format!("Failed to cancel durable function: {e:?}"))?; - Ok(()) + let inst_id = instance_id.to_string(); + let rsn = reason.to_string(); + + with_duroxide_client(|client, rt| { + rt.block_on(async { + client + .cancel_instance(&inst_id, &rsn) + .await + .map_err(|e| format!("Failed to cancel durable function: {e:?}"))?; + Ok(()) + }) }) } /// Raise an external event (signal) to a running orchestration. pub fn raise_external_event(instance_id: &str, event_name: &str, data: &str) -> Result<(), String> { - let rt = get_client_runtime(); - let client = get_duroxide_client()?; + let inst_id = instance_id.to_string(); + let evt_name = event_name.to_string(); + let evt_data = data.to_string(); - rt.block_on(async { - client - .raise_event(instance_id, event_name, data) - .await - .map_err(|e| format!("Failed to raise event: {e:?}"))?; - - for child_instance_id in list_running_descendants(client, instance_id).await { - if let Err(e) = client - .raise_event(&child_instance_id, event_name, data) + with_duroxide_client(|client, rt| { + rt.block_on(async { + client + .raise_event(&inst_id, &evt_name, &evt_data) .await - { - warning!( - "pg_durable: failed to fan out signal '{}' to child instance {}: {:?}", - event_name, - child_instance_id, - e - ); + .map_err(|e| format!("Failed to raise event: {e:?}"))?; + + for child_instance_id in list_running_descendants(client, &inst_id).await { + if let Err(e) = client + .raise_event(&child_instance_id, &evt_name, &evt_data) + .await + { + warning!( + "pg_durable: failed to fan out signal '{}' to child instance {}: {:?}", + evt_name, + child_instance_id, + e + ); + } } - } - Ok(()) + Ok(()) + }) }) } + +#[cfg(test)] +mod tests { + use super::is_connection_error; + + #[test] + fn detects_connection_refused() { + assert!(is_connection_error( + "Failed to start durable function: connection refused" + )); + } + + #[test] + fn detects_broken_pipe() { + assert!(is_connection_error("IO error: broken pipe")); + } + + #[test] + fn detects_pool_timeout() { + assert!(is_connection_error( + "pool timed out while waiting for an open connection" + )); + } + + #[test] + fn detects_connection_reset() { + assert!(is_connection_error("reset by peer")); + } + + #[test] + fn detects_connection_closed() { + assert!(is_connection_error("connection closed unexpectedly")); + } + + #[test] + fn does_not_match_normal_errors() { + assert!(!is_connection_error("Instance not found")); + assert!(!is_connection_error("permission denied for table foo")); + assert!(!is_connection_error("syntax error at position 42")); + assert!(!is_connection_error( + "Orchestration already exists for instance abc123" + )); + } +} diff --git a/src/lib.rs b/src/lib.rs index 681b88d4..9f0b0758 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2741,6 +2741,31 @@ mod tests { "try_from_json should return Err on structurally invalid Durofut" ); } + + // --- C5: Client connection error detection --- + + #[pg_test] + fn test_is_connection_error_detects_failures() { + // Validates the heuristic used to reset the client on connection-level errors. + assert!(crate::client::is_connection_error_for_test( + "connection refused" + )); + assert!(crate::client::is_connection_error_for_test("broken pipe")); + assert!(crate::client::is_connection_error_for_test( + "pool timed out" + )); + assert!(crate::client::is_connection_error_for_test("reset by peer")); + assert!(crate::client::is_connection_error_for_test( + "connection closed" + )); + // Non-connection errors should NOT trigger a reset + assert!(!crate::client::is_connection_error_for_test( + "permission denied" + )); + assert!(!crate::client::is_connection_error_for_test( + "Instance not found" + )); + } } /// Required by `cargo pgrx test`