Reconnect workers after peer stream cancellation

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 282b8096-753c-41be-928d-2b53604825d3

Author: torosent

CHANGELOG.md

@@ -20,6 +20,9 @@ FIXED
 - 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`.
+- Fixed workers stopping permanently when the server reset the `GetWorkItems`
+  stream. Workers now reconnect after a peer cancellation and still exit normally
+  when stopped.
 
 ## v1.7.2
 

durabletask-azuremanaged/CHANGELOG.md

@@ -8,6 +8,8 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 ## Unreleased
 
 - Added `rewind_orchestration()` to `DurableTaskSchedulerClient` and `AsyncDurableTaskSchedulerClient` (inherited from the base clients) to rewind a failed orchestration instance to its last known good state.
+- Fixed Durable Task Scheduler workers stopping permanently when the service reset
+  the `GetWorkItems` stream.
 
 ## v1.7.2
 

durabletask/worker.py

@@ -1063,7 +1063,7 @@ def stream_reader() -> None:
                     invalidate_connection(recreate_channel=recreate_channel)
                 error_details = str(rpc_error)
 
-                if error_code == grpc.StatusCode.CANCELLED:
+                if error_code == grpc.StatusCode.CANCELLED and self._shutdown.is_set():
                     self._logger.info(f"Disconnected from {self._host_address}")
                     break
                 elif error_code == grpc.StatusCode.UNAVAILABLE:

tests/durabletask/test_worker_resiliency.py

@@ -34,6 +34,14 @@ def __str__(self):
         return self._details
 
 
+def _cancel_after_shutdown(worker):
+    def cancel(*_args, **_kwargs):
+        worker._shutdown.set()
+        raise FakeRpcError(grpc.StatusCode.CANCELLED, "stop")
+
+    return cancel
+
+
 class FakeResponseStream:
     def __init__(self, items=(), error: grpc.RpcError | None = None):
         self._items = list(items)
@@ -225,7 +233,7 @@ async def test_worker_applies_configured_hello_timeout(monkeypatch):
     monkeypatch.setattr(worker._shutdown, "wait", lambda timeout: False)
 
     stub = MagicMock()
-    stub.GetWorkItems.side_effect = FakeRpcError(grpc.StatusCode.CANCELLED, "stop")
+    stub.GetWorkItems.side_effect = _cancel_after_shutdown(worker)
 
     monkeypatch.setattr("durabletask.worker.shared.get_grpc_channel", lambda *args, **kwargs: MagicMock())
     monkeypatch.setattr("durabletask.worker.stubs.TaskHubSidecarServiceStub", lambda channel: stub)
@@ -235,6 +243,32 @@ async def test_worker_applies_configured_hello_timeout(monkeypatch):
     assert stub.Hello.call_args.kwargs["timeout"] == 12.5
 
 
+@pytest.mark.asyncio
+async def test_worker_reconnects_after_peer_cancelled_stream(monkeypatch):
+    worker = TaskHubGrpcWorker(channel=MagicMock())
+    worker._async_worker_manager = DummyWorkerManager()
+    monkeypatch.setattr(worker._shutdown, "wait", lambda timeout: False)
+
+    first_stub = MagicMock()
+    first_stub.GetWorkItems.side_effect = FakeRpcError(
+        grpc.StatusCode.CANCELLED,
+        "peer reset",
+    )
+    second_stub = MagicMock()
+    second_stub.GetWorkItems.side_effect = _cancel_after_shutdown(worker)
+    stub_factory = MagicMock(side_effect=[first_stub, second_stub])
+
+    monkeypatch.setattr(
+        "durabletask.worker.stubs.TaskHubSidecarServiceStub",
+        stub_factory,
+    )
+
+    await worker._async_run_loop()
+
+    first_stub.GetWorkItems.assert_called_once()
+    second_stub.GetWorkItems.assert_called_once()
+
+
 @pytest.mark.asyncio
 async def test_worker_does_not_recreate_sdk_owned_channel_for_non_transport_setup_errors(monkeypatch):
     worker = TaskHubGrpcWorker(
@@ -254,7 +288,7 @@ def get_grpc_channel(*args, **kwargs):
     first_stub.Hello.side_effect = RuntimeError("boom")
 
     second_stub = MagicMock()
-    second_stub.GetWorkItems.side_effect = FakeRpcError(grpc.StatusCode.CANCELLED, "stop")
+    second_stub.GetWorkItems.side_effect = _cancel_after_shutdown(worker)
 
     stubs = [first_stub, second_stub]
 
@@ -293,10 +327,7 @@ def get_grpc_channel(*args, **kwargs):
             grpc.StatusCode.UNAVAILABLE,
             "second transport failure",
         )))),
-        MagicMock(GetWorkItems=MagicMock(side_effect=FakeRpcError(
-            grpc.StatusCode.CANCELLED,
-            "stop",
-        ))),
+        MagicMock(GetWorkItems=MagicMock(side_effect=_cancel_after_shutdown(worker))),
     ]
     stub_channels = []
 
@@ -352,10 +383,7 @@ def get_grpc_channel(*args, **kwargs):
     stub_channels = []
     stubs = [
         MagicMock(GetWorkItems=MagicMock(return_value=blocking_stream)),
-        MagicMock(GetWorkItems=MagicMock(side_effect=FakeRpcError(
-            grpc.StatusCode.CANCELLED,
-            "stop",
-        ))),
+        MagicMock(GetWorkItems=MagicMock(side_effect=_cancel_after_shutdown(worker))),
     ]
 
     def create_stub(channel):
@@ -393,10 +421,7 @@ def get_grpc_channel(*args, **kwargs):
     stub_channels = []
     stubs = [
         MagicMock(GetWorkItems=MagicMock(return_value=FakeResponseStream())),
-        MagicMock(GetWorkItems=MagicMock(side_effect=FakeRpcError(
-            grpc.StatusCode.CANCELLED,
-            "stop",
-        ))),
+        MagicMock(GetWorkItems=MagicMock(side_effect=_cancel_after_shutdown(worker))),
     ]
 
     def create_stub(channel):
@@ -440,10 +465,7 @@ def complete_activity(response):
     first_stub.CompleteActivityTask.side_effect = complete_activity
 
     second_stub = MagicMock()
-    second_stub.GetWorkItems.side_effect = FakeRpcError(
-        grpc.StatusCode.CANCELLED,
-        "stop",
-    )
+    second_stub.GetWorkItems.side_effect = _cancel_after_shutdown(worker)
 
     stubs = [first_stub, second_stub]
     stub_channels = []
@@ -495,10 +517,7 @@ def get_grpc_channel(*args, **kwargs):
     first_stub.GetWorkItems.return_value = FakeResponseStream(items=[_make_activity_work_item()])
 
     second_stub = MagicMock()
-    second_stub.GetWorkItems.side_effect = FakeRpcError(
-        grpc.StatusCode.CANCELLED,
-        "stop",
-    )
+    second_stub.GetWorkItems.side_effect = _cancel_after_shutdown(worker)
 
     stubs = [first_stub, second_stub]
 
@@ -561,10 +580,7 @@ def complete_activity(response):
     first_stub.CompleteActivityTask.side_effect = complete_activity
 
     second_stub = MagicMock()
-    second_stub.GetWorkItems.side_effect = FakeRpcError(
-        grpc.StatusCode.CANCELLED,
-        "stop",
-    )
+    second_stub.GetWorkItems.side_effect = _cancel_after_shutdown(worker)
 
     stubs = [first_stub, second_stub]
     stub_channels = []
@@ -637,10 +653,7 @@ def complete_activity(response):
     first_stub.CompleteActivityTask.side_effect = complete_activity
 
     second_stub = MagicMock()
-    second_stub.GetWorkItems.side_effect = FakeRpcError(
-        grpc.StatusCode.CANCELLED,
-        "stop",
-    )
+    second_stub.GetWorkItems.side_effect = _cancel_after_shutdown(worker)
 
     stubs = [first_stub, second_stub]
     stub_channels = []
@@ -684,10 +697,7 @@ def complete_activity(response):
     first_stub.CompleteActivityTask.side_effect = complete_activity
 
     second_stub = MagicMock()
-    second_stub.GetWorkItems.side_effect = FakeRpcError(
-        grpc.StatusCode.CANCELLED,
-        "stop",
-    )
+    second_stub.GetWorkItems.side_effect = _cancel_after_shutdown(worker)
 
     stubs = [first_stub, second_stub]
     stub_channels = []
@@ -751,10 +761,7 @@ def fake_delay(attempt, *, base_seconds, cap_seconds):
         "connect failed",
     )
     second_stub = MagicMock()
-    second_stub.GetWorkItems.side_effect = FakeRpcError(
-        grpc.StatusCode.CANCELLED,
-        "stop",
-    )
+    second_stub.GetWorkItems.side_effect = _cancel_after_shutdown(worker)
     stubs = [first_stub, second_stub]
 
     def create_stub(current_channel):
@@ -789,10 +796,7 @@ async def test_worker_never_replaces_caller_owned_channel_during_transport_failu
             grpc.StatusCode.UNAVAILABLE,
             "transport failure",
         )))),
-        MagicMock(GetWorkItems=MagicMock(side_effect=FakeRpcError(
-            grpc.StatusCode.CANCELLED,
-            "stop",
-        ))),
+        MagicMock(GetWorkItems=MagicMock(side_effect=_cancel_after_shutdown(worker))),
     ]
 
     def create_stub(channel):
@@ -846,10 +850,7 @@ def get_grpc_channel(*args, **kwargs):
             items=[work_item],
             error=FakeRpcError(grpc.StatusCode.UNAVAILABLE, "second transport failure"),
         ))),
-        MagicMock(GetWorkItems=MagicMock(side_effect=FakeRpcError(
-            grpc.StatusCode.CANCELLED,
-            "stop",
-        ))),
+        MagicMock(GetWorkItems=MagicMock(side_effect=_cancel_after_shutdown(worker))),
     ]
 
     monkeypatch.setattr("durabletask.worker.shared.get_grpc_channel", get_grpc_channel)