Support resilient conditional loops

Author: tjgreen42

src/orchestrations/execute_function_graph.rs

@@ -344,15 +344,26 @@ fn contextualize_subtree_failure(context: &str, error: String) -> NodeError {
 }
 
 #[derive(Debug, PartialEq, Eq)]
-enum IsolatedIterationFailure {
-    Continue(String),
-    Propagate(String),
+enum FailureIsolatedBodyOutcome {
+    Succeeded(String),
+    ApplicationFailed,
+    Break(String),
 }
 
-fn classify_isolated_iteration_failure(error: String) -> IsolatedIterationFailure {
-    match decode_subtree_application_failure(&error) {
-        Some(message) => IsolatedIterationFailure::Continue(message),
-        None => IsolatedIterationFailure::Propagate(error),
+fn classify_isolated_body_result(
+    result: Result<String, String>,
+    results: &mut HashMap<String, String>,
+) -> Result<FailureIsolatedBodyOutcome, NodeError> {
+    match result {
+        Ok(raw) => match parse_subtree_envelope(&raw, "LOOP iteration", results) {
+            Ok(body_result) => Ok(FailureIsolatedBodyOutcome::Succeeded(body_result)),
+            Err(NodeError::Break(value)) => Ok(FailureIsolatedBodyOutcome::Break(value)),
+            Err(error) => Err(error),
+        },
+        Err(error) => match decode_subtree_application_failure(&error) {
+            Some(_) => Ok(FailureIsolatedBodyOutcome::ApplicationFailed),
+            None => Err(NodeError::Failure(error)),
+        },
     }
 }
 
@@ -1120,7 +1131,27 @@ async fn run_loop_iteration(
             Err(NodeError::Failure(error)) => return Err(NodeError::Failure(error)),
         };
 
-    // While-condition: if present and false, exit the loop.
+    Box::pin(evaluate_loop_condition(
+        ctx,
+        graph,
+        node,
+        condition_node,
+        body_result,
+        results,
+        exec_ctx,
+    ))
+    .await
+}
+
+async fn evaluate_loop_condition(
+    ctx: &OrchestrationContext,
+    graph: &FunctionGraph,
+    node: &FunctionNode,
+    condition_node: Option<&str>,
+    body_result: String,
+    results: &mut HashMap<String, String>,
+    exec_ctx: &ExecutionContext,
+) -> Result<Option<String>, NodeError> {
     if let Some(condition_node_id) = condition_node {
         ctx.trace_info("Evaluating loop condition");
         let condition_result =
@@ -1152,48 +1183,32 @@ async fn run_loop_iteration(
     Ok(None)
 }
 
-async fn run_failure_isolated_iteration(
+async fn run_failure_isolated_body(
     ctx: &OrchestrationContext,
     graph: &FunctionGraph,
-    node: &FunctionNode,
     loop_node_id: &str,
     body_id: &str,
     results: &mut HashMap<String, String>,
     exec_ctx: &ExecutionContext,
-) -> Result<Option<String>, NodeError> {
+) -> Result<FailureIsolatedBodyOutcome, NodeError> {
     let child_input = build_subtree_input(graph, body_id, results, exec_ctx)?;
     let child_id = subtree_instance_id(ctx, body_id);
 
     ctx.trace_info(format!(
         "Starting failure-isolated iteration for LOOP node {loop_node_id} as child {child_id}"
     ));
 
-    let raw = match ctx
+    let child_result = ctx
         .schedule_sub_orchestration_with_id(SUBTREE_NAME, child_id, child_input)
-        .await
-    {
-        Ok(raw) => raw,
-        Err(error) => match classify_isolated_iteration_failure(error) {
-            IsolatedIterationFailure::Continue(error) => {
-                ctx.trace_warn(format!(
-                    "LOOP node {loop_node_id} iteration application failure; continuing: {error}"
-                ));
-                return Ok(None);
-            }
-            IsolatedIterationFailure::Propagate(error) => {
-                return Err(NodeError::Failure(error));
-            }
-        },
-    };
+        .await;
 
-    match parse_subtree_envelope(&raw, "LOOP iteration", results) {
-        Ok(_) => Ok(None),
-        Err(NodeError::Break(value)) => {
-            store_named_result(ctx, node, &value, results, "LOOP");
-            Ok(Some(value))
-        }
-        Err(error) => Err(error),
+    let outcome = classify_isolated_body_result(child_result, results)?;
+    if outcome == FailureIsolatedBodyOutcome::ApplicationFailed {
+        ctx.trace_warn(format!(
+            "LOOP node {loop_node_id} iteration application failure; continuing"
+        ));
     }
+    Ok(outcome)
 }
 
 /// Execute a loop node inline, driving the *current* orchestration's `continue_as_new`.
@@ -1237,10 +1252,29 @@ async fn execute_loop_node(
     ctx.trace_info("Executing loop iteration");
 
     let final_result = if config.continue_on_failure {
-        Box::pin(run_failure_isolated_iteration(
-            ctx, graph, node, node_id, body_id, results, exec_ctx,
+        match Box::pin(run_failure_isolated_body(
+            ctx, graph, node_id, body_id, results, exec_ctx,
         ))
-        .await
+        .await?
+        {
+            FailureIsolatedBodyOutcome::Succeeded(body_result) => {
+                Box::pin(evaluate_loop_condition(
+                    ctx,
+                    graph,
+                    node,
+                    config.condition_node.as_deref(),
+                    body_result,
+                    results,
+                    exec_ctx,
+                ))
+                .await?
+            }
+            FailureIsolatedBodyOutcome::ApplicationFailed => None,
+            FailureIsolatedBodyOutcome::Break(value) => {
+                store_named_result(ctx, node, &value, results, "LOOP");
+                Some(value)
+            }
+        }
     } else {
         Box::pin(run_loop_iteration(
             ctx,
@@ -1251,10 +1285,10 @@ async fn execute_loop_node(
             results,
             exec_ctx,
         ))
-        .await
+        .await?
     };
 
-    if let Some(final_result) = final_result? {
+    if let Some(final_result) = final_result {
         return Ok(final_result);
     }
 
@@ -2161,6 +2195,31 @@ mod tests {
         assert!(MAX_LOOP_ITERATIONS.is_power_of_two());
     }
 
+    #[test]
+    fn isolated_iteration_continues_typed_application_failure() {
+        let encoded = encode_subtree_application_failure("division by zero");
+        assert_eq!(
+            encoded,
+            r#"{"pg_durable_subtree_failure":"application","message":"division by zero"}"#
+        );
+        let mut results = HashMap::new();
+        assert_eq!(
+            classify_isolated_body_result(Err(encoded), &mut results).unwrap(),
+            FailureIsolatedBodyOutcome::ApplicationFailed
+        );
+    }
+
+    #[test]
+    fn isolated_iteration_propagates_unrecognized_runtime_failure() {
+        let runtime_error =
+            "sub-orchestration instance id 'child' already exists and is terminal".to_string();
+        let mut results = HashMap::new();
+        match classify_isolated_body_result(Err(runtime_error.clone()), &mut results) {
+            Err(NodeError::Failure(error)) => assert_eq!(error, runtime_error),
+            other => panic!("expected runtime failure, got {other:?}"),
+        }
+    }
+
     #[test]
     fn subtree_instance_ids_are_stable_and_generation_scoped() {
         assert_eq!(
@@ -2174,25 +2233,57 @@ mod tests {
     }
 
     #[test]
-    fn isolated_iteration_continues_typed_application_failure() {
-        let encoded = encode_subtree_application_failure("division by zero");
+    fn isolated_body_success_returns_result_and_merges_results() {
+        let raw = envelope_json(
+            Some("Normal"),
+            "42",
+            serde_json::json!({"body_result": "stored"}),
+        );
+        let mut results = HashMap::new();
+
         assert_eq!(
-            encoded,
-            r#"{"pg_durable_subtree_failure":"application","message":"division by zero"}"#
+            classify_isolated_body_result(Ok(raw), &mut results).unwrap(),
+            FailureIsolatedBodyOutcome::Succeeded("42".to_string())
         );
         assert_eq!(
-            classify_isolated_iteration_failure(encoded),
-            IsolatedIterationFailure::Continue("division by zero".to_string())
+            results.get("body_result").map(String::as_str),
+            Some("stored")
         );
     }
 
     #[test]
-    fn isolated_iteration_propagates_unrecognized_runtime_failure() {
-        let runtime_error =
-            "sub-orchestration instance id 'child' already exists and is terminal".to_string();
+    fn isolated_body_consumes_typed_application_failure() {
+        let mut results = HashMap::new();
+
+        assert_eq!(
+            classify_isolated_body_result(
+                Err(encode_subtree_application_failure("boom")),
+                &mut results,
+            )
+            .unwrap(),
+            FailureIsolatedBodyOutcome::ApplicationFailed
+        );
+    }
+
+    #[test]
+    fn isolated_body_propagates_unrecognized_runtime_failure() {
+        let mut results = HashMap::new();
+
+        match classify_isolated_body_result(Err("instance id collision".to_string()), &mut results)
+        {
+            Err(NodeError::Failure(error)) => assert_eq!(error, "instance id collision"),
+            other => panic!("expected runtime failure, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn isolated_body_preserves_break_value() {
+        let raw = envelope_json(Some("Break"), r#"{"status":"done"}"#, serde_json::json!({}));
+        let mut results = HashMap::new();
+
         assert_eq!(
-            classify_isolated_iteration_failure(runtime_error.clone()),
-            IsolatedIterationFailure::Propagate(runtime_error)
+            classify_isolated_body_result(Ok(raw), &mut results).unwrap(),
+            FailureIsolatedBodyOutcome::Break(r#"{"status":"done"}"#.to_string())
         );
     }
 

tests/e2e/sql/70_loop_continue_on_failure.sql

@@ -183,6 +183,126 @@ END $$;
 DROP TABLE _nested_continue_instance;
 DROP TABLE test_nested_continue;
 
+DROP FUNCTION IF EXISTS test_conditional_continue_body();
+DROP FUNCTION IF EXISTS test_conditional_continue_condition();
+DROP FUNCTION IF EXISTS test_conditional_continue_failing_condition();
+DROP SEQUENCE IF EXISTS test_conditional_continue_attempt_seq;
+DROP TABLE IF EXISTS test_conditional_continue_state;
+
+CREATE TABLE test_conditional_continue_state (
+    body_attempts INT NOT NULL DEFAULT 0,
+    condition_checks INT NOT NULL DEFAULT 0
+);
+INSERT INTO test_conditional_continue_state DEFAULT VALUES;
+
+-- Sequence values are not rolled back when the first function call raises, allowing the
+-- second successful call to persist the total attempt count in the state table.
+CREATE SEQUENCE test_conditional_continue_attempt_seq;
+
+CREATE FUNCTION test_conditional_continue_body() RETURNS INT
+LANGUAGE plpgsql AS $$
+DECLARE
+    attempts INT;
+BEGIN
+    attempts := nextval('test_conditional_continue_attempt_seq');
+    UPDATE test_conditional_continue_state
+    SET body_attempts = attempts;
+
+    IF attempts = 1 THEN
+        RAISE EXCEPTION 'transient body failure';
+    END IF;
+    RETURN attempts;
+END
+$$;
+
+CREATE FUNCTION test_conditional_continue_condition() RETURNS BOOLEAN
+LANGUAGE plpgsql AS $$
+BEGIN
+    UPDATE test_conditional_continue_state
+    SET condition_checks = condition_checks + 1;
+    RETURN false;
+END
+$$;
+
+CREATE FUNCTION test_conditional_continue_failing_condition() RETURNS BOOLEAN
+LANGUAGE plpgsql AS $$
+BEGIN
+    RAISE EXCEPTION 'fatal condition failure';
+END
+$$;
+
+CREATE TEMP TABLE _conditional_continue_instances (
+    scenario TEXT PRIMARY KEY,
+    instance_id TEXT NOT NULL
+);
+
+INSERT INTO _conditional_continue_instances
+SELECT 'body-recovery',
+       df.start(
+           df.loop(
+               'SELECT test_conditional_continue_body()',
+               'SELECT test_conditional_continue_condition()',
+               continue_on_failure => true
+           ),
+           'test-conditional-loop-continues-after-body-failure'
+       );
+
+INSERT INTO _conditional_continue_instances
+SELECT 'condition-failure',
+       df.start(
+           df.loop(
+               'SELECT 42',
+               'SELECT test_conditional_continue_failing_condition()',
+               continue_on_failure => true
+           ),
+           'test-conditional-loop-condition-failure-is-fatal'
+       );
+
+DO $$
+DECLARE
+    recovered_status TEXT;
+    condition_failure_status TEXT;
+    body_attempts INT;
+    condition_checks INT;
+BEGIN
+    SELECT df.await_instance(instance_id, 30)
+    INTO recovered_status
+    FROM _conditional_continue_instances
+    WHERE scenario = 'body-recovery';
+
+    SELECT df.await_instance(instance_id, 30)
+    INTO condition_failure_status
+    FROM _conditional_continue_instances
+    WHERE scenario = 'condition-failure';
+
+    SELECT s.body_attempts, s.condition_checks
+    INTO body_attempts, condition_checks
+    FROM test_conditional_continue_state s;
+
+    IF recovered_status IS DISTINCT FROM 'completed' THEN
+        RAISE EXCEPTION
+            'TEST FAILED [conditional continue]: expected completed, got %',
+            recovered_status;
+    END IF;
+    IF body_attempts <> 2 OR condition_checks <> 1 THEN
+        RAISE EXCEPTION
+            'TEST FAILED [conditional continue]: expected body/condition = 2/1, got %/%',
+            body_attempts, condition_checks;
+    END IF;
+    IF condition_failure_status IS DISTINCT FROM 'failed' THEN
+        RAISE EXCEPTION
+            'TEST FAILED [condition failure]: expected failed, got %',
+            condition_failure_status;
+    END IF;
+END $$;
+
+DROP TABLE _conditional_continue_instances;
+DROP FUNCTION test_conditional_continue_body();
+DROP FUNCTION test_conditional_continue_condition();
+DROP FUNCTION test_conditional_continue_failing_condition();
+DROP SEQUENCE test_conditional_continue_attempt_seq;
+DROP TABLE test_conditional_continue_state;
+
 CREATE TEMP TABLE _cancel_continue_instance (instance_id TEXT);
 
 INSERT INTO _cancel_continue_instance(instance_id)