Skip to content

← Concurrency with Threads and asyncio step 56 of 56

Hard End-to-End

Capstone: Build an Async Worker Pool with Retries and Clean Shutdown

Build WorkerPool, which runs a handler on items from an async source, several at a time, and yields each result as soon as it is ready. It combines the whole track in one class, with a bounded queue, a TaskGroup, timeouts, retries, cancellation and a clean shutdown.

async def numbers() -> AsyncIterator[int]:
    for n in range(3):
        yield n


async def slow_square(n: int) -> int:
    await asyncio.sleep(0.1 * (3 - n))  # later numbers finish sooner
    return n * n


async def main() -> None:
    pool: WorkerPool[int, int] = WorkerPool(
        slow_square, workers=3, queue_size=3, item_timeout=1.0
    )
    async for square in pool.run(numbers()):
        print(square)  # 4, then 1, then 0

The three calls overlap, so main takes about 0.3 seconds instead of 0.6, and the squares print in the order they finish.

What to write

Fill in WorkerPool with this interface. Keep the type parameters T for inputs and R for results, and make the file pass mypy --strict. The exercise on typing async code covers Awaitable and AsyncIterator, and the aclosing exercise covers AsyncGenerator.

class WorkerPool[T, R]:
    def __init__(
        self,
        handle: Callable[[T], Awaitable[R]],
        *,
        workers: int,
        queue_size: int,
        item_timeout: float,
        retryable: tuple[type[Exception], ...] = (),
        max_retries: int = 0,
        metrics: Callable[[str, str], None] | None = None,
    ) -> None: ...

    def request_stop(self) -> None: ...
    async def run(self, source: AsyncIterator[T]) -> AsyncGenerator[R, None]: ...
  • __init__ accepts only positive workers and queue_size, a finite positive item_timeout, and a max_retries of zero or more. Anything else raises ValueError.
  • A pool runs once. Iterating run a second time raises RuntimeError("pool is single-use").

Take in work through a bounded queue

This section reuses the producer and consumers from the async queue exercise.

  • Use one asyncio.Queue(maxsize=queue_size) for input and a fixed set of exactly workers consumer tasks.
  • One producer task pulls items from source and queues them. At any moment, the producer has pulled at most queue_size + workers + 1 items that are not finished yet. The + 1 is the item it holds while it waits for room in the queue.
  • Each item the source yields is separate work, even if two items are equal.
  • Wrap every input and result in the supplied _Item before it goes on a queue. Then a value that is None or a Sentinel stays distinct from your end marker.

Yield results as they finish

  • Consumers put each successful result on a second, unbounded asyncio.Queue, and run yields from it. Consumers keep working while the caller pauses between results.
  • Yield results in the order consumers finish them, not in input order. Two items that finish close together may swap places between runs.

Give each attempt a deadline

This combines the timeouts exercise and the cancellation exercise.

  • Each call to the handler is an attempt. Run each one in a new task, so it gets its own copy of the context, as in the contextvars exercise.
  • If the attempt is still running after item_timeout, it fails with TimeoutError("item timed out"). The attempt fails this way even if the handler returns a value while the pool cancels it.
  • To stop a late attempt, cancel its task and wait up to item_timeout for it. Cancel and wait at most three times. After that, wait for the task to finish, however long it takes.
  • Once the task finishes, read its exception with task.exception(), unless task.cancelled() returns True. Otherwise asyncio logs "Task exception was never retrieved".
  • If the handler cancels itself, its task ends cancelled before the deadline. The attempt then fails with RuntimeError("handler cancelled itself").
  • If the caller cancels run, stop the attempt the same way, wait for it, and then re-raise the CancelledError.

Retry the failures you choose

  • If an attempt's exception is an instance of a class in retryable, run the handler again, up to max_retries more times. Start each retry only after the previous attempt has finished.
  • Retry a timeout only if TimeoutError is in retryable.
  • When the retries run out, the item fails for good with the exception from its last attempt. Record that exception object as it is.

