fix(worker): preserve recovery signals and stacked lock-release paths

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 974213a5-9101-41b4-9bde-aa86a72f25e1

Author: andystaples

CHANGELOG.md

@@ -17,8 +17,8 @@ now takes precedence over hypothetical earlier storage failures.
 - Orchestrations that exceed the configured external payload storage limit now
 fail without retrying and release their entity locks, allowing other callers to
 progress. Transient payload-store failures remain retryable.
-- Unexpected non-cancellation orchestration exits, including generator teardown
-failures, now fail with a clear reason instead of leaving an accepted dispatch
+- Previously unhandled non-cancellation `BaseException` exits, including generator
+teardown failures, now fail with a clear reason instead of leaving an accepted dispatch
 unfinished. Cancellation and process-control signals retain their semantics.
 The gRPC worker abandons work that exits before an outcome is sent and resets
 the originating work-item stream after an uncertain completion or abandonment

azure-functions-durable/CHANGELOG.md

@@ -24,8 +24,8 @@ FIXED
 - With the corresponding core `durabletask` SDK fix, oversized messages fail
 large-payload validation before any uploads or payload replacements. Size
 errors now take precedence over hypothetical earlier storage failures.
-- With the corresponding core `durabletask` SDK fix, unexpected non-cancellation
-orchestration execution or generator teardown exits now produce a failed response
+- With the corresponding core `durabletask` SDK fix, previously unhandled non-cancellation
+`BaseException` exits from orchestration execution or teardown now produce a failed response
 with the exception's type, message, and stack trace instead of losing the response.
 - With the corresponding core `durabletask` SDK fix, orchestrator generators
 are closed at the end of each dispatch, including while waiting for durable work.

durabletask-azuremanaged/CHANGELOG.md

@@ -15,8 +15,8 @@ errors now take precedence over hypothetical earlier storage failures.
 - With the corresponding core `durabletask` SDK fix, orchestrations that exceed
 the configured external payload storage limit fail without retrying and release
 their entity locks. Transient payload-store failures remain retryable.
-- With the corresponding core `durabletask` SDK fix, unexpected non-cancellation
-orchestration exits now fail with a clear reason. Cancellation and process-control
+- With the corresponding core `durabletask` SDK fix, previously unhandled
+non-cancellation `BaseException` exits now fail with a clear reason. Cancellation and process-control
 exits recover dispatch ownership, and uncertain completion or abandonment RPCs
 reset the originating work-item stream rather than retaining unfinished work.
 - With the corresponding core `durabletask` SDK fix, orchestrator generators

durabletask/worker.py

@@ -1014,7 +1014,10 @@ def stream_reader(response_stream: Any = response_stream) -> None:
                                     self._execute_orchestrator,
                                     _on_outcome_uncertain=response_stream.cancel,
                                 ),
-                                self._cancel_orchestrator,
+                                partial(
+                                    self._cancel_orchestrator,
+                                    _on_outcome_uncertain=response_stream.cancel,
+                                ),
                                 work_item.orchestratorRequest,
                                 stub,
                                 work_item.completionToken,
