diff --git a/src/client.rs b/src/client.rs index 625ea738..bb8a48bc 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,82 @@ 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 !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 { + // SAFETY: Each PostgreSQL backend is a separate process (fork model). + // This code runs in a single-threaded tokio runtime with no worker + // threads. No concurrent thread can be reading env simultaneously. + unsafe { + std::env::set_var("DUROXIDE_PG_POOL_MAX", "1"); + } + + let store = new_backend_provider(&pg_conn_str, schema).await?; + Ok::(Client::new(store)) + })?; - if !is_worker_ready() { - return Err( - "pg_durable background worker not yet initialized — try again in a moment".to_string(), - ); + 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 +202,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/dsl.rs b/src/dsl.rs index 8db09104..1a6a0889 100644 --- a/src/dsl.rs +++ b/src/dsl.rs @@ -903,6 +903,7 @@ pub fn start( instance_id: instance_id.clone(), label: label.map(|s| s.to_string()), vars, + loop_iteration: 0, }; let input_json = serde_json::to_string(&input).unwrap_or(instance_id.clone()); diff --git a/src/lib.rs b/src/lib.rs index 72ddad44..71f5c2c4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2549,6 +2549,163 @@ mod tests { // Cancel immediately so the BGW does not attempt to execute this instance. Spi::run(&format!("SELECT df.cancel('{instance_id}')")).unwrap(); } + + // ======================================================================== + // Regression Tests - Reliability Hardening + // ======================================================================== + + // --- 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" + )); + } + + // --- H6: CGNAT SSRF blocklist --- + + #[pg_test] + fn test_ssrf_blocks_cgnat_range() { + use std::net::{IpAddr, Ipv4Addr}; + // 100.64.0.0/10 must be blocked + assert!( + crate::ssrf::check_blocked_ip(IpAddr::V4(Ipv4Addr::new(100, 64, 0, 1))).is_some(), + "100.64.0.1 (CGNAT) should be blocked" + ); + assert!( + crate::ssrf::check_blocked_ip(IpAddr::V4(Ipv4Addr::new(100, 127, 255, 254))).is_some(), + "100.127.255.254 (CGNAT) should be blocked" + ); + // Outside CGNAT range should be allowed + assert!( + crate::ssrf::check_blocked_ip(IpAddr::V4(Ipv4Addr::new(100, 128, 0, 1))).is_none(), + "100.128.0.1 (NOT CGNAT) should be allowed" + ); + } + + // --- M1: Row-set expansion limit --- + + #[pg_test] + fn test_row_set_expansion_limit_via_dsl() { + // The row-set expansion limit (10,000 rows) is enforced inside + // expand_row_set(). This is tested thoroughly in the unit test + // types::tests::test_row_set_expansion_rejects_oversized_result. + // Here we just verify the types module is accessible and the limit works + // at the substitution layer by checking a small expansion works. + use crate::types::substitute_all; + use std::collections::HashMap; + + let mut results = HashMap::new(); + let json = r#"{"rows":[{"id":1},{"id":2}],"row_count":2}"#; + results.insert("batch".to_string(), json.to_string()); + + let sys = crate::types::SystemVars { + instance_id: "test1234".to_string(), + label: None, + }; + let vars = HashMap::new(); + let result = substitute_all("SELECT * FROM $batch.*", &results, &vars, &sys); + assert!(result.is_ok(), "Small row-set should expand successfully"); + assert!( + result.unwrap().contains("VALUES"), + "Should produce a VALUES clause" + ); + } + + // --- M7: Loop iteration counter persisted across continue_as_new --- + + #[pg_test] + fn test_function_input_loop_iteration_serialization() { + use crate::types::FunctionInput; + + // Verify loop_iteration is preserved through serialization + let input = FunctionInput { + instance_id: "test123".to_string(), + label: Some("test".to_string()), + vars: std::collections::HashMap::new(), + loop_iteration: 42, + }; + let json = serde_json::to_string(&input).unwrap(); + let deserialized: FunctionInput = serde_json::from_str(&json).unwrap(); + assert_eq!( + deserialized.loop_iteration, 42, + "loop_iteration must survive serialization round-trip" + ); + } + + #[pg_test] + fn test_function_input_loop_iteration_defaults_to_zero() { + use crate::types::FunctionInput; + + // Verify backward compat: old FunctionInput JSON without loop_iteration + // deserializes with loop_iteration = 0 + let json = r#"{"instance_id":"abc12345","label":"test","vars":{}}"#; + let input: FunctionInput = serde_json::from_str(json).unwrap(); + assert_eq!( + input.loop_iteration, 0, + "Missing loop_iteration should default to 0 for backward compatibility" + ); + } + + // --- M8: Malformed loop condition config detection --- + + #[pg_test] + fn test_malformed_loop_condition_detected_at_validate() { + // A LOOP node whose condition_node is a plain string (not a Durofut object) + // should be rejected by validate_recursive because for_each_config_child + // requires condition_node to deserialize as a valid Durofut. + let node = Durofut { + node_type: "LOOP".to_string(), + left_node: Some(Box::new(Durofut { + node_type: "SQL".to_string(), + query: Some("SELECT 1".to_string()), + ..Default::default() + })), + // Malformed config: valid JSON but condition_node is a string, not a Durofut object. + query: Some(r#"{"condition_node": "nonexist"}"#.to_string()), + ..Default::default() + }; + // Validate should fail because condition_node is not a valid Durofut object + let err = node.validate_recursive().unwrap_err(); + assert!( + err.contains("condition_node"), + "Error should mention condition_node, got: {err}" + ); + + // But if the config is totally not JSON, for_each_config_child skips it + // (it's treated as a plain query string, not a config object). + let non_json_node = Durofut { + node_type: "LOOP".to_string(), + left_node: Some(Box::new(Durofut { + node_type: "SQL".to_string(), + query: Some("SELECT 1".to_string()), + ..Default::default() + })), + query: Some("this is not json at all!!!".to_string()), + ..Default::default() + }; + assert!( + non_json_node.validate_recursive().is_ok(), + "LOOP with non-JSON config passes DSL validation (caught at execution time)" + ); + } } /// Required by `cargo pgrx test` diff --git a/src/orchestrations/execute_function_graph.rs b/src/orchestrations/execute_function_graph.rs index 815d70da..26fff1f8 100644 --- a/src/orchestrations/execute_function_graph.rs +++ b/src/orchestrations/execute_function_graph.rs @@ -30,6 +30,8 @@ pub const SUBTREE_NAME: &str = "pg_durable::orchestration::execute-subtree"; struct ExecutionContext { vars: HashMap, label: Option, + /// Loop iteration counter (persisted across continue_as_new generations). + loop_iteration: u64, } /// Envelope returned by `execute_subtree` containing the SQL result and the updated @@ -116,6 +118,7 @@ pub async fn execute(ctx: OrchestrationContext, input_json: String) -> Result bool { serde_json::from_str::(result) @@ -498,35 +510,53 @@ async fn execute_loop_node( // Check while-condition if present if let Some(ref config_str) = node.query { - if let Ok(config) = serde_json::from_str::(config_str) { - if let Some(condition_node_id) = config["condition_node"].as_str() { - ctx.trace_info("Evaluating loop condition"); - let condition_result = Box::pin(execute_function_node_with_vars( - ctx, - graph, - condition_node_id, - results, - exec_ctx, - )) - .await?; - - // Parse condition result to check truthiness (uses evaluate_condition to extract boolean from SQL result) - let should_continue = evaluate_condition(&condition_result).unwrap_or(false); - ctx.trace_info(format!( - "Loop condition evaluated to: {condition_result} (continue={should_continue})" - )); + match serde_json::from_str::(config_str) { + Ok(config) => { + if let Some(condition_node_id) = config["condition_node"].as_str() { + ctx.trace_info("Evaluating loop condition"); + let condition_result = Box::pin(execute_function_node_with_vars( + ctx, + graph, + condition_node_id, + results, + exec_ctx, + )) + .await?; + + // Parse condition result to check truthiness (uses evaluate_condition to extract boolean from SQL result) + let should_continue = evaluate_condition(&condition_result).unwrap_or(false); + ctx.trace_info(format!( + "Loop condition evaluated to: {condition_result} (continue={should_continue})" + )); - if !should_continue { - ctx.trace_info("Loop condition false, exiting loop"); - store_named_result(ctx, node, &body_result, results, "LOOP"); - return Ok(body_result); + if !should_continue { + ctx.trace_info("Loop condition false, exiting loop"); + store_named_result(ctx, node, &body_result, results, "LOOP"); + return Ok(body_result); + } } } + Err(e) => { + // M8: Malformed condition config should fail the loop rather than + // silently creating an infinite loop without exit condition. + return Err(format!( + "LOOP node {node_id}: failed to parse condition config: {e}" + )); + } } } ctx.trace_info("Continuing as new for next loop iteration"); + // M7: Enforce maximum iteration count to prevent runaway infinite loops + let next_iteration = exec_ctx.loop_iteration + 1; + if next_iteration >= MAX_LOOP_ITERATIONS { + return Err(format!( + "Loop exceeded maximum iteration count of {MAX_LOOP_ITERATIONS}. \ + Use df.break() to exit the loop or restructure the workflow." + )); + } + // Enforce a minimum per-iteration wall-clock duration to prevent // busy-looping (e.g. `df.loop(df.sleep(0))`). Compute the elapsed time // from the deterministic clock; if the iteration finished faster than @@ -551,6 +581,7 @@ async fn execute_loop_node( instance_id: graph.instance_id.clone(), label: exec_ctx.label.clone(), vars: exec_ctx.vars.clone(), + loop_iteration: next_iteration, }; // duroxide 0.1.1: continue_as_new returns an awaitable future - return it directly diff --git a/src/ssrf.rs b/src/ssrf.rs index ee4c1fae..7ceab950 100644 --- a/src/ssrf.rs +++ b/src/ssrf.rs @@ -104,6 +104,7 @@ fn check_blocked_ipv4(ip: Ipv4Addr) -> Option<&'static str> { match octets { [0, ..] => Some("reserved (0.0.0.0/8)"), [10, ..] => Some("private (10.0.0.0/8)"), + [100, b, ..] if (64..=127).contains(&b) => Some("shared/CGNAT (100.64.0.0/10)"), [127, ..] => Some("loopback (127.0.0.0/8)"), [169, 254, ..] => Some("link-local (169.254.0.0/16)"), [172, b, ..] if (16..=31).contains(&b) => Some("private (172.16.0.0/12)"), @@ -236,7 +237,10 @@ pub fn validate_url_allowlist(url: &str) -> Result<(), String> { /// Extract the hostname (without port or brackets) from a URL. /// /// Returns `None` for malformed URLs or URLs without a `://` scheme separator. -#[cfg(not(feature = "http-allow-all"))] +#[cfg(any( + feature = "http-allow-azure-domains", + feature = "http-allow-test-domains" +))] fn extract_host(url: &str) -> Option { // Strip scheme let after_scheme = url.find("://").map(|i| &url[i + 3..])?; @@ -420,6 +424,21 @@ mod tests { assert!(check_blocked_ip(IpAddr::V4(Ipv4Addr::new(0, 255, 255, 255))).is_some()); } + #[cfg(not(feature = "http-allow-all"))] + #[test] + fn blocks_cgnat_rfc6598() { + // 100.64.0.0/10 — Carrier-Grade NAT (RFC 6598) + // Used by cloud providers for internal routing / metadata + assert!(check_blocked_ip(IpAddr::V4(Ipv4Addr::new(100, 64, 0, 0))).is_some()); + assert!(check_blocked_ip(IpAddr::V4(Ipv4Addr::new(100, 64, 0, 1))).is_some()); + assert!(check_blocked_ip(IpAddr::V4(Ipv4Addr::new(100, 100, 100, 100))).is_some()); + assert!(check_blocked_ip(IpAddr::V4(Ipv4Addr::new(100, 127, 255, 255))).is_some()); + // Edge: 100.63.x.x is NOT CGNAT + assert!(check_blocked_ip(IpAddr::V4(Ipv4Addr::new(100, 63, 255, 255))).is_none()); + // Edge: 100.128.x.x is NOT CGNAT + assert!(check_blocked_ip(IpAddr::V4(Ipv4Addr::new(100, 128, 0, 0))).is_none()); + } + // --- IPv4 allowed (public) --- #[test] @@ -539,7 +558,10 @@ mod tests { // --- extract_host helper --- - #[cfg(not(feature = "http-allow-all"))] + #[cfg(any( + feature = "http-allow-azure-domains", + feature = "http-allow-test-domains" + ))] #[test] fn extract_host_basic() { assert_eq!( @@ -555,7 +577,10 @@ mod tests { assert_eq!(extract_host("http://user:pass@host/p"), Some("host".into())); } - #[cfg(not(feature = "http-allow-all"))] + #[cfg(any( + feature = "http-allow-azure-domains", + feature = "http-allow-test-domains" + ))] #[test] fn extract_host_query_and_fragment() { // Query-only URL (no path slash after authority) @@ -585,7 +610,10 @@ mod tests { ); } - #[cfg(not(feature = "http-allow-all"))] + #[cfg(any( + feature = "http-allow-azure-domains", + feature = "http-allow-test-domains" + ))] #[test] fn extract_host_none_cases() { assert_eq!(extract_host("no-scheme"), None); diff --git a/src/types.rs b/src/types.rs index eb646969..a98cd33f 100644 --- a/src/types.rs +++ b/src/types.rs @@ -590,6 +590,10 @@ fn extract_column_value( /// Expand `$name.*` into an inline `VALUES` subquery (SQL) or JSON array (raw). fn expand_row_set(name: &str, json_str: &str, for_sql: bool) -> Result { + /// Maximum number of rows allowed in `$name.*` expansion to prevent + /// unbounded SQL string allocation from large result sets. + const MAX_ROWSET_EXPANSION: usize = 10_000; + let json: serde_json::Value = serde_json::from_str(json_str) .map_err(|e| format!("${name}.* — invalid result JSON: {e}"))?; @@ -598,6 +602,15 @@ fn expand_row_set(name: &str, json_str: &str, for_sql: bool) -> Result MAX_ROWSET_EXPANSION { + return Err(format!( + "${name}.* — result has {} rows, exceeding the maximum of {} for row-set expansion. \ + Use pagination or intermediate tables for large result sets.", + rows.len(), + MAX_ROWSET_EXPANSION + )); + } + if !for_sql { return Ok(serde_json::to_string(rows).unwrap()); } @@ -858,6 +871,10 @@ pub struct FunctionInput { pub label: Option, #[serde(default)] pub vars: std::collections::HashMap, + /// Loop iteration counter, incremented on each `continue_as_new`. + /// Used to enforce a maximum iteration safeguard. + #[serde(default)] + pub loop_iteration: u64, } /// Configuration for HTTP requests @@ -1511,4 +1528,50 @@ mod tests { substitute_all_raw("Hello $doc.name", &results, &empty_vars(), &sys_vars()).unwrap(); assert_eq!(out, "Hello Alice"); } + + #[test] + fn test_row_set_expansion_rejects_oversized_result() { + // Build a JSON result with more than 10,000 rows + let mut rows = Vec::new(); + for i in 0..10_001 { + rows.push(serde_json::json!({"id": i})); + } + let json_str = serde_json::json!({"rows": rows, "row_count": 10_001}).to_string(); + let results = make_results(&[("big", &json_str)]); + + let result = substitute_all("SELECT * FROM $big.*", &results, &empty_vars(), &sys_vars()); + assert!( + result.is_err(), + "Should reject row-set expansion > 10,000 rows" + ); + let err = result.unwrap_err(); + assert!( + err.contains("exceeding the maximum"), + "Error should mention the limit, got: {err}" + ); + } + + #[test] + fn test_row_set_expansion_accepts_within_limit() { + // Build a JSON result with exactly 100 rows (well within limit) + let mut rows = Vec::new(); + for i in 0..100 { + rows.push(serde_json::json!({"id": i, "name": format!("item_{i}")})); + } + let json_str = serde_json::json!({"rows": rows, "row_count": 100}).to_string(); + let results = make_results(&[("batch", &json_str)]); + + let result = substitute_all( + "SELECT * FROM $batch.*", + &results, + &empty_vars(), + &sys_vars(), + ); + assert!( + result.is_ok(), + "Should accept row-set expansion within limit" + ); + let sql = result.unwrap(); + assert!(sql.contains("VALUES"), "Should produce VALUES clause"); + } } diff --git a/src/worker.rs b/src/worker.rs index 4f6d66ad..ac4fc2f9 100644 --- a/src/worker.rs +++ b/src/worker.rs @@ -125,8 +125,7 @@ async fn run_duroxide_runtime() { ); // Management pool: consolidates former polling and activity pools into one. - // Used for extension-existence polling, epoch sentinels, worker-ready writes, - // graph loading, and status updates. Sized by the max_management_connections GUC. + // Used for graph loading and status updates. Sized by the max_management_connections GUC. // Retry in a loop so the worker survives the target database not yet existing // (e.g. pg_regress creates `contrib_regression` after PostgreSQL starts). let mgmt_pool = loop { @@ -150,13 +149,38 @@ async fn run_duroxide_runtime() { } }; + // Dedicated polling pool: a separate 1-connection pool used exclusively for + // extension-existence checks and epoch sentinel heartbeats. This isolation + // prevents activity work (graph loading, status updates) from starving the + // health-check loop and causing spurious runtime shutdowns under high load. + let poll_pool = loop { + if is_shutdown_requested() { + log!("pg_durable: shutdown requested before poll pool created, exiting"); + return; + } + match sqlx::postgres::PgPoolOptions::new() + .max_connections(1) + .connect(&pg_conn_str) + .await + { + Ok(pool) => break pool, + Err(e) => { + log!( + "pg_durable: failed to create poll pool (will retry in 5s): {}", + e + ); + tokio::time::sleep(Duration::from_secs(5)).await; + } + } + }; + loop { if is_shutdown_requested() { log!("pg_durable: shutdown requested, exiting"); break; } - if !wait_for_extension_creation(&mgmt_pool, WAIT_FOR_EXTENSION_POLL_INTERVAL).await { + if !wait_for_extension_creation(&poll_pool, WAIT_FOR_EXTENSION_POLL_INTERVAL).await { break; } @@ -191,7 +215,7 @@ async fn run_duroxide_runtime() { // Write a sentinel so we can detect drop+recreate even if the // extension is always present in pg_extension between polls. - let epoch_id = match write_epoch_sentinel(&mgmt_pool).await { + let epoch_id = match write_epoch_sentinel(&poll_pool).await { Ok(id) => { log!("pg_durable: epoch sentinel written ({})", id); Some(id) @@ -203,7 +227,7 @@ async fn run_duroxide_runtime() { }; run_until_extension_dropped_or_shutdown( - &mgmt_pool, + &poll_pool, duroxide_runtime, EXTENSION_DROP_POLL_INTERVAL, SHUTDOWN_CHECK_INTERVAL, @@ -431,12 +455,17 @@ async fn initialize_duroxide_runtime( log!("pg_durable: initializing duroxide runtime..."); // Control duroxide provider pool size via env var (the only mechanism - // without modifying duroxide-pg). BGW is single-threaded so no - // concurrent readers. Note: std::env::set_var becomes unsafe in Rust 2024 edition. - std::env::set_var( - "DUROXIDE_PG_POOL_MAX", - get_max_duroxide_connections().to_string(), - ); + // without modifying duroxide-pg). + // + // SAFETY: The BGW tokio runtime uses new_current_thread() — no additional + // OS threads are spawned. PostgreSQL's fork model means this process has no + // other threads that could be reading the environment concurrently. + unsafe { + std::env::set_var( + "DUROXIDE_PG_POOL_MAX", + get_max_duroxide_connections().to_string(), + ); + } // Create the user-execution semaphore once — the GUC is Postmaster-context // so the value never changes within a worker lifetime.