All samples

Research sample · Python 3.14

Graceful Python asyncio worker shutdown

Research question

Audit this Python asyncio worker against the declared runtime. Focus on cancellation propagation, SIGTERM shutdown, task ownership, retry races, queue draining, and current asyncio techniques. Ground every high-priority finding in the supplied code.

Topics covered: asyncio graceful shutdown, Python task cancellation, SIGTERM workers

Default research

Python 3.14 · 16,938 characters

Research complete

Question

Audit this Python asyncio worker against the declared runtime. Focus on cancellation propagation, SIGTERM shutdown, task ownership, retry races, queue draining, and current asyncio techniques. Ground every high-priority finding in the supplied code.

Submitted code · asyncio-worker.pyPython
"""Discovery example: a bounded asyncio worker with graceful shutdown.

Declared runtime: Python 3.14.

The sample is deliberately self-contained. Persistence and transport are
protocols so lifecycle, cancellation, retry, and ownership behavior remain
visible without requiring a particular database or message broker.
"""

from __future__ import annotations

import asyncio
import contextlib
import dataclasses
import datetime as dt
import enum
import json
import logging
import random
import signal
import sys
import uuid
from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine
from typing import Any, Protocol, TypeVar


Json = None | bool | int | float | str | list["Json"] | dict[str, "Json"]
T = TypeVar("T")
UTC = dt.UTC


class JobState(enum.StrEnum):
    PENDING = "pending"
    RUNNING = "running"
View all 554 lines
"""Discovery example: a bounded asyncio worker with graceful shutdown.

Declared runtime: Python 3.14.

The sample is deliberately self-contained. Persistence and transport are
protocols so lifecycle, cancellation, retry, and ownership behavior remain
visible without requiring a particular database or message broker.
"""

from __future__ import annotations

import asyncio
import contextlib
import dataclasses
import datetime as dt
import enum
import json
import logging
import random
import signal
import sys
import uuid
from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine
from typing import Any, Protocol, TypeVar


Json = None | bool | int | float | str | list["Json"] | dict[str, "Json"]
T = TypeVar("T")
UTC = dt.UTC


class JobState(enum.StrEnum):
    PENDING = "pending"
    RUNNING = "running"
    RETRY_WAIT = "retry_wait"
    SUCCEEDED = "succeeded"
    DEAD = "dead"


@dataclasses.dataclass(frozen=True, slots=True)
class Failure:
    name: str
    message: str
    code: str
    retryable: bool


@dataclasses.dataclass(frozen=True, slots=True)
class Job:
    id: uuid.UUID
    tenant_id: uuid.UUID
    operation_key: str
    payload: Json
    state: JobState
    attempts: int
    max_attempts: int
    available_at: dt.datetime
    lease_owner: str | None = None
    lease_token: uuid.UUID | None = None
    lease_expires_at: dt.datetime | None = None


@dataclasses.dataclass(frozen=True, slots=True)
class Claim:
    job: Job
    owner: str
    token: uuid.UUID


@dataclasses.dataclass(frozen=True, slots=True)
class WorkerPolicy:
    concurrency: int = 4
    input_capacity: int = 32
    lease_seconds: float = 30.0
    heartbeat_seconds: float = 10.0
    attempt_timeout_seconds: float = 20.0
    shutdown_timeout_seconds: float = 25.0
    poll_seconds: float = 0.25
    base_retry_seconds: float = 1.0
    max_retry_seconds: float = 60.0

    def validate(self) -> None:
        if self.concurrency < 1:
            raise ValueError("concurrency must be positive")
        if self.input_capacity < self.concurrency:
            raise ValueError("input capacity must be at least concurrency")
        if self.heartbeat_seconds >= self.lease_seconds / 2:
            raise ValueError("heartbeat must be less than half of lease duration")
        if self.attempt_timeout_seconds <= 0 or self.shutdown_timeout_seconds <= 0:
            raise ValueError("timeouts must be positive")


class WorkerError(Exception):
    code = "WORKER_ERROR"
    retryable = True