@@ -1173,6 +1176,7 @@ def _execute_orchestrator(
             stub.CompleteOrchestratorTask(res)
         except BaseException as error:
             recovery_error: BaseException | None = None
+            outcome_uncertain = outcome_attempted
             self._logger.error(
                 f"Orchestrator dispatch for '{req.instanceId}' exited without a confirmed outcome.",
                 exc_info=True,
@@ -1182,33 +1186,35 @@ def _execute_orchestrator(
                     self._cancel_orchestrator(req, stub, completionToken)
                 except BaseException as abandonment_error:
                     recovery_error = abandonment_error
-                    outcome_attempted = True
-                    self._logger.error(
-                        f"Failed to abandon orchestrator dispatch for '{req.instanceId}'.",
-                        exc_info=True,
-                    )
+                    outcome_uncertain = True
             # A failed RPC may already have been accepted. Do not retry it or
             # send a conflicting outcome; release ownership through its stream.
-            if outcome_attempted and _on_outcome_uncertain is not None:
-                self._logger.error(
-                    f"Resetting the owning work-item stream for '{req.instanceId}' after an uncertain outcome."
-                )
+            if outcome_uncertain and _on_outcome_uncertain is not None:
                 try:
-                    _on_outcome_uncertain()
+                    self._reset_orchestrator_stream(req.instanceId, _on_outcome_uncertain)
                 except BaseException as reset_error:
                     if recovery_error is None or (
                         isinstance(recovery_error, Exception) and not isinstance(reset_error, Exception)
                     ):
                         recovery_error = reset_error
-                    self._logger.error(
-                        f"Failed to reset the owning work-item stream for '{req.instanceId}'.",
-                        exc_info=True,
-                    )
             if not isinstance(error, Exception):
                 raise
             if recovery_error is not None and not isinstance(recovery_error, Exception):
                 raise recovery_error
 
+    def _reset_orchestrator_stream(self, instance_id: str, reset: Callable[[], Any]) -> None:
+        self._logger.error(
+            f"Resetting the owning work-item stream for '{instance_id}' after an uncertain outcome."
+        )
+        try:
+            reset()
+        except BaseException:
+            self._logger.error(
+                f"Failed to reset the owning work-item stream for '{instance_id}'.",
+                exc_info=True,
+            )
+            raise
+
     def _get_orchestrator_response(
             self,
             req: pb.OrchestratorRequest,
@@ -1351,12 +1357,27 @@ def _cancel_orchestrator(
             req: pb.OrchestratorRequest,
             stub: Any,
             completionToken: Any,
+            *,
+            _on_outcome_uncertain: Callable[[], Any] | None = None,
     ) -> None:
-        stub.AbandonTaskOrchestratorWorkItem(
-            pb.AbandonOrchestrationTaskRequest(
-                completionToken=completionToken
+        try:
+            stub.AbandonTaskOrchestratorWorkItem(
+                pb.AbandonOrchestrationTaskRequest(
+                    completionToken=completionToken
+                )
             )
-        )
+        except BaseException as error:
+            self._logger.error(
+                f"Failed to abandon orchestrator dispatch for '{req.instanceId}'.",
+                exc_info=True,
+            )
+            if _on_outcome_uncertain is not None:
+                try:
+                    self._reset_orchestrator_stream(req.instanceId, _on_outcome_uncertain)
+                except BaseException as reset_error:
+                    if isinstance(error, Exception) and not isinstance(reset_error, Exception):
+                        raise
+            raise
         self._logger.info(f"Cancelled orchestration task for invocation ID: {req.instanceId}")
 
     def _on_activity_execution_started(self, req: pb.ActivityRequest) -> None:

tests/azure-functions-durable/test_worker_compat.py

@@ -56,6 +56,34 @@ def test_worker_uses_propagate_only_tracing():
     assert worker.emit_trace_spans is False
 
 
+@pytest.mark.parametrize("use_async", [False, True])
+async def test_worker_returns_failed_response_for_unexpected_teardown(monkeypatch, use_async):
+    monkeypatch.setattr(payloads, "_payload_store", None)
+
+    class TeardownExit(BaseException):
+        pass
+
+    def orchestrator(context):
+        try:
+            yield context.call_activity("work")
+        finally:
+            raise TeardownExit("deterministic teardown failure")
+
+    worker = DurableFunctionsWorker()
+    request = _encode_orchestrator_request("teardown")
+    if use_async:
+        encoded = await worker.execute_orchestration_request_async(orchestrator, request)
+    else:
+        encoded = worker.execute_orchestration_request(orchestrator, request)
+    response = _decode_orchestrator_response(encoded)
+    completion = _get_completion_action(response)
+
+    assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
+    assert completion.failureDetails.errorType.endswith(".TeardownExit")
+    assert completion.failureDetails.errorMessage == "deterministic teardown failure"
+    assert len(response.actions) == 1
+
+
 def test_worker_created_before_configuration_hydrates_and_externalizes(monkeypatch, payload_store_factory):
     monkeypatch.setattr(payloads, "_payload_store", None)
     worker = DurableFunctionsWorker()

tests/durabletask-azuremanaged/test_azuremanaged_grpc_resiliency.py

@@ -1,10 +1,13 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 from collections.abc import Mapping
 from typing import Any
 
+from durabletask import task
+from durabletask.internal import helpers
+from durabletask.internal import orchestrator_service_pb2 as pb
 from durabletask.azuremanaged.client import (
     AsyncDurableTaskSchedulerClient,
     DurableTaskSchedulerClient,
@@ -40,6 +43,34 @@ def test_dts_worker_passes_resiliency_options_to_base_worker():
     assert mock_init.call_args.kwargs["resiliency_options"] is resiliency
 
 
+def test_dts_worker_completes_unexpected_teardown_failure():
+    class TeardownExit(BaseException):
+        pass
+
+    def orchestrator(ctx: task.OrchestrationContext, _: Any):
+        try:
+            yield ctx.call_activity("work")
+        finally:
+            raise TeardownExit("deterministic teardown failure")
+
+    with DurableTaskSchedulerWorker(
+            host_address="localhost:4001", taskhub="hub", token_credential=None) as worker:
+        name = worker.add_orchestrator(orchestrator)
+        request = pb.OrchestratorRequest(instanceId="instance", newEvents=[
+            helpers.new_orchestrator_started_event(),
+            helpers.new_execution_started_event(name, "instance"),
+        ])
+        stub = Mock()
+        worker._execute_orchestrator(request, stub, "original-token")
+
+    stub.CompleteOrchestratorTask.assert_called_once()
+    stub.AbandonTaskOrchestratorWorkItem.assert_not_called()
+    response = stub.CompleteOrchestratorTask.call_args.args[0]
+    assert response.completionToken == "original-token"
+    assert response.actions[0].completeOrchestration.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
+    assert response.actions[0].completeOrchestration.failureDetails.errorMessage == "deterministic teardown failure"
+
+
 def test_dts_worker_passes_exception_properties_provider_to_base_worker():
     class PropertiesProvider:
         def get_exception_properties(self, exception: Exception) -> Mapping[str, Any] | None:

tests/durabletask/test_orchestration_dispatch_recovery.py

@@ -96,16 +96,18 @@ def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Generator[task.Task[
     assert "secondary cleanup failure" not in completion.failureDetails.errorMessage
 
 
-@pytest.mark.parametrize("release_before_wait", [False, True])
-def test_unexpected_teardown_failure_preserves_entity_lock_releases(release_before_wait: bool):
+@pytest.mark.parametrize("wait", ["locked-event", "released-event", "entity-call"])
+def test_unexpected_teardown_failure_preserves_entity_lock_releases(wait: str):
     entity_id = EntityInstanceId("Counter", "key")
 
     def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Generator[task.Task[Any], Any, None]:
         acquired = False
         try:
             with (yield ctx.lock_entities([entity_id])):
                 acquired = True
-                if not release_before_wait:
+                if wait == "entity-call":
+                    yield ctx.call_entity(entity_id, "get")
+                elif wait == "locked-event":
                     yield ctx.wait_for_external_event("ready")
             yield ctx.wait_for_external_event("ready")
         finally:

tests/durabletask/test_worker_resiliency.py

@@ -747,7 +747,8 @@ def create_stub(channel):
 
 
 @pytest.mark.asyncio
-async def test_orchestrator_uncertain_outcome_resets_only_originating_stream(monkeypatch):
+@pytest.mark.parametrize("operation", ["complete", "abandon"])
+async def test_orchestrator_uncertain_outcome_resets_only_originating_stream(monkeypatch, operation):
     channel = MagicMock()
     worker = TaskHubGrpcWorker(channel=channel)
     manager = DummyWorkerManager()
@@ -762,6 +763,7 @@ async def test_orchestrator_uncertain_outcome_resets_only_originating_stream(mon
     first_stub = MagicMock()
     first_stub.GetWorkItems.return_value = first_stream
     first_stub.CompleteOrchestratorTask.side_effect = RuntimeError("completion outcome unknown")
+    first_stub.AbandonTaskOrchestratorWorkItem.side_effect = RuntimeError("abandonment outcome unknown")
     next_stub = MagicMock()
     next_stub.GetWorkItems.return_value = next_stream
     stop_stub = MagicMock()
@@ -775,12 +777,18 @@ async def test_orchestrator_uncertain_outcome_resets_only_originating_stream(mon
     next_stream.cancelled = False
     worker._response_stream = next_stream
     _, submission = manager.submissions[0]
-    handler, _, request, original_stub, token = submission
-    handler(request, original_stub, token)
-
-    first_stub.CompleteOrchestratorTask.assert_called_once()
-    assert first_stub.CompleteOrchestratorTask.call_args.args[0].completionToken == "original-token"
-    first_stub.AbandonTaskOrchestratorWorkItem.assert_not_called()
+    handler, cancel, request, original_stub, token = submission
+    if operation == "complete":
+        handler(request, original_stub, token)
+        first_stub.CompleteOrchestratorTask.assert_called_once()
+        assert first_stub.CompleteOrchestratorTask.call_args.args[0].completionToken == "original-token"
+        first_stub.AbandonTaskOrchestratorWorkItem.assert_not_called()
+    else:
+        with pytest.raises(RuntimeError, match="abandonment outcome unknown"):
+            cancel(request, original_stub, token)
+        first_stub.CompleteOrchestratorTask.assert_not_called()
+        first_stub.AbandonTaskOrchestratorWorkItem.assert_called_once()
+        assert first_stub.AbandonTaskOrchestratorWorkItem.call_args.args[0].completionToken == "original-token"
     assert first_stream.cancelled
     assert not next_stream.cancelled
     channel.close.assert_not_called()