diff --git a/CHANGELOG.md b/CHANGELOG.md index 03a7ff22..52072dae 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,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 diff --git a/durabletask/worker.py b/durabletask/worker.py index 1770b758..1cf7c61c 100644 --- a/durabletask/worker.py +++ b/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): diff --git a/tests/durabletask/test_worker_concurrency_loop_async.py b/tests/durabletask/test_worker_concurrency_loop_async.py index 686e12b8..c05d4674 100644 --- a/tests/durabletask/test_worker_concurrency_loop_async.py +++ b/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