From 0277bea15f51b9115dfe30a777697ce86e94f836 Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Sun, 26 Jul 2026 23:09:04 -0700 Subject: [PATCH 1/2] Bound async worker task creation to the concurrency limit `_AsyncWorkerManager._consume_queue()` created an `asyncio.Task` for every dequeued work item and only acquired the concurrency semaphore later, inside `_process_work_item()`. The semaphore therefore throttled execution but not allocation, so during a burst with slow handlers the number of pending task objects grew with the queue depth rather than with the configured concurrency. The loop also rescanned the entire `running_tasks` set on every iteration. The consumer now acquires a semaphore permit *before* dequeuing an item and hands that permit off to the task it creates, which releases it once the work item is done. The queue stays the bounded backlog and in-flight work item tasks never exceed the configured `ConcurrencyOptions` limits. Completed tasks are removed from `running_tasks` via `task.add_done_callback(set.discard)` instead of an O(n) sweep. Semantics are preserved: there is still no await between `queue.get()` returning and `asyncio.create_task()`, so an item can never be dropped by cancellation; the permit is released if the dequeue times out or raises; and `queue.task_done()` is still called exactly once per item, before the permit is released, on both the success and failure paths. Fixes #193 Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 010b6832-802b-425d-957c-cf29b7143392 --- CHANGELOG.md | 4 + durabletask/worker.py | 54 +++++--- .../test_worker_concurrency_loop_async.py | 130 ++++++++++++++++++ 3 files changed, 168 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 99372609..1f4d286b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,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. +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 diff --git a/durabletask/worker.py b/durabletask/worker.py index 7efb732c..3897f4fc 100644 --- a/durabletask/worker.py +++ b/durabletask/worker.py @@ -3434,43 +3434,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): diff --git a/tests/durabletask/test_worker_concurrency_loop_async.py b/tests/durabletask/test_worker_concurrency_loop_async.py index 686e12b8..155b3cf4 100644 --- a/tests/durabletask/test_worker_concurrency_loop_async.py +++ b/tests/durabletask/test_worker_concurrency_loop_async.py @@ -95,3 +95,133 @@ 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 wait until its queues are bound to this loop.""" + 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, wait_for_queues = _start_manager_and_wait_for_queues(manager) + await wait_for_queues + 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, wait_for_queues = _start_manager_and_wait_for_queues(manager) + await wait_for_queues + 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 From f75b20e9ddae2b2d8da583c6dab109b95d86050e Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Mon, 27 Jul 2026 09:41:00 -0700 Subject: [PATCH 2/2] Clarify that the queue-readiness helper returns a coroutine The second element of the tuple returned by _start_manager_and_wait_for_queues is an already-created coroutine object, not a callable. Naming the local wait_for_queues made wait wait_for_queues read like a missing call -- it misled an automated reviewer and would mislead a human too. Rename the local to queues_ready and say so in the docstring. No behavior change. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 010b6832-802b-425d-957c-cf29b7143392 --- .../test_worker_concurrency_loop_async.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/tests/durabletask/test_worker_concurrency_loop_async.py b/tests/durabletask/test_worker_concurrency_loop_async.py index 155b3cf4..c05d4674 100644 --- a/tests/durabletask/test_worker_concurrency_loop_async.py +++ b/tests/durabletask/test_worker_concurrency_loop_async.py @@ -98,7 +98,11 @@ async def run_test(): def _start_manager_and_wait_for_queues(manager): - """Start the manager loop and wait until its queues are bound to this loop.""" + """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(): @@ -160,8 +164,8 @@ async def cancel_work(idx): async def run_test(): manager.prepare_for_run() - worker_task, wait_for_queues = _start_manager_and_wait_for_queues(manager) - await wait_for_queues + 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) @@ -204,8 +208,8 @@ async def cancel_work(idx): async def run_test(): manager.prepare_for_run() - worker_task, wait_for_queues = _start_manager_and_wait_for_queues(manager) - await wait_for_queues + 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