Stop, drain and report failures

This follows the stages of the graceful shutdown exercise. To drain means to finish every item already queued or in progress and yield its result.

  • request_stop() is a plain def method, and it is safe to call more than once.
  • If you call request_stop() before iterating, run ends at once and pulls zero items.
  • After a stop, the producer stops pulling items. Wake it if it is waiting on the source or on a full queue. An item it already pulled may still be handled.
  • Then drain the queue, and run yields the drained results.
  • An item that fails for good also requests a stop.
  • If the source raises an Exception, record it and stop. Retries apply only to the handler.
  • If pulling from the source raises CancelledError before any stop was requested, record RuntimeError("source cancelled itself") and stop.
  • After the drain, run raises an ExceptionGroup that holds the original recorded exception objects. If the pool recorded no exceptions, run ends normally.

Finish every task the pool starts

The TaskGroup exercise and the exercise on owning background tasks cover these pieces.

  • Run the producer and the consumers in an asyncio.TaskGroup inside a separate driver task. run only reads the result queue and yields. A yield inside the async with block would let a failing child task cancel the caller's code between results.
  • On a normal drain, once the producer has finished, put one SENTINEL per consumer on the input queue. Each consumer exits when it takes one.
  • If the driver is cancelled, let the TaskGroup cancel the consumers. Don't wait in a finally to queue SENTINEL, because no consumer is left to take it.
  • If the caller cancels the task iterating run, or calls aclose() on the generator that run returns, cancel the driver. Wait for every pool task to finish, then re-raise. Keep waiting through a second cancellation.
  • When run ends for any reason, every task the pool started has finished.
  • Leave the source open. The caller owns the source and decides when to close it.

Report metrics

  • If metrics is set, call metrics(event, detail) exactly once for each of these item events, and only for these:
    • "ok" when an item succeeds, with repr() of the input item.
    • "retry" before each retry, with repr() of the exception.
    • "failure" when an item fails for good, with repr() of the exception.
  • Ignore any Exception from repr() or from metrics, so the pool works the same with a broken metrics backend. Catch Exception, not BaseException, so CancelledError and KeyboardInterrupt still get through.

What the driver checks

Keep solve and the helpers below WorkerPool as they are. The test handler can succeed, fail a set number of times, fail for good, hang, return late, cancel itself, or ignore its first cancellation. stop_after=0 calls request_stop() before iterating, a positive value calls it after that many results, and a negative value skips it.

  • results holds every yielded result, sorted to keep the output stable.
  • attempts counts handler calls per item.
  • max_lead is the largest count of items produced but not yet handled successfully. Failed items stay in that count. lead_ok is True when max_lead <= queue_size + workers + 1.
  • outcome is "ok", "ExceptionGroup", or the name of another exception run raised. group_size, group_types and group_messages give the group's size and its sorted type names and messages.
  • ctx_ok is True when every handler read back its own ContextVar value.
  • In the tests where the metrics callback always raises, metrics_raised counts its calls.
  • task_delta counts tasks still running after run ends. It must be 0.

Going further

Test run yourself for result order, early close, a producer stuck on a full queue, and repeated cancellation. The report shows only the final outcome.

  • Only the input queue is bounded, so a slow caller lets results pile up in memory.
  • A retry repeats the handler's side effects. Make handlers safe to run twice.
  • A deadline can cancel a handler only at an await. A blocking call delays every task on the event loop, and a handler that ignores cancellation keeps the pool waiting.
  • Breaking out of async for leaves the pool's tasks running until the event loop closes the generator later. contextlib.aclosing(...) closes it at once. Its argument needs an aclose() method in its type, so the return type of run is AsyncGenerator.

The capstone track extends this pool in Capstone: Extend the Worker Pool to Keep Going After Failures, where a failed item no longer stops the pool from taking new work.