From bbcd2cc94bb192b5004c7604e9b73a5cf3216bce Mon Sep 17 00:00:00 2001 From: Haider Zahid Date: Mon, 31 Aug 2026 14:28:02 +0200 Subject: [PATCH 1/3] Fix caller transaction graph handoff Track the originating PostgreSQL transaction while the worker waits for a newly started graph. Use durable, bounded probes so long commits resume, rollbacks fail cleanly, transient database errors retry, and wait history remains bounded. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- CHANGELOG.md | 1 + USER_GUIDE.md | 7 + docs/ARCHITECTURE.md | 14 +- docs/E2E_TESTING.md | 1 + docs/upgrade-testing.md | 9 + src/activities/load_function_graph.rs | 423 +++++++++++++++---- src/dsl.rs | 10 + src/lib.rs | 2 + src/orchestrations/execute_function_graph.rs | 150 ++++++- src/registry.rs | 10 + src/types.rs | 57 +++ tests/e2e/sql/68_long_caller_transaction.sql | 308 ++++++++++++++ 12 files changed, 909 insertions(+), 83 deletions(-) create mode 100644 tests/e2e/sql/68_long_caller_transaction.sql diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c7e7400..3a47d62d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,7 @@ Pre-1.0 note: while `pg_durable` is in major version `0`, minor releases may inc ### Fixed +- **Caller-transaction handoff:** `df.start()` now tracks the originating transaction until it commits or aborts, so legal caller transactions lasting more than five seconds no longer leave a `pending` `df.instances` row paired with a failed engine execution. Graph admission uses durable backoff and bounded-history compaction rather than holding a worker connection while it waits. - **Worker connection role names (#364):** catalog role names are now passed verbatim when opening workflow connections, preventing quote-wrapped names from being reinterpreted as a different role. - **Restricted HTTP transport (#342, #363):** restricted allow-list builds now require HTTPS so credentials and request bodies cannot be sent over plaintext HTTP; development-only `http-allow-all` builds continue to permit HTTP. diff --git a/USER_GUIDE.md b/USER_GUIDE.md index 7d6581cf..6f1c65a8 100644 --- a/USER_GUIDE.md +++ b/USER_GUIDE.md @@ -169,6 +169,13 @@ runs in. It changes nothing about the durable function that gets started: | `'caller'` (default) | Joins the caller's transaction; a `ROLLBACK` discards the durable function. | | `'new'` | Runs in its own transaction on a separate session; **survives a rollback of the caller's transaction**. | +The caller transaction may remain open for an arbitrary amount of time after +`df.start()` returns. The worker follows that transaction's outcome without +holding an execution connection: it begins the workflow after commit and +terminates the engine record without executing SQL after rollback. Rolling back +only the savepoint that contains `df.start()` is also treated as a rollback even +if the enclosing transaction later commits. + `'new'` provides the same rollback-survival outcome as Oracle autonomous transactions and `REQUIRES_NEW` propagation for **asynchronously started work**. It is not a synchronous autonomous routine: only the durable launch has diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index ff62aa49..74b8168d 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -452,6 +452,19 @@ pub async fn execute( } ``` +`df.start()` commits the duroxide start independently while its `df.instances` +and `df.nodes` writes remain in the caller's transaction. New orchestration +inputs therefore carry the originating top-level transaction ID. Graph +admission uses a single-shot probe: load immediately when the graph is visible, +otherwise inspect `pg_xact_status()`, then wait with deterministic durable +timers while the transaction is in progress. An abort terminates the engine +record without executing SQL; a committed transaction whose graph is still +absent identifies a savepoint rollback. The wait periodically uses +`continue_as_new` to bound replay history and never holds a management +connection between probes. Historical orchestration inputs omit the transaction +ID and continue scheduling the original load activity with its original input, +preserving in-flight replay compatibility. + ### Node Execution Internal node handlers return `NodeResult`, a `Result` whose error arm is a typed @@ -809,4 +822,3 @@ SELECT df.start( 2. **Phase 2 (Execution)**: Background worker's duroxide runtime picks up the orchestration. `LoadFunctionGraph` activity loads the graph. Orchestration walks the graph, scheduling activities for each step. Results flow between nodes via `$variable` substitution. Loops use `continue_as_new` for durability. The key insight is that **graph construction is synchronous** (in user transaction) while **execution is asynchronous and durable** (in background worker via duroxide replay). - diff --git a/docs/E2E_TESTING.md b/docs/E2E_TESTING.md index bbf88d2e..bf6b4d6a 100644 --- a/docs/E2E_TESTING.md +++ b/docs/E2E_TESTING.md @@ -72,6 +72,7 @@ The test suite is organized into 23 files. Files `01`–`09` open with `SET SESS | `15_rls.sql` | RLS on `df.instances` / `df.nodes` / `df.vars` — per-user visibility, cross-user cancel/signal denied, column-level UPDATE, superuser bypass, per-user variable isolation | | `16_heartbeat.sql` | Worker heartbeat liveness — `df._worker_epoch.last_seen_at` advances over time | | `52_node_id_collision_across_instances.sql` | Cross-instance node-ID collision — two instances own the same 8-hex node id; asserts composite-PK coexistence, that `(instance_id, id)` addresses exactly one row, `df.result()` is instance-scoped, and a scoped `update_node_status`-style UPDATE affects exactly one row (issue #129) | +| `68_long_caller_transaction.sql` | Caller-transaction handoff beyond five seconds, transient graph-probe failure, whole rollback, and savepoint rollback | ### Build-Phase Specific diff --git a/docs/upgrade-testing.md b/docs/upgrade-testing.md index 294fcae0..a06efd5a 100644 --- a/docs/upgrade-testing.md +++ b/docs/upgrade-testing.md @@ -203,6 +203,15 @@ gate, so they never need to be added to the exclude list. Each schema-changing PR should add a section here documenting what changed, what the upgrade script handles, and any backward compatibility considerations. +### v0.2.6 → v0.2.7 + +#### Transaction-aware graph admission +- **Runtime change (no DDL):** New caller-mode starts include the top-level PostgreSQL transaction ID in the root orchestration input. A versioned single-shot activity probes graph visibility and `pg_xact_status()`; the deterministic orchestration waits with capped backoff and periodically `continue_as_new`s to bound replay history. +- **Rollback behavior:** A whole-transaction abort fails the df-less engine record without executing SQL. A committed origin transaction with no visible graph is reported distinctly as a likely savepoint rollback. Transient graph/pool errors return a retry state rather than terminally failing the orchestration. +- **Replay compatibility:** Historical `FunctionInput` payloads deserialize with no origin transaction ID and schedule the original `pg_durable::activity::load-function-graph` activity with the same raw instance-ID input. Existing in-flight history therefore retains its operation name, order, and input bytes. The new activity name and input shape are used only for starts created by the new binary. +- **Scenario A/B2 considerations:** No extension schema or persisted `df` data changes; no upgrade DDL is required. +- **Scenario B1 considerations:** The new binary uses PostgreSQL's built-in `pg_current_xact_id()` / `pg_xact_status()` functions and existing `df.instances` / `df.nodes` columns, all available across the supported v0.2.2+ provider line. New starts work against old extension schemas without runtime schema detection. + ### v0.2.5 → v0.2.6 #### Remove `df.ensure_durofut()` diff --git a/src/activities/load_function_graph.rs b/src/activities/load_function_graph.rs index 14d2239e..6ccfc42e 100644 --- a/src/activities/load_function_graph.rs +++ b/src/activities/load_function_graph.rs @@ -9,6 +9,7 @@ use duroxide::ActivityContext; use sqlx::{PgPool, Row}; use std::sync::Arc; +use std::time::Duration; use crate::types::{ is_role_superuser_name, superuser_instances_enabled, FunctionGraph, FunctionNode, @@ -16,88 +17,91 @@ use crate::types::{ /// Activity name for registration and scheduling pub const NAME: &str = "pg_durable::activity::load-function-graph"; +/// Transaction-aware graph probe used only by inputs created by the current binary. +pub const TRANSACTION_AWARE_NAME: &str = + "pg_durable::activity::probe-function-graph-transaction-v1"; +const TRANSACTION_PROBE_QUERY_TIMEOUT: Duration = Duration::from_secs(2); /// Retry configuration for waiting on uncommitted transactions pub const MAX_WAIT_SECS: u64 = 5; pub const POLL_INTERVAL_MS: u64 = 100; -/// Load a function graph from the database, with retry logic for transaction visibility -pub async fn execute( - ctx: ActivityContext, - pool: Arc, - instance_id: String, -) -> Result { - ctx.trace_info(format!( - "Loading function graph for instance: {instance_id}" - )); +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub struct TransactionAwareLoadInput { + pub instance_id: String, + pub origin_xid: String, +} - let instance_query = "SELECT root_node, r.rolname AS submitted_by - FROM df.instances i - LEFT JOIN pg_catalog.pg_roles r ON r.oid = i.submitted_by::oid - WHERE i.id = $1"; +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[serde(tag = "state", rename_all = "snake_case")] +pub enum TransactionGraphProbe { + Ready { graph: String }, + InProgress, + Retry, + Aborted, + CommittedMissing, +} - // Retry loop: wait for instance data to appear - let start_time = std::time::Instant::now(); - let (root_node_id, instance_submitted_by): (String, String) = loop { - match sqlx::query(instance_query) - .bind(&instance_id) - .fetch_optional(pool.as_ref()) - .await - { - Ok(Some(row)) => { - let submitted_by: Option = row.get("submitted_by"); - match submitted_by { - Some(name) => break (row.get("root_node"), name), - None => { - return Err(format!( - "Instance {instance_id}: submitted_by role no longer exists in pg_roles" - )) - } - } - } - Ok(None) => { - let elapsed = start_time.elapsed(); - if elapsed.as_secs() >= MAX_WAIT_SECS { - return Err(format!( - "Instance {instance_id} not found after {MAX_WAIT_SECS}s (transaction may have been rolled back)" - )); - } - if elapsed.as_millis() < POLL_INTERVAL_MS as u128 * 2 { - ctx.trace_info(format!( - "Instance {instance_id} not yet visible, waiting for transaction commit..." - )); - } - tokio::time::sleep(std::time::Duration::from_millis(POLL_INTERVAL_MS)).await; - } - Err(e) => { - let elapsed = start_time.elapsed(); - if elapsed.as_secs() >= MAX_WAIT_SECS { - return Err(format!( - "Instance {instance_id} not found after {MAX_WAIT_SECS}s: {e}" - )); - } - tokio::time::sleep(std::time::Duration::from_millis(POLL_INTERVAL_MS)).await; - } +#[derive(Debug)] +enum LoadGraphError { + Retryable(String), + Permanent(String), +} + +impl LoadGraphError { + fn into_message(self) -> String { + match self { + Self::Retryable(message) | Self::Permanent(message) => message, } - }; + } +} + +const INSTANCE_QUERY: &str = "SELECT root_node, r.rolname AS submitted_by + FROM df.instances i + LEFT JOIN pg_catalog.pg_roles r ON r.oid = i.submitted_by::oid + WHERE i.id = $1"; + +async fn find_visible_instance( + pool: &PgPool, + instance_id: &str, +) -> Result, sqlx::Error> { + sqlx::query(INSTANCE_QUERY) + .bind(instance_id) + .fetch_optional(pool) + .await +} + +async fn load_visible_graph( + ctx: &ActivityContext, + pool: &PgPool, + instance_id: String, + instance_row: sqlx::postgres::PgRow, +) -> Result { + let root_node_id: String = instance_row.get("root_node"); + let instance_submitted_by: Option = instance_row.get("submitted_by"); + let instance_submitted_by = instance_submitted_by.ok_or_else(|| { + LoadGraphError::Permanent(format!( + "Instance {instance_id}: submitted_by role no longer exists in pg_roles" + )) + })?; // Worker-side superuser guard: reject before executing any user SQL. // This closes the forgery path where a BYPASSRLS role inserts rows with // submitted_by = directly, bypassing the df.start() check. if !superuser_instances_enabled() { - match is_role_superuser_name(pool.as_ref(), &instance_submitted_by).await { + match is_role_superuser_name(pool, &instance_submitted_by).await { Ok(true) => { - return Err(format!( + return Err(LoadGraphError::Permanent(format!( "pg_durable blocked instance {instance_id}: submitted_by role \ \"{instance_submitted_by}\" is a superuser, but \ pg_durable.enable_superuser_instances is off" - )); + ))); } Ok(false) => {} Err(e) => { - return Err(format!( + return Err(LoadGraphError::Retryable(format!( "pg_durable: superuser check failed for instance {instance_id}: {e}" - )); + ))); } } } @@ -110,33 +114,26 @@ pub async fn execute( LEFT JOIN pg_catalog.pg_roles r ON r.oid = n.submitted_by::oid WHERE n.instance_id = $1"#; - let rows = match sqlx::query(nodes_query) + let rows = sqlx::query(nodes_query) .bind(&instance_id) - .fetch_all(pool.as_ref()) + .fetch_all(pool) .await - { - Ok(rows) => rows, - Err(e) => return Err(format!("Failed to load function nodes: {e}")), - }; + .map_err(|e| LoadGraphError::Retryable(format!("Failed to load function nodes: {e}")))?; let mut nodes = std::collections::BTreeMap::new(); for row in rows { let id: String = row.get("id"); let submitted_by: Option = row.get("submitted_by"); - let submitted_by = match submitted_by { - Some(name) => name, - None => { - return Err(format!( + let submitted_by = submitted_by.ok_or_else(|| { + LoadGraphError::Permanent(format!( "Instance {instance_id}: node {id} submitted_by role no longer exists in pg_roles" )) - } - }; + })?; // No per-node superuser check needed: a composite FK // (instance_id, submitted_by) REFERENCES df.instances (id, submitted_by) // guarantees every node shares the instance's submitted_by. // The instance-level check above already covers the superuser case. - let node = FunctionNode { id: id.clone(), node_type: row.get("node_type"), @@ -161,5 +158,281 @@ pub async fn execute( graph.nodes.len() )); - serde_json::to_string(&graph).map_err(|e| format!("Failed to serialize graph: {e}")) + serde_json::to_string(&graph) + .map_err(|e| LoadGraphError::Permanent(format!("Failed to serialize graph: {e}"))) +} + +/// Load a function graph from the database, with retry logic for transaction visibility +pub async fn execute( + ctx: ActivityContext, + pool: Arc, + instance_id: String, +) -> Result { + ctx.trace_info(format!( + "Loading function graph for instance: {instance_id}" + )); + + // Retry loop: wait for instance data to appear + let start_time = std::time::Instant::now(); + let instance_row = loop { + match find_visible_instance(pool.as_ref(), &instance_id).await { + Ok(Some(row)) => break row, + Ok(None) => { + let elapsed = start_time.elapsed(); + if elapsed.as_secs() >= MAX_WAIT_SECS { + return Err(format!( + "Instance {instance_id} not found after {MAX_WAIT_SECS}s (transaction may have been rolled back)" + )); + } + if elapsed.as_millis() < POLL_INTERVAL_MS as u128 * 2 { + ctx.trace_info(format!( + "Instance {instance_id} not yet visible, waiting for transaction commit..." + )); + } + tokio::time::sleep(std::time::Duration::from_millis(POLL_INTERVAL_MS)).await; + } + Err(e) => { + let elapsed = start_time.elapsed(); + if elapsed.as_secs() >= MAX_WAIT_SECS { + return Err(format!( + "Instance {instance_id} not found after {MAX_WAIT_SECS}s: {e}" + )); + } + tokio::time::sleep(std::time::Duration::from_millis(POLL_INTERVAL_MS)).await; + } + } + }; + + load_visible_graph(&ctx, pool.as_ref(), instance_id, instance_row) + .await + .map_err(LoadGraphError::into_message) +} + +fn serialize_probe(probe: &TransactionGraphProbe) -> Result { + serde_json::to_string(probe).map_err(|e| format!("Failed to serialize graph probe result: {e}")) +} + +fn retry_probe( + ctx: &ActivityContext, + operation: &str, + instance_id: &str, + error: impl std::fmt::Display, +) -> Result { + ctx.trace_info(format!( + "Transient {operation} failure for instance {instance_id}; graph probe will retry: {error}" + )); + serialize_probe(&TransactionGraphProbe::Retry) +} + +/// Probe graph visibility without pinning an activity task while the caller's +/// transaction remains open. The orchestration schedules a durable timer and +/// invokes this single-shot activity again when the xid is still in progress. +pub async fn probe_transaction( + ctx: ActivityContext, + pool: Arc, + input_json: String, +) -> Result { + let input: TransactionAwareLoadInput = serde_json::from_str(&input_json) + .map_err(|e| format!("Invalid transaction-aware graph probe input: {e}"))?; + if input.origin_xid.parse::().is_err() { + return Err(format!( + "Invalid origin transaction id \"{}\" for instance {}", + input.origin_xid, input.instance_id + )); + } + + ctx.trace_info(format!( + "Probing graph visibility for instance {} from origin transaction {}", + input.instance_id, input.origin_xid + )); + + let visible = match tokio::time::timeout( + TRANSACTION_PROBE_QUERY_TIMEOUT, + find_visible_instance(pool.as_ref(), &input.instance_id), + ) + .await + { + Ok(result) => result, + Err(_) => { + return retry_probe( + &ctx, + "graph visibility check", + &input.instance_id, + "timed out after 2s", + ); + } + }; + + match visible { + Ok(Some(row)) => { + let loaded = match tokio::time::timeout( + TRANSACTION_PROBE_QUERY_TIMEOUT, + load_visible_graph(&ctx, pool.as_ref(), input.instance_id.clone(), row), + ) + .await + { + Ok(result) => result, + Err(_) => { + return retry_probe( + &ctx, + "graph load", + &input.instance_id, + "timed out after 2s", + ); + } + }; + return match loaded { + Ok(graph) => serialize_probe(&TransactionGraphProbe::Ready { graph }), + Err(LoadGraphError::Retryable(error)) => { + retry_probe(&ctx, "graph load", &input.instance_id, error) + } + Err(LoadGraphError::Permanent(error)) => Err(error), + }; + } + Ok(None) => {} + Err(e) => { + return retry_probe(&ctx, "graph visibility check", &input.instance_id, e); + } + } + + let transaction_status_query = tokio::time::timeout( + TRANSACTION_PROBE_QUERY_TIMEOUT, + sqlx::query_scalar("SELECT pg_catalog.pg_xact_status($1::text::xid8)::text") + .bind(&input.origin_xid) + .fetch_one(pool.as_ref()), + ) + .await; + let transaction_status: Option = match transaction_status_query { + Ok(Ok(status)) => status, + Ok(Err(e)) => { + return retry_probe(&ctx, "origin transaction check", &input.instance_id, e); + } + Err(_) => { + return retry_probe( + &ctx, + "origin transaction check", + &input.instance_id, + "timed out after 2s", + ); + } + }; + + let probe = match transaction_status.as_deref() { + Some("in progress") => TransactionGraphProbe::InProgress, + Some("aborted") => TransactionGraphProbe::Aborted, + Some("committed") => { + // The status and graph reads use separate READ COMMITTED statements. + // Re-read once after observing commit so a commit between the first + // graph query and pg_xact_status cannot be misclassified as missing. + let visible = match tokio::time::timeout( + TRANSACTION_PROBE_QUERY_TIMEOUT, + find_visible_instance(pool.as_ref(), &input.instance_id), + ) + .await + { + Ok(result) => result, + Err(_) => { + return retry_probe( + &ctx, + "post-commit graph visibility check", + &input.instance_id, + "timed out after 2s", + ); + } + }; + match visible { + Ok(Some(row)) => { + let loaded = match tokio::time::timeout( + TRANSACTION_PROBE_QUERY_TIMEOUT, + load_visible_graph(&ctx, pool.as_ref(), input.instance_id.clone(), row), + ) + .await + { + Ok(result) => result, + Err(_) => { + return retry_probe( + &ctx, + "post-commit graph load", + &input.instance_id, + "timed out after 2s", + ); + } + }; + match loaded { + Ok(graph) => TransactionGraphProbe::Ready { graph }, + Err(LoadGraphError::Retryable(error)) => { + return retry_probe( + &ctx, + "post-commit graph load", + &input.instance_id, + error, + ); + } + Err(LoadGraphError::Permanent(error)) => return Err(error), + } + } + Ok(None) => TransactionGraphProbe::CommittedMissing, + Err(e) => { + return retry_probe( + &ctx, + "post-commit graph visibility check", + &input.instance_id, + e, + ); + } + } + } + Some(other) => { + return Err(format!( + "Origin transaction {} for instance {} has unknown status \"{}\"", + input.origin_xid, input.instance_id, other + )) + } + None => { + return Err(format!( + "Origin transaction {} for instance {} has no available status", + input.origin_xid, input.instance_id + )) + } + }; + + serialize_probe(&probe) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn transaction_probe_input_round_trips() { + let input = TransactionAwareLoadInput { + instance_id: "deadbeef".to_string(), + origin_xid: "12345".to_string(), + }; + let json = serde_json::to_string(&input).unwrap(); + assert_eq!( + serde_json::from_str::(&json).unwrap(), + input + ); + } + + #[test] + fn transaction_probe_results_have_stable_tags() { + assert_eq!( + serde_json::to_string(&TransactionGraphProbe::InProgress).unwrap(), + r#"{"state":"in_progress"}"# + ); + assert_eq!( + serde_json::to_string(&TransactionGraphProbe::Aborted).unwrap(), + r#"{"state":"aborted"}"# + ); + assert_eq!( + serde_json::to_string(&TransactionGraphProbe::Retry).unwrap(), + r#"{"state":"retry"}"# + ); + assert_eq!( + serde_json::to_string(&TransactionGraphProbe::CommittedMissing).unwrap(), + r#"{"state":"committed_missing"}"# + ); + } } diff --git a/src/dsl.rs b/src/dsl.rs index a8926544..5212054f 100644 --- a/src/dsl.rs +++ b/src/dsl.rs @@ -1271,6 +1271,14 @@ fn start_in_caller_transaction(fut: &str, label: Option<&str>, database: Option< vars }); + // df.start() hands the orchestration to duroxide over a separately committed + // connection before this caller transaction commits. Carry the owning + // top-level xid so the worker can wait for its actual outcome instead of + // guessing that a transaction lasting more than a fixed timeout rolled back. + let origin_xid = Spi::get_one::("SELECT pg_catalog.pg_current_xact_id()::text") + .unwrap_or_else(|e| pgrx::error!("failed to capture df.start() transaction id: {e}")) + .unwrap_or_else(|| pgrx::error!("df.start() transaction id is unavailable")); + // Start the orchestration via duroxide let input = FunctionInput { instance_id: instance_id.clone(), @@ -1280,6 +1288,8 @@ fn start_in_caller_transaction(fut: &str, label: Option<&str>, database: Option< // Generation 0 loads the graph from df.nodes; only a root loop continuing as new // carries it inline. graph: None, + origin_xid: Some(origin_xid), + graph_wait_attempt: 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 08826a7e..5afef48f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3339,6 +3339,8 @@ mod tests { vars: std::collections::HashMap::new(), loop_iteration: 42, graph: None, + origin_xid: None, + graph_wait_attempt: 0, }; let json = serde_json::to_string(&input).unwrap(); let deserialized: FunctionInput = serde_json::from_str(&json).unwrap(); diff --git a/src/orchestrations/execute_function_graph.rs b/src/orchestrations/execute_function_graph.rs index d2e42dbb..66854048 100644 --- a/src/orchestrations/execute_function_graph.rs +++ b/src/orchestrations/execute_function_graph.rs @@ -17,6 +17,7 @@ use cron::Schedule as CronSchedule; use duroxide::OrchestrationContext; use crate::activities; +use crate::activities::load_function_graph::{TransactionAwareLoadInput, TransactionGraphProbe}; use crate::types::{ evaluate_condition, string_map_to_json, substitute_all, substitute_all_raw, FunctionGraph, FunctionInput, FunctionNode, SystemVars, @@ -121,6 +122,105 @@ impl From<&str> for NodeError { /// Result type for node handlers: `Ok` value string, or a typed control-flow/failure error. type NodeResult = Result; +const INITIAL_TRANSACTION_POLL_MS: u64 = 100; +const MAX_TRANSACTION_POLL_MS: u64 = 5_000; +const GRAPH_WAIT_POLLS_PER_EXECUTION: u32 = 64; + +fn transaction_poll_delay(attempt: u32) -> Duration { + let multiplier = 1u64 << attempt.min(6); + Duration::from_millis( + INITIAL_TRANSACTION_POLL_MS + .saturating_mul(multiplier) + .min(MAX_TRANSACTION_POLL_MS), + ) +} + +fn should_compact_graph_wait(polls_in_execution: u32) -> bool { + polls_in_execution >= GRAPH_WAIT_POLLS_PER_EXECUTION +} + +fn graph_wait_continuation( + input: &FunctionInput, + next_attempt: u32, +) -> Result { + let mut continuation = input.clone(); + continuation.graph_wait_attempt = next_attempt; + serde_json::to_string(&continuation) +} + +async fn load_initial_graph( + ctx: &OrchestrationContext, + input: &FunctionInput, +) -> Result { + let Some(origin_xid) = input.origin_xid.as_ref() else { + // Replay compatibility: historical FunctionInput payloads have no xid. + // They must schedule the original activity name with the exact original + // raw instance-id input bytes. + return ctx + .schedule_activity( + activities::load_function_graph::NAME, + input.instance_id.clone(), + ) + .await; + }; + + let probe_input = serde_json::to_string(&TransactionAwareLoadInput { + instance_id: input.instance_id.clone(), + origin_xid: origin_xid.clone(), + }) + .map_err(|e| format!("Failed to serialize transaction-aware graph probe: {e}"))?; + + let mut attempt = input.graph_wait_attempt; + let mut polls_in_execution = 0u32; + loop { + let raw = ctx + .schedule_activity( + activities::load_function_graph::TRANSACTION_AWARE_NAME, + probe_input.clone(), + ) + .await?; + let probe: TransactionGraphProbe = serde_json::from_str(&raw) + .map_err(|e| format!("Failed to parse transaction-aware graph probe result: {e}"))?; + + match probe { + TransactionGraphProbe::Ready { graph } => return Ok(graph), + TransactionGraphProbe::InProgress | TransactionGraphProbe::Retry => { + let delay = transaction_poll_delay(attempt); + ctx.trace_info(format!( + "Instance {} graph is waiting on origin transaction {}; retrying in {:?}", + input.instance_id, origin_xid, delay + )); + ctx.schedule_timer(delay).await; + attempt = attempt.saturating_add(1); + polls_in_execution = polls_in_execution.saturating_add(1); + + if should_compact_graph_wait(polls_in_execution) { + let continuation_json = graph_wait_continuation(input, attempt) + .map_err(|e| format!("Failed to serialize graph-wait continuation: {e}"))?; + ctx.trace_info(format!( + "Compacting graph-admission history for instance {} after {} polls", + input.instance_id, polls_in_execution + )); + return ctx.continue_as_new(continuation_json).await; + } + } + TransactionGraphProbe::Aborted => { + return Err(format!( + "Instance {} origin transaction {} aborted before its graph became visible", + input.instance_id, origin_xid + )) + } + TransactionGraphProbe::CommittedMissing => { + return Err(format!( + "Instance {} origin transaction {} committed but its instance graph is absent \ + (the start may have been rolled back to a savepoint)", + input.instance_id, origin_xid + )) + } + } + } +} + /// Distinguishes a normal subtree result from one that unwound via `df.break()`. /// /// Stored as `Option` in the envelope (see `SubtreeEnvelope::control`): a @@ -190,13 +290,7 @@ pub async fn execute(ctx: OrchestrationContext, input_json: String) -> Result json, - None => match ctx - .schedule_activity( - activities::load_function_graph::NAME, - input.instance_id.clone(), - ) - .await - { + None => match load_initial_graph(&ctx, &input).await { Ok(json) => json, Err(e) => { // load_function_graph failed (e.g., superuser blocked). @@ -992,6 +1086,8 @@ async fn execute_loop_node( vars: exec_ctx.vars.clone(), loop_iteration: next_iteration, graph: Some(graph_json), + origin_xid: None, + graph_wait_attempt: 0, }; serde_json::to_string(&new_input) .map_err(|e| format!("Failed to serialize loop input: {e}"))? @@ -1803,6 +1899,46 @@ mod tests { } } + #[test] + fn transaction_poll_backoff_is_deterministic_and_capped() { + assert_eq!(transaction_poll_delay(0), Duration::from_millis(100)); + assert_eq!(transaction_poll_delay(1), Duration::from_millis(200)); + assert_eq!(transaction_poll_delay(5), Duration::from_millis(3_200)); + assert_eq!(transaction_poll_delay(6), Duration::from_millis(5_000)); + assert_eq!( + transaction_poll_delay(u32::MAX), + Duration::from_millis(5_000) + ); + assert!(!should_compact_graph_wait( + GRAPH_WAIT_POLLS_PER_EXECUTION - 1 + )); + assert!(should_compact_graph_wait(GRAPH_WAIT_POLLS_PER_EXECUTION)); + assert!(should_compact_graph_wait(u32::MAX)); + } + + #[test] + fn graph_wait_compaction_preserves_admission_state_only() { + let input = FunctionInput { + instance_id: "deadbeef".to_string(), + label: Some("waiting".to_string()), + vars: HashMap::from([("key".to_string(), "value".to_string())]), + loop_iteration: 7, + graph: None, + origin_xid: Some("12345".to_string()), + graph_wait_attempt: 63, + }; + + let json = graph_wait_continuation(&input, 64).unwrap(); + let continued: FunctionInput = serde_json::from_str(&json).unwrap(); + assert_eq!(continued.instance_id, input.instance_id); + assert_eq!(continued.label, input.label); + assert_eq!(continued.vars, input.vars); + assert_eq!(continued.loop_iteration, 7); + assert_eq!(continued.graph, None); + assert_eq!(continued.origin_xid.as_deref(), Some("12345")); + assert_eq!(continued.graph_wait_attempt, 64); + } + #[test] fn parse_legacy_break_sentinel_decodes_string_value() { // A JSON string value round-trips as the quoted JSON string, matching the old diff --git a/src/registry.rs b/src/registry.rs index 8bcc79b7..1eee1517 100644 --- a/src/registry.rs +++ b/src/registry.rs @@ -16,6 +16,7 @@ use crate::orchestrations; pub fn create_activity_registry(pool: Arc, semaphore: Arc) -> ActivityRegistry { let sql_semaphore = semaphore; let graph_pool = pool.clone(); + let transaction_graph_pool = pool.clone(); let status_pool = pool.clone(); let node_status_pool = pool.clone(); let http_pool = pool.clone(); @@ -30,6 +31,15 @@ pub fn create_activity_registry(pool: Arc, semaphore: Arc) -> let pool = graph_pool.clone(); async move { activities::load_function_graph::execute(ctx, pool, instance_id).await } }) + .register( + activities::load_function_graph::TRANSACTION_AWARE_NAME, + move |ctx: ActivityContext, input_json: String| { + let pool = transaction_graph_pool.clone(); + async move { + activities::load_function_graph::probe_transaction(ctx, pool, input_json).await + } + }, + ) .register(activities::update_instance_status::NAME, move |ctx: ActivityContext, input_json: String| { let pool = status_pool.clone(); async move { activities::update_instance_status::execute(ctx, pool, input_json).await } diff --git a/src/types.rs b/src/types.rs index f967fe62..2ad67ad2 100644 --- a/src/types.rs +++ b/src/types.rs @@ -1249,6 +1249,10 @@ pub struct FunctionGraph { pub nodes: std::collections::BTreeMap, } +fn is_zero_u32(value: &u32) -> bool { + *value == 0 +} + /// Input structure passed to duroxide functions #[derive(Debug, Clone, Serialize, Deserialize)] pub struct FunctionInput { @@ -1268,6 +1272,19 @@ pub struct FunctionInput { /// once no matter how many iterations it runs. #[serde(default, skip_serializing_if = "Option::is_none")] pub graph: Option, + /// Top-level transaction that owns the initial df.instances/df.nodes writes. + /// + /// New starts carry this so graph loading can distinguish a still-open caller + /// transaction from a rollback. Historical inputs omit it and retain the + /// legacy bounded graph-load activity for replay compatibility. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub origin_xid: Option, + /// Number of graph-admission polls already made for `origin_xid`. + /// + /// Kept separate from `loop_iteration`: admission can compact its history + /// with continue_as_new before user graph execution begins. + #[serde(default, skip_serializing_if = "is_zero_u32")] + pub graph_wait_attempt: u32, } pub(crate) fn serialize_string_map( @@ -2196,6 +2213,8 @@ mod tests { vars: forward, loop_iteration: 0, graph: None, + origin_xid: None, + graph_wait_attempt: 0, }; let reverse_input = FunctionInput { instance_id: "instance".to_string(), @@ -2203,6 +2222,8 @@ mod tests { vars: reverse, loop_iteration: 0, graph: None, + origin_xid: None, + graph_wait_attempt: 0, }; assert_eq!( serde_json::to_string(&forward_input).unwrap(), @@ -2210,6 +2231,42 @@ mod tests { ); } + #[test] + fn function_input_origin_xid_is_backward_compatible() { + let legacy_json = r#"{"instance_id":"abc12345","label":null,"vars":{},"loop_iteration":0}"#; + let legacy: FunctionInput = serde_json::from_str(legacy_json).unwrap(); + assert_eq!(legacy.origin_xid, None); + assert!(!serde_json::to_string(&legacy) + .unwrap() + .contains("origin_xid")); + + let current = FunctionInput { + instance_id: "abc12345".to_string(), + label: None, + vars: std::collections::HashMap::new(), + loop_iteration: 0, + graph: None, + origin_xid: Some("123456".to_string()), + graph_wait_attempt: 0, + }; + let current_json = serde_json::to_string(¤t).unwrap(); + assert!(current_json.ends_with(r#","origin_xid":"123456"}"#)); + assert_eq!( + serde_json::from_str::(¤t_json) + .unwrap() + .origin_xid + .as_deref(), + Some("123456") + ); + + let compacted = FunctionInput { + graph_wait_attempt: 64, + ..current + }; + let compacted_json = serde_json::to_string(&compacted).unwrap(); + assert!(compacted_json.ends_with(r#","graph_wait_attempt":64}"#)); + } + fn sys_vars() -> SystemVars { SystemVars { instance_id: "test-id".to_string(), diff --git a/tests/e2e/sql/68_long_caller_transaction.sql b/tests/e2e/sql/68_long_caller_transaction.sql new file mode 100644 index 00000000..f995ad1f --- /dev/null +++ b/tests/e2e/sql/68_long_caller_transaction.sql @@ -0,0 +1,308 @@ +-- Copyright (c) Microsoft Corporation. +-- Licensed under the PostgreSQL License. + +-- A caller-mode df.start() persists its graph in the caller's transaction but +-- hands the orchestration to duroxide on a separate committed connection. The +-- worker must follow the originating transaction rather than treating a legal +-- transaction lasting more than five seconds as a rollback. + +CREATE TEMP TABLE _long_tx_state (instance_id TEXT, scenario TEXT); +DROP TABLE IF EXISTS long_tx_effects; +CREATE TABLE long_tx_effects ( + instance_id TEXT NOT NULL, + scenario TEXT NOT NULL, + PRIMARY KEY (instance_id, scenario) +); + +CREATE OR REPLACE FUNCTION pg_temp.engine_info(p_instance_id TEXT) +RETURNS TABLE(status TEXT, output TEXT) +LANGUAGE plpgsql +AS $$ +DECLARE + provider_schema TEXT := df.duroxide_schema(); +BEGIN + RETURN QUERY EXECUTE format( + 'SELECT i.status, i.output FROM %I.get_instance_info($1) i', + provider_schema + ) USING p_instance_id; +END +$$; + +-- A long transaction that commits must still execute exactly once. +BEGIN; +INSERT INTO _long_tx_state +SELECT df.start( + 'INSERT INTO long_tx_effects(instance_id, scenario) + VALUES (''{sys_instance_id}'', ''committed'')', + 'long-caller-transaction' +), 'committed'; + +DO $$ +DECLARE + v_instance_id TEXT; + engine_status TEXT; + attempts INT := 0; +BEGIN + SELECT s.instance_id INTO v_instance_id + FROM _long_tx_state s + WHERE s.scenario = 'committed'; + LOOP + SELECT i.status INTO engine_status + FROM pg_temp.engine_info(v_instance_id) i; + EXIT WHEN lower(COALESCE(engine_status, '')) = 'running' OR attempts >= 300; + PERFORM pg_sleep(0.1); + attempts := attempts + 1; + END LOOP; + + IF lower(COALESCE(engine_status, '')) != 'running' THEN + RAISE EXCEPTION + 'TEST FAILED [long commit]: engine did not start while caller transaction was open (status=%)', + engine_status; + END IF; + + -- Exceed the legacy fixed graph-visibility timeout after the engine is + -- definitely running. + PERFORM pg_sleep(6); +END +$$; +COMMIT; + +DO $$ +DECLARE + v_instance_id TEXT; + final_status TEXT; + effect_count INT; +BEGIN + SELECT s.instance_id INTO v_instance_id + FROM _long_tx_state s + WHERE s.scenario = 'committed'; + SELECT df.await_instance(v_instance_id, 30) INTO final_status; + SELECT count(*) INTO effect_count + FROM long_tx_effects e + WHERE e.instance_id = v_instance_id + AND e.scenario = 'committed'; + + IF final_status IS DISTINCT FROM 'completed' OR effect_count != 1 THEN + RAISE EXCEPTION + 'TEST FAILED [long commit]: status=%, effects=% (expected completed/1)', + final_status, effect_count; + END IF; +END +$$; + +-- A transient management-connection failure while probing an uncommitted graph +-- must be retried durably rather than recorded as a terminal activity failure. +BEGIN; +LOCK TABLE df.instances IN ACCESS EXCLUSIVE MODE; +INSERT INTO _long_tx_state +SELECT df.start( + 'INSERT INTO long_tx_effects(instance_id, scenario) + VALUES (''{sys_instance_id}'', ''transient-retry'')', + 'graph-probe-transient-retry' +), 'transient-retry'; + +DO $$ +DECLARE + v_instance_id TEXT; + engine_status TEXT; + target_pid INT; + attempts INT := 0; +BEGIN + SELECT s.instance_id INTO v_instance_id + FROM _long_tx_state s + WHERE s.scenario = 'transient-retry'; + + LOOP + SELECT i.status INTO engine_status + FROM pg_temp.engine_info(v_instance_id) i; + EXIT WHEN lower(COALESCE(engine_status, '')) = 'running' OR attempts >= 300; + PERFORM pg_sleep(0.1); + attempts := attempts + 1; + END LOOP; + IF lower(COALESCE(engine_status, '')) != 'running' THEN + RAISE EXCEPTION + 'TEST FAILED [transient retry]: engine did not start (status=%)', + engine_status; + END IF; + + attempts := 0; + LOOP + SELECT pid INTO target_pid + FROM pg_stat_activity + WHERE application_name = 'pg_durable:worker:management' + AND wait_event_type = 'Lock' + ORDER BY pid + LIMIT 1; + EXIT WHEN target_pid IS NOT NULL OR attempts >= 300; + PERFORM pg_sleep(0.05); + attempts := attempts + 1; + END LOOP; + IF target_pid IS NULL THEN + RAISE EXCEPTION + 'TEST FAILED [transient retry]: no blocked graph-probe backend found'; + END IF; + IF NOT pg_terminate_backend(target_pid) THEN + RAISE EXCEPTION + 'TEST FAILED [transient retry]: could not terminate graph-probe backend %', + target_pid; + END IF; + + -- Leave time for the killed query to return Retry and the next lock-blocked + -- probe to hit its bounded query timeout before releasing the table lock. + PERFORM pg_sleep(3); +END +$$; +COMMIT; + +DO $$ +DECLARE + v_instance_id TEXT; + final_status TEXT; + effect_count INT; +BEGIN + SELECT s.instance_id INTO v_instance_id + FROM _long_tx_state s + WHERE s.scenario = 'transient-retry'; + SELECT df.await_instance(v_instance_id, 30) INTO final_status; + SELECT count(*) INTO effect_count + FROM long_tx_effects e + WHERE e.instance_id = v_instance_id + AND e.scenario = 'transient-retry'; + + IF final_status IS DISTINCT FROM 'completed' OR effect_count != 1 THEN + RAISE EXCEPTION + 'TEST FAILED [transient retry]: status=%, effects=% (expected completed/1)', + final_status, effect_count; + END IF; +END +$$; + +-- A whole-transaction rollback has an aborted originating xid. It must fail +-- without running SQL and without leaving a df.instances row. +BEGIN; +INSERT INTO _long_tx_state +SELECT df.start( + 'INSERT INTO long_tx_effects(instance_id, scenario) + VALUES (''{sys_instance_id}'', ''whole-rollback'')', + 'whole-transaction-rollback' +) AS instance_id, 'whole-rollback-open' +RETURNING instance_id AS whole_rollback_id \gset + +DO $$ +DECLARE + v_instance_id TEXT; + engine_status TEXT; + attempts INT := 0; +BEGIN + SELECT s.instance_id INTO v_instance_id + FROM _long_tx_state s + WHERE s.scenario = 'whole-rollback-open'; + + LOOP + SELECT i.status INTO engine_status + FROM pg_temp.engine_info(v_instance_id) i; + EXIT WHEN lower(COALESCE(engine_status, '')) = 'running' OR attempts >= 300; + PERFORM pg_sleep(0.1); + attempts := attempts + 1; + END LOOP; + + IF lower(COALESCE(engine_status, '')) != 'running' THEN + RAISE EXCEPTION + 'TEST FAILED [whole rollback]: engine did not start while caller transaction was open (status=%)', + engine_status; + END IF; + + -- The transaction remains legal and in progress past the old timeout. The + -- transaction-aware path must keep waiting until the rollback below. + PERFORM pg_sleep(6); +END +$$; +ROLLBACK; + +INSERT INTO _long_tx_state VALUES (:'whole_rollback_id', 'whole-rollback'); + +DO $$ +DECLARE + v_instance_id TEXT; + engine_status TEXT; + engine_output TEXT; + attempts INT := 0; +BEGIN + SELECT s.instance_id INTO v_instance_id + FROM _long_tx_state s + WHERE s.scenario = 'whole-rollback'; + + LOOP + SELECT i.status, i.output INTO engine_status, engine_output + FROM pg_temp.engine_info(v_instance_id) i; + EXIT WHEN lower(COALESCE(engine_status, '')) = 'failed' OR attempts >= 300; + PERFORM pg_sleep(0.1); + attempts := attempts + 1; + END LOOP; + + IF lower(COALESCE(engine_status, '')) != 'failed' + OR engine_output NOT LIKE '%origin transaction%aborted%' THEN + RAISE EXCEPTION + 'TEST FAILED [whole rollback]: engine status=%, output=%', + engine_status, engine_output; + END IF; + IF EXISTS (SELECT 1 FROM df.instances i WHERE i.id = v_instance_id) + OR EXISTS (SELECT 1 FROM long_tx_effects e WHERE e.instance_id = v_instance_id) THEN + RAISE EXCEPTION + 'TEST FAILED [whole rollback]: rolled-back df row or side effect exists for %', + v_instance_id; + END IF; +END +$$; + +-- Rolling back only the start's savepoint leaves the top-level xid committed +-- but the graph absent. This must be distinguished from an in-progress commit. +BEGIN; +SAVEPOINT before_start; +SELECT df.start( + 'INSERT INTO long_tx_effects(instance_id, scenario) + VALUES (''{sys_instance_id}'', ''savepoint-rollback'')', + 'savepoint-rollback' +) AS savepoint_rollback_id \gset +ROLLBACK TO SAVEPOINT before_start; +COMMIT; + +INSERT INTO _long_tx_state VALUES (:'savepoint_rollback_id', 'savepoint-rollback'); + +DO $$ +DECLARE + v_instance_id TEXT; + engine_status TEXT; + engine_output TEXT; + attempts INT := 0; +BEGIN + SELECT s.instance_id INTO v_instance_id + FROM _long_tx_state s + WHERE s.scenario = 'savepoint-rollback'; + + LOOP + SELECT i.status, i.output INTO engine_status, engine_output + FROM pg_temp.engine_info(v_instance_id) i; + EXIT WHEN lower(COALESCE(engine_status, '')) = 'failed' OR attempts >= 300; + PERFORM pg_sleep(0.1); + attempts := attempts + 1; + END LOOP; + + IF lower(COALESCE(engine_status, '')) != 'failed' + OR engine_output NOT LIKE '%committed%graph%absent%' THEN + RAISE EXCEPTION + 'TEST FAILED [savepoint rollback]: engine status=%, output=%', + engine_status, engine_output; + END IF; + IF EXISTS (SELECT 1 FROM df.instances i WHERE i.id = v_instance_id) + OR EXISTS (SELECT 1 FROM long_tx_effects e WHERE e.instance_id = v_instance_id) THEN + RAISE EXCEPTION + 'TEST FAILED [savepoint rollback]: rolled-back df row or side effect exists for %', + v_instance_id; + END IF; +END +$$; + +DROP TABLE _long_tx_state; +DROP TABLE long_tx_effects; +SELECT 'TEST PASSED: long caller transaction handoff' AS result; From 1c34e84f8423f0656f500a33fa91ee39e0489973 Mon Sep 17 00:00:00 2001 From: Haider Zahid Date: Tue, 1 Sep 2026 10:01:29 +0200 Subject: [PATCH 2/3] Address review feedback on caller transaction graph handoff Resolves all 5 line comments from the automated review of the caller-mode df.start() transaction-handoff fix: 1. HIGH: snapshot-visibility race after pg_xact_status() reports "committed". Postgres can mark a transaction committed before ProcArrayEndTransaction removes it from the running set used by a fresh snapshot, so an immediate re-read could misclassify a valid graph as CommittedMissing. Now checks pg_visible_in_snapshot(origin_xid, pg_current_snapshot()) first and retries until the xid is snapshot-visible before trusting a re-read. 2. HIGH: fixed 2s deadline could livelock a valid graph load (role validation + up to MAX_GRAPH_NODES rows + serialization). Graph loading now uses its own GRAPH_LOAD_QUERY_TIMEOUT (20s client-side) plus Postgres-side statement_timeout/lock_timeout backstops, decoupled from the short TRANSACTION_PROBE_QUERY_TIMEOUT used for cheap visibility and transaction-status probes. 3. HIGH: permanent database errors retried forever. Added SQLSTATE classification (classify_sqlstate/classify_sqlx_error): an allowlist of genuinely transient codes (connection failures, serialization/deadlock, our own statement/lock timeouts, admin-cancel/recovery conflicts) retries as before; everything else (insufficient_privilege, undefined_table, undefined_function, schema/decode errors, ...) now fails the activity immediately with its SQLSTATE embedded. Also added a bounded transient retry budget (MAX_GRAPH_RETRY_ATTEMPTS) tracked independently from the caller-transaction wait counter, so transient DB failures can no longer poll indefinitely even when individually "retryable". 4. MEDIUM: the regression test didn't prove the timeout-retry path executed and could pick an unrelated backend by pid alone. Rewrote 68_long_caller_transaction.sql to identify the graph-probe backend by its exact query text, and to require a strictly later query_start (rather than a different pid, since the management pool may legitimately reuse a connection) for both the post-kill retry and a further, independent timeout-driven retry before the lock is released. Also fixes a real bug found while validating this test: pg_stat_activity's snapshot is cached for the lifetime of the enclosing transaction, so a long-lived DO block must call pg_stat_clear_snapshot() before every poll to observe live backend state. 5. LOW: false OR NULL evaluates to NULL in PL/pgSQL, so a NULL engine_output on failure bypassed the rollback assertions. Added engine_output IS NULL OR ... to both rollback checks. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/activities/load_function_graph.rs | 297 +++++++++++++++++-- src/dsl.rs | 1 + src/lib.rs | 1 + src/orchestrations/execute_function_graph.rs | 79 ++++- src/types.rs | 27 ++ tests/e2e/sql/68_long_caller_transaction.sql | 119 +++++++- 6 files changed, 485 insertions(+), 39 deletions(-) diff --git a/src/activities/load_function_graph.rs b/src/activities/load_function_graph.rs index 6ccfc42e..6151636a 100644 --- a/src/activities/load_function_graph.rs +++ b/src/activities/load_function_graph.rs @@ -20,7 +20,25 @@ pub const NAME: &str = "pg_durable::activity::load-function-graph"; /// Transaction-aware graph probe used only by inputs created by the current binary. pub const TRANSACTION_AWARE_NAME: &str = "pg_durable::activity::probe-function-graph-transaction-v1"; +/// Deadline for cheap visibility/transaction-status probes only (a single +/// primary-key lookup or `pg_xact_status()` call). Graph *loading* uses its own, +/// much longer policy below — see `GRAPH_LOAD_QUERY_TIMEOUT`. const TRANSACTION_PROBE_QUERY_TIMEOUT: Duration = Duration::from_secs(2); +/// Deadline for the graph-loading path: role validation, fetching up to +/// `MAX_GRAPH_NODES` (10,000) node rows, constructing the graph, and +/// serializing it. Kept well above `TRANSACTION_PROBE_QUERY_TIMEOUT` so a +/// valid-but-large graph load is never mistaken for a stuck probe. +const GRAPH_LOAD_QUERY_TIMEOUT: Duration = Duration::from_secs(20); +/// Postgres-side `statement_timeout` applied while fetching node rows, so a +/// runaway query is cancelled by the server with a proper SQLSTATE instead of +/// relying solely on the client-side `GRAPH_LOAD_QUERY_TIMEOUT` dropping the +/// connection. +const GRAPH_LOAD_STATEMENT_TIMEOUT_MS: u64 = 15_000; +/// Postgres-side `lock_timeout` applied while fetching node rows, so a load +/// blocked behind a conflicting lock on `df.nodes`/`pg_roles` fails fast with +/// `55P03` (classified transient, see `classify_sqlstate`) rather than +/// consuming the whole statement timeout waiting to even start. +const GRAPH_LOAD_LOCK_TIMEOUT_MS: u64 = 5_000; /// Retry configuration for waiting on uncommitted transactions pub const MAX_WAIT_SECS: u64 = 5; @@ -56,6 +74,72 @@ impl LoadGraphError { } } +/// Whether a database error observed while probing/loading a graph is worth +/// retrying (bounded, see `MAX_GRAPH_RETRY_ATTEMPTS` in the orchestration) or +/// should fail the workflow immediately. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ErrorClass { + Transient, + Permanent, +} + +/// Classify a Postgres SQLSTATE as transient (worth a bounded retry) or +/// permanent (fail immediately). Deliberately conservative: any code not on +/// the transient allowlist - including codes we don't recognize - is treated +/// as permanent, since silently retrying an unrecognized error forever is +/// exactly the bug class this classification exists to prevent. +fn classify_sqlstate(code: &str) -> ErrorClass { + match code { + // Connection Exception class: transport/connection-level failures + // that are expected to be transient. + "08000" | "08001" | "08003" | "08004" | "08006" | "08007" | "08P01" => { + ErrorClass::Transient + } + // Concurrency conflicts that are expected to clear on retry. + "40001" /* serialization_failure */ | "40P01" /* deadlock_detected */ => { + ErrorClass::Transient + } + // Our own statement/lock timeouts (see GRAPH_LOAD_STATEMENT_TIMEOUT_MS / + // GRAPH_LOAD_LOCK_TIMEOUT_MS): the query was cancelled by policy, not + // because it can never succeed. + "57014" /* query_canceled */ | "55P03" /* lock_not_available */ => ErrorClass::Transient, + // Admin-initiated cancellation / hot-standby conflicts: transient by + // nature (server restart, failover, recovery conflict). + "57P01" | "57P02" | "57P03" => ErrorClass::Transient, + _ => ErrorClass::Permanent, + } +} + +/// Extract the SQLSTATE code from a sqlx error, if any. +fn sqlstate_of(error: &sqlx::Error) -> Option { + error + .as_database_error() + .and_then(|db| db.code()) + .map(|code| code.into_owned()) +} + +/// Classify a sqlx error as transient or permanent. Errors without a SQLSTATE +/// (pool exhaustion, IO/TLS/protocol failures - i.e. connection failures that +/// never reached the server) are treated as transient. +fn classify_sqlx_error(error: &sqlx::Error) -> ErrorClass { + match sqlstate_of(error) { + Some(code) => classify_sqlstate(&code), + None => ErrorClass::Transient, + } +} + +/// Wrap a sqlx error observed while loading a graph into the appropriately +/// classified `LoadGraphError`, embedding the SQLSTATE (or "unknown") so +/// terminal failures are diagnosable. +fn classified_load_graph_error(operation: &str, error: sqlx::Error) -> LoadGraphError { + let code = sqlstate_of(&error).unwrap_or_else(|| "unknown".to_string()); + let message = format!("{operation} failed (SQLSTATE {code}): {error}"); + match classify_sqlx_error(&error) { + ErrorClass::Transient => LoadGraphError::Retryable(message), + ErrorClass::Permanent => LoadGraphError::Permanent(message), + } +} + const INSTANCE_QUERY: &str = "SELECT root_node, r.rolname AS submitted_by FROM df.instances i LEFT JOIN pg_catalog.pg_roles r ON r.oid = i.submitted_by::oid @@ -71,6 +155,50 @@ async fn find_visible_instance( .await } +/// Fetch a graph's node rows with a real Postgres-side timeout backstop. +/// +/// Runs inside its own transaction so `SET LOCAL statement_timeout` / +/// `SET LOCAL lock_timeout` apply only to this read and are automatically +/// discarded when the transaction ends - no risk of leaking a modified +/// timeout onto a pooled connection reused by unrelated work. The read never +/// writes anything, so the transaction is always rolled back regardless of +/// outcome (rollback vs. commit makes no observable difference here; rollback +/// avoids depending on the connection's default transaction characteristics). +async fn fetch_node_rows( + pool: &PgPool, + instance_id: &str, +) -> Result, sqlx::Error> { + const NODES_QUERY: &str = r#"SELECT n.id, n.node_type, n.query, n.result_name, + n.left_node, n.right_node, + r.rolname AS submitted_by, + n.database + FROM df.nodes n + LEFT JOIN pg_catalog.pg_roles r ON r.oid = n.submitted_by::oid + WHERE n.instance_id = $1"#; + + let mut tx = pool.begin().await?; + sqlx::query(&format!( + "SET LOCAL statement_timeout = '{GRAPH_LOAD_STATEMENT_TIMEOUT_MS}ms'" + )) + .execute(&mut *tx) + .await?; + sqlx::query(&format!( + "SET LOCAL lock_timeout = '{GRAPH_LOAD_LOCK_TIMEOUT_MS}ms'" + )) + .execute(&mut *tx) + .await?; + + let rows = sqlx::query(NODES_QUERY) + .bind(instance_id) + .fetch_all(&mut *tx) + .await; + + // Best-effort: this is a read-only transaction, so a rollback failure + // (e.g. connection already dropped) doesn't change the outcome we report. + let _ = tx.rollback().await; + rows +} + async fn load_visible_graph( ctx: &ActivityContext, pool: &PgPool, @@ -106,19 +234,9 @@ async fn load_visible_graph( } } - let nodes_query = r#"SELECT n.id, n.node_type, n.query, n.result_name, - n.left_node, n.right_node, - r.rolname AS submitted_by, - n.database - FROM df.nodes n - LEFT JOIN pg_catalog.pg_roles r ON r.oid = n.submitted_by::oid - WHERE n.instance_id = $1"#; - - let rows = sqlx::query(nodes_query) - .bind(&instance_id) - .fetch_all(pool) + let rows = fetch_node_rows(pool, &instance_id) .await - .map_err(|e| LoadGraphError::Retryable(format!("Failed to load function nodes: {e}")))?; + .map_err(|e| classified_load_graph_error("Failed to load function nodes", e))?; let mut nodes = std::collections::BTreeMap::new(); for row in rows { @@ -224,6 +342,52 @@ fn retry_probe( serialize_probe(&TransactionGraphProbe::Retry) } +/// Route a sqlx error encountered while probing (visibility/transaction-status +/// checks) through classification: transient errors become a bounded `Retry` +/// probe result, permanent errors fail the activity immediately with the +/// SQLSTATE embedded so the workflow doesn't poll forever on an unrecoverable +/// condition (e.g. insufficient_privilege, undefined_table). +fn probe_error_outcome( + ctx: &ActivityContext, + operation: &str, + instance_id: &str, + error: sqlx::Error, +) -> Result { + match classify_sqlx_error(&error) { + ErrorClass::Transient => retry_probe(ctx, operation, instance_id, error), + ErrorClass::Permanent => { + let code = sqlstate_of(&error).unwrap_or_else(|| "unknown".to_string()); + Err(format!( + "Instance {instance_id}: {operation} failed permanently (SQLSTATE {code}): {error}" + )) + } + } +} + +/// Decision point for the exact race window this handles: PostgreSQL can +/// record a transaction as committed before `ProcArrayEndTransaction` removes +/// it from the running-transactions set used to build a fresh snapshot. In +/// that window, a graph re-read can spuriously return no rows even though the +/// transaction is genuinely committed and the graph exists. Extracted as a +/// pure function so the decision itself has direct unit coverage. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum CommittedProbeStep { + /// The origin xid is committed but not yet snapshot-visible: retry + /// (bounded) instead of trusting a re-read. + AwaitSnapshotVisibility, + /// The origin xid is committed and snapshot-visible: a re-read that finds + /// no graph can be trusted as genuine `CommittedMissing`. + ReadyToReread, +} + +fn committed_probe_step(snapshot_visible: bool) -> CommittedProbeStep { + if snapshot_visible { + CommittedProbeStep::ReadyToReread + } else { + CommittedProbeStep::AwaitSnapshotVisibility + } +} + /// Probe graph visibility without pinning an activity task while the caller's /// transaction remains open. The orchestration schedules a durable timer and /// invokes this single-shot activity again when the xid is still in progress. @@ -266,7 +430,7 @@ pub async fn probe_transaction( match visible { Ok(Some(row)) => { let loaded = match tokio::time::timeout( - TRANSACTION_PROBE_QUERY_TIMEOUT, + GRAPH_LOAD_QUERY_TIMEOUT, load_visible_graph(&ctx, pool.as_ref(), input.instance_id.clone(), row), ) .await @@ -277,7 +441,7 @@ pub async fn probe_transaction( &ctx, "graph load", &input.instance_id, - "timed out after 2s", + format!("timed out after {}s", GRAPH_LOAD_QUERY_TIMEOUT.as_secs()), ); } }; @@ -291,7 +455,7 @@ pub async fn probe_transaction( } Ok(None) => {} Err(e) => { - return retry_probe(&ctx, "graph visibility check", &input.instance_id, e); + return probe_error_outcome(&ctx, "graph visibility check", &input.instance_id, e); } } @@ -305,7 +469,7 @@ pub async fn probe_transaction( let transaction_status: Option = match transaction_status_query { Ok(Ok(status)) => status, Ok(Err(e)) => { - return retry_probe(&ctx, "origin transaction check", &input.instance_id, e); + return probe_error_outcome(&ctx, "origin transaction check", &input.instance_id, e); } Err(_) => { return retry_probe( @@ -321,9 +485,56 @@ pub async fn probe_transaction( Some("in progress") => TransactionGraphProbe::InProgress, Some("aborted") => TransactionGraphProbe::Aborted, Some("committed") => { + // pg_xact_status() can report "committed" before + // ProcArrayEndTransaction removes the xid from the running set + // used to build a fresh snapshot. Check snapshot visibility + // explicitly before trusting a re-read as proof the graph is + // absent - otherwise a valid committed graph can be permanently + // misclassified as CommittedMissing during this narrow window. + let snapshot_visible_query = tokio::time::timeout( + TRANSACTION_PROBE_QUERY_TIMEOUT, + sqlx::query_scalar::<_, bool>( + "SELECT pg_catalog.pg_visible_in_snapshot($1::text::xid8, pg_catalog.pg_current_snapshot())", + ) + .bind(&input.origin_xid) + .fetch_one(pool.as_ref()), + ) + .await; + let snapshot_visible = match snapshot_visible_query { + Ok(Ok(visible)) => visible, + Ok(Err(e)) => { + return probe_error_outcome( + &ctx, + "post-commit snapshot visibility check", + &input.instance_id, + e, + ); + } + Err(_) => { + return retry_probe( + &ctx, + "post-commit snapshot visibility check", + &input.instance_id, + "timed out after 2s", + ); + } + }; + + if committed_probe_step(snapshot_visible) == CommittedProbeStep::AwaitSnapshotVisibility + { + return retry_probe( + &ctx, + "post-commit snapshot visibility", + &input.instance_id, + "origin transaction committed but not yet snapshot-visible", + ); + } + // The status and graph reads use separate READ COMMITTED statements. // Re-read once after observing commit so a commit between the first // graph query and pg_xact_status cannot be misclassified as missing. + // Snapshot visibility is now confirmed above, so a `None` here is a + // genuine CommittedMissing, not a visibility race. let visible = match tokio::time::timeout( TRANSACTION_PROBE_QUERY_TIMEOUT, find_visible_instance(pool.as_ref(), &input.instance_id), @@ -343,7 +554,7 @@ pub async fn probe_transaction( match visible { Ok(Some(row)) => { let loaded = match tokio::time::timeout( - TRANSACTION_PROBE_QUERY_TIMEOUT, + GRAPH_LOAD_QUERY_TIMEOUT, load_visible_graph(&ctx, pool.as_ref(), input.instance_id.clone(), row), ) .await @@ -354,7 +565,7 @@ pub async fn probe_transaction( &ctx, "post-commit graph load", &input.instance_id, - "timed out after 2s", + format!("timed out after {}s", GRAPH_LOAD_QUERY_TIMEOUT.as_secs()), ); } }; @@ -373,7 +584,7 @@ pub async fn probe_transaction( } Ok(None) => TransactionGraphProbe::CommittedMissing, Err(e) => { - return retry_probe( + return probe_error_outcome( &ctx, "post-commit graph visibility check", &input.instance_id, @@ -435,4 +646,52 @@ mod tests { r#"{"state":"committed_missing"}"# ); } + + #[test] + fn committed_probe_step_awaits_snapshot_visibility_until_confirmed() { + // The exact race this covers: pg_xact_status() already reports + // "committed" but the xid isn't snapshot-visible yet - must not be + // treated as ready to re-read (that would risk a spurious + // CommittedMissing classification for a genuinely committed graph). + assert_eq!( + committed_probe_step(false), + CommittedProbeStep::AwaitSnapshotVisibility + ); + assert_eq!( + committed_probe_step(true), + CommittedProbeStep::ReadyToReread + ); + } + + #[test] + fn classify_sqlstate_allows_connection_and_our_own_timeout_codes() { + for code in [ + "08000", "08001", "08003", "08004", "08006", "08007", "08P01", "40001", "40P01", + "57014", "55P03", "57P01", "57P02", "57P03", + ] { + assert_eq!( + classify_sqlstate(code), + ErrorClass::Transient, + "expected {code} to classify as transient" + ); + } + } + + #[test] + fn classify_sqlstate_treats_privilege_and_schema_errors_as_permanent() { + for code in ["42501", "42883", "42P01", "22P02", "23505"] { + assert_eq!( + classify_sqlstate(code), + ErrorClass::Permanent, + "expected {code} to classify as permanent" + ); + } + } + + #[test] + fn classify_sqlstate_defaults_unknown_codes_to_permanent() { + // Deliberate: an unrecognized SQLSTATE must not be silently retried + // forever - that's exactly the bug class being fixed. + assert_eq!(classify_sqlstate("XXUNK"), ErrorClass::Permanent); + } } diff --git a/src/dsl.rs b/src/dsl.rs index 5212054f..7e399bef 100644 --- a/src/dsl.rs +++ b/src/dsl.rs @@ -1290,6 +1290,7 @@ fn start_in_caller_transaction(fut: &str, label: Option<&str>, database: Option< graph: None, origin_xid: Some(origin_xid), graph_wait_attempt: 0, + graph_retry_attempt: 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 5afef48f..985ce7c0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3341,6 +3341,7 @@ mod tests { graph: None, origin_xid: None, graph_wait_attempt: 0, + graph_retry_attempt: 0, }; let json = serde_json::to_string(&input).unwrap(); let deserialized: FunctionInput = serde_json::from_str(&json).unwrap(); diff --git a/src/orchestrations/execute_function_graph.rs b/src/orchestrations/execute_function_graph.rs index 66854048..3ac382df 100644 --- a/src/orchestrations/execute_function_graph.rs +++ b/src/orchestrations/execute_function_graph.rs @@ -125,6 +125,14 @@ type NodeResult = Result; const INITIAL_TRANSACTION_POLL_MS: u64 = 100; const MAX_TRANSACTION_POLL_MS: u64 = 5_000; const GRAPH_WAIT_POLLS_PER_EXECUTION: u32 = 64; +/// Bound on how many transient `Retry` outcomes (DB errors, query timeouts, +/// snapshot-visibility lag - see `probe_transaction`) a single graph admission +/// will absorb before failing the orchestration outright. Unlike waiting on +/// the caller's own open transaction (`InProgress`, legitimately unbounded), +/// these retries indicate the worker's own machinery is unhealthy, and +/// polling forever would hide a permanent problem behind misleading +/// "waiting on transaction" logs. +const MAX_GRAPH_RETRY_ATTEMPTS: u32 = 20; fn transaction_poll_delay(attempt: u32) -> Duration { let multiplier = 1u64 << attempt.min(6); @@ -139,12 +147,21 @@ fn should_compact_graph_wait(polls_in_execution: u32) -> bool { polls_in_execution >= GRAPH_WAIT_POLLS_PER_EXECUTION } +/// Whether the bounded transient-retry budget for graph admission has been +/// exhausted. Pure so the boundary (`== MAX_GRAPH_RETRY_ATTEMPTS` vs `>`) has +/// direct unit coverage. +fn graph_retry_budget_exceeded(retry_attempt: u32) -> bool { + retry_attempt > MAX_GRAPH_RETRY_ATTEMPTS +} + fn graph_wait_continuation( input: &FunctionInput, - next_attempt: u32, + next_wait_attempt: u32, + next_retry_attempt: u32, ) -> Result { let mut continuation = input.clone(); - continuation.graph_wait_attempt = next_attempt; + continuation.graph_wait_attempt = next_wait_attempt; + continuation.graph_retry_attempt = next_retry_attempt; serde_json::to_string(&continuation) } @@ -170,7 +187,8 @@ async fn load_initial_graph( }) .map_err(|e| format!("Failed to serialize transaction-aware graph probe: {e}"))?; - let mut attempt = input.graph_wait_attempt; + let mut wait_attempt = input.graph_wait_attempt; + let mut retry_attempt = input.graph_retry_attempt; let mut polls_in_execution = 0u32; loop { let raw = ctx @@ -184,19 +202,51 @@ async fn load_initial_graph( match probe { TransactionGraphProbe::Ready { graph } => return Ok(graph), - TransactionGraphProbe::InProgress | TransactionGraphProbe::Retry => { - let delay = transaction_poll_delay(attempt); + TransactionGraphProbe::InProgress => { + let delay = transaction_poll_delay(wait_attempt); ctx.trace_info(format!( "Instance {} graph is waiting on origin transaction {}; retrying in {:?}", input.instance_id, origin_xid, delay )); ctx.schedule_timer(delay).await; - attempt = attempt.saturating_add(1); + wait_attempt = wait_attempt.saturating_add(1); + polls_in_execution = polls_in_execution.saturating_add(1); + + if should_compact_graph_wait(polls_in_execution) { + let continuation_json = + graph_wait_continuation(input, wait_attempt, retry_attempt).map_err( + |e| format!("Failed to serialize graph-wait continuation: {e}"), + )?; + ctx.trace_info(format!( + "Compacting graph-admission history for instance {} after {} polls", + input.instance_id, polls_in_execution + )); + return ctx.continue_as_new(continuation_json).await; + } + } + TransactionGraphProbe::Retry => { + retry_attempt = retry_attempt.saturating_add(1); + if graph_retry_budget_exceeded(retry_attempt) { + return Err(format!( + "Instance {} graph admission for origin transaction {} failed: \ + exceeded {MAX_GRAPH_RETRY_ATTEMPTS} transient retries", + input.instance_id, origin_xid + )); + } + let delay = transaction_poll_delay(retry_attempt); + ctx.trace_info(format!( + "Instance {} graph admission for origin transaction {} hit a transient \ + failure ({}/{MAX_GRAPH_RETRY_ATTEMPTS}); retrying in {:?}", + input.instance_id, origin_xid, retry_attempt, delay + )); + ctx.schedule_timer(delay).await; polls_in_execution = polls_in_execution.saturating_add(1); if should_compact_graph_wait(polls_in_execution) { - let continuation_json = graph_wait_continuation(input, attempt) - .map_err(|e| format!("Failed to serialize graph-wait continuation: {e}"))?; + let continuation_json = + graph_wait_continuation(input, wait_attempt, retry_attempt).map_err( + |e| format!("Failed to serialize graph-wait continuation: {e}"), + )?; ctx.trace_info(format!( "Compacting graph-admission history for instance {} after {} polls", input.instance_id, polls_in_execution @@ -1088,6 +1138,7 @@ async fn execute_loop_node( graph: Some(graph_json), origin_xid: None, graph_wait_attempt: 0, + graph_retry_attempt: 0, }; serde_json::to_string(&new_input) .map_err(|e| format!("Failed to serialize loop input: {e}"))? @@ -1926,9 +1977,10 @@ mod tests { graph: None, origin_xid: Some("12345".to_string()), graph_wait_attempt: 63, + graph_retry_attempt: 2, }; - let json = graph_wait_continuation(&input, 64).unwrap(); + let json = graph_wait_continuation(&input, 64, 3).unwrap(); let continued: FunctionInput = serde_json::from_str(&json).unwrap(); assert_eq!(continued.instance_id, input.instance_id); assert_eq!(continued.label, input.label); @@ -1937,6 +1989,15 @@ mod tests { assert_eq!(continued.graph, None); assert_eq!(continued.origin_xid.as_deref(), Some("12345")); assert_eq!(continued.graph_wait_attempt, 64); + assert_eq!(continued.graph_retry_attempt, 3); + } + + #[test] + fn graph_retry_budget_exceeded_allows_exactly_max_attempts() { + assert!(!graph_retry_budget_exceeded(0)); + assert!(!graph_retry_budget_exceeded(MAX_GRAPH_RETRY_ATTEMPTS)); + assert!(graph_retry_budget_exceeded(MAX_GRAPH_RETRY_ATTEMPTS + 1)); + assert!(graph_retry_budget_exceeded(u32::MAX)); } #[test] diff --git a/src/types.rs b/src/types.rs index 2ad67ad2..aa0b9ee6 100644 --- a/src/types.rs +++ b/src/types.rs @@ -1285,6 +1285,15 @@ pub struct FunctionInput { /// with continue_as_new before user graph execution begins. #[serde(default, skip_serializing_if = "is_zero_u32")] pub graph_wait_attempt: u32, + /// Number of transient graph-admission `Retry` outcomes already observed + /// for `origin_xid` (DB errors, query timeouts, snapshot-visibility lag). + /// + /// Tracked independently from `graph_wait_attempt`: waiting on the + /// caller's own open transaction (`InProgress`) is legitimately + /// unbounded, but waiting on the worker's own machinery is not, so it is + /// bounded separately (see `MAX_GRAPH_RETRY_ATTEMPTS`). + #[serde(default, skip_serializing_if = "is_zero_u32")] + pub graph_retry_attempt: u32, } pub(crate) fn serialize_string_map( @@ -2215,6 +2224,7 @@ mod tests { graph: None, origin_xid: None, graph_wait_attempt: 0, + graph_retry_attempt: 0, }; let reverse_input = FunctionInput { instance_id: "instance".to_string(), @@ -2224,6 +2234,7 @@ mod tests { graph: None, origin_xid: None, graph_wait_attempt: 0, + graph_retry_attempt: 0, }; assert_eq!( serde_json::to_string(&forward_input).unwrap(), @@ -2248,6 +2259,7 @@ mod tests { graph: None, origin_xid: Some("123456".to_string()), graph_wait_attempt: 0, + graph_retry_attempt: 0, }; let current_json = serde_json::to_string(¤t).unwrap(); assert!(current_json.ends_with(r#","origin_xid":"123456"}"#)); @@ -2265,6 +2277,21 @@ mod tests { }; let compacted_json = serde_json::to_string(&compacted).unwrap(); assert!(compacted_json.ends_with(r#","graph_wait_attempt":64}"#)); + + let retry_compacted = FunctionInput { + graph_retry_attempt: 3, + ..compacted + }; + let retry_compacted_json = serde_json::to_string(&retry_compacted).unwrap(); + assert!( + retry_compacted_json.ends_with(r#","graph_wait_attempt":64,"graph_retry_attempt":3}"#) + ); + assert_eq!( + serde_json::from_str::(&retry_compacted_json) + .unwrap() + .graph_retry_attempt, + 3 + ); } fn sys_vars() -> SystemVars { diff --git a/tests/e2e/sql/68_long_caller_transaction.sql b/tests/e2e/sql/68_long_caller_transaction.sql index f995ad1f..abcb46cf 100644 --- a/tests/e2e/sql/68_long_caller_transaction.sql +++ b/tests/e2e/sql/68_long_caller_transaction.sql @@ -105,7 +105,12 @@ DO $$ DECLARE v_instance_id TEXT; engine_status TEXT; - target_pid INT; + first_pid INT; + first_query_start TIMESTAMPTZ; + second_pid INT; + second_query_start TIMESTAMPTZ; + third_pid INT; + third_query_start TIMESTAMPTZ; attempts INT := 0; BEGIN SELECT s.instance_id INTO v_instance_id @@ -125,31 +130,121 @@ BEGIN engine_status; END IF; + -- Identify the graph-probe backend by the query it is actually running + -- (not just any lock-waiting management-pool backend, which could match + -- an unrelated query) so we can recognize a genuinely new attempt. + -- Match on the graph-probe's exact leading text ("SELECT root_node, ... + -- AS submitted_by") rather than the generic "FROM df.instances" + -- substring: the worker's periodic terminal-instance retention/cleanup + -- query also reads df.instances and also blocks on this same lock, so a + -- loose substring match can pick that unrelated backend instead. + -- + -- We identify a genuinely NEW retry attempt by query_start advancing, + -- not by the backend pid changing: the management pool may legitimately + -- reuse the same physical connection across retries (a client-side + -- timeout abandons a query without necessarily evicting/reopening the + -- pooled connection), so requiring a different pid would wrongly fail + -- even when the retry path is working correctly. query_start is set + -- fresh each time a backend begins a new query - even one that + -- immediately blocks acquiring a lock - so an advancing query_start on a + -- matching row unambiguously proves a new attempt was submitted. + -- + -- We key off state = 'active' rather than wait_event_type = 'Lock': a + -- backend parked in the heavyweight lock manager only reports + -- wait_event_type = 'Lock' for the brief instant it first enters the + -- wait - PostgreSQL's periodic deadlock-check wakeup clears it back to + -- NULL for most of the actual wait even though the backend is still + -- blocked making no progress (state stays 'active', the query text is + -- unchanged). Since our own transaction holds an ACCESS EXCLUSIVE lock + -- on df.instances for this whole scenario, no backend anywhere can be + -- genuinely, successfully executing this query against df.instances + -- while we hold it, so "state = active" plus matching application_name + -- and query text is an unambiguous, race-free signal that a backend is + -- blocked behind our lock. + -- pg_stat_activity's backing function takes a snapshot of all backend + -- statuses that is cached for the lifetime of the current transaction; + -- repeated reads within the same transaction (which this whole DO block + -- runs in, since it holds the ACCESS EXCLUSIVE lock throughout) would + -- otherwise silently return the same frozen point-in-time view forever. + -- pg_stat_clear_snapshot() must be called before every poll so each + -- iteration observes genuinely live backend state. attempts := 0; LOOP - SELECT pid INTO target_pid + PERFORM pg_stat_clear_snapshot(); + SELECT pid, query_start INTO first_pid, first_query_start FROM pg_stat_activity WHERE application_name = 'pg_durable:worker:management' - AND wait_event_type = 'Lock' - ORDER BY pid + AND state = 'active' + AND query ILIKE 'SELECT root_node,%submitted_by%FROM df.instances%' + ORDER BY query_start LIMIT 1; - EXIT WHEN target_pid IS NOT NULL OR attempts >= 300; + EXIT WHEN first_pid IS NOT NULL OR attempts >= 400; PERFORM pg_sleep(0.05); attempts := attempts + 1; END LOOP; - IF target_pid IS NULL THEN + IF first_pid IS NULL THEN RAISE EXCEPTION 'TEST FAILED [transient retry]: no blocked graph-probe backend found'; END IF; - IF NOT pg_terminate_backend(target_pid) THEN + IF NOT pg_terminate_backend(first_pid) THEN RAISE EXCEPTION 'TEST FAILED [transient retry]: could not terminate graph-probe backend %', - target_pid; + first_pid; END IF; - -- Leave time for the killed query to return Retry and the next lock-blocked - -- probe to hit its bounded query timeout before releasing the table lock. - PERFORM pg_sleep(3); + -- Prove the connection-failure retry path actually ran: a fresh attempt + -- with a later query_start must appear. Nothing else can make this + -- query progress while the ACCESS EXCLUSIVE lock is held, so a new + -- query_start can only appear because the terminated connection's error + -- was classified transient and retried by the orchestration - not + -- treated as a terminal activity failure. + attempts := 0; + LOOP + PERFORM pg_stat_clear_snapshot(); + SELECT pid, query_start INTO second_pid, second_query_start + FROM pg_stat_activity + WHERE application_name = 'pg_durable:worker:management' + AND state = 'active' + AND query ILIKE 'SELECT root_node,%submitted_by%FROM df.instances%' + AND query_start > first_query_start + ORDER BY query_start + LIMIT 1; + EXIT WHEN second_pid IS NOT NULL OR attempts >= 400; + PERFORM pg_sleep(0.05); + attempts := attempts + 1; + END LOOP; + IF second_pid IS NULL THEN + RAISE EXCEPTION + 'TEST FAILED [transient retry]: no retry attempt observed after terminating %; the connection-failure retry path did not execute', + first_pid; + END IF; + + -- Prove the *timeout* path also runs, independent of our explicit kill: + -- nothing else can terminate the second attempt, so a third, later + -- query_start can only appear once that attempt's own bounded + -- client-side query timeout elapsed and the orchestration rescheduled + -- it - proof that the timeout-retry path (not just the kill-retry path) + -- actually executes before we release the lock. + attempts := 0; + LOOP + PERFORM pg_stat_clear_snapshot(); + SELECT pid, query_start INTO third_pid, third_query_start + FROM pg_stat_activity + WHERE application_name = 'pg_durable:worker:management' + AND state = 'active' + AND query ILIKE 'SELECT root_node,%submitted_by%FROM df.instances%' + AND query_start > second_query_start + ORDER BY query_start + LIMIT 1; + EXIT WHEN third_pid IS NOT NULL OR attempts >= 400; + PERFORM pg_sleep(0.05); + attempts := attempts + 1; + END LOOP; + IF third_pid IS NULL THEN + RAISE EXCEPTION + 'TEST FAILED [transient retry]: no timeout-driven retry observed after %; the bounded query-timeout path did not execute', + second_pid; + END IF; END $$; COMMIT; @@ -241,6 +336,7 @@ BEGIN END LOOP; IF lower(COALESCE(engine_status, '')) != 'failed' + OR engine_output IS NULL OR engine_output NOT LIKE '%origin transaction%aborted%' THEN RAISE EXCEPTION 'TEST FAILED [whole rollback]: engine status=%, output=%', @@ -289,6 +385,7 @@ BEGIN END LOOP; IF lower(COALESCE(engine_status, '')) != 'failed' + OR engine_output IS NULL OR engine_output NOT LIKE '%committed%graph%absent%' THEN RAISE EXCEPTION 'TEST FAILED [savepoint rollback]: engine status=%, output=%', From 2a4478dec5a1d440a2bf01ef8cdb45a94d3961a2 Mon Sep 17 00:00:00 2001 From: Pino de Candia Date: Wed, 2 Sep 2026 15:47:45 +0000 Subject: [PATCH 3/3] Harden caller-transaction handoff per follow-up review - Move the caller-transaction handoff changelog entry to the 0.2.8 Unreleased section after rebasing onto origin/main. - Apply server-side statement_timeout/lock_timeout to every cheap probe (visibility, transaction-status, snapshot) and the role superuser check so a blocked probe is cancelled by PostgreSQL and its connection returned cleanly to the pool instead of pinning a management connection. - Durably retry terminal df.instances status writes (completed/failed) with a bounded budget so a transient management-plane outage cannot leave the row non-terminal while the engine execution is terminal. --- CHANGELOG.md | 5 +- src/activities/load_function_graph.rs | 87 ++++++++++++++--- src/orchestrations/execute_function_graph.rs | 99 ++++++++++++++------ src/types.rs | 31 +++++- 4 files changed, 174 insertions(+), 48 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a47d62d..6967c191 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,10 @@ Pre-1.0 note: while `pg_durable` is in major version `0`, minor releases may inc ## [0.2.8] - Unreleased +### Fixed + +- **Caller-transaction handoff:** `df.start()` now tracks the originating transaction until it commits or aborts, so legal caller transactions lasting more than five seconds no longer leave a `pending` `df.instances` row paired with a failed engine execution. Graph admission uses durable backoff and bounded-history compaction rather than holding a worker connection while it waits. + ## [0.2.7] - 2026-08-31 ### Added @@ -18,7 +22,6 @@ Pre-1.0 note: while `pg_durable` is in major version `0`, minor releases may inc ### Fixed -- **Caller-transaction handoff:** `df.start()` now tracks the originating transaction until it commits or aborts, so legal caller transactions lasting more than five seconds no longer leave a `pending` `df.instances` row paired with a failed engine execution. Graph admission uses durable backoff and bounded-history compaction rather than holding a worker connection while it waits. - **Worker connection role names (#364):** catalog role names are now passed verbatim when opening workflow connections, preventing quote-wrapped names from being reinterpreted as a different role. - **Restricted HTTP transport (#342, #363):** restricted allow-list builds now require HTTPS so credentials and request bodies cannot be sent over plaintext HTTP; development-only `http-allow-all` builds continue to permit HTTP. diff --git a/src/activities/load_function_graph.rs b/src/activities/load_function_graph.rs index 6151636a..6ca33bab 100644 --- a/src/activities/load_function_graph.rs +++ b/src/activities/load_function_graph.rs @@ -39,6 +39,19 @@ const GRAPH_LOAD_STATEMENT_TIMEOUT_MS: u64 = 15_000; /// `55P03` (classified transient, see `classify_sqlstate`) rather than /// consuming the whole statement timeout waiting to even start. const GRAPH_LOAD_LOCK_TIMEOUT_MS: u64 = 5_000; +/// Postgres-side `statement_timeout` applied to every cheap probe query +/// (visibility, transaction-status, snapshot). Set below the 2s client-side +/// `TRANSACTION_PROBE_QUERY_TIMEOUT` so the *server* cancels a stuck probe and +/// hands the connection back to the pool cleanly, before the client-side +/// `tokio::time::timeout` would drop the in-flight future and force the pool to +/// drain (or discard) the connection — the failure mode that can otherwise +/// exhaust the small management pool under a conflicting lock. +const PROBE_STATEMENT_TIMEOUT_MS: u64 = 1_500; +/// Postgres-side `lock_timeout` applied to every cheap probe query, so a probe +/// blocked behind a conflicting lock (e.g. on `df.instances`) fails fast with +/// `55P03` (classified transient) instead of occupying a management connection +/// until the client deadline. +const PROBE_LOCK_TIMEOUT_MS: u64 = 1_500; /// Retry configuration for waiting on uncommitted transactions pub const MAX_WAIT_SECS: u64 = 5; @@ -145,14 +158,72 @@ const INSTANCE_QUERY: &str = "SELECT root_node, r.rolname AS submitted_by LEFT JOIN pg_catalog.pg_roles r ON r.oid = i.submitted_by::oid WHERE i.id = $1"; +/// Begin a transaction with server-side `statement_timeout` / `lock_timeout` +/// applied via `SET LOCAL`, so every cheap probe query is cancelled by +/// PostgreSQL before the client-side deadline fires. The timeouts are scoped to +/// the transaction and discarded on rollback, so they never leak onto a pooled +/// connection reused by unrelated work. Callers run their read against the +/// returned transaction and roll it back (the probes are read-only, so the +/// rollback vs. commit distinction is not observable). +async fn begin_probe_tx( + pool: &PgPool, +) -> Result, sqlx::Error> { + let mut tx = pool.begin().await?; + sqlx::query(&format!( + "SET LOCAL statement_timeout = '{PROBE_STATEMENT_TIMEOUT_MS}ms'" + )) + .execute(&mut *tx) + .await?; + sqlx::query(&format!( + "SET LOCAL lock_timeout = '{PROBE_LOCK_TIMEOUT_MS}ms'" + )) + .execute(&mut *tx) + .await?; + Ok(tx) +} + async fn find_visible_instance( pool: &PgPool, instance_id: &str, ) -> Result, sqlx::Error> { - sqlx::query(INSTANCE_QUERY) + let mut tx = begin_probe_tx(pool).await?; + let row = sqlx::query(INSTANCE_QUERY) .bind(instance_id) - .fetch_optional(pool) - .await + .fetch_optional(&mut *tx) + .await; + // Read-only transaction: a rollback failure cannot change the result. + let _ = tx.rollback().await; + row +} + +/// Probe the caller's origin transaction status (`pg_xact_status`) under +/// server-enforced probe timeouts, so a blocked catalog read is cancelled by +/// PostgreSQL rather than pinning a management connection. +async fn probe_origin_transaction_status( + pool: &PgPool, + origin_xid: &str, +) -> Result, sqlx::Error> { + let mut tx = begin_probe_tx(pool).await?; + let status = sqlx::query_scalar("SELECT pg_catalog.pg_xact_status($1::text::xid8)::text") + .bind(origin_xid) + .fetch_one(&mut *tx) + .await; + let _ = tx.rollback().await; + status +} + +/// Probe whether the origin transaction is visible in a fresh snapshot, under +/// server-enforced probe timeouts (see `begin_probe_tx`). +async fn probe_snapshot_visible(pool: &PgPool, origin_xid: &str) -> Result { + let mut tx = begin_probe_tx(pool).await?; + let visible = sqlx::query_scalar::<_, bool>( + "SELECT pg_catalog.pg_visible_in_snapshot($1::text::xid8, pg_catalog.pg_current_snapshot())", + ) + .bind(origin_xid) + .fetch_one(&mut *tx) + .await; + let _ = tx.rollback().await; + visible } /// Fetch a graph's node rows with a real Postgres-side timeout backstop. @@ -461,9 +532,7 @@ pub async fn probe_transaction( let transaction_status_query = tokio::time::timeout( TRANSACTION_PROBE_QUERY_TIMEOUT, - sqlx::query_scalar("SELECT pg_catalog.pg_xact_status($1::text::xid8)::text") - .bind(&input.origin_xid) - .fetch_one(pool.as_ref()), + probe_origin_transaction_status(pool.as_ref(), &input.origin_xid), ) .await; let transaction_status: Option = match transaction_status_query { @@ -493,11 +562,7 @@ pub async fn probe_transaction( // misclassified as CommittedMissing during this narrow window. let snapshot_visible_query = tokio::time::timeout( TRANSACTION_PROBE_QUERY_TIMEOUT, - sqlx::query_scalar::<_, bool>( - "SELECT pg_catalog.pg_visible_in_snapshot($1::text::xid8, pg_catalog.pg_current_snapshot())", - ) - .bind(&input.origin_xid) - .fetch_one(pool.as_ref()), + probe_snapshot_visible(pool.as_ref(), &input.origin_xid), ) .await; let snapshot_visible = match snapshot_visible_query { diff --git a/src/orchestrations/execute_function_graph.rs b/src/orchestrations/execute_function_graph.rs index 3ac382df..89af983b 100644 --- a/src/orchestrations/execute_function_graph.rs +++ b/src/orchestrations/execute_function_graph.rs @@ -134,6 +134,15 @@ const GRAPH_WAIT_POLLS_PER_EXECUTION: u32 = 64; /// "waiting on transaction" logs. const MAX_GRAPH_RETRY_ATTEMPTS: u32 = 20; +/// Bound on how many times a *terminal* `df.instances` status write +/// (`completed`/`failed`) is durably retried before the orchestration gives up. +/// The engine execution is already terminal at these sites, so a dropped write +/// would leave `df.status()` / `df.await_instance()` trusting a stale +/// non-terminal row indefinitely. Retrying across durable timers lets a +/// transient management-plane outage self-heal once the database recovers, +/// while the bound still prevents an unhealthy worker from spinning forever. +const MAX_STATUS_FINALIZE_ATTEMPTS: u32 = 20; + fn transaction_poll_delay(attempt: u32) -> Duration { let multiplier = 1u64 << attempt.min(6); Duration::from_millis( @@ -300,6 +309,58 @@ struct SubtreeEnvelope { results: HashMap, } +/// Durably drive a *terminal* `df.instances` status write to completion. +/// +/// The engine execution is already terminal when this is called, so a +/// best-effort write that is silently dropped on a transient database/pool +/// failure would leave `df.status()` / `df.await_instance()` trusting a stale +/// non-terminal (`pending`/`running`) row while the engine is finished — the +/// two status surfaces diverging indefinitely. Retry the update activity across +/// durable timers so a transient management-plane outage self-heals once the +/// database recovers. The retry budget is bounded (`MAX_STATUS_FINALIZE_ATTEMPTS`) +/// so a persistently unhealthy worker cannot spin forever; on exhaustion the +/// divergence is at least surfaced in the trace rather than hidden. +/// +/// This is deterministic: activity results and timer fires are recorded in +/// history, so replay follows the same attempt sequence. +async fn finalize_instance_status(ctx: &OrchestrationContext, instance_id: &str, status: &str) { + let status_input = serde_json::json!({ + "instance_id": instance_id, + "status": status, + }) + .to_string(); + + let mut attempt = 0u32; + loop { + match ctx + .schedule_activity( + activities::update_instance_status::NAME, + status_input.clone(), + ) + .await + { + Ok(_) => return, + Err(e) => { + if attempt >= MAX_STATUS_FINALIZE_ATTEMPTS { + ctx.trace_info(format!( + "Instance {instance_id}: giving up finalizing status to '{status}' after \ + {MAX_STATUS_FINALIZE_ATTEMPTS} attempts; df.instances may remain \ + non-terminal while the engine execution is terminal: {e}" + )); + return; + } + let delay = transaction_poll_delay(attempt); + ctx.trace_info(format!( + "Instance {instance_id}: finalizing status to '{status}' failed \ + (attempt {attempt}/{MAX_STATUS_FINALIZE_ATTEMPTS}); retrying in {delay:?}: {e}" + )); + ctx.schedule_timer(delay).await; + attempt = attempt.saturating_add(1); + } + } + } +} + /// Execute a complete function graph — the entry point for a durable function. /// /// # Control flow @@ -344,17 +405,11 @@ pub async fn execute(ctx: OrchestrationContext, input_json: String) -> Result json, Err(e) => { // load_function_graph failed (e.g., superuser blocked). - // Mark the instance as failed before propagating. - let status_input = serde_json::json!({ - "instance_id": input.instance_id, - "status": "failed" - }); - let _ = ctx - .schedule_activity( - activities::update_instance_status::NAME, - status_input.to_string(), - ) - .await; + // Mark the instance as failed before propagating. The engine + // execution is about to end terminally, so this terminal write + // is retried durably rather than dropped — otherwise the row + // would stay non-terminal while the engine is failed. + finalize_instance_status(&ctx, &input.instance_id, "failed").await; return Err(e); } }, @@ -414,29 +469,11 @@ pub async fn execute(ctx: OrchestrationContext, input_json: String) -> Result { ctx.trace_info(format!("Function completed with result: {result}")); - let status_input = serde_json::json!({ - "instance_id": input.instance_id, - "status": "completed" - }); - let _ = ctx - .schedule_activity( - activities::update_instance_status::NAME, - status_input.to_string(), - ) - .await; + finalize_instance_status(&ctx, &input.instance_id, "completed").await; } Err(err) => { ctx.trace_info(format!("Function failed with error: {err}")); - let status_input = serde_json::json!({ - "instance_id": input.instance_id, - "status": "failed" - }); - let _ = ctx - .schedule_activity( - activities::update_instance_status::NAME, - status_input.to_string(), - ) - .await; + finalize_instance_status(&ctx, &input.instance_id, "failed").await; } } diff --git a/src/types.rs b/src/types.rs index aa0b9ee6..58d43906 100644 --- a/src/types.rs +++ b/src/types.rs @@ -112,13 +112,34 @@ pub fn is_role_superuser_oid(role_oid: pgrx::pg_sys::Oid) -> Result Result { - sqlx::query_scalar::<_, bool>("SELECT rolsuper FROM pg_catalog.pg_roles WHERE rolname = $1") - .bind(role_name) - .fetch_optional(pool) + let err = |e: sqlx::Error| format!("superuser check failed for role '{}': {}", role_name, e); + let mut tx = pool.begin().await.map_err(err)?; + sqlx::query("SET LOCAL statement_timeout = '1500ms'") + .execute(&mut *tx) + .await + .map_err(err)?; + sqlx::query("SET LOCAL lock_timeout = '1500ms'") + .execute(&mut *tx) .await - .map_err(|e| format!("superuser check failed for role '{}': {}", role_name, e)) - .and_then(|opt| opt.ok_or_else(|| format!("role '{}' not found in pg_roles", role_name))) + .map_err(err)?; + let result = sqlx::query_scalar::<_, bool>( + "SELECT rolsuper FROM pg_catalog.pg_roles WHERE rolname = $1", + ) + .bind(role_name) + .fetch_optional(&mut *tx) + .await + .map_err(err) + .and_then(|opt| opt.ok_or_else(|| format!("role '{}' not found in pg_roles", role_name))); + // Read-only transaction: a rollback failure cannot change the result. + let _ = tx.rollback().await; + result } /// Maximum nesting depth for workflow graphs. Bounds recursive graph walkers