Capstone: Extend the Worker Pool to Keep Going After Failures

This capstone extends the concurrency track's capstone, "Build an Async Worker Pool with Retries and Clean Shutdown". The pool still takes items from an async source through a bounded queue, runs a handler on several at a time, and yields each result as soon as it is ready. Three rules change:

An item that fails for good no longer stops the pool. The pool keeps taking items, and raises every failure together in an ExceptionGroup at the end.
A timeout is always worth a retry, whether or not TimeoutError is in retryable.
metrics receives the original input item instead of a repr() string, and the event for a failed item is "failed".

Every other rule is the same. If you solved the earlier capstone, start from your WorkerPool and make those three changes. The starter here names the end marker _DONE instead of SENTINEL.

async def numbers() -> AsyncGenerator[int, None]:
    for n in range(10):
        yield n


async def times_ten(n: int) -> int:
    await asyncio.sleep(0.01)
    if n == 2:
        raise ValueError("bad item")
    return n * 10


async def main() -> None:
    pool: WorkerPool[int, int] = WorkerPool(
        times_ten, workers=2, queue_size=2, item_timeout=1.0
    )
    results: list[int] = []
    try:
        async with aclosing(pool.run(numbers())) as output:
            async for value in output:
                results.append(value)
    except ExceptionGroup as group:
        print(sorted(results))  # [0, 10, 30, 40, 50, 60, 70, 80, 90]
        print(group.exceptions)  # (ValueError('bad item'),)

Item 2 fails, and the pool still handles items 3 to 9. The earlier pool would stop taking items after item 2 failed, so the last few would never run.

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.

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, T], 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

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 most queue_size + workers + 1 pulled items wait for their handler to start. The + 1 is the item the producer 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 even a _Done, 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

Each call to the handler is an attempt. Run each one in a new task, so it gets its own copy of the context. Setting a ContextVar in one attempt then never changes the value another attempt sees.
handle may return any Awaitable, and asyncio.create_task needs a coroutine. Call handle(item) inside a small async def that awaits it, so the whole call runs in the new task's context.
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 timeouts and the failures you choose

If an attempt fails with TimeoutError, or with 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.
Any other Exception fails the item at once, with no retry.
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 and drain

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.

Keep going, then report every failure

An item that fails for good does not stop the pool. Record its exception, and keep taking and handling items.
If the source raises an Exception, record it, stop taking items, and drain. Retries apply only to the handler.
If pulling from the source raises CancelledError before any stop was requested, record RuntimeError("source cancelled itself"), stop taking items, and drain.
When run has yielded every result, it 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

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 _DONE 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 _DONE, 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, item) with the original input item, exactly once for each of these events, and only for these:
- "ok" when an item succeeds.
- "retry" before each retry.
- "failed" when an item fails for good.
Ignore any Exception 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 driver's items are distinct integers, and its handler returns item * 2. Per item, behaviour makes the handler succeed, sleep for 30 seconds ("slow"), raise Permanent ("fatal"), or raise Transient flaky_failures times and then succeed ("flaky"). The driver closes both the source and run with aclosing.

With stop_after unset, the report has these keys:

mode is "full".
results holds every yielded result, sorted.
failures lists the type names in the ExceptionGroup, sorted.
retried lists the items whose handler ran more than once.
bounded is True when the source never got more than queue_size + workers + 1 items ahead of the handlers that had started.
isolated is True when every attempt started with the default ContextVar value and kept its own value throughout.
tasks_clean is True when no task the pool started is still running after run ends.
emitted lists the distinct metrics events, sorted.

A stop_after of zero or less calls request_stop() before iterating. A positive value calls it as the source resumes after yielding that many items. With stop_after set, the report has these keys:

mode is "stop".
no_exception is True when run raised no ExceptionGroup.
exactly_once is True when run yielded no result twice.
subset is True when every result came from an input item.
tasks_clean and bounded mean the same as above.

Going further

Test run yourself for result order, early close, a producer stuck on a full queue, and repeated cancellation. The report sorts results and 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 new task copies the context, not the objects in it. Two attempts that read the same mutable object from a ContextVar still share it.
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.