class PermanentJobError(WorkerError):
    code = "PERMANENT_JOB_ERROR"
    retryable = False


class LeaseLostError(WorkerError):
    code = "LEASE_LOST"


class WorkerStoppingError(WorkerError):
    code = "WORKER_STOPPING"


def serialize_failure(error: BaseException) -> Failure:
    if isinstance(error, asyncio.CancelledError):
        return Failure("CancelledError", "Job attempt was cancelled.", "CANCELLED", True)
    if isinstance(error, WorkerError):
        return Failure(type(error).__name__, str(error), error.code, error.retryable)
    return Failure(type(error).__name__, str(error), "UNEXPECTED_ERROR", True)


def utc_now() -> dt.datetime:
    return dt.datetime.now(UTC)


def retry_delay(policy: WorkerPolicy, attempt: int) -> float:
    exponential = min(
        policy.base_retry_seconds * (2 ** max(attempt - 1, 0)),
        policy.max_retry_seconds,
    )
    return exponential * random.uniform(0.8, 1.2)


class JobStore(Protocol):
    async def enqueue(
        self,
        *,
        tenant_id: uuid.UUID,
        operation_key: str,
        payload: Json,
        max_attempts: int,
    ) -> tuple[Job, bool]: ...

    async def claim_next(
        self,
        *,
        tenant_id: uuid.UUID,
        owner: str,
        lease_seconds: float,
        now: dt.datetime,
    ) -> Claim | None: ...

    async def renew(
        self,
        *,
        claim: Claim,
        lease_seconds: float,
        now: dt.datetime,
    ) -> bool: ...

    async def complete(self, *, claim: Claim, result: Json, now: dt.datetime) -> bool: ...

    async def fail(
        self,
        *,
        claim: Claim,
        failure: Failure,
        available_at: dt.datetime | None,
        now: dt.datetime,
    ) -> bool: ...

    async def release_expired(
        self, *, tenant_id: uuid.UUID, now: dt.datetime, limit: int
    ) -> int: ...


class Transport(Protocol):
    async def deliver(self, job: Job) -> Json: ...


@dataclasses.dataclass(slots=True)
class LogContext:
    logger: logging.Logger
    tenant_id: uuid.UUID
    worker_id: str

    def write(self, level: int, event: str, **fields: Json) -> None:
        self.logger.log(
            level,
            json.dumps(
                {
                    "event": event,
                    "tenant_id": str(self.tenant_id),
                    "worker_id": self.worker_id,
                    **fields,
                },
                sort_keys=True,
            ),
        )


class TaskSupervisor:
    """Own tasks and collect failures without losing references."""

    def __init__(self, log: LogContext) -> None:
        self._log = log
        self._tasks: set[asyncio.Task[Any]] = set()
        self._closed = False

    def create(self, awaitable: Coroutine[Any, Any, T], *, name: str) -> asyncio.Task[T]:
        if self._closed:
            awaitable.close()
            raise RuntimeError("task supervisor is closed")
        task = asyncio.create_task(awaitable, name=name)
        self._tasks.add(task)
        task.add_done_callback(self._done)
        return task

    def _done(self, task: asyncio.Task[Any]) -> None:
        self._tasks.discard(task)
        if task.cancelled():
            return
        try:
            error = task.exception()
        except asyncio.CancelledError:
            return
        if error is not None:
            self._log.write(
                logging.ERROR,
                "background_task_failed",
                task=task.get_name(),
                error=repr(error),
            )

    async def cancel_and_wait(self, reason: str) -> None:
        self._closed = True
        tasks = tuple(self._tasks)
        for task in tasks:
            task.cancel(reason)
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)

    async def wait(self) -> None:
        tasks = tuple(self._tasks)
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)

    @property
    def count(self) -> int:
        return len(self._tasks)


