Merge branch 'main' into andystaples-client-parity-features

Author: andystaples

CHANGELOG.md

@@ -21,6 +21,10 @@ CHANGED
 - **Breaking:** `FailureDetails.error_type` — and the `errorType` value sent over the wire — is now the fully-qualified type name (`module.ClassName`, e.g. `builtins.ValueError`, `durabletask.task.TaskFailedError`) instead of the bare class name, matching the .NET and Java SDKs. Code that compared `error_type` against a bare name (for example `== "ValueError"`) must be updated to the qualified name or, preferably, switched to `FailureDetails.is_caused_by()`. Because this value is persisted and crosses the orchestration boundary, failures produced by older workers may still carry a bare name; `is_caused_by()` accepts both.
 - **Breaking:** Retired the process-global history-export context. `durabletask.extensions.history_export.bind_context()` and `clear_context()` have been removed; the export activities now resolve their `HistoryExportContext` per invocation via a resolver captured at registration. `ExportHistoryClient.register_worker()` continues to wire this up automatically, so most callers are unaffected. Code that called `bind_context(HistoryExportContext(client, writer))` directly should instead register the activities with `durabletask.extensions.history_export.activities.register(worker, lambda: HistoryExportContext(client, writer))` (or a lazier resolver), or use `ExportHistoryClient.register_worker()`.
 
+FIXED
+
+- Fixed the worker allocating one `asyncio` task per queued work item before applying the concurrency limit, which made memory use and event-loop scheduling overhead grow with the queue backlog during bursts. In-flight work item tasks are now bounded by the configured `ConcurrencyOptions` limits.
+
 ## v1.8.0
 
 ADDED

durabletask/worker.py

@@ -3440,43 +3440,57 @@ async def run(self) -> None:
                 self._pool_is_shutdown = True
 
     async def _consume_queue(self, queue: asyncio.Queue[_WorkItem], semaphore: asyncio.Semaphore) -> None:
-        # List to track running tasks
+        # Set of running tasks. Completed tasks remove themselves through a done
+        # callback so that this loop never has to scan the whole set.
         running_tasks: set[asyncio.Task[Any]] = set()
 
         while True:
-            # Clean up completed tasks
-            done_tasks = {task for task in running_tasks if task.done()}
-            running_tasks -= done_tasks
-
             # Exit if shutdown is set and the queue is empty and no tasks are running
             if self._shutdown and queue.empty() and not running_tasks:
                 break
 
+            # Reserve concurrency capacity *before* dequeuing a work item so that
+            # the number of allocated tasks stays bounded by the configured
+            # concurrency limit instead of growing with the queue depth. The
+            # permit is handed off to the task created below, which releases it
+            # once the work item has been processed.
+            await semaphore.acquire()
+            permit_owned = True
             try:
-                work = await asyncio.wait_for(queue.get(), timeout=1.0)
-            except asyncio.TimeoutError:
-                continue
+                try:
+                    work = await asyncio.wait_for(queue.get(), timeout=1.0)
+                except asyncio.TimeoutError:
+                    continue
 
-            func, cancellation_func, args, kwargs = work
-            # Create a concurrent task for processing
-            task = asyncio.create_task(
-                self._process_work_item(semaphore, queue, func, cancellation_func, args, kwargs)
-            )
-            running_tasks.add(task)
+                func, cancellation_func, args, kwargs = work
+                # Create a concurrent task for processing
+                task = asyncio.create_task(
+                    self._process_work_item(semaphore, queue, func, cancellation_func, args, kwargs)
+                )
+                permit_owned = False
+                running_tasks.add(task)
+                task.add_done_callback(running_tasks.discard)
+            finally:
+                if permit_owned:
+                    semaphore.release()
 
     async def _process_work_item(
             self, semaphore: asyncio.Semaphore, queue: asyncio.Queue[_WorkItem],
             func: Callable[..., Any], cancellation_func: Callable[..., Any],
             args: tuple[Any, ...], kwargs: dict[str, Any]
     ) -> None:
-        async with semaphore:
+        # The semaphore permit is acquired by ``_consume_queue`` before this task
+        # is created and is released here once the work item is done.
+        try:
+            await self._run_func(func, *args, **kwargs)
+        except Exception as work_exception:
+            self._logger.error(f"Uncaught error while processing work item, item will be abandoned: {work_exception}")
+            await self._run_func(cancellation_func, *args, **kwargs)
+        finally:
             try:
-                await self._run_func(func, *args, **kwargs)
-            except Exception as work_exception:
-                self._logger.error(f"Uncaught error while processing work item, item will be abandoned: {work_exception}")
-                await self._run_func(cancellation_func, *args, **kwargs)
-            finally:
                 queue.task_done()
