Fix when_all to wait for all tasks before surfacing failures (#173)

* Fix when_all to wait for all tasks before surfacing failures

when_all previously failed fast on the first child task failure. It now waits for every child task to complete before surfacing the first failure, matching the semantics of .NET's Task.WhenAll.

* Fix is_failed guard by storing exception

* Constructor ordering to prevent AttributeError

* Update CHANGELOG.md

Author: andystaples

CHANGELOG.md

@@ -19,6 +19,7 @@ FIXED
 - Fixed the orchestration version being set from an unset `executionStarted.version` field. Presence is now checked with `HasField`, so `OrchestrationContext` no longer reports a version when none was provided.
 - Fixed orchestrations failing with `OrchestrationStateError` when a `genericEvent` history event was replayed (for example, the marker the Durable Functions extension appends when rewinding an orchestration). Such informational events are now ignored during replay, matching the .NET worker.
 - Fixed a lock-granted entity response over the legacy entity protocol raising while trying to deserialize an empty operation result. Lock-granted events no longer attempt result deserialization.
+- Fixed `task.when_all()` failing fast when one of its child tasks failed. It now waits for every child task to complete before surfacing the first failure, matching the semantics of .NET's `Task.WhenAll`.
 
 ## v1.7.2
 

durabletask/task.py

@@ -531,6 +531,8 @@ def get_exception(self) -> TaskFailedError:
 class CompositeTask(Task[T]):
     """A task that is composed of other tasks."""
     _tasks: list[Task[Any]]
+    _completed_tasks: int
+    _failed_tasks: int
 
     def __init__(self, tasks: list[Task[Any]]):
         super().__init__()
@@ -554,9 +556,14 @@ class WhenAllTask(CompositeTask[list[T]]):
     """A task that completes when all of its child tasks complete."""
 
     def __init__(self, tasks: list[Task[T]]):
+        # Initialize state that on_child_completed() reads BEFORE invoking the
+        # base constructor: CompositeTask.__init__ calls on_child_completed()
+        # for any children that are already complete, so `_pending_exception`
+        # must exist first. The base constructor also initializes
+        # `_completed_tasks`/`_failed_tasks` to 0 and then accounts for
+        # pre-completed children, so they must not be reset afterwards.
+        self._pending_exception: TaskFailedError | None = None
         super().__init__(cast(list[Task[Any]], tasks))
-        self._completed_tasks = 0
-        self._failed_tasks = 0
 
     @property
     def pending_tasks(self) -> int:
@@ -567,13 +574,26 @@ def on_child_completed(self, task: Task[Any]) -> None:
         if self.is_complete:
             raise ValueError('The task has already completed.')
         self._completed_tasks += 1
-        if task.is_failed and self._exception is None:
-            self._exception = task.get_exception()
-            self._is_complete = True
+        if task.is_failed:
+            self._failed_tasks += 1
+            if self._pending_exception is None:
+                # Stage the first failure but do NOT expose it via `_exception`
+                # yet. Exposing it now would make `is_failed` return True while
+                # `is_complete` is still False, diverging from .NET's
+                # Task.WhenAll (which does not fault until all children finish).
+                self._pending_exception = task.get_exception()
         if self._completed_tasks == len(self._tasks):
-            # The order of the result MUST match the order of the tasks provided to the constructor.
-            self._result = [child.get_result() for child in self._tasks]
+            # Only complete once every child task has completed. This matches the
+            # semantics of .NET's Task.WhenAll: the composite task waits for all
+            # children to finish and, if any failed, surfaces the first failure
+            # rather than failing fast on the first error.
             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 = [child.get_result() for child in self._tasks]
 
     def get_completed_tasks(self) -> int:
         return self._completed_tasks

tests/durabletask/test_orchestration_executor.py

@@ -1615,15 +1615,27 @@ def orchestrator(ctx: task.OrchestrationContext, _):
             i + 1, activity_name, encoded_input=str(i)))
 
     # 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 must NOT fail fast: it waits for every child task to complete
+    # before surfacing a failure (matching .NET's Task.WhenAll semantics), 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, JsonDataConverter())
+    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, JsonDataConverter())
     result = executor.execute(TEST_INSTANCE_ID, old_events, new_events)
     actions = result.actions
@@ -1634,6 +1646,67 @@ def orchestrator(ctx: task.OrchestrationContext, _):
     assert str(ex) in complete_action.failureDetails.errorMessage
 
 
+def test_when_all_defers_failure_until_all_children_complete():
+    """when_all must not report is_failed until every child task completes.
+
+    This mirrors .NET's Task.WhenAll, where the returned task does not fault
+    (failed => complete) until all children have finished. A composite that
+    exposed is_failed while still incomplete would surprise consumers that
+    assume failed implies complete.
+    """
+    t1 = task.CompletableTask()
+    t2 = task.CompletableTask()
+    t3 = task.CompletableTask()
+    when_all = task.when_all([t1, t2, t3])
+
+    assert not when_all.is_complete
+    assert not when_all.is_failed
+
+    # First child fails: the composite must stay pending and must NOT report
+    # failure yet.
+    t1.fail("boom", Exception("boom"))
+    assert not when_all.is_complete
+    assert not when_all.is_failed
+
+    # A later child completes successfully: still pending, still not failed.
+    t2.complete("ok")
+    assert not when_all.is_complete
+    assert not when_all.is_failed
+
+    # Once the final child completes, the composite completes and only then
+    # surfaces the first failure.
+    t3.complete("ok")
+    assert when_all.is_complete
+    assert when_all.is_failed
+
+
+def test_when_all_handles_pre_completed_children():
+    """when_all must account for children that are already complete/failed.
+
+    CompositeTask.__init__ invokes on_child_completed() for children that are
+    already complete, so when_all must be constructible from an already-failed
+    child without raising (regression test for AttributeError on
+    `_pending_exception`).
+    """
+    # An already-failed child plus a still-pending child.
+    failed = task.CompletableTask()
+    failed.fail("boom", Exception("boom"))
+    pending = task.CompletableTask()
+
+    when_all = task.when_all([failed, pending])
+
+    # The pre-completed failure is accounted for but not yet surfaced.
+    assert when_all.get_completed_tasks() == 1
+    assert not when_all.is_complete
+    assert not when_all.is_failed
+
+    # Completing the remaining child finishes the composite and surfaces the
+    # first failure.
+    pending.complete("ok")
+    assert when_all.is_complete
+    assert when_all.is_failed
+
+
 def test_when_any():
     """Tests that a when_any pattern works correctly"""
     def hello(_, name: str):