class LeaseHeartbeat:
    def __init__(
        self,
        *,
        store: JobStore,
        claim: Claim,
        policy: WorkerPolicy,
        log: LogContext,
    ) -> None:
        self._store = store
        self._claim = claim
        self._policy = policy
        self._log = log
        self._task: asyncio.Task[None] | None = None
        self._lost = asyncio.Event()
        self._stopping = asyncio.Event()

    async def __aenter__(self) -> LeaseHeartbeat:
        if self._task is not None:
            raise RuntimeError("heartbeat already started")
        self._task = asyncio.create_task(
            self._run(),
            name=f"heartbeat:{self._claim.job.id}",
        )
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        traceback: Any,
    ) -> None:
        self._stopping.set()
        task = self._task
        if task is None:
            return
        task.cancel("heartbeat context closed")
        with contextlib.suppress(asyncio.CancelledError):
            await task
        self._task = None

    async def _run(self) -> None:
        try:
            while not self._stopping.is_set():
                await asyncio.sleep(self._policy.heartbeat_seconds)
                renewed = await self._store.renew(
                    claim=self._claim,
                    lease_seconds=self._policy.lease_seconds,
                    now=utc_now(),
                )
                if not renewed:
                    self._lost.set()
                    self._log.write(
                        logging.WARNING,
                        "lease_lost",
                        job_id=str(self._claim.job.id),
                    )
                    return
        except asyncio.CancelledError:
            raise
        except Exception as error:
            self._lost.set()
            self._log.write(
                logging.ERROR,
                "heartbeat_failed",
                job_id=str(self._claim.job.id),
                error=repr(error),
            )

    async def wait_lost(self) -> None:
        await self._lost.wait()


async def race_lease_and_delivery(
    *,
    heartbeat: LeaseHeartbeat,
    deliver: Awaitable[Json],
) -> Json:
    delivery_task = asyncio.create_task(deliver, name="job-delivery")
    lease_task = asyncio.create_task(heartbeat.wait_lost(), name="lease-loss")
    try:
        done, _ = await asyncio.wait(
            {delivery_task, lease_task},
            return_when=asyncio.FIRST_COMPLETED,
        )
        if lease_task in done:
            delivery_task.cancel("lease lost")
            with contextlib.suppress(asyncio.CancelledError):
                await delivery_task
            raise LeaseLostError("lease was lost during delivery")
        lease_task.cancel("delivery completed")
        with contextlib.suppress(asyncio.CancelledError):
            await lease_task
        return await delivery_task
    finally:
        for task in (delivery_task, lease_task):
            if not task.done():
                task.cancel("race cleanup")


