Build a Service That Shuts Down in Stages

Build the core of a worker service. It pulls items from an async source and runs a handler task for each one. When the source runs out or someone asks it to stop, it shuts down in stages.

service: Service[int] = Service(drain_timeout=5.0)
loop.add_signal_handler(signal.SIGTERM, service.request_stop)
await service.serve(jobs(), handle)  # returns once every handler has finished

What to write

type State = Literal["idle", "running", "draining", "stopped"]

def describe(state: State) -> str: ...

class Service[T]:
    def __init__(self, *, drain_timeout: float) -> None: ...
    def state(self) -> State: ...
    def request_stop(self) -> None: ...
    def pending(self) -> tuple[asyncio.Task[None], ...]: ...
    async def serve(self, work: AsyncIterator[T],
                    handle: Callable[[T], Awaitable[None]]) -> None: ...

describe returns "not started", "accepting work", "finishing in-flight work" and "done" for the four states in order. End its match with case _: assert_never(state).

Set up the service

• A new service is idle. A negative, infinite or NaN drain_timeout raises ValueError.
• A second call to serve raises RuntimeError.
• request_stop() is a plain function. A second call changes nothing, and it works before serve starts.
• pending() returns the service's unfinished tasks, in any order.

Pull work while running

• Set the state to running. If a stop was requested before serve, pull nothing.
• Check the stop flag before every pull and right after it. Once it is set, launch no more handlers.
• Run each handler in its own task. handle may return any Awaitable, and create_task needs a coroutine, so wrap the call in a small async def.
• After each launch, await asyncio.sleep(0) once. The tests count items on exactly this schedule.
• Run the intake loop as one task, and a stop-wait task that awaits the stop request as another. Wait for whichever finishes first. Only cancelling intake can wake a source that is stuck waiting.

Drain in three windows

When the source runs out, a stop is requested, the source raises, or the caller cancels serve, set the state to draining. Then:

1. Cancel the intake task and the stop-wait task.
2. Wait up to drain_timeout for the remaining tasks.
3. Cancel the tasks still running, and wait up to drain_timeout again.
4. Cancel them again, and wait a third time.

Stop early when no tasks remain. asyncio.wait raises ValueError for an empty set.

Finish

• If no tasks remain, set the state to stopped and return.
• If tasks survive the third wait, keep the state draining, keep them in pending(), and raise TimeoutError. When the last of them finishes, set the state to stopped.
• Read the exception of every finished task that was not cancelled, for example in a done callback. A failing handler does not make serve fail.
• Run the drain as its own task and await it through asyncio.shield. If the caller cancels serve, even repeatedly, keep waiting for the drain, then re-raise the CancelledError.
• If the source raised, re-raise its error after the drain. If the drain also timed out, raise the caller's cancellation or the source error instead of TimeoutError.

What the driver reports

• produced: how many items the source yielded.
• started and handled: sorted IDs of the handlers that began, and of those that finished without an error.
• cancelled: one ID for each CancelledError a handler caught. The stubborn and incomplete handlers each catch more than one.
• returned_normally and outcome: whether serve returned, or else the exception's type name.
• final_state, description and pending: the state, its description, and how many other tasks were alive when serve ended. A clean shutdown gives "stopped", "done" and 0.

In blocked_source, the stop arrives while the source waits forever. In incomplete, one handler ignores every cancellation, so serve must raise TimeoutError in state draining with one pending task. The driver then releases that handler.