Fix worker shutdown channel draining

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

Author: Bernd Verst

CHANGELOG.md

@@ -33,8 +33,9 @@ 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.
+- Fixed `TaskHubGrpcWorker` so in-flight and queued work item completions keep
+  draining across graceful gRPC stream resets and worker shutdown 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

@@ -690,6 +690,8 @@ def start(self):
         if self._auto_generate_work_item_filters:
             self._work_item_filters = WorkItemFilters._from_registry(self._registry)
 
+        self._shutdown.clear()
+
         def run_loop():
             loop = asyncio.new_event_loop()
             asyncio.set_event_loop(loop)
@@ -701,6 +703,7 @@ def run_loop():
         self._is_running = True
 
     async def _async_run_loop(self):
+        self._async_worker_manager.prepare_for_run()
         worker_task = asyncio.create_task(self._async_worker_manager.run())
         current_channel = self._channel
         current_stub = None
@@ -1060,6 +1063,11 @@ def stop(self):
             self._response_stream.cancel()
         if self._runLoop is not None:
             self._runLoop.join(timeout=30)
+            if self._runLoop.is_alive():
+                self._logger.info(
+                    "Waiting for pending work items to finish before completing shutdown..."
+                )
+                self._runLoop.join()
         self._async_worker_manager.shutdown()
         self._logger.info("Worker shutdown completed")
         self._is_running = False