class Worker:
    def __init__(
        self,
        *,
        tenant_id: uuid.UUID,
        store: JobStore,
        transport: Transport,
        policy: WorkerPolicy,
        logger: logging.Logger,
    ) -> None:
        policy.validate()
        self.tenant_id = tenant_id
        self.store = store
        self.transport = transport
        self.policy = policy
        self.worker_id = f"worker-{uuid.uuid4()}"
        self.log = LogContext(logger, tenant_id, self.worker_id)
        self.supervisor = TaskSupervisor(self.log)
        self.queue: asyncio.Queue[Claim] = asyncio.Queue(maxsize=policy.input_capacity)
        self.stopping = asyncio.Event()
        self._stop_reason = "shutdown requested"

    def request_stop(self, reason: str) -> None:
        if self.stopping.is_set():
            return
        self._stop_reason = reason
        self.stopping.set()
        self.log.write(logging.INFO, "shutdown_requested", reason=reason)

    async def run(self) -> None:
        await self.store.release_expired(
            tenant_id=self.tenant_id,
            now=utc_now(),
            limit=100,
        )
        try:
            async with asyncio.TaskGroup() as group:
                group.create_task(self._poll(), name="poller")
                for index in range(self.policy.concurrency):
                    group.create_task(self._consume(index), name=f"consumer:{index}")
                await self.stopping.wait()
                await self._drain_for_shutdown()
        except* WorkerStoppingError:
            pass
        finally:
            await self.supervisor.cancel_and_wait(self._stop_reason)

    async def _poll(self) -> None:
        while not self.stopping.is_set():
            if self.queue.full():
                await asyncio.sleep(self.policy.poll_seconds)
                continue
            claim = await self.store.claim_next(
                tenant_id=self.tenant_id,
                owner=self.worker_id,
                lease_seconds=self.policy.lease_seconds,
                now=utc_now(),
            )
            if claim is None:
                await asyncio.sleep(self.policy.poll_seconds)
                continue
            await self.queue.put(claim)
            self.log.write(
                logging.INFO,
                "job_claimed",
                job_id=str(claim.job.id),
                attempt=claim.job.attempts,
            )
        raise WorkerStoppingError(self._stop_reason)

    async def _consume(self, index: int) -> None:
        while True:
            if self.stopping.is_set() and self.queue.empty():
                raise WorkerStoppingError(self._stop_reason)
            try:
                claim = await asyncio.wait_for(self.queue.get(), timeout=0.25)
            except TimeoutError:
                continue
            try:
                await self._execute(claim, index)
            finally:
                self.queue.task_done()

    async def _execute(self, claim: Claim, consumer: int) -> None:
        self.log.write(
            logging.INFO,
            "job_started",
            job_id=str(claim.job.id),
            consumer=consumer,
        )
        try:
            async with LeaseHeartbeat(
                store=self.store,
                claim=claim,
                policy=self.policy,
                log=self.log,
            ) as heartbeat:
                async with asyncio.timeout(self.policy.attempt_timeout_seconds):
                    result = await race_lease_and_delivery(
                        heartbeat=heartbeat,
                        deliver=self.transport.deliver(claim.job),
                    )
            completed = await self.store.complete(
                claim=claim,
                result=result,
                now=utc_now(),
            )
            if not completed:
                raise LeaseLostError(f"completion rejected for {claim.job.id}")
            self.log.write(logging.INFO, "job_succeeded", job_id=str(claim.job.id))
        except asyncio.CancelledError:
            await self._record_failure(
                claim,
                Failure("CancelledError", "worker stopped", "CANCELLED", True),
            )
            raise
        except TimeoutError:
            await self._record_failure(
                claim,
                Failure("TimeoutError", "attempt timed out", "ATTEMPT_TIMEOUT", True),
            )
        except Exception as error:
            await self._record_failure(claim, serialize_failure(error))

    async def _record_failure(self, claim: Claim, failure: Failure) -> None:
        exhausted = claim.job.attempts >= claim.job.max_attempts
        retryable = failure.retryable and not exhausted
        available_at = (
            utc_now() + dt.timedelta(seconds=retry_delay(self.policy, claim.job.attempts))
            if retryable
            else None
        )
        recorded = await self.store.fail(
            claim=claim,
            failure=dataclasses.replace(failure, retryable=retryable),
            available_at=available_at,
            now=utc_now(),
        )
        self.log.write(
            logging.WARNING,
            "job_retry_scheduled" if retryable else "job_dead",
            job_id=str(claim.job.id),
            code=failure.code,
            recorded=recorded,
            available_at=available_at.isoformat() if available_at else None,
        )

    async def _drain_for_shutdown(self) -> None:
        try:
            async with asyncio.timeout(self.policy.shutdown_timeout_seconds):
                await self.queue.join()
                await self.supervisor.wait()
        except TimeoutError:
            self.log.write(
                logging.WARNING,
                "shutdown_deadline_exceeded",
                queued=self.queue.qsize(),
                background=self.supervisor.count,
            )
            await self.supervisor.cancel_and_wait("shutdown deadline exceeded")


@contextlib.asynccontextmanager
async def installed_signal_handlers(worker: Worker) -> AsyncIterator[None]:
    loop = asyncio.get_running_loop()
    installed: list[signal.Signals] = []

    def stop_for(sig: signal.Signals) -> None:
        worker.request_stop(sig.name)

    if sys.platform != "win32":
        for sig in (signal.SIGTERM, signal.SIGINT):
            loop.add_signal_handler(sig, stop_for, sig)
            installed.append(sig)
    try:
        yield
    finally:
        for sig in installed:
            loop.remove_signal_handler(sig)