+            finally:
+                semaphore.release()
 
     async def _run_func(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
         if inspect.iscoroutinefunction(func):

tests/durabletask/test_client.py

@@ -214,9 +214,10 @@ def install_resilient_test_stubs(client):
     call so newly created stubs continue to participate in failure tracking.
 
     The interceptor's ``_on_recreate`` callback is intentionally left alone:
-    it is the client's ``_schedule_recreate`` (fire-and-forget), and tests
-    that need to observe the recreate's completion can wait on
-    ``client._recreate_done_event``.
+    it is the client's ``_schedule_recreate`` (fire-and-forget). Synchronous
+    tests that chain recreates must observe completion via
+    ``await_sync_recreate``; waiting on ``client._recreate_done_event`` alone
+    is not sufficient (see that helper's docstring).
     """
     is_async = inspect.iscoroutinefunction(client._maybe_recreate_channel)
     wrapper_cls = _ResilientAsyncTestStub if is_async else _ResilientSyncTestStub
@@ -241,6 +242,28 @@ def wrapped_recreate():
     client._maybe_recreate_channel = wrapped_recreate
 
 
+def await_sync_recreate(client: TaskHubGrpcClient, timeout: float = 5.0) -> None:
+    """Block until the synchronous client's recreate thread has fully exited.
+
+    ``_recreate_done_event`` is set by ``_run_recreate`` from *inside* the
+    recreate thread, so the event is signalled while that thread is still
+    alive. ``_schedule_recreate`` single-flights on ``is_alive()``, so a
+    trigger issued in that window is dropped *before* it reaches
+    ``_recreate_done_event.clear()`` -- leaving the event set from the
+    previous recreate. A later ``wait()`` would then return immediately off
+    that stale signal, and assertions about the next recreate would fail.
+
+    Tests that chain recreates must therefore wait on the same condition the
+    single-flight guard reads. Joining the thread makes ``is_alive()``
+    deterministically ``False`` before the next trigger is issued.
+    """
+    assert client._recreate_done_event.wait(timeout=timeout)
+    recreate_thread = client._recreate_thread
+    if recreate_thread is not None:
+        recreate_thread.join(timeout=timeout)
+        assert not recreate_thread.is_alive()
+
+
 class FakePayloadStore(PayloadStore):
     TOKEN_PREFIX = 'fake://'
 
@@ -669,10 +692,10 @@ def test_sync_client_close_closes_all_retired_sdk_channels_immediately():
         # Wait for the first fire-and-forget recreate to complete so the
         # single-flight guard in _schedule_recreate does not drop the second
         # trigger.
-        assert client._recreate_done_event.wait(timeout=5.0)
+        await_sync_recreate(client)
         with pytest.raises(FakeRpcError):
             client.get_orchestration_state("abc")
-        assert client._recreate_done_event.wait(timeout=5.0)
+        await_sync_recreate(client)
 
         client.close()
 
@@ -912,21 +935,21 @@ def test_sync_client_recreate_cooldown_prevents_immediate_repeated_recreation():
             client.get_orchestration_state("abc")
         # Wait for the fire-and-forget recreate to complete before asserting
         # the channel was swapped.
-        assert client._recreate_done_event.wait(timeout=5.0)
+        await_sync_recreate(client)
         assert client._channel is second_channel
         assert mock_get_channel.call_count == 2
 
         with pytest.raises(FakeRpcError):
             client.get_orchestration_state("abc")
         # Cooldown should fire-and-forget but exit without recreating; wait
         # for the no-op recreate to complete so the assertion is deterministic.
-        assert client._recreate_done_event.wait(timeout=5.0)
+        await_sync_recreate(client)
         assert client._channel is second_channel
         assert mock_get_channel.call_count == 2
 
         with pytest.raises(FakeRpcError):
             client.get_orchestration_state("abc")
-        assert client._recreate_done_event.wait(timeout=5.0)
+        await_sync_recreate(client)
         assert client._channel is third_channel
 
     expected_channel_call = call(

tests/durabletask/test_worker_concurrency_loop_async.py

@@ -95,3 +95,137 @@ async def run_test():
         await worker_task
     asyncio.run(run_test())
     asyncio.run(run_test())
+
+
+def _start_manager_and_wait_for_queues(manager):
+    """Start the manager loop and return an awaitable that resolves once its queues exist.
+
+    The second element of the returned tuple is an already-created coroutine
+    object (a readiness awaitable), not a callable -- await it directly.
+    """
+    worker_task = asyncio.create_task(manager.run())
+
+    async def wait_for_queues():
+        for _ in range(100):
+            if manager.activity_queue is not None:
+                return
+            await asyncio.sleep(0.01)
+        raise RuntimeError("Worker manager queues were never initialized")
+
+    return worker_task, wait_for_queues()
+
+
+def test_async_worker_manager_bounds_task_allocation():
+    """Work item tasks must be allocated only up to the concurrency limit.
+
+    The concurrency semaphore is acquired before a task is created, so a burst
+    of queued work items must not allocate one asyncio.Task per queued item.
+    """
+    limit = 2
+    total_items = 25
+    options = ConcurrencyOptions(
+        maximum_concurrent_activity_work_items=limit,
+        maximum_concurrent_orchestration_work_items=limit,
+        maximum_concurrent_entity_work_items=limit,
+        maximum_thread_pool_workers=limit,
+    )
+    manager = TaskHubGrpcWorker(concurrency_options=options)._async_worker_manager
+
+    state = {"allocated": 0, "peak_allocated": 0, "running": 0, "peak_running": 0}
+    completed = []
+    original_process_work_item = manager._process_work_item
+
+    def counting_process_work_item(*args, **kwargs):
+        # The coroutine object is created synchronously by asyncio.create_task,
+        # so this counts allocated work item tasks, not just started ones.
+        state["allocated"] += 1
+        state["peak_allocated"] = max(state["peak_allocated"], state["allocated"])
+        inner = original_process_work_item(*args, **kwargs)
+
+        async def tracked():
+            try:
+                return await inner
+            finally:
+                state["allocated"] -= 1
+
+        return tracked()
+
+    manager._process_work_item = counting_process_work_item
+
+    async def slow_work(idx):
+        state["running"] += 1
+        state["peak_running"] = max(state["peak_running"], state["running"])
+        await asyncio.sleep(0.02)
+        state["running"] -= 1
+        completed.append(idx)
+
+    async def cancel_work(idx):
+        pass
+
+    async def run_test():
+        manager.prepare_for_run()
+        worker_task, queues_ready = _start_manager_and_wait_for_queues(manager)
+        await queues_ready
+        assert manager.activity_queue is not None
+        for i in range(total_items):
+            manager.submit_activity(slow_work, cancel_work, i)
+        await asyncio.wait_for(manager.activity_queue.join(), timeout=60)
+        manager.shutdown()
+        await asyncio.wait_for(worker_task, timeout=60)
+
+    asyncio.run(run_test())
+
+    assert sorted(completed) == list(range(total_items))
+    assert state["peak_allocated"] <= limit, (
+        f"Expected at most {limit} concurrently allocated work item tasks, "
+        f"got {state['peak_allocated']}"
+    )
+    # The concurrency limit must still be fully usable.
+    assert state["peak_running"] == limit, (
+        f"Expected {limit} concurrently running work items, got {state['peak_running']}"
+    )
+    assert state["allocated"] == 0
+
+
+def test_async_worker_manager_calls_task_done_exactly_once_on_failure():
+    """Failing work items must still mark the queue item done exactly once."""
+    total_items = 6
+    options = ConcurrencyOptions(
+        maximum_concurrent_activity_work_items=2,
+        maximum_concurrent_orchestration_work_items=2,
+        maximum_concurrent_entity_work_items=2,
+        maximum_thread_pool_workers=2,
+    )
+    manager = TaskHubGrpcWorker(concurrency_options=options)._async_worker_manager
+    cancelled = []
+    task_done_calls = []
+
+    async def failing_work(idx):
+        raise RuntimeError(f"boom {idx}")
+
+    async def cancel_work(idx):
+        cancelled.append(idx)
+
+    async def run_test():
+        manager.prepare_for_run()
+        worker_task, queues_ready = _start_manager_and_wait_for_queues(manager)
+        await queues_ready
+        queue = manager.activity_queue
+        assert queue is not None
+        original_task_done = queue.task_done
+
+        def counting_task_done():
+            task_done_calls.append(1)
+            original_task_done()
+
+        queue.task_done = counting_task_done  # type: ignore[method-assign]
+        for i in range(total_items):
+            manager.submit_activity(failing_work, cancel_work, i)
+        await asyncio.wait_for(queue.join(), timeout=60)
+        manager.shutdown()
+        await asyncio.wait_for(worker_task, timeout=60)
+
+    asyncio.run(run_test())
+
+    assert sorted(cancelled) == list(range(total_items))
+    assert len(task_done_calls) == total_items