Fix worker channel retirement for in-flight completions

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

Author: Bernd Verst

CHANGELOG.md

@@ -33,6 +33,8 @@ FIXED
   so configured hello timeouts apply on fresh connections, received work resets
   failure tracking, SDK-owned channels are refreshed and cleaned up safely, and
   caller-owned channels are never recreated or closed during reconnects.
+- Fixed `TaskHubGrpcWorker` so in-flight work item completions can finish after
+  a graceful gRPC stream reset before the worker retires an SDK-owned channel.
 - Improved sync and async gRPC clients so repeated transport failures recreate
   SDK-owned channels, while long-poll deadlines, successful replies, and
   application-level RPC errors do not trigger unnecessary channel replacement.

durabletask/worker.py

@@ -10,7 +10,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta, timezone
-from threading import Event, Thread
+from threading import Event, Lock, Thread
 from types import GeneratorType
 from enum import Enum
 from typing import Any, Generator, Optional, Sequence, Tuple, TypeVar, Union, overload
@@ -130,6 +130,73 @@ class _WorkItemStreamOutcome(Enum):
     SILENT_DISCONNECT = "silent_disconnect"
 
 
+@dataclass
+class _TrackedChannelState:
+    channel: Any
+    ref_count: int = 0
+    close_when_released: bool = False
+
+
+class _InFlightChannelTracker:
+    def __init__(self):
+        self._lock = Lock()
+        self._states: dict[int, _TrackedChannelState] = {}
+
+    def acquire(self, channel: Any):
+        channel_key = id(channel)
+        with self._lock:
+            state = self._states.get(channel_key)
+            if state is None:
+                state = _TrackedChannelState(channel=channel)
+                self._states[channel_key] = state
+            state.ref_count += 1
+
+        released = False
+
+        def release() -> None:
+            nonlocal released
+            if released:
+                return
+            released = True
+
+            channel_to_close = None
+            with self._lock:
+                state = self._states.get(channel_key)
+                if state is None:
+                    return
+
+                state.ref_count -= 1
+                if state.ref_count == 0:
+                    if state.close_when_released:
+                        channel_to_close = state.channel
+                    del self._states[channel_key]
+
+            if channel_to_close is not None:
+                self._close_channel(channel_to_close)
+
+        return release
+
+    def retire(self, channel: Any) -> None:
+        channel_key = id(channel)
+        channel_to_close = None
+        with self._lock:
+            state = self._states.get(channel_key)
+            if state is None:
+                channel_to_close = channel
+            else:
+                state.close_when_released = True
+
+        if channel_to_close is not None:
+            self._close_channel(channel_to_close)
+
+    @staticmethod
+    def _close_channel(channel: Any) -> None:
+        try:
+            channel.close()
+        except Exception:
+            pass
+
+
 class VersioningOptions:
     """Configuration options for orchestrator and activity versioning.
 
@@ -642,6 +709,7 @@ async def _async_run_loop(self):
         failure_tracker = FailureTracker(
             threshold=self._resiliency_options.channel_recreate_failure_threshold,
         )
+        in_flight_channel_tracker = _InFlightChannelTracker()
 
         def get_reconnect_delay_seconds() -> float:
             return get_full_jitter_delay_seconds(
@@ -671,6 +739,45 @@ def create_fresh_connection():
                 current_stub = None
                 raise
 
+        def wrap_execution(handler, release):
+            def wrapped(*args, **kwargs):
+                result = handler(*args, **kwargs)
+                release()
+                return result
+
+            return wrapped
+
+        def wrap_cancellation(handler, release):
+            def wrapped(*args, **kwargs):
+                try:
+                    return handler(*args, **kwargs)
+                finally:
+                    release()
+
+            return wrapped
+
+        def submit_work_item(
+                submit_func,
+                handler,
+                cancellation_handler,
+                request,
+                stub,
+                completion_token,
+                channel,
+        ):
+            release = in_flight_channel_tracker.acquire(channel)
+            try:
+                submit_func(
+                    wrap_execution(handler, release),
+                    wrap_cancellation(cancellation_handler, release),
+                    request,
+                    stub,
+                    completion_token,
+                )
+            except Exception:
+                release()
+                raise
+
         def invalidate_connection(
                 *,
                 recreate_channel: bool = False,
@@ -700,10 +807,7 @@ def invalidate_connection(
                     and self._can_recreate_channel()
                     and (recreate_channel or close_channel)
             ):
-                try:
-                    current_channel.close()
-                except Exception:
-                    pass
+                in_flight_channel_tracker.retire(current_channel)
                 current_channel = None
             current_stub = None
 
@@ -742,7 +846,9 @@ def should_invalidate_connection(rpc_error):
                     continue
             try:
                 assert current_stub is not None
+                assert current_channel is not None
                 stub = current_stub
+                channel = current_channel
                 capabilities = []
                 if self._payload_store is not None:
                     capabilities.append(pb.WORKER_CAPABILITY_LARGE_PAYLOADS)
@@ -822,36 +928,44 @@ def stream_reader():
 
                         failure_tracker.record_success()
                         if work_item.HasField("orchestratorRequest"):
-                            self._async_worker_manager.submit_orchestration(
+                            submit_work_item(
+                                self._async_worker_manager.submit_orchestration,
                                 self._execute_orchestrator,
                                 self._cancel_orchestrator,
                                 work_item.orchestratorRequest,
                                 stub,
                                 work_item.completionToken,
+                                channel,
                             )
                         elif work_item.HasField("activityRequest"):
-                            self._async_worker_manager.submit_activity(
+                            submit_work_item(
+                                self._async_worker_manager.submit_activity,
                                 self._execute_activity,
                                 self._cancel_activity,
                                 work_item.activityRequest,
                                 stub,
                                 work_item.completionToken,
+                                channel,
                             )
                         elif work_item.HasField("entityRequest"):
-                            self._async_worker_manager.submit_entity_batch(
+                            submit_work_item(
+                                self._async_worker_manager.submit_entity_batch,
                                 self._execute_entity_batch,
                                 self._cancel_entity_batch,
                                 work_item.entityRequest,
                                 stub,
                                 work_item.completionToken,
+                                channel,
                             )
                         elif work_item.HasField("entityRequestV2"):
-                            self._async_worker_manager.submit_entity_batch(
+                            submit_work_item(
+                                self._async_worker_manager.submit_entity_batch,
                                 self._execute_entity_batch,
                                 self._cancel_entity_batch,
                                 work_item.entityRequestV2,
                                 stub,
-                                work_item.completionToken
+                                work_item.completionToken,
+                                channel,
                             )
                         else:
                             self._logger.warning(

tests/durabletask/test_worker_resiliency.py

@@ -78,6 +78,16 @@ def shutdown(self):
         self._shutdown_event.set()
 
 
+def _complete_activity_request(req, stub, completion_token):
+    stub.CompleteActivityTask(
+        pb.ActivityResponse(
+            instanceId=req.orchestrationInstance.instanceId,
+            taskId=req.taskId,
+            completionToken=completion_token,
+        )
+    )
+
+
 def _make_activity_work_item() -> pb.WorkItem:
     return pb.WorkItem(
         activityRequest=pb.ActivityRequest(
@@ -331,6 +341,119 @@ def create_stub(channel):
     created_channels[1].close.assert_called_once()
 
 
+@pytest.mark.asyncio
+async def test_worker_defers_sdk_owned_channel_close_until_inflight_completion_finishes(monkeypatch):
+    worker = TaskHubGrpcWorker()
+    worker_manager = DummyWorkerManager()
+    worker._async_worker_manager = worker_manager
+    worker._execute_activity = _complete_activity_request
+    monkeypatch.setattr(worker._shutdown, "wait", lambda timeout: False)
+
+    created_channels = []
+
+    def get_grpc_channel(*args, **kwargs):
+        channel = MagicMock(name=f"channel-{len(created_channels) + 1}")
+        created_channels.append(channel)
+        return channel
+
+    completed_responses = []
+    first_stub = MagicMock()
+    first_stub.GetWorkItems.return_value = FakeResponseStream(items=[_make_activity_work_item()])
+
+    def complete_activity(response):
+        assert created_channels[0].close.call_count == 0
+        completed_responses.append(response)
+
+    first_stub.CompleteActivityTask.side_effect = complete_activity
+
+    second_stub = MagicMock()
+    second_stub.GetWorkItems.side_effect = FakeRpcError(
+        grpc.StatusCode.CANCELLED,
+        "stop",
+    )
+
+    stubs = [first_stub, second_stub]
+    stub_channels = []
+
+    def create_stub(channel):
+        stub_channels.append(channel)
+        return stubs.pop(0)
+
+    monkeypatch.setattr("durabletask.worker.shared.get_grpc_channel", get_grpc_channel)
+    monkeypatch.setattr("durabletask.worker.stubs.TaskHubSidecarServiceStub", create_stub)
+
+    await worker._async_run_loop()
+
+    assert len(worker_manager.submissions) == 1
+    assert len(created_channels) == 2
+    assert stub_channels == created_channels
+    created_channels[0].close.assert_not_called()
+    created_channels[1].close.assert_called_once()
+
+    _, submission = worker_manager.submissions[0]
+    func, _, req, stub, completion_token = submission
+    func(req, stub, completion_token)
+
+    assert len(completed_responses) == 1
+    assert completed_responses[0].completionToken == "token"
+    created_channels[0].close.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_worker_never_closes_caller_owned_channel_after_graceful_reset(monkeypatch):
+    provided_channel = MagicMock(name="provided-channel")
+    worker = TaskHubGrpcWorker(channel=provided_channel)
+    worker_manager = DummyWorkerManager()
+    worker._async_worker_manager = worker_manager
+    worker._execute_activity = _complete_activity_request
+    monkeypatch.setattr(worker._shutdown, "wait", lambda timeout: False)
+
+    completed_responses = []
+    first_stub = MagicMock()
+    first_stub.GetWorkItems.return_value = FakeResponseStream(items=[_make_activity_work_item()])
+
+    def complete_activity(response):
+        assert provided_channel.close.call_count == 0
+        completed_responses.append(response)
+
+    first_stub.CompleteActivityTask.side_effect = complete_activity
+
+    second_stub = MagicMock()
+    second_stub.GetWorkItems.side_effect = FakeRpcError(
+        grpc.StatusCode.CANCELLED,
+        "stop",
+    )
+
+    stubs = [first_stub, second_stub]
+    stub_channels = []
+
+    def create_stub(channel):
+        stub_channels.append(channel)
+        return stubs.pop(0)
+
+    monkeypatch.setattr(
+        "durabletask.worker.shared.get_grpc_channel",
+        lambda *args, **kwargs: pytest.fail(
+            "SDK channel factory should not run for caller-owned channels"
+        ),
+    )
+    monkeypatch.setattr("durabletask.worker.stubs.TaskHubSidecarServiceStub", create_stub)
+
+    await worker._async_run_loop()
+
+    assert len(worker_manager.submissions) == 1
+    assert stub_channels == [provided_channel, provided_channel]
+    provided_channel.close.assert_not_called()
+
+    _, submission = worker_manager.submissions[0]
+    func, _, req, stub, completion_token = submission
+    func(req, stub, completion_token)
+
+    assert len(completed_responses) == 1
+    assert completed_responses[0].completionToken == "token"
+    provided_channel.close.assert_not_called()
+
+
 @pytest.mark.asyncio
 async def test_worker_uses_reconnect_backoff_helper_after_connection_failure(monkeypatch):
     worker = TaskHubGrpcWorker(