@@ -2883,11 +2891,22 @@ def __init__(self, concurrency_options: ConcurrencyOptions, logger: logging.Logg
         self._pending_activity_work: list = []
         self._pending_orchestration_work: list = []
         self._pending_entity_batch_work: list = []
-        self.thread_pool = ThreadPoolExecutor(
-            max_workers=concurrency_options.maximum_thread_pool_workers,
+        self.thread_pool = self._create_thread_pool()
+        self._shutdown = False
+
+    def _create_thread_pool(self) -> ThreadPoolExecutor:
+        return ThreadPoolExecutor(
+            max_workers=self.concurrency_options.maximum_thread_pool_workers,
             thread_name_prefix="DurableTask",
         )
+
+    def _ensure_thread_pool(self) -> None:
+        if getattr(self.thread_pool, "_shutdown", False):
+            self.thread_pool = self._create_thread_pool()
+
+    def prepare_for_run(self) -> None:
         self._shutdown = False
+        self._ensure_thread_pool()
 
     def _ensure_queues_for_current_loop(self):
         """Ensure queues are bound to the current event loop."""
@@ -2962,8 +2981,7 @@ def _ensure_queues_for_current_loop(self):
         self._pending_entity_batch_work.clear()
 
     async def run(self):
-        # Reset shutdown flag in case this manager is being reused
-        self._shutdown = False
+        self._ensure_thread_pool()
 
         # Ensure queues are properly bound to the current event loop
         self._ensure_queues_for_current_loop()
@@ -3025,6 +3043,9 @@ async def run(self):
                 except Exception as cancellation_exception:
                     self._logger.error(f"Uncaught error while cancelling entity batch work item: {cancellation_exception}")
             self.shutdown()
+        finally:
+            if not getattr(self.thread_pool, "_shutdown", False):
+                self.thread_pool.shutdown(wait=True)
 
     async def _consume_queue(self, queue: asyncio.Queue, semaphore: asyncio.Semaphore):
         # List to track running tasks
@@ -3068,12 +3089,7 @@ async def _run_func(self, func, *args, **kwargs):
             return await func(*args, **kwargs)
         else:
             loop = asyncio.get_running_loop()
-            # Avoid submitting to executor after shutdown
-            if (
-                    getattr(self, "_shutdown", False) and getattr(self, "thread_pool", None) and getattr(
-                        self.thread_pool, "_shutdown", False)
-            ):
-                return None
+            self._ensure_thread_pool()
             return await loop.run_in_executor(
                 self.thread_pool, lambda: func(*args, **kwargs)
             )
@@ -3113,11 +3129,10 @@ def submit_entity_batch(self, func, cancellation_func, *args, **kwargs):
 
     def shutdown(self):
         self._shutdown = True
-        self.thread_pool.shutdown(wait=True)
 
     async def reset_for_new_run(self):
         """Reset the manager state for a new run."""
-        self._shutdown = False
+        self.prepare_for_run()
         # Clear any existing queues - they'll be recreated when needed
         if self.activity_queue is not None:
             # Clear existing queue by creating a new one

tests/durabletask/test_worker_concurrency_loop.py

@@ -73,6 +73,7 @@ def cancel_dummy_activity(req, stub, completionToken):
 
     async def run_test():
         # Start the worker manager's run loop in the background
+        worker._async_worker_manager.prepare_for_run()
         worker_task = asyncio.create_task(worker._async_worker_manager.run())
         for req in orchestrator_requests:
             worker._async_worker_manager.submit_orchestration(dummy_orchestrator, cancel_dummy_orchestrator, req, stub, DummyCompletionToken())
@@ -133,6 +134,7 @@ def fn(*args, **kwargs):
 
     # Run the manager loop in a thread (sync context)
     def run_manager():
+        manager.prepare_for_run()
         asyncio.run(manager.run())
 
     t = threading.Thread(target=run_manager)

tests/durabletask/test_worker_concurrency_loop_async.py

@@ -72,6 +72,7 @@ async def cancel_dummy_activity(req, stub, completionToken):
     async def run_test():
         # Clear stub state before each run
         stub.completed.clear()
+        grpc_worker._async_worker_manager.prepare_for_run()
         worker_task = asyncio.create_task(grpc_worker._async_worker_manager.run())
         # Need to yield to that thread in order to let it start up on the second run
         startup_attempts = 0

tests/durabletask/test_worker_resiliency.py

@@ -1,13 +1,18 @@
 import asyncio
 import grpc
-from threading import Event
+from threading import Event, Timer
 from unittest.mock import MagicMock
 
 import pytest
 
 from durabletask.grpc_options import GrpcWorkerResiliencyOptions
 from durabletask.internal import orchestrator_service_pb2 as pb
-from durabletask.worker import TaskHubGrpcWorker, _WorkItemStreamOutcome
+from durabletask.worker import (
+    _AsyncWorkerManager,
+    ConcurrencyOptions,
+    TaskHubGrpcWorker,
+    _WorkItemStreamOutcome,
+)
 
 
 class FakeRpcError(grpc.RpcError):
@@ -62,6 +67,9 @@ def __init__(self):
         self._shutdown_event = asyncio.Event()
         self.submissions: list[tuple[str, tuple]] = []
 
+    def prepare_for_run(self):
+        self._shutdown_event = asyncio.Event()
+
     async def run(self):
         await self._shutdown_event.wait()
 
@@ -88,16 +96,58 @@ def _complete_activity_request(req, stub, completion_token):
     )
 
 
-def _make_activity_work_item() -> pb.WorkItem:
+def _make_activity_work_item(
+    task_id: int = 1,
+    completion_token: str = "token",
+    instance_id: str = "instance-id",
+) -> pb.WorkItem:
     return pb.WorkItem(
         activityRequest=pb.ActivityRequest(
             name="test_activity",
-            taskId=1,
-            orchestrationInstance=pb.OrchestrationInstance(instanceId="instance-id"),
+            taskId=task_id,
+            orchestrationInstance=pb.OrchestrationInstance(instanceId=instance_id),
         ),
-        completionToken="token",
+        completionToken=completion_token,
+    )
+
+
+async def _wait_for_condition(predicate, *, timeout: float = 2.0):
+    loop = asyncio.get_running_loop()
+    deadline = loop.time() + timeout
+    while not predicate():
+        if loop.time() >= deadline:
+            raise AssertionError("condition was not met before timeout")
+        await asyncio.sleep(0.01)
+
+
+@pytest.mark.asyncio
+async def test_async_worker_manager_honors_shutdown_requested_before_run():
+    manager = _AsyncWorkerManager(
+        ConcurrencyOptions(maximum_thread_pool_workers=1),
+        MagicMock(),
     )
 
+    manager.shutdown()
+    await asyncio.wait_for(manager.run(), timeout=1.0)
+
+
+def test_worker_start_clears_prior_shutdown_request():
+    worker = TaskHubGrpcWorker()
+    worker._shutdown.set()
+    run_started = Event()
+
+    async def fake_run_loop():
+        run_started.set()
+
+    worker._async_run_loop = fake_run_loop
+    worker.start()
+    worker._runLoop.join(timeout=1.0)
+
+    assert run_started.is_set() is True
+    assert worker._shutdown.is_set() is False
+
+    worker.stop()
+
 
 def test_worker_classifies_graceful_close_before_first_message():
     worker = TaskHubGrpcWorker(
@@ -399,6 +449,148 @@ def create_stub(channel):
     created_channels[0].close.assert_called_once()
 
 
+@pytest.mark.asyncio
+async def test_worker_shutdown_drains_real_manager_work_before_closing_retired_sdk_channel(monkeypatch):
+    worker = TaskHubGrpcWorker(
+        concurrency_options=ConcurrencyOptions(
+            maximum_concurrent_activity_work_items=1,
+            maximum_thread_pool_workers=1,
+        )
+    )
+    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
+
+    allow_first_completion = Event()
+    first_completion_started = Event()
+    completed_task_ids = []
+    first_stub = MagicMock()
+    first_stub.GetWorkItems.return_value = FakeResponseStream(items=[
+        _make_activity_work_item(task_id=1, completion_token="token-1"),
+        _make_activity_work_item(task_id=2, completion_token="token-2"),
+    ])
+
+    def complete_activity(response):
+        completed_task_ids.append(response.taskId)
+        if response.taskId == 1:
+            first_completion_started.set()
+            Timer(0.2, allow_first_completion.set).start()
+            assert allow_first_completion.wait(timeout=5.0)
+        elif response.taskId == 2:
+            assert created_channels[0].close.call_count == 0
+
+    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)
+
+    run_task = asyncio.create_task(worker._async_run_loop())
+    await asyncio.wait_for(run_task, timeout=2.0)
+
+    assert first_completion_started.is_set() is True
+    assert len(created_channels) == 2
+    assert stub_channels == created_channels
+    assert completed_task_ids == [1, 2]
+    created_channels[0].close.assert_called_once()
+    created_channels[1].close.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_worker_shutdown_runs_real_manager_cancellation_wrapper_before_closing_retired_sdk_channel(monkeypatch):
+    worker = TaskHubGrpcWorker(
+        concurrency_options=ConcurrencyOptions(
+            maximum_concurrent_activity_work_items=1,
+            maximum_thread_pool_workers=1,
+        )
+    )
+    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
+
+    allow_first_completion = Event()
+    first_completion_started = Event()
+    completed_task_ids = []
+    cancelled_task_ids = []
+
+    def execute_activity(req, stub, completion_token):
+        if req.taskId == 1:
+            _complete_activity_request(req, stub, completion_token)
+        else:
+            raise RuntimeError("boom")
+
+    def cancel_activity(req, stub, completion_token):
+        cancelled_task_ids.append(req.taskId)
+        assert created_channels[0].close.call_count == 0
+
+    worker._execute_activity = execute_activity
+    worker._cancel_activity = cancel_activity
+
+    first_stub = MagicMock()
+    first_stub.GetWorkItems.return_value = FakeResponseStream(items=[
+        _make_activity_work_item(task_id=1, completion_token="token-1"),
+        _make_activity_work_item(task_id=2, completion_token="token-2"),
+    ])
+
+    def complete_activity(response):
+        completed_task_ids.append(response.taskId)
+        Timer(0.2, allow_first_completion.set).start()
+        first_completion_started.set()
+        assert allow_first_completion.wait(timeout=5.0)
+
+    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)
+
+    run_task = asyncio.create_task(worker._async_run_loop())
+    await asyncio.wait_for(run_task, timeout=2.0)
+
+    assert first_completion_started.is_set() is True
+    assert len(created_channels) == 2
+    assert stub_channels == created_channels
+    assert completed_task_ids == [1]
+    assert cancelled_task_ids == [2]
+    created_channels[0].close.assert_called_once()
+    created_channels[1].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")