async def serve(worker: Worker) -> None:
    async with installed_signal_handlers(worker):
        await worker.run()


def configure_logging() -> logging.Logger:
    logger = logging.getLogger("discovery.worker")
    logger.setLevel(logging.INFO)
    if not logger.handlers:
        handler = logging.StreamHandler()
        handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
        logger.addHandler(handler)
    return logger


async def main(store: JobStore, transport: Transport, tenant_id: uuid.UUID) -> None:
    worker = Worker(
        tenant_id=tenant_id,
        store=store,
        transport=transport,
        policy=WorkerPolicy(),
        logger=configure_logging(),
    )
    await serve(worker)

Result

Python3 priority findings2 displayed sources

The drain timer is not a shutdown deadline

The supervisor uses TaskGroup and asyncio.timeout correctly in isolation, but the enclosing TaskGroup can continue waiting after the drain timer expires. Cancellation cleanup and race participants also need stronger task ownership.

Priority findings

03
  1. 01

    Make the SIGTERM deadline bound the whole worker

    Fix now

    The timeout covers queue draining, not TaskGroup exit. The poller and consumers remain owned by the group and can keep shutdown alive after the warning is logged.

    asyncio-worker.py:391

    await self._drain_for_shutdown()
    S7 · WEB-1
  2. 02

    Await every cancelled race participant

    Fix now

    The final cleanup cancels unfinished delivery and lease tasks without awaiting them, allowing their cancellation or exception handling to outlive the race helper.

    asyncio-worker.py:345-347

    task.cancel("race cleanup")
    S7
  3. 03

    Define durable failure behavior under cancellation

    Fix next

    A second cancellation can interrupt failure persistence. Choose whether cancellation must finish recording a retry or rely on lease expiry, then test that contract.

    asyncio-worker.py:460-464

    except asyncio.CancelledError: await self._record_failure(
    S7

Close and drain the queue explicitly

Python 3.14 can reject new puts and wake blocked queue operations while allowing existing work to drain.

Modernized techniquePython
async def stop_workers(queue: asyncio.Queue[Claim]) -> None:
    queue.shutdown(False)
    await queue.join()

async def consume(queue: asyncio.Queue[Claim]) -> None:
    try:
        while True:
            claim = await queue.get()
            try:
                await execute(claim)
            finally:
                queue.task_done()
    except asyncio.QueueShutDown:
        return

Authoritative sources

S7Task cancellationWEB-1Queue.shutdown

What this Python asyncio shutdown audit covers

This Python 3.14 example models a long-running asyncio worker with a bounded queue, TaskGroup supervision, lease renewal, retry scheduling, delivery races, and SIGTERM handling. It is the kind of code that can appear correct in a happy-path test while still hanging during deployment or losing state during cancellation.

Hattrick validates the declared runtime before evaluating syntax or API availability. The research then traces cancellation through the actual task graph and checks the recommendations against the Python 3.14 asyncio documentation, including the newer Queue.shutdown behavior.

Engineering questions answered

  • Does asyncio.timeout bound the full TaskGroup shutdown path?
  • Must cancelled race tasks be awaited before a helper returns?
  • How should a worker persist retry state when CancelledError interrupts a job?

How Hattrick approaches the question

  1. 01

    Runtime-aware validation

    Parses the sample with Python 3.14 before making language-version or asyncio API claims.

  2. 02

    Cancellation-path analysis

    Follows ownership across TaskGroup children, delivery tasks, lease tasks, and durable failure recording.

  3. 03

    Current migration guidance

    Identifies where Queue.shutdown(False) can replace polling and make graceful draining explicit.

Practical takeaway

A drain timeout must own the tasks it intends to stop

Graceful shutdown is not just waiting for queue.join(). A production worker needs a clear producer stop, queue closure, bounded drain, explicit child cancellation, and an awaited outcome for every task that was created.