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>

Author: Haider Zahid

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<String> {
+    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<Vec<sqlx::postgres::PgRow>, 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<String, String> {
+    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<String> = 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);
+    }
 }