when_all waits for all tasks before surfacing the first failure (#1166) (#1175)

Previously a WhenAllTask completed as soon as one child failed, while
sibling tasks were still running and their later completions were
dropped. Now the composite stages the first failure and keeps waiting;
it completes only once every child has finished, then fails with the
first recorded exception (or returns the ordered results if none
failed).



(cherry picked from commit 47f0f23)

Signed-off-by: Albert Callarisa <albert@diagrid.io>
Signed-off-by: dapr-bot <dapr-bot@users.noreply.github.com>
Co-authored-by: Albert Callarisa <albert@diagrid.io>
Co-authored-by: Yaron Schneider <schneider.yaron@live.com>

Author: 3 people

ext/dapr-ext-workflow/dapr/ext/workflow/_durabletask/task.py

@@ -405,11 +405,17 @@ def on_child_completed(self, task: Task[T]):
 
 
 class WhenAllTask(CompositeTask[list[T]]):
-    """A task that completes when all of its child tasks complete."""
+    """A task that completes once all of its child tasks complete.
+
+    If any child fails, the composite still waits for every child to finish
+    and then fails with the first failure.
+    """
 
     def __init__(self, tasks: list[Task[T]]):
-        # CompositeTask.__init__ already counted pre-completed children; do not reset
+        # CompositeTask.__init__ replays pre-completed children into on_child_completed,
+        # so _pending_exception must exist before it runs; likewise do not reset
         # _completed_tasks after it, or a deferred when_all over them hangs forever.
+        self._pending_exception: Optional[Exception] = None
         super().__init__(tasks)
         # If there are no child tasks, this composite should complete immediately
         if len(self._tasks) == 0:
@@ -423,20 +429,22 @@ def pending_tasks(self) -> int:
 
     def on_child_completed(self, task: Task[T]):
         if self.is_complete:
-            # Already completed (e.g. a previous child failed), ignore late arrivals
             return
         self._completed_tasks += 1
-        if task.is_failed and self._exception is None:
-            self._exception = task.get_exception()
-            self._is_complete = True
-            if self._parent is not None:
-                self._parent.on_child_completed(self)
-        elif self._completed_tasks == len(self._tasks):
+        if task.is_failed and self._pending_exception is None:
+            # Stage the first failure without exposing it via _exception yet:
+            # that would make is_failed True while is_complete is still False.
+            self._pending_exception = task.get_exception()
+        if self._completed_tasks < len(self._tasks):
+            return
+        self._is_complete = True
+        if self._pending_exception is not None:
+            self._exception = self._pending_exception
+        else:
             # The order of the result MUST match the order of the tasks provided to the constructor.
             self._result = [task.get_result() for task in self._tasks]
-            self._is_complete = True
-            if self._parent is not None:
-                self._parent.on_child_completed(self)
+        if self._parent is not None:
+            self._parent.on_child_completed(self)
 
     def get_completed_tasks(self) -> int:
         return self._completed_tasks
@@ -629,7 +637,8 @@ def on_child_completed(self, completed_task: Task):
 
 
 def when_all(tasks: list[Task[T]]) -> WhenAllTask[T]:
-    """Returns a task that completes when all of the provided tasks complete or when one of the tasks fail."""
+    """Returns a task that completes once all of the provided tasks complete,
+    surfacing the first failure (if any) only after every task has finished."""
     return WhenAllTask(tasks)
 
 

ext/dapr-ext-workflow/dapr/ext/workflow/dapr_workflow_context.py

@@ -172,8 +172,8 @@ def is_patched(self, patch_name: str) -> bool:
 
 
 def when_all(tasks: List[task.Task[T]]) -> task.WhenAllTask[T]:
-    """Returns a task that completes when all of the provided tasks complete or when one of the
-    tasks fail."""
+    """Returns a task that completes once all of the provided tasks complete,
+    surfacing the first failure (if any) only after every task has finished."""
     return task.when_all(tasks)
 
 

ext/dapr-ext-workflow/tests/durabletask/test_orchestration_executor.py

@@ -1204,16 +1204,29 @@ def orchestrator(ctx: task.OrchestrationContext, _):
         )
 
     # 5 of the tasks complete successfully, 1 fails, and 4 are still running.
-    # The expectation is that the orchestration will fail immediately.
-    new_events = []
+    # when_all waits for every child task to complete before surfacing the
+    # failure, so the orchestration is expected to still be running with zero
+    # new actions.
+    ex = Exception('Kah-BOOOOM!!!')
+    partial_events = []
     for i in range(5):
+        partial_events.append(
+            helpers.new_task_completed_event(i + 1, encoded_output=print_int(None, i))
+        )
+    partial_events.append(helpers.new_task_failed_event(6, ex))
+
+    executor = worker._OrchestrationExecutor(registry, TEST_LOGGER)
+    result = executor.execute(TEST_INSTANCE_ID, old_events, partial_events)
+    assert len(result.actions) == 0
+
+    # Once the remaining 4 tasks also complete, the orchestration fails and
+    # surfaces the first task failure.
+    new_events = list(partial_events)
+    for i in range(6, 10):
         new_events.append(
             helpers.new_task_completed_event(i + 1, encoded_output=print_int(None, i))
         )
-    ex = Exception('Kah-BOOOOM!!!')
-    new_events.append(helpers.new_task_failed_event(6, ex))
 
-    # Now test with the full set of new events. We expect the orchestration to complete.
     executor = worker._OrchestrationExecutor(registry, TEST_LOGGER)
     result = executor.execute(TEST_INSTANCE_ID, old_events, new_events)
     actions = result.actions
@@ -1277,13 +1290,8 @@ def orchestrator(ctx: task.OrchestrationContext, _):
 
 
 def test_when_all_success_after_failure_does_not_crash():
-    """Tests that task completions arriving after when_all already failed
-    do not crash the orchestration.
-
-    This is a regression test: previously a ValueError was raised when
-    a successful task completed after the WhenAllTask was already marked
-    complete due to a prior child failure.
-    """
+    """Tests that a success arriving after a failure completes the set and
+    surfaces the failure to the orchestrator where it can be caught."""
 
     def dummy_activity(ctx, _):
         pass

ext/dapr-ext-workflow/tests/durabletask/test_task.py

@@ -267,34 +267,75 @@ def test_when_all_failure_after_success_still_reports_failure():
         all_task.get_result()
 
 
-def test_when_all_failure_before_success_still_reports_failure():
-    """When a child fails before the other children succeed,
-    the WhenAllTask must complete with the failure immediately."""
+def test_when_all_defers_failure_until_all_children_complete():
+    """After a child failure the WhenAllTask keeps waiting for the remaining
+    children; the first failure is surfaced once every child has completed."""
     c1 = task.CompletableTask()
     c2 = task.CompletableTask()
+    c3 = task.CompletableTask()
 
-    all_task = task.when_all([c1, c2])
+    all_task = task.when_all([c1, c2, c3])
 
-    # c1 fails first
     c1.fail('activity failed', _make_failure_details('activity failed'))
 
+    assert not all_task.is_complete
+    assert not all_task.is_failed
+    assert all_task.get_completed_tasks() == 1
+
+    c2.complete('two')
+
+    assert not all_task.is_complete
+    assert not all_task.is_failed
+    assert all_task.get_completed_tasks() == 2
+
+    c3.complete('three')
+
     assert all_task.is_complete
     assert all_task.is_failed
+    assert all_task.get_completed_tasks() == 3
     with pytest.raises(task.TaskFailedError):
         all_task.get_result()
 
-    # c2 succeeds after — must not raise ValueError
-    c2.complete('two')
 
-    # WhenAllTask should still be in the same failed state
+def test_when_all_surfaces_first_failure_when_multiple_children_fail():
+    """When several children fail, the WhenAllTask reports the first failure."""
+    c1 = task.CompletableTask()
+    c2 = task.CompletableTask()
+
+    all_task = task.when_all([c1, c2])
+
+    c1.fail('first error', _make_failure_details('first error'))
+    c2.fail('second error', _make_failure_details('second error'))
+
+    assert all_task.is_complete
+    assert all_task.is_failed
+    with pytest.raises(task.TaskFailedError, match='first error'):
+        all_task.get_result()
+
+
+def test_when_all_with_pre_failed_child_waits_for_remaining():
+    """A child that already failed before construction is staged, not surfaced,
+    until the remaining children complete."""
+    failed_child = task.CompletableTask()
+    failed_child.fail('activity failed', _make_failure_details('activity failed'))
+    pending_child = task.CompletableTask()
+
+    all_task = task.when_all([failed_child, pending_child])
+
+    assert not all_task.is_complete
+    assert not all_task.is_failed
+    assert all_task.get_completed_tasks() == 1
+
+    pending_child.complete('ok')
+
     assert all_task.is_complete
     assert all_task.is_failed
     with pytest.raises(task.TaskFailedError):
         all_task.get_result()
 
 
 def test_when_all_failure_propagates_to_parent():
-    """When a WhenAllTask fails due to a child failure,
+    """When a WhenAllTask fails after all children complete,
     it should notify its parent composite task."""
     c1 = task.CompletableTask()
     c2 = task.CompletableTask()
@@ -306,6 +347,12 @@ def test_when_all_failure_propagates_to_parent():
 
     c1.fail('activity failed', _make_failure_details('activity failed'))
 
+    # The failure is staged until c2 completes, so neither composite is done yet
+    assert not all_task.is_complete
+    assert not any_task.is_complete
+
+    c2.complete('two')
+
     assert all_task.is_complete
     assert all_task.is_failed
     # The parent WhenAnyTask should also have completed