Worker

Worker process internals: dispatcher, consumer, heartbeat, leader election, shutdown.

taskq.worker itself is a lazy __getattr__ shim (re-exports selected names from its submodules on first access) and renders no useful members on its own. The directives below target the concrete submodules instead. See also the Workers guide for a task-oriented walkthrough.

Process entry point, registration, producer/consumer loops

run

Worker-runtime wiring helpers and process entry point.

This module forms the production seam between WorkerDeps assembly (open_worker_deps, deps.py) and the heartbeat task spawn. It also hosts the worker process entry point worker_main and the _main bootstrap coroutine that wires the full TaskGroup of long-lived siblings.

Deviations from the original sketch

• Signal handlers go through install_signal_handlers, not inline lambdas.
• orchestrate_shutdown takes no tg parameter.
• ProcessScope/ThreadScope/LoopScope are bootstrapped inside open_worker_deps rather than before it (M3 single-process deployment , process exit and deps exit are coterminal; see _main comment block).
• The _local_queue_seed keyword-only parameter is a test seam, not public API.
• local_queue maxsize uses max_concurrency (not batch_size , no batch_size field exists on WorkerSettings).

M1 stub consumers accept a stub_work_timeout keyword-only parameter (default 60.0 s) that controls the sentinel sleep duration. Integration tests may pass a shorter override (e.g. stub_work_timeout=2.0) when seeding jobs that must complete naturally during a short test run. The bootstrap accepts the default; M2 replaces the stubs with the real producer/consumer.

__all__ module-attribute

__all__ = [
    "_emit_resolved_capacity_startup_lines",
    "_main",
    "consumer_loop_stub",
    "deregister_worker",
    "di_consumer_loop",
    "producer_loop",
    "producer_loop_stub",
    "register_worker",
    "worker_main",
    "worker_main_async",
]

worker_main

worker_main(
    settings: WorkerSettings,
    *,
    actor_registry: Mapping[str, ActorRef[Any, Any]]
    | None = None,
    di_registry: ProviderRegistry | None = None,
    cron_registry: list[CronScheduleSpec] | None = None,
    connections: WorkerConnections | None = None,
    pg_credential_provider: PgCredentialProvider
    | None = None,
    rate_limit_registry: RateLimitRegistry | None = None,
    until_idle: bool = False,
    idle_settle_window: float | None = None,
    idle_poll_interval: float | None = None,
    idle_max_runtime: float | None = None,
) -> int

Worker process entry point (synchronous).

Thin wrapper: runs :func:worker_main_async under an asyncio.Runner and returns its int result. Uses Runner (not asyncio.run) for finer control over teardown.

Because this function owns the loop, an application that must build loop-bound dependencies (asyncpg pools) before constructing its actors cannot use it - await :func:worker_main_async from its own coroutine instead. Both take exactly these parameters.

actor_registry is a mapping from short name to :class:ActorRef containing every @actor-decorated handler this worker intends to run. Forwarded to :func:_main for the bootstrap config sync.

pg_credential_provider is the resolved Postgres credential provider for the worker-internal per-slot transaction pool (see :func:_main). Pass the same provider used to build connections, the slot pool does not read WorkerConnections. None authenticates it from the direct DSN.

di_registry is an optional pre-configured :class:ProviderRegistry containing application-specific provider registrations (database pools, HTTP clients, etc.). When supplied, the worker uses it instead of creating a fresh registry, callers must NOT call validate() before passing it here; the worker calls validate() as part of its bootstrap sequence. WorkerSettings and Clock are registered automatically if not already present.

rate_limit_registry is an optional owned :class:RateLimitRegistry for this worker (e.g. one instance per process in a multi-process deployment). When None, resolution falls back to a RateLimitRegistry value provider at Scope.LOOP in di_registry, then to the module singleton, import-time .register() on the singleton keeps working exactly as before. Co-present with a RateLimitRegistry provider in di_registry this raises TypeError (ambiguous, bootstrap and dispatch would diverge); pass one or the other. Forwarded to :func:_main.

cron_registry is an optional list of :class:CronScheduleSpec objects to auto-register at startup. When None (the default), get_registered_crons() is used instead, schedules declared via the @cron decorator are auto-discovered. When an explicit list is passed (even empty []), only those schedules are registered; decorator-registered schedules are skipped. For each spec, a direct INSERT INTO … cron_schedules is executed inside try/except asyncpg.UniqueViolationError: pass, the DB (actor, name) UNIQUE constraint prevents duplicates, so concurrent worker replicas can safely race. Startup auto-discovery is create-only, skip-on-conflict: existing cron_schedules rows are never modified by the registration pass. If a @cron decorator's parameters change after the schedule was first registered, the operator must manually update or delete and recreate the schedule.

until_idle enables drain mode (see :func:_main for details). idle_settle_window, idle_poll_interval, and idle_max_runtime override the corresponding settings when until_idle is True; they are ignored otherwise.

Returns the exit code from :func:_main, 0 on clean shutdown, 3 if any jobs failed in drain mode, 4 if idle_max_runtime was exceeded.

Source code in src/taskq/worker/_bootstrap.py

def worker_main(
    settings: WorkerSettings,
    *,
    actor_registry: Mapping[str, ActorRef[Any, Any]] | None = None,
    di_registry: ProviderRegistry | None = None,
    cron_registry: list[CronScheduleSpec] | None = None,
    connections: WorkerConnections | None = None,
    pg_credential_provider: PgCredentialProvider | None = None,
    rate_limit_registry: RateLimitRegistry | None = None,
    until_idle: bool = False,
    idle_settle_window: float | None = None,
    idle_poll_interval: float | None = None,
    idle_max_runtime: float | None = None,
) -> int:
    """Worker process entry point (synchronous).

    Thin wrapper: runs :func:`worker_main_async` under an ``asyncio.Runner``
    and returns its int result. Uses ``Runner`` (not ``asyncio.run``) for
    finer control over teardown.

    Because this function owns the loop, an application that must build
    loop-bound dependencies (asyncpg pools) before constructing its actors
    cannot use it - await :func:`worker_main_async` from its own coroutine
    instead. Both take exactly these parameters.

    ``actor_registry`` is a mapping from short name to :class:`ActorRef`
    containing every ``@actor``-decorated handler this worker intends to
    run. Forwarded to :func:`_main` for the  bootstrap config sync.

    ``pg_credential_provider`` is the resolved Postgres credential
    provider for the worker-internal per-slot transaction pool (see
    :func:`_main`). Pass the same provider used to build
    ``connections``, the slot pool does not read WorkerConnections.
    ``None`` authenticates it from the direct DSN.

    ``di_registry`` is an optional pre-configured :class:`ProviderRegistry`
    containing application-specific provider registrations (database pools,
    HTTP clients, etc.).  When supplied, the worker uses it instead of
    creating a fresh registry, callers must NOT call ``validate()`` before
    passing it here; the worker calls ``validate()`` as part of its bootstrap
    sequence.  ``WorkerSettings`` and ``Clock`` are registered automatically
    if not already present.

    ``rate_limit_registry`` is an optional owned :class:`RateLimitRegistry`
    for this worker (e.g. one instance per process in a multi-process
    deployment).  When ``None``, resolution falls back to a
    ``RateLimitRegistry`` value provider at ``Scope.LOOP`` in
    ``di_registry``, then to the module singleton, import-time
    ``.register()`` on the singleton keeps working exactly as before.
    Co-present with a ``RateLimitRegistry`` provider in ``di_registry``
    this raises ``TypeError`` (ambiguous, bootstrap and dispatch would
    diverge); pass one or the other.
    Forwarded to :func:`_main`.

    ``cron_registry`` is an optional list of :class:`CronScheduleSpec`
    objects to auto-register at startup.  When ``None`` (the default),
    ``get_registered_crons()`` is used instead, schedules declared via
    the ``@cron`` decorator are auto-discovered.  When an explicit list
    is passed (even empty ``[]``), only those schedules are registered;
    decorator-registered schedules are skipped.  For each spec, a direct
    ``INSERT INTO … cron_schedules`` is executed inside
    ``try/except asyncpg.UniqueViolationError: pass``, the DB ``(actor, name)``
    UNIQUE constraint prevents duplicates, so concurrent worker replicas
    can safely race.  Startup auto-discovery is **create-only,
    skip-on-conflict**: existing ``cron_schedules`` rows are never
    modified by the registration pass.  If a ``@cron`` decorator's
    parameters change after the schedule was first registered, the
    operator must manually update or delete and recreate the schedule.

    ``until_idle`` enables drain mode (see :func:`_main` for details).
    ``idle_settle_window``, ``idle_poll_interval``, and
    ``idle_max_runtime`` override the corresponding settings when
    ``until_idle`` is True; they are ignored otherwise.

    Returns the exit code from :func:`_main`, 0 on clean shutdown,
    3 if any jobs failed in drain mode, 4 if ``idle_max_runtime`` was
    exceeded.
    """
    with asyncio.Runner() as runner:
        code = runner.run(
            worker_main_async(
                settings,
                actor_registry=actor_registry,
                di_registry=di_registry,
                connections=connections,
                cron_registry=cron_registry,
                pg_credential_provider=pg_credential_provider,
                rate_limit_registry=rate_limit_registry,
                until_idle=until_idle,
                idle_settle_window=idle_settle_window,
                idle_poll_interval=idle_poll_interval,
                idle_max_runtime=idle_max_runtime,
            )
        )
    # The drain's verdict is guarded by the ENTRYPOINTS (the taskq worker
    # command and the e2e/system-e2e worker entries), not here. This must
    # not end with ``signal.signal(SIGTERM, SIG_IGN)`` - correct for a
    # process about to ``sys.exit``, poison for any other caller: an
    # ignored disposition is inherited by every process this host forks
    # from then on (a child spawned after an in-process ``worker_main``
    # ignores SIGTERM and dies to SIGKILL, rc -9). A library function
    # returns control to its caller and must leave the host's signal state
    # exactly as it found it; the verdict guard belongs where the process
    # is genuinely about to die. SIGINT is left alone everywhere: its
    # default disposition is KeyboardInterrupt, a live interactive
    # contract, not a silent status eraser.
    return code

worker_main_async async

worker_main_async(
    settings: WorkerSettings,
    *,
    actor_registry: Mapping[str, ActorRef[Any, Any]]
    | None = None,
    di_registry: ProviderRegistry | None = None,
    cron_registry: list[CronScheduleSpec] | None = None,
    connections: WorkerConnections | None = None,
    pg_credential_provider: PgCredentialProvider
    | None = None,
    rate_limit_registry: RateLimitRegistry | None = None,
    until_idle: bool = False,
    idle_settle_window: float | None = None,
    idle_poll_interval: float | None = None,
    idle_max_runtime: float | None = None,
) -> int

Run a worker to completion on the caller's event loop.

Same parameters, same semantics and same exit codes as :func:worker_main - which is now the thin synchronous wrapper that drives this coroutine under its own asyncio.Runner. See :func:worker_main for what each parameter means.

Why this exists as public API: worker_main owns the loop, and asyncpg pools are loop-bound. An application whose actors are closures over its own dependencies has to build those dependencies (pools, HTTP clients) on the loop the worker will run on, before the actor registry is constructed - and there is no hook in worker_main to do that. Pre-creating a pool in a throwaway loop and handing it to worker_main's Runner is not an option: the pool would belong to a dead loop. Awaiting this from inside the caller's own coroutine puts dependency construction, actor construction, and the worker on one loop::

async def main() -> int:
    pool = await asyncpg.create_pool(dsn)          # this loop
    registry = build_actors(pool)                  # closures over it
    try:
        return await worker_main_async(settings, actor_registry=registry)
    finally:
        await pool.close()

raise SystemExit(asyncio.run(main()))

Signal handling, logging setup and cron auto-discovery are identical to :func:worker_main; nothing about running here changes the worker's behaviour. The caller owns whatever it built - TaskQ closes only the resources it created itself (see :class:~taskq.connections.WorkerConnections for the ownership rule that governs anything passed via connections).

Source code in src/taskq/worker/_bootstrap.py

async def worker_main_async(
    settings: WorkerSettings,
    *,
    actor_registry: Mapping[str, ActorRef[Any, Any]] | None = None,
    di_registry: ProviderRegistry | None = None,
    cron_registry: list[CronScheduleSpec] | None = None,
    connections: WorkerConnections | None = None,
    pg_credential_provider: PgCredentialProvider | None = None,
    rate_limit_registry: RateLimitRegistry | None = None,
    until_idle: bool = False,
    idle_settle_window: float | None = None,
    idle_poll_interval: float | None = None,
    idle_max_runtime: float | None = None,
) -> int:
    """Run a worker to completion on the **caller's** event loop.

    Same parameters, same semantics and same exit codes as
    :func:`worker_main` - which is now the thin synchronous wrapper that
    drives this coroutine under its own ``asyncio.Runner``. See
    :func:`worker_main` for what each parameter means.

    Why this exists as public API: ``worker_main`` owns the loop, and
    asyncpg pools are **loop-bound**. An application whose actors are
    closures over its own dependencies has to build those dependencies
    (pools, HTTP clients) on the loop the worker will run on, *before* the
    actor registry is constructed - and there is no hook in ``worker_main``
    to do that. Pre-creating a pool in a throwaway loop and handing it to
    ``worker_main``'s Runner is not an option: the pool would belong to a
    dead loop. Awaiting this from inside the caller's own coroutine puts
    dependency construction, actor construction, and the worker on one
    loop::

        async def main() -> int:
            pool = await asyncpg.create_pool(dsn)          # this loop
            registry = build_actors(pool)                  # closures over it
            try:
                return await worker_main_async(settings, actor_registry=registry)
            finally:
                await pool.close()

        raise SystemExit(asyncio.run(main()))

    Signal handling, logging setup and cron auto-discovery are identical to
    :func:`worker_main`; nothing about running here changes the worker's
    behaviour. The caller owns whatever it built - TaskQ closes only the
    resources it created itself (see :class:`~taskq.connections.WorkerConnections`
    for the ownership rule that governs anything passed via ``connections``).
    """
    from taskq.scheduler import get_registered_crons

    schedule_specs = cron_registry if cron_registry is not None else get_registered_crons()
    setup_logging(
        level=settings.log_level,
        log_format=settings.log_format,
        events_level=settings.log_events_level,
    )
    return await _main(
        settings,
        actor_registry=actor_registry,
        _registry=di_registry,
        _cron_registry=schedule_specs,
        connections=connections,
        pg_credential_provider=pg_credential_provider,
        rate_limit_registry=rate_limit_registry,
        until_idle=until_idle,
        idle_settle_window=idle_settle_window,
        idle_poll_interval=idle_poll_interval,
        idle_max_runtime=idle_max_runtime,
    )

__getattr__

__getattr__(name: str) -> object

Source code in src/taskq/worker/run.py

def __getattr__(name: str) -> object:
    if name in (
        "_main",
        "_emit_sub_enqueue_startup_warnings",
        "_emit_unconsumed_queue_startup_warnings",
    ):
        from taskq.worker._bootstrap import (
            _emit_sub_enqueue_startup_warnings,
            _emit_unconsumed_queue_startup_warnings,
            _main,
        )

        return {
            "_main": _main,
            "_emit_sub_enqueue_startup_warnings": _emit_sub_enqueue_startup_warnings,
            "_emit_unconsumed_queue_startup_warnings": _emit_unconsumed_queue_startup_warnings,
        }[name]
    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

make_heartbeat_kwargs

make_heartbeat_kwargs(
    deps: WorkerDeps,
    worker_id: UUID,
    backend: Backend,
    cancel_wake_event: Event | None = None,
) -> dict[str, object]

Return keyword arguments that wire a cancel controller and optional cancel-wake event into heartbeat_loop for a given worker.

Usage (production, owned by the orchestration layer)::

kwargs = make_heartbeat_kwargs(deps, worker_id, backend, cancel_wake_event)
await heartbeat_loop(deps, worker_id, shutdown, **kwargs)

Returns:

Type	Description
dict[str, object]	{"cancel_controller": ..., "cancel_wake_event": ...}.

Source code in src/taskq/worker/run.py

def make_heartbeat_kwargs(
    deps: WorkerDeps,
    worker_id: UUID,
    backend: Backend,
    cancel_wake_event: asyncio.Event | None = None,
) -> dict[str, object]:
    """Return keyword arguments that wire a cancel controller and optional
    cancel-wake event into heartbeat_loop for a given worker.

    Usage (production, owned by the orchestration layer)::

        kwargs = make_heartbeat_kwargs(deps, worker_id, backend, cancel_wake_event)
        await heartbeat_loop(deps, worker_id, shutdown, **kwargs)

    Returns:
        ``{"cancel_controller": ..., "cancel_wake_event": ...}``.
    """
    return {
        "cancel_controller": make_cancel_controller(deps, worker_id, backend),
        "cancel_wake_event": cancel_wake_event,
    }

producer_loop async

producer_loop(
    deps: WorkerDeps,
    local_queue: Queue[JobRow],
    shutdown_event: Event,
    producer_stop_event: Event,
    *,
    backend: Backend,
    worker_id: UUID,
    slot_freed_event: Event | None = None,
    rng: Random | None = None,
) -> None

Dispatch pending jobs from the database and feed them into local_queue.

On each iteration the producer:

1. Waits for a wake signal (NOTIFY-driven asyncio.Event) or the poll_interval fallback timer, whichever fires first.
2. Calls backend.dispatch_batch() to atomically claim up to the worker's genuinely free slots, local_queue.maxsize - local_queue.qsize() - active_jobs.count() - active_jobs.intent_count() pending jobs (pending → running) using FOR UPDATE SKIP LOCKED.
3. Puts each returned :class:JobRow onto local_queue for the consumer tasks.

A round that returns fewer rows than it asked for arms a short jittered cooldown (:data:_CLAIM_COOLDOWN_SECONDS) before the next round, so a burst of wakes or freed slots costs one round rather than one per trigger; a full round re-claims immediately.

Slot accounting: a worker with max_concurrency slots may hold at most that many rows locked at once, the claim sizes by the slots actually free (queue emptiness minus jobs actively running), not by queue emptiness alone, which allowed up to 2x max_concurrency rows locked (a full local queue while every consumer was busy still looked like max_concurrency free slots: double reclaim exposure on a crash, and head-of-line latency behind long jobs while peer workers idle). The get()-to-register window is counted through the claim intent the consumer loop records at the take, so a consumer paused in DI resolution or a denied-admission local retry is still a held slot to this arithmetic; the completion-side slot frees at the consumer's deregister, and the wake the consumer sets there re-arms this producer the moment accounting settles, no claim waits for the next poll tick.

Exits cleanly when either shutdown_event or producer_stop_event is set.

slot_freed_event is set by the consumer loops at the two slot-release points, when a local_queue.get() drains a queue slot, and when a job's active_jobs registration ends; the bootstrap wires one shared event into this loop and every consumer. rng supplies the fallback-poll jitter (a test seam; production uses the module RNG seeded per process). Both default to standalone behaviour: a private event nobody sets degrades the slot-refill wait to the fallback cadence, and the module RNG jitters as in production.

Source code in src/taskq/worker/run.py

async def producer_loop(
    deps: WorkerDeps,
    local_queue: asyncio.Queue[JobRow],
    shutdown_event: asyncio.Event,
    producer_stop_event: asyncio.Event,
    *,
    backend: Backend,
    worker_id: UUID,
    slot_freed_event: asyncio.Event | None = None,
    rng: random.Random | None = None,
) -> None:
    """Dispatch pending jobs from the database and feed them into ``local_queue``.

    On each iteration the producer:

    1. Waits for a wake signal (NOTIFY-driven ``asyncio.Event``) or the
       ``poll_interval`` fallback timer, whichever fires first.
    2. Calls ``backend.dispatch_batch()`` to atomically claim up to the
       worker's genuinely free slots, ``local_queue.maxsize -
       local_queue.qsize() - active_jobs.count() -
       active_jobs.intent_count()`` pending jobs
       (pending → running) using ``FOR UPDATE SKIP LOCKED``.
    3. Puts each returned :class:`JobRow` onto ``local_queue`` for the
       consumer tasks.

    A round that returns fewer rows than it asked for arms a short
    jittered cooldown (:data:`_CLAIM_COOLDOWN_SECONDS`) before the next
    round, so a burst of wakes or freed slots costs one round rather than
    one per trigger; a full round re-claims immediately.

    Slot accounting: a worker with ``max_concurrency`` slots may
    hold at most that many rows locked at once, the claim sizes by the
    slots actually free (queue emptiness minus jobs actively running),
    not by queue emptiness alone, which allowed up to 2x
    ``max_concurrency`` rows locked (a full local queue while every
    consumer was busy still looked like ``max_concurrency`` free slots:
    double reclaim exposure on a crash, and head-of-line latency behind
    long jobs while peer workers idle). The get()-to-register window is
    counted through the claim intent the consumer loop records at the
    take, so a consumer paused in DI resolution or a denied-admission
    local retry is still a held slot to this arithmetic; the
    completion-side slot frees at the consumer's ``deregister``, and
    the wake the consumer sets there re-arms this producer the moment
    accounting settles, no claim waits for the next poll tick.

    Exits cleanly when either ``shutdown_event`` or ``producer_stop_event``
    is set.

    ``slot_freed_event`` is set by the consumer loops at the two
    slot-release points, when a ``local_queue.get()`` drains a queue
    slot, and when a job's ``active_jobs`` registration ends; the
    bootstrap wires one shared event into this loop and every consumer.
    ``rng`` supplies the
    fallback-poll jitter (a test seam; production uses the module RNG
    seeded per process). Both default to standalone behaviour: a private
    event nobody sets degrades the slot-refill wait to the fallback
    cadence, and the module RNG jitters as in production.
    """
    settings = deps.settings
    queues = settings.queues
    lock_lease_td = timedelta(seconds=settings.lock_lease)
    notify_enabled = getattr(settings, "notify_enabled", False)
    # Same getattr-with-default convention as notify_enabled above: the
    # producer-loop unit tests drive SimpleNamespace stand-ins for
    # WorkerSettings, and the pooled flag defaults to False (direct DSN).
    pooled = bool(getattr(settings, "pg_is_pooled", False))
    poll_interval = settings.notify_poll_interval if notify_enabled else settings.poll_interval
    rng_source = rng if rng is not None else _PRODUCER_RNG
    # Wakes this producer at the consumers' two slot-release points: a
    # local_queue.get() draining a queue slot, and a job's active_jobs
    # deregister freeing an active slot (see the saturation branch
    # below). None keeps the loop standalone: a private event nobody
    # sets degrades the bounded wait to exactly the fixed-cadence poll
    # it replaces.
    slot_freed = slot_freed_event if slot_freed_event is not None else asyncio.Event()

    _producer_log.info(
        "producer-loop-start",
        queues=queues,
        poll_interval=poll_interval,
        notify_enabled=notify_enabled,
        max_concurrency=settings.max_concurrency,
        worker_id=str(worker_id),
    )
    # Set once any claim round commits rows. Guards the exit hand-back
    # below: a producer that never claimed cannot hold a locked row, so
    # its exit owes the fleet no write (the common idle-shutdown shape).
    made_a_claim = False
    # The unexpected-error backstop, the same one every leader
    # maintenance loop carries (see taskq.worker._transient): a
    # NON-transient dispatch failure (a revoked grant, a code bug, a
    # data error) is logged loudly each round but must not retry
    # forever into a zombie that keeps ticking, keeps its liveness
    # registration fresh and stays ready while claiming nothing. After
    # ``max_consecutive`` consecutive unexpected rounds the ORIGINAL
    # error propagates and tears the worker down deliberately;
    # transient failures never feed it, and only a fully successful
    # round resets the streak, exactly the leader loops' semantics.
    guard = UnexpectedLoopErrorGuard("worker.producer")
    # Monotonic deadline before which no claim round may start, armed by
    # a short round (see _CLAIM_COOLDOWN_SECONDS), so the triggers that
    # land while it runs (wakes, freed slots) coalesce into one round.
    claim_not_before = 0.0

    async with contextlib.AsyncExitStack() as stack:
        wake_event: asyncio.Event | None = None
        if notify_enabled:
            _subscribe_wake = getattr(backend, "subscribe_wake", None)
            if callable(_subscribe_wake):
                wake_event = await stack.enter_async_context(
                    cast(
                        "contextlib.AbstractAsyncContextManager[asyncio.Event]",
                        _subscribe_wake(queues=queues),
                    )
                )
                _producer_log.info("producer-subscribed-wake", worker_id=str(worker_id))
            else:
                _producer_log.warning(
                    "producer-no-wake-subscribe",
                    note="subscribe_wake not available; falling back to poll-only",
                    worker_id=str(worker_id),
                )

        while not (shutdown_event.is_set() or producer_stop_event.is_set()):
            deps.liveness.tick("producer", period=poll_interval)
            # The worker's genuinely free slots: every slot is
            # either empty, lent to a queued row (qsize), claimed by a
            # consumer whose registration has not absorbed the take yet
            # (intent_count: the claim-to-register window where DI
            # resolution and the admission acquire run), or occupied by
            # a running job (active_jobs). Sizing the claim by queue
            # emptiness alone counted a fully-busy worker's slots as
            # free whenever its queue had drained, up to 2x
            # max_concurrency rows locked fleet-wide, 2x the reclaim
            # exposure on a crash, and pending work locked behind long
            # jobs while peer workers idled. The get()-to-register
            # window (a row taken from the queue but not yet in
            # active_jobs) is one scheduler step wide and bounded by the
            # consumer count; the reverse, a row finished but not yet
            # deregistered, delays only its own slot's re-claim until
            # the deregister-side wake, never a poll tick.
            available = (
                local_queue.maxsize
                - local_queue.qsize()
                - deps.active_jobs.count()
                - deps.active_jobs.intent_count()
            )
            if available <= 0:
                # All consumer slots busy and the local queue full. A
                # consumer's get() frees a queue slot and a job's
                # completion (deregister) frees an active slot, the
                # consumer loops set slot_freed at exactly those two
                # points, so the next claim begins the moment either
                # lands instead of on the next poll tick.
                # Bounded, not bare: an event that is never set (broken
                # wiring, a consumer-less worker) must still leave this
                # loop re-checking on the fallback cadence.
                with contextlib.suppress(TimeoutError):
                    await asyncio.wait_for(slot_freed.wait(), timeout=_SLOT_REFILL_POLL_SECONDS)
                slot_freed.clear()
                continue

            cooldown_remaining = claim_not_before - time.monotonic()
            if cooldown_remaining > 0:
                # Every trigger that lands during this wait, more wakes,
                # more freed slots, is answered by the single round that
                # follows, sized to the slots free by then. Re-entering
                # the loop re-reads availability and the stop flags.
                await asyncio.sleep(cooldown_remaining)
                continue

            # Cleared BEFORE the round, not after: a NOTIFY that lands
            # while the round runs may announce a row the round's
            # snapshot predates, and must survive as exactly one
            # follow-up round. Everything that arrived before this point
            # is answered by the round itself.
            if wake_event is not None:
                wake_event.clear()
            round_started = time.monotonic()
            try:
                jobs = await backend.dispatch_batch(
                    worker_id=worker_id,
                    queues=queues,
                    limit=available,
                    lock_lease=lock_lease_td,
                )
            except Exception as exc:
                # Transient shapes degrade quietly: log at warning and
                # retry next tick (the TRANSIENT_PG_ERRORS doctrine; the
                # pooled extension covers a transaction-mode pooler
                # remapping a prepared statement's server connection out
                # from under a cached name, SQLSTATE 26000/42P05, declared
                # via TASKQ_PG_IS_POOLED). Everything else keeps the loud,
                # exception-level record: a real bug must stay visible.
                if is_transient_pg_error(exc, pooled=pooled):
                    _producer_log.warning(
                        "dispatch-batch-transient",
                        worker_id=str(worker_id),
                        error_class=type(exc).__name__,
                        error=str(exc),
                    )
                else:
                    _producer_log.exception("dispatch-batch-error", worker_id=str(worker_id))
                    # Outside the transient set: feed the backstop. At the
                    # consecutive-failure cap this re-raises the original
                    # error, ending the zombie-tick state the loop would
                    # otherwise never leave (see the guard construction
                    # above and taskq.worker._transient).
                    guard.unexpected(exc)
                with contextlib.suppress(asyncio.CancelledError):
                    await asyncio.sleep(poll_interval)
                continue

            # The round completed without error: reset the backstop's
            # streak. An empty round counts, a claimed round counts, and
            # neither a transient nor an unexpected failure reaches here.
            guard.ok()

            # DEBUG_ONLY (TASKQ_LOG_EVENTS_LEVEL=debug; see
            # taskq.obs._structlog's classification table): the poll
            # cadence detail, one line per round. stdlib INFO deliberately
            # - the events-level filter, not the root logger level, owns
            # this line's visibility, and it is dropped at
            # info/warning/off (the default stream stays byte-identical).
            _producer_log.info(
                "poll-cadence",
                kind="debug",
                worker_id=str(worker_id),
                claimed=len(jobs),
                available=available,
            )

            if len(jobs) < available:
                # A short round: the backlog is drained or a peer won it.
                # A round on its heels would only re-run the claim CTE and
                # the loser's window expansions for nothing. A full round
                # arms no floor, there is backlog to drain, and the next
                # freed slot claims immediately.
                claim_not_before = round_started + _jittered_poll_interval(
                    _CLAIM_COOLDOWN_SECONDS, rng_source
                )

            if jobs:
                made_a_claim = True
                for job in jobs:
                    # A row this worker disowned, the sweep re-pended and
                    # this claim took back is a live job of ours again:
                    # its lease must be renewed from here on.
                    deps.disowned_jobs.discard(job.id)
                    # Fence the claim before the first await: the put is
                    # this loop's first yield since the claim committed,
                    # and the lost-claim probe (the heartbeat tick) must
                    # see the row held from the instant it is running and
                    # locked here, through the queue residence, until the
                    # consumer's take moves the coverage to the intent
                    # map.
                    deps.active_jobs.mark_enqueued(job.id)
                    await local_queue.put(job)
                continue

            wake_wait = asyncio.create_task(wake_event.wait()) if wake_event is not None else None
            # Jittered per wait, not fixed: an idle fleet sharing one
            # interval polls in phase, and any transient event
            # re-synchronizes it into periodic DB load spikes (see
            # _jittered_poll_interval for the vendor precedent).
            poll_wait = asyncio.create_task(
                asyncio.sleep(_jittered_poll_interval(poll_interval, rng_source))
            )
            stop_wait = asyncio.create_task(producer_stop_event.wait())
            shutdown_wait = asyncio.create_task(shutdown_event.wait())
            # A consumer's slot release (get() or deregister) is itself a
            # wake: an empty round caused by a closed admission gate (the
            # queue-cap damper admits nothing while a running cohort holds
            # the cap) self-corrects the moment that cohort completes, and
            # the completion's release lands here as this event. Without
            # it the producer sleeps the full fallback poll (5s shipped)
            # with pending work, free slots, and an idle DB: the burst
            # cadence collapses to one cohort per poll. An empty round
            # on a genuinely drained queue never sees the event set and
            # still waits out the poll.
            slot_wait = asyncio.create_task(slot_freed.wait())

            all_waits = [
                w
                for w in (
                    wake_wait,
                    poll_wait,
                    stop_wait,
                    shutdown_wait,
                    slot_wait,
                )
                if w is not None
            ]

            try:
                _done, pending = await asyncio.wait(
                    all_waits,
                    return_when=asyncio.FIRST_COMPLETED,
                )
                for task in pending:
                    task.cancel()
                    with contextlib.suppress(asyncio.CancelledError):
                        await task
            finally:
                for task in all_waits:
                    if not task.done():
                        task.cancel()
                        with contextlib.suppress(asyncio.CancelledError):
                            await task

            # Cleared after the wait, not before: a release landing while
            # the round ran is covered by the round's own claim (the loop
            # continues past a full round without waiting); one landing
            # between this clear and the next round's snapshot re-arms the
            # event, and the next empty round's wait returns to it
            # immediately, a bounded re-check, never a hot loop (the wait
            # is re-entered only after a round that claimed nothing).
            slot_freed.clear()

            if poll_wait in _done:
                # The fallback poll is its own cadence, already jittered
                # and never shorter than a burst: a poll-timed round owes
                # no further wait, whatever the previous round returned.
                claim_not_before = 0.0

    # Exit hand-back, on the DRAINING path only (producer_stop_event is
    # what the orchestrator sets at DRAINING entry): the DRAINING pass
    # re-pends what this worker held when it ran, but a claim round already
    # in flight at that moment commits AFTER it, the producer only
    # observes the stop event between rounds. Those rows are locked to a
    # process on its way out and no later phase sees them (they never reach
    # the active-jobs registry), so the producer hands back whatever it
    # still holds as its own last act on that path: by loop exit no further
    # claim of this worker's can commit, which is exactly the ordering the
    # single DRAINING pass could not give. The statement is the same
    # bounded, registry-excluding, attempt-refunding one (idempotent, a
    # row the first pass already released no longer matches), and its
    # failure mode is the helper's own (log + return 0; the lease-expiry
    # sweep remains the backstop). A bare shutdown_event exit (the
    # external-stop path, no orchestration) keeps the lease-reclaim shape
    # it has always had.
    if producer_stop_event.is_set() and made_a_claim:
        handed_back = await drain_local_queue_to_pending(deps, worker_id)
        if handed_back:
            _producer_log.info(
                "producer-exit-handback",
                worker_id=str(worker_id),
                rows_re_pended=handed_back,
            )

    reason = "producer_stop_event" if producer_stop_event.is_set() else "shutdown_event"
    _producer_log.info("producer-loop-exit", reason=reason)
    deps.liveness.forget("producer")

producer_loop_stub async

producer_loop_stub(
    deps: WorkerDeps,
    local_queue: Queue[JobRow],
    shutdown_event: Event,
    producer_stop_event: Event,
    *,
    backend: Backend,
    worker_id: UUID,
) -> None

Observe producer_stop_event and shutdown_event; exit cleanly.

Why the unused parameters: this is a drop-in stand-in for :func:producer_loop (used when a worker runs consume-only), so it takes that function's signature verbatim and the caller can swap one for the other without a shim. deps, local_queue, backend and worker_id are exactly the arguments it does not need because it never dispatches, a divergent signature here would move that branch into every call site.

Outer loop: while not (producer_stop_event.is_set() or shutdown_event.is_set()). Body races producer_stop_event.wait() against shutdown_event.wait() via asyncio.wait(..., return_when=FIRST_COMPLETED); the loser is cancelled to avoid a pending-task leak.

Source code in src/taskq/worker/run.py

async def producer_loop_stub(
    deps: WorkerDeps,
    local_queue: asyncio.Queue[JobRow],
    shutdown_event: asyncio.Event,
    producer_stop_event: asyncio.Event,
    *,
    backend: Backend,
    worker_id: UUID,
) -> None:
    """Observe producer_stop_event and shutdown_event; exit cleanly.

    Why the unused parameters: this is a drop-in stand-in for
    :func:`producer_loop` (used when a worker runs consume-only), so it
    takes that function's signature verbatim and the caller can swap one
    for the other without a shim. ``deps``, ``local_queue``, ``backend``
    and ``worker_id`` are exactly the arguments it does not need because
    it never dispatches, a divergent signature here would move that
    branch into every call site.

    Outer loop: ``while not (producer_stop_event.is_set() or shutdown_event.is_set())``.
    Body races ``producer_stop_event.wait()`` against ``shutdown_event.wait()`` via
    ``asyncio.wait(..., return_when=FIRST_COMPLETED)``; the loser is cancelled
    to avoid a pending-task leak.
    """
    while not (producer_stop_event.is_set() or shutdown_event.is_set()):
        stop_wait = asyncio.create_task(producer_stop_event.wait())
        shutdown_wait = asyncio.create_task(shutdown_event.wait())
        try:
            _done, pending = await asyncio.wait(
                [stop_wait, shutdown_wait],
                return_when=asyncio.FIRST_COMPLETED,
            )
            for task in pending:
                task.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await task
        finally:
            for task in [stop_wait, shutdown_wait]:
                if not task.done():
                    task.cancel()

    reason = "producer_stop_event" if producer_stop_event.is_set() else "shutdown_event"
    _producer_log.info("producer-loop-exit", reason=reason)

consumer_loop_stub async

consumer_loop_stub(
    deps: WorkerDeps,
    local_queue: Queue[JobRow],
    shutdown_event: Event,
    *,
    backend: Backend,
    worker_id: UUID,
    stub_work_timeout: float = 60.0,
    slot_freed_event: Event | None = None,
) -> None

Pull one job per iteration, register, sleep sentinel, write terminal status.

Outer loop: while not shutdown_event.is_set(). Races local_queue.get() against shutdown_event.wait(); on shutdown win the queue waiter is cancelled and the stub returns cleanly.

On job get the stub registers in deps.active_jobs, awaits a cancellable sentinel, writes terminal state via backend (shielded), and deregisters in finally.

Source code in src/taskq/worker/run.py

async def consumer_loop_stub(
    deps: WorkerDeps,
    local_queue: asyncio.Queue[JobRow],
    shutdown_event: asyncio.Event,
    *,
    backend: Backend,
    worker_id: UUID,
    stub_work_timeout: float = 60.0,
    slot_freed_event: asyncio.Event | None = None,
) -> None:
    """Pull one job per iteration, register, sleep sentinel, write terminal status.

    Outer loop: ``while not shutdown_event.is_set()``.
    Races ``local_queue.get()`` against ``shutdown_event.wait()``; on shutdown
    win the queue waiter is cancelled and the stub returns cleanly.

    On job get the stub registers in ``deps.active_jobs``, awaits a cancellable
    sentinel, writes terminal state via ``backend`` (shielded), and deregisters
    in ``finally``.
    """
    while not (shutdown_event.is_set() or deps.producer_stop_event.is_set()):
        q_get = asyncio.create_task(local_queue.get())
        shut_wait = asyncio.create_task(shutdown_event.wait())
        stop_wait = asyncio.create_task(deps.producer_stop_event.wait())
        try:
            _done, pending = await asyncio.wait(
                [q_get, shut_wait, stop_wait],
                return_when=asyncio.FIRST_COMPLETED,
            )
            for task in pending:
                task.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await task
            if q_get not in _done:
                # A stop signal won the race and nothing was taken.
                return
            if deps.producer_stop_event.is_set():
                # DRAINING owns every row this worker holds: the get TAKEN
                # a row on the same turn the drain began. The hand-back
                # re-pends it to the fleet (the take moved only the
                # in-process queue entry, never the DB row), so running
                # the stale local copy here would execute the job body
                # twice. Return without dispatching.
                return
            # Why fall through on a shutdown_event-only both-done turn:
            # the get has already TAKEN the job out of local_queue and no
            # drain owns it (producer_stop_event is unset), returning here
            # would discard it with no consumer run, no terminal write, and
            # no release, leaving recovery to lock-lease expiry. The taken
            # job runs this final iteration; the outer while's shutdown
            # check then exits the loop.
        finally:
            for task in (q_get, shut_wait, stop_wait):
                if not task.done():
                    task.cancel()

        job: JobRow = q_get.result()

        # Record the claim intent BEFORE any await: between this take and
        # register() the DB row is running, locked here, and invisible to
        # active_jobs, and a hand-back pass running in that window would
        # re-pend it to the fleet while this loop is about to execute it.
        # The token is what scopes this iteration's resolve_claim below
        # to ITS OWN claim (the issue-461 class): a same-worker re-claim
        # can re-mark the key before this iteration unwinds.
        _claim = deps.active_jobs.mark_claimed(job.id)

        # Slot-release point #1 of 2: the get() above dropped
        # qsize by one, so a producer held up on queue capacity can
        # claim again, wake it now. Point #2 is the deregister at the
        # end of this iteration (the active-side slot), because the
        # producer's availability subtracts active jobs too.
        if slot_freed_event is not None:
            slot_freed_event.set()

        current_task = asyncio.current_task()
        if current_task is None:
            raise RuntimeError("consumer_loop_stub must run inside a TaskGroup")

        # The ONE parent context: tag inheritance AND the fan-out ledger
        # stamp (LIB-2) — every child enqueue in the actor body carries
        # this job's id as parent_id, the exact pending-children
        # accounting the client's backpressure read serves.
        with parent_tags(tuple(job.tags), job_id=job.id):
            ctx: JobContext[_StubPayload] = JobContext(
                job_id=job.id,
                actor=job.actor,
                queue=job.queue,
                attempt=job.attempt,
                claim_epoch=job.claim_epoch,
                snooze_count=job.snooze_count,
                worker_id=worker_id,
                payload=_StubPayload(),
                jobs=SubJobEnqueuer(
                    loop_scope_resolved=None,
                    worker_pool=None,
                    backend=backend,
                ),
                log=bind_job_context(
                    _consumer_log,
                    job_id=job.id,
                    actor=job.actor,
                    queue=job.queue,
                    attempt=job.attempt,
                    identity_key=job.identity_key,
                    trace_id=job.trace_id or "",
                ),
            )

            # The registration is captured for the exit path below (issue
            # 461): deregister must pop THIS attempt's entry, not a
            # bare-id pop of whatever the key holds if a re-claim
            # overwrote it mid-run.
            _stub_entry = await deps.active_jobs.register(job.id, current_task, ctx)  # type: ignore[arg-type]  # Why: JobContext[_StubPayload] is a JobContext[BaseModel]; pyright cannot widen Generic[TChild] to Generic[TParent] without explicit covariance.

            try:
                try:
                    await asyncio.wait_for(
                        ctx.cancel_event.wait(),
                        timeout=stub_work_timeout,
                    )
                except asyncio.CancelledError:
                    # shield_with_retrieval, not plain asyncio.shield: a second
                    # cancel landing while this write is detached must not
                    # strand its outcome unretrieved (see taskq._shield).
                    with contextlib.suppress(asyncio.CancelledError):
                        await _stub_terminal_write(backend, job, worker_id, cancelled=True)
                    raise
                except TimeoutError:
                    pass

                await _stub_terminal_write(
                    backend, job, worker_id, cancelled=ctx.cancellation_requested
                )
                # fallback_result_ttl is not forwarded here: the stub path has
                # no actor registry and therefore no @actor(result_ttl=...)
                # literal to supply. If the stored actor_config.result_ttl is
                # cleared (NULL), the enqueue-pinned expiry survives, the
                # original bug, on the one path that structurally cannot fix
                # itself. This is acceptable because the stub is a test/dev
                # sentinel, not a production consumer. The real consumer
                # (_consumer.py) resolves the literal from the registry and
                # passes it through mark_succeeded_with_conn.

            except asyncio.CancelledError:
                with contextlib.suppress(asyncio.CancelledError):
                    await _stub_terminal_write(backend, job, worker_id, cancelled=True)
                raise

            except Exception:
                _consumer_log.exception("consumer-stub-error", job_id=str(job.id))

            finally:
                deps.active_jobs.resolve_claim(job.id, _claim)
                await deps.active_jobs.deregister(job.id, _stub_entry)
                # Slot-release point #2: the producer's
                # availability subtracts active jobs, so this slot
                # frees at the deregister above, not at the get() that
                # only lent the queue slot, wake the producer the
                # moment accounting settles.
                if slot_freed_event is not None:
                    slot_freed_event.set()

di_consumer_loop async

di_consumer_loop(
    deps: WorkerDeps,
    local_queue: Queue[JobRow],
    shutdown_event: Event,
    *,
    backend: Backend,
    worker_id: UUID,
    registry: ProviderRegistry,
    process_scope: ProcessScope,
    thread_scope: ThreadScope,
    loop_scope: LoopScope,
    actor_registry: Mapping[str, ActorRef[Any, Any]],
    enqueuer: SubJobEnqueuer,
    slot_freed_event: Event | None = None,
) -> None

Pull one job per iteration and dispatch via dispatch_one_job.

Outer loop: while not shutdown_event.is_set(). Races local_queue.get() against shutdown_event.wait(); on shutdown win the queue waiter is cancelled and the loop returns cleanly.

Each job is dispatched through dispatch_one_job which composes build_actor_scope + consume_one_job, providing DI-aware actor invocation with per-invocation TRANSIENT scope teardown.

Source code in src/taskq/worker/run.py

async def di_consumer_loop(
    deps: WorkerDeps,
    local_queue: asyncio.Queue[JobRow],
    shutdown_event: asyncio.Event,
    *,
    backend: Backend,
    worker_id: UUID,
    registry: ProviderRegistry,
    process_scope: ProcessScope,
    thread_scope: ThreadScope,
    loop_scope: LoopScope,
    actor_registry: Mapping[str, ActorRef[Any, Any]],
    enqueuer: SubJobEnqueuer,
    slot_freed_event: asyncio.Event | None = None,
) -> None:
    """Pull one job per iteration and dispatch via dispatch_one_job.

    Outer loop: ``while not shutdown_event.is_set()``.
    Races ``local_queue.get()`` against ``shutdown_event.wait()``; on shutdown
    win the queue waiter is cancelled and the loop returns cleanly.

    Each job is dispatched through dispatch_one_job which composes
    build_actor_scope + consume_one_job, providing DI-aware actor
    invocation with per-invocation TRANSIENT scope teardown.
    """
    clock_obj = process_scope.get(Clock)
    if clock_obj is None or not isinstance(clock_obj, Clock):
        raise MissingProvider(
            type_name="Clock",
            required_by="worker.di_consumer_loop (ProcessScope must have a cached Clock after bootstrap)",
        )
    clock: Clock = clock_obj

    while not (shutdown_event.is_set() or deps.producer_stop_event.is_set()):
        _fork_at = take_parent_fork_event()
        if _fork_at is not None:
            # The heartbeat loop carries the same report; this one gets the
            # job's log context on a fork that landed mid-dispatch. Report
            # and continue: the parent's connections were not written by
            # the child unless the child also used them, and a protocol
            # -dead pooled connection is discarded and rebuilt by the pool
            # (the in-flight job sees a transient error, not silence).
            _consumer_log.error(
                "fork-detected-in-worker-process",
                fork_age_secs=round(time.monotonic() - _fork_at, 3),
                detail=(
                    "a fork() happened in this worker process while its "
                    "connections were live; the child inherits every socket "
                    "(pools, LISTEN, redis, the loop's pipes). A job body or "
                    "library must not fork: use subprocess.Popen (default "
                    "close_fds=True drops the inherited descriptors), or "
                    "open fresh resources in the child after the fork."
                ),
            )
        q_get = asyncio.create_task(local_queue.get())
        shut_wait = asyncio.create_task(shutdown_event.wait())
        stop_wait = asyncio.create_task(deps.producer_stop_event.wait())
        try:
            _done, pending = await asyncio.wait(
                [q_get, shut_wait, stop_wait],
                return_when=asyncio.FIRST_COMPLETED,
            )
            for task in pending:
                task.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await task
            if q_get not in _done:
                # A stop signal won the race and nothing was taken.
                return
            if deps.producer_stop_event.is_set():
                # DRAINING owns every row this worker holds: the get TAKEN
                # a row on the same turn the drain began. Running it would
                # double-execute, the row is locked to this worker and
                # never reached the active-jobs registry, so the hand-back
                # (the orchestrator's DRAINING pass for rows claimed before
                # it ran, the producer's exit pass for a claim round that
                # committed after) re-pends it to the fleet while this
                # loop's stale local copy would also run here. Return
                # without dispatching. The taken row needs no release from
                # THIS loop: the take moved only the in-process queue
                # entry, never the DB row, so whichever hand-back pass
                # runs finds it exactly as claimed.
                return
            # Why fall through on a shutdown_event-only both-done turn:
            # the get has already TAKEN the job out of local_queue and no
            # drain owns it (producer_stop_event is unset, this is the
            # external-exit path, not a graceful-shutdown orchestration) ,
            # returning here would discard it with no dispatch, no
            # terminal write, and no release, leaving recovery to
            # lock-lease expiry. The taken job runs this final iteration;
            # the outer while's shutdown check then exits the loop.
        finally:
            for task in (q_get, shut_wait, stop_wait):
                if not task.done():
                    task.cancel()

        job: JobRow = q_get.result()

        # Record the claim intent BEFORE any await: between this take and
        # register() the DB row is running, locked here, and invisible to
        # active_jobs, and a hand-back pass running in that window would
        # re-pend it to the fleet while this loop is about to execute it.
        # The token scopes this iteration's resolve_claim below to ITS
        # OWN claim (the issue-461 class): a same-worker re-claim's take
        # can re-mark the key while this iteration's attempt unwinds.
        _claim = deps.active_jobs.mark_claimed(job.id)

        # Slot-release point #1 of 2: the get() above dropped
        # qsize by one, so a producer held up on queue capacity can
        # claim again, wake it now. Point #2 is the finally around
        # dispatch_one_job below (the active-side slot), because the
        # producer's availability subtracts active jobs too. Fires on
        # every iteration, every exit path (success, failure, snooze,
        # not-found release) passes it.
        if slot_freed_event is not None:
            slot_freed_event.set()

        if job.actor not in actor_registry:
            _consumer_log.error(
                "dispatch-actor-not-found",
                job_id=str(job.id),
                actor=job.actor,
            )
            # Release the claimed job instead of leaving it 'running' until
            # lease expiry, a worker whose registry has the actor can then
            # pick it up. The short delay keeps this worker from re-claiming
            # it in a hot loop.
            #
            # Contract: an unregistered actor parks the job at the snooze
            # cadence, budget-free, mark_snoozed's default 'snoozed'
            # outcome refunds the claim's attempt increment, so a job whose
            # actor is missing (through no fault of its own) never burns
            # retry budget while it waits for a worker that can run it;
            # the stranded-jobs detector surfaces it. This is not an
            # actor-requested deferral semantically, but the snooze write
            # is the closest bounded outcome, a delay, a release, and a
            # released_reason marker in one transition.
            try:
                release_outcome = await backend.mark_snoozed(
                    job.id,
                    worker_id,
                    timedelta(seconds=10),
                    metadata_update={"released_reason": "actor-not-found"},
                    attempt=job.attempt,
                    claim_epoch=job.claim_epoch,
                )
            except Exception:
                _consumer_log.exception(
                    "dispatch-actor-not-found-release-failed",
                    job_id=str(job.id),
                    actor=job.actor,
                )
                # The row is still running under this worker's lock with
                # nothing left to move it: disown it so the heartbeat stops
                # renewing the lease and the reclaim sweep hands it back.
                deps.disowned_jobs.add(job.id)
            else:
                # The snooze tri-state is the write's fence and must be
                # read, not dropped: "scheduled" released the row,
                # "failed" terminalised it on the deadline arm, and
                # "noop" means the row stopped being this worker's to
                # move between claim and release (a reclaim re-pended
                # it, or a racing cancel terminalised it), the new
                # owner holds it, so there is nothing to release, but
                # a discarded outcome would make the fenced-out write
                # indistinguishable from a landed one.
                if release_outcome == "noop":
                    _consumer_log.debug(
                        "dispatch-actor-not-found-release-noop",
                        job_id=str(job.id),
                        actor=job.actor,
                    )
                elif release_outcome == "failed":
                    _consumer_log.info(
                        "dispatch-actor-not-found-release-deadline-exceeded",
                        job_id=str(job.id),
                        actor=job.actor,
                    )
            # This exit path bypasses the dispatch finally below: drop the
            # claim intent here, or a stale id would fence a future claim
            # of the same row out of every hand-back pass.
            deps.active_jobs.resolve_claim(job.id, _claim)
            continue

        actor_ref = actor_registry[job.actor]
        # The attempt runs in a per-job child task, and the child is what
        # consume_one_job registers as the job's inflight attempt
        # (asyncio.current_task() inside the child). Registering THIS loop
        # task instead (awaiting dispatch_one_job inline) would let the
        # cancel ladder's phase-2 escalation
        # (cancel.py's active.task.cancel()) cancel the whole loop:
        # CancelledError is a BaseException, the handlers below catch only
        # SlotPoolAcquireError and Exception, so it would propagate out of the
        # while, the TaskGroup would discard the cancelled sibling, nothing
        # would respawn it, and the slot would never serve another row (the producer's
        # availability counts queue depth minus active jobs, never live
        # consumers, so the worker stranded max_concurrency rows and parked).
        # The child scopes every cancel of a registered attempt task to the
        # one job: phase-2 escalation, the shutdown orchestrator's
        # CANCELLING/FORCING phases, and the isolate path all cancel the
        # child, this loop survives, and the cancel's delivery semantics are
        # the inline ones (a cancel of this loop while it awaits the child
        # propagates into the child first, through the same frames an inline
        # cancel hits).
        dispatch_task = asyncio.create_task(
            dispatch_one_job(
                backend=backend,
                deps=deps,
                job=job,
                worker_id=worker_id,
                registry=registry,
                process_scope=process_scope,
                thread_scope=thread_scope,
                loop_scope=loop_scope,
                actor_ref=actor_ref,  # type: ignore[arg-type]  # Why: ActorRef[Any, Any] is not ActorRef[BaseModel, BaseModel | None]; pyright cannot widen the generic parameters, but the runtime contract is sound, actor_ref carries the correct payload_type and fn.
                actor_config=actor_ref.config,
                clock=clock,
                active_jobs=deps.active_jobs,
                max_retry_backoff=deps.settings.max_retry_backoff,
                enqueuer=enqueuer,
            ),
            name=f"job-attempt:{job.id}",
        )
        try:
            try:
                outcome = await dispatch_task
            except asyncio.CancelledError:
                # Route by WHO was cancelled. This loop task's own
                # cancellation (the TaskGroup's teardown cancel, count
                # elevated) or any shutdown signal means the child was killed
                # as part of the worker coming down: re-raise, the loop dies
                # with its job. The
                # remaining case is the child ALONE being cancelled while
                # this loop lives: an operator cancel's phase-2 escalation
                # (or an actor that raised CancelledError itself). The child
                # already terminalised the job through its own CancelledError
                # handler (mark_cancelled, or mark_interrupted on a shutdown
                # origin); the CancelledError out of the await is the job's
                # outcome surfacing, not this loop's death. Absorb it as the
                # cancelled outcome and keep serving: the loop's survival is
                # the invariant, one stuck job must not take the slot with
                # it. No uncancel() is owed: this task was never cancelled,
                # its count is untouched on this path.
                current = asyncio.current_task()
                if (
                    (current is not None and current.cancelling() > 0)
                    or shutdown_event.is_set()
                    or deps.producer_stop_event.is_set()
                    or deps.shutdown_phase is not ShutdownPhase.NONE
                ):
                    raise
                outcome = "cancelled"
                _consumer_log.info(
                    "dispatch-force-cancelled-slot-continues",
                    job_id=str(job.id),
                    actor=job.actor,
                )
            if outcome == "failed":
                deps.drain_failures += 1
        except SlotPoolAcquireError:
            # Infrastructure, not a job outcome: counting this as a drain
            # failure would make a Kubernetes Job / CI drain step report
            # job failures that never happened. The acquire was recorded
            # (counter) and logged (per-occurrence cause, job id) at the
            # raise site. The row though is still running under this
            # worker's lock with no runner left to move it: disown it so
            # the heartbeat stops renewing the lease and the reclaim
            # sweep hands it back - the same treatment the
            # actor-not-found release-failure arm above gets (leaving
            # the job to lease reclaim would be a no-op: the tick's
            # renewal keeps that lease alive forever, so the recovery
            # it names never comes and the row stays lost).
            _disown_job(deps.disowned_jobs, job)
            continue
        except Exception:
            _consumer_log.exception("dispatch-failed", job_id=str(job.id))
            deps.drain_failures += 1
        finally:
            # The claim intent's life ends here on every path: register()
            # absorbed it when the job ran, and every non-running path
            # (skip, release, acquire failure, error) is past the window
            # the intent exists to fence. Resolving in the loop's own
            # finally cannot straddle a drain: the row is no longer
            # running-and-unowned by the time this runs. Identity-scoped
            # to THIS iteration's token (the issue-461 class): a
            # same-worker re-claim's take can have re-marked the key
            # while this iteration's attempt was unwinding, and this
            # resolve must not erase the live claim's intent.
            deps.active_jobs.resolve_claim(job.id, _claim)
            # Slot-release point #2: the producer's availability
            # subtracts active jobs, so the slot this job held frees at
            # the deregister dispatch_one_job's own finally has run by
            # every path that reaches here, not at the get() that only
            # lent the queue slot. The wake is what re-arms the producer
            # the moment accounting settles; without it a finished job's
            # replacement claim would wait for the fallback poll tick.
            # (The SlotPoolAcquireError path never registered, so its
            # wake is redundant with the get()-point one, bounded, and
            # the price of one unconditional release point.)
            if slot_freed_event is not None:
                slot_freed_event.set()

register_worker async

register_worker(
    pool: Pool, settings: WorkerSettings
) -> UUID

Register the current worker in taskq.workers and return its UUID.

Generates a UUIDv7, inserts a row into {schema}.workers, and returns the new UUID. Acquires from pool with a 2.0 s timeout; on timeout or connection error the failure is logged and re-raised (registering is fatal).

If settings.worker_label or settings.workgroup_instance are set, they are stored directly for cross-process correlation and health checking.

The row's metadata records the worker's runtime facts: whether NOTIFY dispatch is enabled, and max_concurrency, the capacity the worker runs at, which sizes local_queue and bounds every dispatch.

Source code in src/taskq/worker/run.py

async def register_worker(pool: asyncpg.Pool, settings: WorkerSettings) -> UUID:
    """Register the current worker in ``taskq.workers`` and return its UUID.

    Generates a UUIDv7, inserts a row into ``{schema}.workers``, and returns
    the new UUID.  Acquires from *pool* with a 2.0 s timeout; on timeout or
    connection error the failure is logged and re-raised (registering is
    fatal).

    If ``settings.worker_label`` or ``settings.workgroup_instance`` are set,
    they are stored directly for cross-process correlation and health checking.

    The row's metadata records the worker's runtime facts: whether NOTIFY
    dispatch is enabled, and ``max_concurrency``, the capacity the worker
    runs at, which sizes ``local_queue`` and bounds every dispatch.
    """
    worker_id = new_uuid()
    schema = settings.schema_name
    require_schema(schema)
    hostname = socket.gethostname()
    pid = os.getpid()
    queues = settings.queues

    maybe_label: str | None = settings.worker_label
    maybe_instance_raw = settings.workgroup_instance
    maybe_instance: UUID | None = UUID(maybe_instance_raw) if maybe_instance_raw else None

    notify_enabled = getattr(settings, "notify_enabled", False)
    # The workers row carries the capacity this worker actually runs at:
    # ``max_concurrency`` sizes ``local_queue`` and bounds every dispatch,
    # so a fleet's effective parallelism stays queryable from the database
    # for monitoring and coordination.
    metadata: dict[str, object] = {
        "notify_enabled": notify_enabled,
        "max_concurrency": settings.max_concurrency,
    }

    sql = (
        f'INSERT INTO "{schema}".workers '  # noqa: S608  # Why: schema validated against _IDENT_RE before interpolation; asyncpg cannot bind identifiers as parameters.
        "(id, hostname, pid, queues, worker_label, workgroup_instance, metadata) "
        "VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb)"
    )
    metadata_json = jsonb_param(metadata)

    try:
        async with pool.acquire(timeout=2.0) as conn:
            await conn.execute(
                sql, worker_id, hostname, pid, queues, maybe_label, maybe_instance, metadata_json
            )
    except TRANSIENT_PG_ERRORS as e:
        _reg_log.error("register-worker-failed", error=str(e))
        raise

    return worker_id

deregister_worker async

deregister_worker(
    pool: Pool, settings: WorkerSettings, worker_id: UUID
) -> bool

Remove the worker row from {schema}.workers (best-effort).

Acquires from pool with a 2.0 s timeout. On timeout or connection error, logs a structured warning deregister_worker_failed and reports the failure WITHOUT raising: shutdown MUST NOT block on this cleanup, and the caller's backstop (the deps-teardown retry, then the recovery sweep) decides what a still-unknown row fate means.

Returns the DELETE's outcome: True when the row is confirmed gone (the DELETE executed and reported rowcount > 0), False when the delete is NOT confirmed, either because a transient failure was swallowed here (the dominant class: pool-acquire timeout, conn loss; the row's fate is unknown) or because the DELETE matched no row (the sweep or a peer already removed it, so there was nothing to delete). A non-transient failure still raises: it would reproduce identically in any retry, so reporting it as transient-shaped False would livelock the backstop while logging like a blip.

Source code in src/taskq/worker/run.py

async def deregister_worker(pool: asyncpg.Pool, settings: WorkerSettings, worker_id: UUID) -> bool:
    """Remove the worker row from ``{schema}.workers`` (best-effort).

    Acquires from *pool* with a 2.0 s timeout.  On timeout or connection
    error, logs a structured warning ``deregister_worker_failed`` and
    reports the failure WITHOUT raising: shutdown MUST NOT block on this
    cleanup, and the caller's backstop (the deps-teardown retry, then the
    recovery sweep) decides what a still-unknown row fate means.

    Returns the DELETE's outcome: ``True`` when the row is confirmed gone
    (the DELETE executed and reported ``rowcount > 0``), ``False`` when the
    delete is NOT confirmed, either because a transient failure was
    swallowed here (the dominant class: pool-acquire timeout, conn loss;
    the row's fate is unknown) or because the DELETE matched no row (the
    sweep or a peer already removed it, so there was nothing to delete).
    A non-transient failure still raises: it would reproduce identically
    in any retry, so reporting it as transient-shaped ``False`` would
    livelock the backstop while logging like a blip.
    """
    schema = settings.schema_name
    require_schema(schema)

    sql = f'DELETE FROM "{schema}".workers WHERE id = $1'  # noqa: S608  # Why: schema validated against _IDENT_RE before interpolation; asyncpg cannot bind identifiers as parameters.

    try:
        async with pool.acquire(timeout=2.0) as conn:
            status: object = await conn.execute(sql, worker_id)
    except TRANSIENT_PG_ERRORS as e:
        _reg_log.warning(
            "deregister_worker_failed",
            worker_id=worker_id,
            error=str(e),
        )
        return False
    return _delete_status_confirmed(status)

Maintenance leader (election, sweeps, cron, prune/archive)

leader

Maintenance leader: election, watchdog, and recovery sweeps. A single elected leader per cluster runs cooperative loops inside one asyncio.TaskGroup: election, watchdog, scheduled-wake (sweep 3), cron, sweep (sweeps 1/2/4), prune (sweep 5), archive expiry (sweep 6), stale worker cleanup, queue depth, reservation slots, and backlog detection. Non-leader pods retry election periodically and skip the gated work , the backlog detector is the deliberate exception (every worker samples; see _backlog_detection_loop). Leadership is a row lease in maintenance_leader: the holder writes an expires_at it chose from leader_lease and renews it every heartbeat interval, and any pod may take the row over once that instant has passed. The row alone decides the election in every state, held, lapsed, or absent, so nothing about the role depends on the liveness of a connection: a holder that dies without a FIN is replaced on a horizon TaskQ controls rather than the server's connection-reaping schedule, and the replacement needs no privilege beyond UPDATE on that row. The schema-qualified advisory lock is still taken after winning, as a courtesy to pods from releases that only understand the lock; it is never required, never waited on, and its absence never gates or costs the role.

Failover SLA

Graceful stop sub-second when the wake lands: the resign is the FIRST act of the shutdown (the shutdown-ordering contract below: the handover happens at shutdown START, not after the drain, so draining adds no latency to the handover), it deletes the row and then broadcasts on the leadership wake channel, a HINT (the row and the fence stay the authority) that makes a follower re-run the fenced elect after at most leader_wake_jitter of damping. The wake is never the guarantee: a lost, duplicated, malformed, or disabled broadcast, and a killed leader (that can notify no one), degrades to the plain bound ≤ heartbeat_interval + one round trip (the resign deletes the row; the next election wins it - the bound the hint degrades to, and the crash bound)

Worker killed ≤ leader_lease + heartbeat_interval + one round trip Silent leader ≤ leader_lease + heartbeat_interval + one round trip Won, unassumable ≤ leader_lease + heartbeat_interval + one failing cycle's connection attempts (each bounded by reload_factory_timeout; the trust-spent hand-back deletes the row: the lapse backstop cannot fire while sub-lease re-wins keep refreshing it; see _hand_back_unassumable_lease) Partition detect ≤ watchdog_interval + heartbeat_interval + 2 s PG failover ≤ heartbeat_interval Watchdog detect ≤ watchdog_interval + heartbeat_interval

Shutdown ordering contract (the no-leadership-while-stopping invariant): a worker that is shutting down NEVER wins, holds, or renews the leader lease, and a leader that begins shutting down hands the lease over at shutdown START, so a rolling deploy never has a leaderless window and never sees a dying pod win or hold leadership.

1. orchestrate_shutdown sets deps.shutdown_start_event BEFORE the DRAINING phase touches any row. That event is the stop signal this runtime obeys, earlier and narrower than the worker-wide shutdown_event, which does not fire until every phase has run.
2. The election loop parks the moment the stop signal is observed: no attempts and no renewals for the rest of the shutdown, however long the drain runs. It drops its liveness registration when it parks (detector 2 must not trip on a loop that stopped by design) and ignores the resign-broadcast wake (:meth:wake_election no-ops).
3. A pod that is LEADING when the stop lands hands the lease over at shutdown START: demote first (the flag drops before any await), then the fenced resign over leader_conn, which is still open - the orchestrator closes it only after the phases. The successor's election tick (or the broadcast wake) elects while THIS pod is still draining its jobs. The teardown resign in run()'s finally remains as the backstop for exits that never run the orchestrator (a sibling crash, a bare cancel of _main).
4. An elect statement that was in flight when the stop landed and WON anyway is handed straight back: fenced resign over the conn the elect just used, and the assume path never runs - the monitor/cron conns do not open, lead() never sets the flag, so no leader-gated loop, sweep or cron tick included, can begin.
5. The drain does not depend on this pod remaining leader: every phase of orchestrate_shutdown writes through the dispatcher pool and the backend with no leadership premise, and the successor covers the maintenance sweeps from its first tick.

__all__ module-attribute

__all__ = [
    "ArchiveExpiryResult",
    "MaintenanceLeader",
    "PruneResult",
    "_build_retention_per_status",
    "_load_actor_retention_overrides",
    "_schedule_utc_to_cron",
    "archive_expiry_sweep",
    "cleanup_stale_workers",
    "complete_stale_batches",
    "prune_terminal_jobs",
]

log module-attribute

log: BoundLogger = get_logger(__name__)

ArchiveExpiryResult dataclass

ArchiveExpiryResult(
    total_deleted: int,
    by_status: dict[str, int],
    expire_before: datetime,
    duration_ms: int,
)

total_deleted instance-attribute

total_deleted: int

by_status instance-attribute

by_status: dict[str, int]

expire_before instance-attribute

expire_before: datetime

duration_ms instance-attribute

duration_ms: int

PruneResult dataclass

PruneResult(
    total_deleted: int,
    archived: int,
    by_actor: dict[str, int],
    by_status: dict[str, int],
    cutoffs: dict[str, datetime],
    duration_ms: int,
)

total_deleted instance-attribute

total_deleted: int

archived instance-attribute

archived: int

by_actor instance-attribute

by_actor: dict[str, int]

by_status instance-attribute

by_status: dict[str, int]

cutoffs instance-attribute

cutoffs: dict[str, datetime]

duration_ms instance-attribute

duration_ms: int

MaintenanceLeader

MaintenanceLeader(
    deps: WorkerDeps,
    worker_id: UUID,
    backend: Backend,
    *,
    clock: Clock,
    rate_limit_registry: RateLimitRegistry | None = None,
    actor_policies: Mapping[str, ActorFirePolicy]
    | None = None,
    cron_specs: Sequence[CronScheduleSpec] | None = None,
)

Elected leader that runs watchdog, sweeps, cron, and prune loops.

Source code in src/taskq/worker/leader.py

def __init__(
    self,
    deps: WorkerDeps,
    worker_id: UUID,
    backend: Backend,
    *,
    clock: Clock,
    rate_limit_registry: RateLimitRegistry | None = None,
    actor_policies: Mapping[str, ActorFirePolicy] | None = None,
    cron_specs: Sequence[CronScheduleSpec] | None = None,
) -> None:
    self._deps = deps
    self._worker_id = worker_id
    self._backend = backend
    self._clock = clock
    # Singleton / max_pending flags the cron tick stamps and enforces
    # on its fires (parity with the client enqueue path); None keeps
    # the tick's no-stamping behavior.
    self._actor_policies = actor_policies
    # The code's declared cron specs, for the takeover recovery. At
    # every leadership assumption the leader re-runs the stale
    # auto-disable recovery over these (the same predicates the boot
    # registration pass applies): during a mixed-version rolling
    # deploy the OLD leader's cron tick can auto-disable a schedule
    # AFTER every new pod has booted, an unmarked (disabled_by=NULL)
    # write the boot pass has already finished matching, so the next
    # leadership change is the last moment the new code can revert it
    # without a full restart. None (tests, spec-less deployments) runs
    # no takeover pass.
    self._cron_specs = list(cron_specs) if cron_specs else []
    self._sweep_ctx = SweepContext(
        deps=deps,
        backend=backend,
        clock=clock,
        worker_id=worker_id,
        rate_limit_registry=rate_limit_registry,
    )
    self._leader_monitor_conn: asyncpg.Connection | None = None
    self._cron_conn: asyncpg.Connection | None = None
    # The (worker_id, elected_at) term this pod last found holding the
    # role while it lost an election, so contention distinguishes a
    # transition from the steady state of following one live peer.
    self._observed_holder: tuple[UUID, datetime] | None = None
    # The last election this pod WON, kept past every demotion for
    # resign()'s fence.  ``stop_leading`` nulls ``deps.leader_term`` on
    # every mid-run step-down, so without this copy a leader that
    # demoted itself (the cron/monitor conn-death paths) and is then
    # torn down cannot fence its resign: the row it left behind lapses
    # only after a whole ``leader_lease``, stalling every successor by
    # exactly the wait resign exists to remove.  The fence stays sound
    # because the resign DELETE matches ``(worker_id, elected_at)``
    # exactly, a successor's takeover rewrites both, so a stale fence
    # deletes nothing.
    self._resign_fence: LeaderTerm | None = None
    # The FIRST term of this pod's current won-but-unassumable episode:
    # a run of election wins whose dedicated-conn opens keep failing
    # Its ``trusted_until`` is the episode's whole trust
    # budget (the same window ``_renew_failed`` gives a renewing
    # leader to recover a transient failure before standing down),
    # because every re-win mints a fresh full window, so nothing
    # else bounds how long the pod can keep a row it cannot use.
    # Cleared by a successful assume (the pod can lead again) and by
    # an OBSERVED ELECTION LOSS (a peer holds the row now, the
    # episode is over, and a later win starts a fresh one with a
    # fresh budget); deliberately NOT cleared by the trust-spent
    # resign itself, so a pod that re-wins while still broken, with
    # no peer having taken the row in between, hands the row
    # straight back instead of buying another window.
    self._unassumable_anchor: LeaderTerm | None = None
    # Log-once latch for a refused courtesy advisory-lock probe: a
    # managed Postgres refusing ``pg_try_advisory_lock`` refuses it on
    # every election win for the life of the grants, so the WARN is
    # emitted on the first refusal and re-armed only when a probe
    # succeeds again (the grant appearing IS a new operational fact).
    self._advisory_lock_refused = False
    # The election-wake seam, ONE handler shape for both callers that
    # arm it: ``wake_election`` (this pod's own resign-broadcast seam)
    # sets it directly, and a PEER's broadcast reaches the same object
    # through run()'s subscription, which re-points this field at the
    # backend registry's event for the run's life (see
    # ``_leader_wake_subscription``). Consumed by the loop's tick wait
    # and damped by ``_follower_park``; a stopping worker ignores it
    # (see ``wake_election``) - the shutdown-ordering contract in this
    # module's doc header.
    self._wake_event = asyncio.Event()

wake_election

wake_election() -> None

Request an immediate election attempt (the resign-broadcast seam).

A hint, never a command: the loop still runs every gate it runs on its ordinary cadence. A worker whose shutdown has begun IGNORES the wake - no attempt, no win, no lease - the shutdown ordering contract in this module's doc header. Safe to call from any task on this loop; the wake is consumed by the election loop's park wait (_follower_park), which damps it by leader_wake_jitter - and refuses it outright while stopping, the same refusal this arm-time gate makes.

Source code in src/taskq/worker/leader.py

def wake_election(self) -> None:
    """Request an immediate election attempt (the resign-broadcast seam).

    A hint, never a command: the loop still runs every gate it runs
    on its ordinary cadence. A worker whose shutdown has begun
    IGNORES the wake - no attempt, no win, no lease - the shutdown
    ordering contract in this module's doc header. Safe to call from
    any task on this loop; the wake is consumed by the election
    loop's park wait (``_follower_park``), which damps it by
    ``leader_wake_jitter`` - and refuses it outright while stopping,
    the same refusal this arm-time gate makes.
    """
    if self._stopping():
        log.debug(
            "election-wake-ignored-stopping",
            kind="election_wake_ignored_stopping",
            worker_id=str(self._worker_id),
        )
        return
    self._wake_event.set()

run async

run(shutdown: Event) -> None

Source code in src/taskq/worker/leader.py

async def run(self, shutdown: asyncio.Event) -> None:
    with _active_leaders_lock:
        _active_leaders.add(self)
    try:
        async with self._leader_wake_subscription(), asyncio.TaskGroup() as tg:
            tg.create_task(self._election_loop(shutdown), name="leader.election")
            tg.create_task(self._watchdog_loop(shutdown), name="leader.watchdog")
            tg.create_task(self._scheduled_wake_loop(shutdown), name="leader.scheduled_wake")
            tg.create_task(self._cron_loop(shutdown), name="leader.cron")
            tg.create_task(self._sweep_loop(shutdown), name="leader.sweep")
            tg.create_task(self._prune_loop(shutdown), name="leader.prune")
            tg.create_task(self._archive_expiry_loop(shutdown), name="leader.archive_expiry")
            tg.create_task(self._queue_depth_loop(shutdown), name="leader.queue_depth")
            tg.create_task(
                self._backlog_detection_loop(shutdown), name="leader.backlog_detection"
            )
            tg.create_task(
                self._reservation_slots_loop(shutdown), name="leader.reservation_slots"
            )
            tg.create_task(self._stranded_jobs_loop(shutdown), name="leader.stranded_jobs")
            tg.create_task(
                self._ratelimit_buckets_loop(shutdown), name="leader.ratelimit_buckets"
            )
            await shutdown.wait()
    finally:
        # Hand the lease back before the conns go: a replacement pod then
        # elects on its next cycle instead of waiting out the lease this
        # one would otherwise leave behind. Best-effort, the lease lapses
        # on its own, so a resign that cannot reach the database costs at
        # most that wait.
        await self.resign()
        # Final teardown, not a mid-run demotion: close with the
        # conn-teardown-close-* family so an ordinary shutdown never
        # pages as an unexpected mid-run close timeout.
        await self._close_leader_owned_conns(mid_run=False)
        with _active_leaders_lock:
            _active_leaders.discard(self)

resign async

resign() -> bool

Hand the lease back so a peer elects on its next cycle.

Fenced on the term, so a resign issued after a takeover cannot delete the successor's row. Best-effort by design: the lease lapses on its own, so a resign that cannot reach the database costs a replacement pod at most one lease rather than correctness.

Runs at run() teardown, AFTER the shutdown orchestrator has already closed and nulled a TaskQ-owned leader_conn, so the write rides the leader-owned monitor conn when the primary is gone. Both conns are idle by then (the loops that used them have exited), and the monitor conn is never the orchestrator's to close.

The fence prefers the live term and falls back to the last election this pod won: a mid-run demotion clears deps.leader_term long before teardown, and without the fallback this pod's own row, never taken over, or the DELETE would fence it out, would sit until the lease lapses while a replacement pod waits on it.

The other caller is _hand_back_unassumable_lease, mid-run and on a pod that never led (leader_term is None, so the fallback fence is the one that applies): it rides the leader conn the winning elect just used, before that conn is dropped.

A resign that DELETED the row then broadcasts the vacancy on the leadership wake channel (_emit_leader_wake), so a successor assumes sub-second instead of on its next tick. The broadcast is best-effort on top of a best-effort resign: it is a HINT, the tick is the guarantee, and a broadcast that never lands costs exactly what the pre-wake design cost.

Returns whether the resign actually DELETED the row. False is the ordinary fenced no-op (a successor's row, or none left) and every best-effort shape (no term to resign, no live conn to ride, the write raising): the no-op and no-term/no-conn shapes are silent by design (at teardown they are normal), while a raised write carries its own leader-resign-failed WARN. A caller whose policy REPORTS a hand-back must not do so on a False here, which is why the trust-spent hand-back WARNs only on a True.

Source code in src/taskq/worker/leader.py

async def resign(self) -> bool:
    """Hand the lease back so a peer elects on its next cycle.

    Fenced on the term, so a resign issued after a takeover cannot delete
    the successor's row. Best-effort by design: the lease lapses on its
    own, so a resign that cannot reach the database costs a replacement
    pod at most one lease rather than correctness.

    Runs at ``run()`` teardown, AFTER the shutdown orchestrator has
    already closed and nulled a TaskQ-owned ``leader_conn``, so the
    write rides the leader-owned monitor conn when the primary is gone.
    Both conns are idle by then (the loops that used them have exited),
    and the monitor conn is never the orchestrator's to close.

    The fence prefers the live term and falls back to the last election
    this pod won: a mid-run demotion clears ``deps.leader_term`` long
    before teardown, and without the fallback this pod's own row, never
    taken over, or the DELETE would fence it out, would sit until the
    lease lapses while a replacement pod waits on it.

    The other caller is ``_hand_back_unassumable_lease``, mid-run and
    on a pod that never led (``leader_term`` is None, so the fallback
    fence is the one that applies): it rides the leader conn the
    winning elect just used, before that conn is dropped.

    A resign that DELETED the row then broadcasts the vacancy on the
    leadership wake channel (``_emit_leader_wake``), so a successor
    assumes sub-second instead of on its next tick. The broadcast is
    best-effort on top of a best-effort resign: it is a HINT, the
    tick is the guarantee, and a broadcast that never lands costs
    exactly what the pre-wake design cost.

    Returns whether the resign actually DELETED the row. ``False`` is
    the ordinary fenced no-op (a successor's row, or none left) and
    every best-effort shape (no term to resign, no live conn to ride,
    the write raising): the no-op and no-term/no-conn shapes are
    silent by design (at teardown they are normal), while a raised
    write carries its own ``leader-resign-failed`` WARN. A caller
    whose policy REPORTS a hand-back must not do so on a ``False``
    here, which is why the trust-spent hand-back WARNs only on
    a ``True``.
    """
    term = self._deps.leader_term or self._resign_fence
    if term is None:
        return False
    conn = self._deps.leader_conn
    if conn is None or conn.is_closed():
        conn = self._leader_monitor_conn
    if conn is None or conn.is_closed():
        return False
    _elect, _renew, resign_sql = build_leader_lease_sql(self._deps.settings.schema_name)
    try:
        async with asyncio.timeout(CLOSE_TIMEOUT_SECS):
            tag = await conn.execute(resign_sql, self._worker_id, term.elected_at)
    except Exception as exc:
        log.warning(
            "leader-resign-failed",
            kind="leader_resign_failed",
            worker_id=str(self._worker_id),
            error=repr(exc),
        )
        return False
    # The command tag is the truthful delete count: "DELETE 1" is a
    # row handed back, "DELETE 0" is the fence correctly matching
    # nothing (a successor's row, or the row already resigned). A
    # caller reporting the hand-back must read this, not the absence
    # of an error: a fenced no-op is a success-shaped failure.
    parts = tag.split()
    deleted = len(parts) == 2 and parts[0].upper() == "DELETE" and parts[1] != "0"
    if not deleted:
        log.info(
            "leader-resign-noop",
            kind="leader_resign_noop",
            worker_id=str(self._worker_id),
            tag=tag,
        )
        return False
    # The broadcast AFTER the delete commits: the DELETE is a single
    # autocommit statement, so this separate statement is sent only
    # once the delete's round trip returned - the commit-gate
    # discipline (never notify inside the tx whose fact the payload
    # claims), without needing the cron tick's self-addressed gate,
    # because there is no surrounding transaction to roll back a
    # committed row out from under the payload.
    await self._emit_leader_wake(term)
    log.info(
        "leader-resigned",
        kind="leader_resigned",
        worker_id=str(self._worker_id),
    )
    return True

archive_expiry_sweep async

archive_expiry_sweep(
    conn: ConnLike,
    *,
    batch_size: int = DEFAULT_PRUNE_BATCH_SIZE,
    schema: str = "taskq",
    drain_gate: Callable[[], bool] | None = None,
    statement_timeout_ms: int = DEFAULT_PRUNE_STATEMENT_TIMEOUT_MS,
    sizer: SweepBatchSizer | None = None,
) -> ArchiveExpiryResult

Hard-delete expired jobs_archive rows, one bounded, self-committing batch at a time.

Each batch is one _EXPIRY_CTE_SQL statement, committed before the next runs, under the same server-side statement_timeout and breaker semantics as :func:prune_terminal_jobs, see that function's docstring and :data:~taskq.constants.DEFAULT_PRUNE_STATEMENT_TIMEOUT_MS for the timeout/breaker contract. drain_gate is called before every batch; a False return stops the drain with the batches already committed.

The POLICY FLOOR (hypertable deployments only): when the connected server has a policy_retention job registered against jobs_archive (:func:taskq.backend._retention_floor.retention_policy_floor, probed ONCE per run — before the batch loop, never per batch — and failing open to None on any probe error), rows older than the policy's own drop_after horizon gain the AND finished_at >= $floor conjunct and are left to the policy, which drops them whole-chunk silently. THE BOUNDARY, stated honestly: OLDER than the floor, the POLICY owns deletion (silent, chunk-granular, watermark-blind — see the pinned boundary in tests/test_timescale_retention_interplay.py); NEWER than it (inside the window), this sweep owns deletion (row-exact on expire_at, the stamp honored exactly inside the window — the young chunk whose rows have passed expire_at while their chunk survives is still this sweep's work, and remains so). The floor's clock is the POLICY's clock (finished_at, the chunk alignment column) and the floor's value is the policy's OWN horizon (parsed from the registered config, not re-derived), so a policy registered with a different interval than the current setting still governs the aged end exactly. A row older than the floor inside a young chunk is dropped when the chunk itself ages past the boundary, at most one chunk interval late — chunk granularity is the trade. On vanilla Postgres the probe fails open to None and this sweep runs full-range, byte-identical to the no-floor behavior.

Source code in src/taskq/worker/_leader_shared.py

async def archive_expiry_sweep(
    conn: ConnLike,
    *,
    batch_size: int = DEFAULT_PRUNE_BATCH_SIZE,
    schema: str = "taskq",
    drain_gate: Callable[[], bool] | None = None,
    statement_timeout_ms: int = DEFAULT_PRUNE_STATEMENT_TIMEOUT_MS,
    sizer: SweepBatchSizer | None = None,
) -> ArchiveExpiryResult:
    """Hard-delete expired ``jobs_archive`` rows, one bounded,
    self-committing batch at a time.

    Each batch is one ``_EXPIRY_CTE_SQL`` statement, committed before the
    next runs, under the same server-side ``statement_timeout`` and
    breaker semantics as :func:`prune_terminal_jobs`, see that
    function's docstring and
    :data:`~taskq.constants.DEFAULT_PRUNE_STATEMENT_TIMEOUT_MS` for the
    timeout/breaker contract. *drain_gate* is called before every batch;
    a ``False`` return stops the drain with the batches already
    committed.

    The POLICY FLOOR (hypertable deployments only): when the connected
    server has a ``policy_retention`` job registered against
    ``jobs_archive`` (:func:`taskq.backend._retention_floor.retention_policy_floor`,
    probed ONCE per run — before the batch loop, never per batch — and
    failing open to None on any probe error), rows older than the
    policy's own ``drop_after`` horizon gain the ``AND finished_at >=
    $floor`` conjunct and are left to the policy, which drops them
    whole-chunk silently. THE BOUNDARY, stated honestly: OLDER than the
    floor, the POLICY owns deletion (silent, chunk-granular,
    watermark-blind — see the pinned boundary in
    ``tests/test_timescale_retention_interplay.py``); NEWER than it
    (inside the window), this sweep owns deletion (row-exact on
    ``expire_at``, the stamp honored exactly inside the window — the
    young chunk whose rows have passed ``expire_at`` while their chunk
    survives is still this sweep's work, and remains so). The floor's
    clock is the POLICY's clock
    (``finished_at``, the chunk alignment column) and the floor's value
    is the policy's OWN horizon (parsed from the registered config, not
    re-derived), so a policy registered with a different interval than
    the current setting still governs the aged end exactly. A row older
    than the floor inside a young chunk is dropped when the chunk itself
    ages past the boundary, at most one chunk interval late — chunk
    granularity is the trade. On vanilla Postgres the probe fails open
    to None and this sweep runs full-range, byte-identical to the
    no-floor behavior.
    """
    require_schema(schema)
    _validate_positive("statement_timeout_ms", statement_timeout_ms)
    if sizer is None:
        _validate_positive("batch_size", batch_size)
    start = time.monotonic()
    total_deleted = 0
    by_status: dict[str, int] = {}
    # Reported on the result for observability; the DELETE predicate itself
    # is server-side (`expire_at < statement_timestamp()` in _EXPIRY_CTE_SQL).
    # Read from the database anyway so the reported instant stays in the
    # predicate's own clock domain, the sibling cutoffs above were
    # documented as display-only and then quietly grew a second consumer.
    expire_before: datetime = await conn.fetchval(_DB_NOW_SQL)
    # ONE probe per sweep run, anchored to the clock domain the reported
    # cutoff above already uses (the database's), never per batch.
    floor = await retention_policy_floor(
        conn, schema, "jobs_archive", "finished_at", now=expire_before
    )
    sql = _compose_expiry_sql(schema, floor)

    while True:
        if drain_gate is not None and not drain_gate():
            break
        size = _effective_prune_batch_size(batch_size, sizer)
        _record_prune_batch_size("archive_expiry", size, sizer)
        args: tuple[object, ...] = (size,) if floor is None else (size, floor)
        rows = await _run_prune_batch(
            conn,
            sql,
            *args,
            statement_timeout_ms=statement_timeout_ms,
            sweep_name="archive_expiry",
            sizer=sizer,
        )
        if not rows:
            break
        batch_total = 0
        for row in rows:
            row_status: str = row["status"]
            cnt: int = row["cnt"]
            batch_total += cnt
            by_status[row_status] = by_status.get(row_status, 0) + cnt
            record_expired_archive_jobs(row_status, cnt)
        total_deleted += batch_total
        if batch_total < size:
            break

    duration_ms = int((time.monotonic() - start) * 1000)
    return ArchiveExpiryResult(
        total_deleted=total_deleted,
        by_status=by_status,
        expire_before=expire_before,
        duration_ms=duration_ms,
    )

cleanup_stale_workers async

cleanup_stale_workers(
    conn: ConnLike,
    *,
    worker_id: UUID,
    staleness: timedelta,
    schema: str = "taskq",
    batch_size: int = DEFAULT_EVENT_WRITER_BATCH_SIZE,
) -> int

Delete worker rows whose last_seen_at exceeds staleness, one bounded batch per call.

One call deletes at most batch_size workers in one short transaction; repeated calls drain the stale set a batch at a time. The caller's worker_id is never deleted. Returns the number of worker rows removed by this call. Worker-level cascade (maintenance_leader, job_attempts) is handled by the DDL ON DELETE clauses, no extra sweeping needed, and the window bounds workers per call, which bounds that per-worker fan-out per transaction: a whole-fleet crash drains in committed batches instead of one transaction rewriting every stale worker's attempt history.

Source code in src/taskq/worker/_leader_shared.py

async def cleanup_stale_workers(
    conn: ConnLike,
    *,
    worker_id: UUID,
    staleness: timedelta,
    schema: str = "taskq",
    batch_size: int = DEFAULT_EVENT_WRITER_BATCH_SIZE,
) -> int:
    """Delete worker rows whose ``last_seen_at`` exceeds *staleness*, one
    bounded batch per call.

    One call deletes at most ``batch_size`` workers in one short
    transaction; repeated calls drain the stale set a batch at a time.
    The caller's *worker_id* is never deleted. Returns the number of
    worker rows removed by this call. Worker-level cascade
    (``maintenance_leader``, ``job_attempts``) is handled by the DDL
    ``ON DELETE`` clauses, no extra sweeping needed, and the window
    bounds workers per call, which bounds that per-worker fan-out per
    transaction: a whole-fleet crash drains in committed batches instead
    of one transaction rewriting every stale worker's attempt history.
    """
    require_schema(schema)
    sql = _CLEANUP_STALE_WORKERS_SQL.format(schema=schema)
    tag = await conn.execute(sql, staleness, worker_id, batch_size)
    return int(tag.rsplit(" ", 1)[-1]) if tag else 0

complete_stale_batches async

complete_stale_batches(
    conn: ConnLike,
    *,
    schema: str = "taskq",
    batch_size: int = DEFAULT_EVENT_WRITER_BATCH_SIZE,
) -> int

Safety net: mark active batches with zero non-terminal jobs as complete.

Covers batches whose completion hook was lost (consumer crash) and intentionally-empty batches (expected_size=0, no jobs at all). One call is one bounded batch: at most batch_size completions, committed, so the sweep loop drains the remainder one call per tick.

Source code in src/taskq/worker/_leader_shared.py

async def complete_stale_batches(
    conn: ConnLike,
    *,
    schema: str = "taskq",
    batch_size: int = DEFAULT_EVENT_WRITER_BATCH_SIZE,
) -> int:
    """Safety net: mark active batches with zero non-terminal jobs as complete.

    Covers batches whose completion hook was lost (consumer crash) and
    intentionally-empty batches (expected_size=0, no jobs at all). One call
    is one bounded batch: at most *batch_size* completions, committed, so
    the sweep loop drains the remainder one call per tick.
    """
    completed: int = await conn.fetchval(complete_stale_batches_sql(schema), batch_size)
    return completed

prune_terminal_jobs async

prune_terminal_jobs(
    conn: ConnLike,
    *,
    retention_per_status: dict[str, timedelta],
    archive_retention: timedelta,
    batch_size: int = DEFAULT_PRUNE_BATCH_SIZE,
    schema: str = "taskq",
    actor_overrides: dict[str, timedelta] | None = None,
    drain_gate: Callable[[], bool] | None = None,
    statement_timeout_ms: int = DEFAULT_PRUNE_STATEMENT_TIMEOUT_MS,
    sizer: SweepBatchSizer | None = None,
) -> PruneResult

Archive-and-delete terminal jobs past their retention, one bounded, self-committing batch at a time.

Each batch is one transaction of two statements (the bounded candidate window, then the lock-bearing _ARCHIVE_CTE_SQL write: jobs → jobs_archive, job_attempts → job_attempts_archive, jobs DELETE), committed before the next batch runs. The batch runs under a server-side statement_timeout (the :data:~taskq.constants.DEFAULT_PRUNE_STATEMENT_TIMEOUT_MS derivation), and when sizer is given its latched tier, not batch_size, sizes every window, so a database that keeps aborting batches is retried at a reduced tier instead of the same one.

drain_gate is called before every batch; a False return stops the drain (committed batches stay committed; the remainder waits for the next attempt). The prune loop passes a gate that returns False on shutdown and ticks detector-2 liveness between batches.

The archive predicate's clock is the database's own (statement_timestamp() in the candidate statement, re-verified at lock time by the write statement's), and the reported cutoffs are anchored to a database-side clock_timestamp() read, see the "Anchored to the database clock" comment in the body below.

Source code in src/taskq/worker/_leader_shared.py

async def prune_terminal_jobs(
    conn: ConnLike,
    *,
    retention_per_status: dict[str, timedelta],
    archive_retention: timedelta,
    batch_size: int = DEFAULT_PRUNE_BATCH_SIZE,
    schema: str = "taskq",
    actor_overrides: dict[str, timedelta] | None = None,
    drain_gate: Callable[[], bool] | None = None,
    statement_timeout_ms: int = DEFAULT_PRUNE_STATEMENT_TIMEOUT_MS,
    sizer: SweepBatchSizer | None = None,
) -> PruneResult:
    """Archive-and-delete terminal jobs past their retention, one bounded,
    self-committing batch at a time.

    Each batch is one transaction of two statements (the bounded
    candidate window, then the lock-bearing ``_ARCHIVE_CTE_SQL`` write:
    jobs → jobs_archive, job_attempts → job_attempts_archive, jobs
    DELETE), committed before the next batch runs. The batch runs under
    a server-side ``statement_timeout`` (the
    :data:`~taskq.constants.DEFAULT_PRUNE_STATEMENT_TIMEOUT_MS`
    derivation), and when *sizer* is given its latched tier, not
    *batch_size*, sizes every window, so a database that keeps aborting
    batches is retried at a reduced tier instead of the same one.

    *drain_gate* is called before every batch; a ``False`` return stops
    the drain (committed batches stay committed; the remainder waits for
    the next attempt). The prune loop passes a gate that returns False on
    shutdown and ticks detector-2 liveness between batches.

    The archive predicate's clock is the database's own
    (``statement_timestamp()`` in the candidate statement, re-verified
    at lock time by the write statement's), and the reported cutoffs are
    anchored to a database-side ``clock_timestamp()`` read, see the
    "Anchored to the database clock" comment in the body below.
    """
    require_schema(schema)
    _validate_positive("statement_timeout_ms", statement_timeout_ms)
    if sizer is None:
        _validate_positive("batch_size", batch_size)
    start = time.monotonic()
    total_deleted = 0
    total_archived = 0
    by_actor: dict[str, int] = {}
    by_status: dict[str, int] = {}
    cutoffs: dict[str, datetime] = {}
    archive_interval = archive_retention

    # Anchored to the database clock, not this process's: the archive
    # predicate is server-side (statement_timestamp() - $2::interval in
    # _ARCHIVE_CTE_SQL), and the caller derives `prune_old_batches`'
    # DELETE cutoff from `max(cutoffs.values())`, so these are NOT
    # display-only, and a Python `now` here would compare an app-clock
    # instant against DB-written `completed_at` values, pruning batches
    # early or late by the skew.
    db_now: datetime = await conn.fetchval(_DB_NOW_SQL)

    for status in TERMINAL_STATUSES:
        retention = retention_per_status.get(status, DEFAULT_PRUNE_RETENTION)
        cutoffs[status] = db_now - retention
        candidate_sql = _ARCHIVE_CANDIDATE_SQL.format(schema=schema)
        write_sql = _ARCHIVE_CTE_SQL.format(schema=schema)

        while True:
            if drain_gate is not None and not drain_gate():
                break
            size = _effective_prune_batch_size(batch_size, sizer)
            _record_prune_batch_size("prune", size, sizer)
            rows = await _run_prune_archive_batch(
                conn,
                candidate_sql=candidate_sql,
                write_sql=write_sql,
                status=status,
                retention=retention,
                size=size,
                archive_interval=archive_interval,
                actor=None,
                statement_timeout_ms=statement_timeout_ms,
                sweep_name="prune",
                sizer=sizer,
            )
            if not rows:
                break
            batch_total = 0
            for row in rows:
                actor_name: str = row["actor"]
                row_status: str = row["status"]
                cnt: int = row["cnt"]
                batch_total += cnt
                by_actor[actor_name] = by_actor.get(actor_name, 0) + cnt
                by_status[row_status] = by_status.get(row_status, 0) + cnt
                record_pruned_jobs(actor_name, row_status, cnt)
                record_archived_jobs(row_status, cnt)
            total_deleted += batch_total
            total_archived += batch_total
            if batch_total < size:
                break

    if actor_overrides:
        for actor_name, actor_retention in actor_overrides.items():
            candidate_sql = _ARCHIVE_CANDIDATE_ACTOR_SQL.format(schema=schema)
            write_sql = _ARCHIVE_CTE_SQL.format(schema=schema)
            for status in TERMINAL_STATUSES:
                if actor_retention >= retention_per_status.get(status, DEFAULT_PRUNE_RETENTION):
                    continue
                while True:
                    if drain_gate is not None and not drain_gate():
                        break
                    size = _effective_prune_batch_size(batch_size, sizer)
                    _record_prune_batch_size("prune", size, sizer)
                    rows = await _run_prune_archive_batch(
                        conn,
                        candidate_sql=candidate_sql,
                        write_sql=write_sql,
                        status=status,
                        retention=actor_retention,
                        size=size,
                        archive_interval=archive_interval,
                        actor=actor_name,
                        statement_timeout_ms=statement_timeout_ms,
                        sweep_name="prune",
                        sizer=sizer,
                    )
                    if not rows:
                        break
                    batch_total = 0
                    for row in rows:
                        a_name: str = row["actor"]
                        r_status: str = row["status"]
                        cnt: int = row["cnt"]
                        batch_total += cnt
                        by_actor[a_name] = by_actor.get(a_name, 0) + cnt
                        by_status[r_status] = by_status.get(r_status, 0) + cnt
                        record_pruned_jobs(a_name, r_status, cnt)
                        record_archived_jobs(r_status, cnt)
                    total_deleted += batch_total
                    total_archived += batch_total
                    if batch_total < size:
                        break

    duration_ms = int((time.monotonic() - start) * 1000)
    return PruneResult(
        total_deleted=total_deleted,
        archived=total_archived,
        by_actor=by_actor,
        by_status=by_status,
        cutoffs=cutoffs,
        duration_ms=duration_ms,
    )

build_leader_lease_sql

build_leader_lease_sql(schema: str) -> tuple[str, str, str]

Render the elect / renew / resign statements for schema.

Validates schema against the canonical identifier regex before formatting; asyncpg cannot bind identifiers as parameters.

Source code in src/taskq/worker/leader.py

def build_leader_lease_sql(schema: str) -> tuple[str, str, str]:
    """Render the elect / renew / resign statements for *schema*.

    Validates *schema* against the canonical identifier regex before
    formatting; asyncpg cannot bind identifiers as parameters.
    """
    require_schema(schema)
    return (
        _LEADER_ELECT_SQL_TEMPLATE.format(schema=schema),
        _LEADER_RENEW_SQL_TEMPLATE.format(schema=schema),
        _LEADER_RESIGN_SQL_TEMPLATE.format(schema=schema),
    )

Dispatch

dispatch

Dispatch per-job DI dispatch.

:func:dispatch_one_job composes :func:build_actor_scope with :func:consume_one_job so the worker's per-job dispatch path is DI-aware: actors with Annotated[T, Scope.X] parameters receive resolved instances at dispatch time, scoped to their effective scope, with TRANSIENT teardown running per invocation.

The dispatch SQL constants and the dispatch_batch asyncpg helper live in :mod:taskq.backend._dispatch_sql (backend layer). This module imports them from there since worker → backend is the correct layer direction.

logger module-attribute

logger: BoundLogger = get_logger(__name__)

SlotPoolAcquireError

SlotPoolAcquireError(
    *, acquire_timeout: float, detail: str | None = None
)

Bases: Exception

A consumer slot could not obtain a usable transaction connection.

Infrastructure, not a job outcome: the job is already claimed and recovers by lock-lease expiry, so the consumer loop must not count this as a job failure. The bounded acquire (or slot-connection init hook) that raised is recorded and logged at the raise site; this type exists so the loop can route it around the failure accounting. The underlying cause stays chained (__cause__).

Source code in src/taskq/worker/dispatch.py

def __init__(self, *, acquire_timeout: float, detail: str | None = None) -> None:
    super().__init__(
        detail
        if detail is not None
        else f"could not acquire a slot-pool connection within {acquire_timeout}s"
    )

dispatch_one_job async

dispatch_one_job(
    *,
    backend: Backend,
    deps: WorkerDeps,
    job: JobRow,
    worker_id: UUID,
    registry: ProviderRegistry,
    process_scope: ProcessScope,
    thread_scope: ThreadScope,
    loop_scope: LoopScope,
    actor_ref: ActorRef[BaseModel, BaseModel | None],
    actor_config: ActorConfigLike,
    clock: Clock,
    active_jobs: ActiveJobRegistry | None = None,
    max_retry_backoff: timedelta = DEFAULT_MAX_RETRY_BACKOFF,
    logger_arg: BoundLogger | None = None,
    enqueuer: SubJobEnqueuer,
) -> AttemptOutcome

Dispatch one job through the DI-resolved actor scope.

1. Acquire the job's transaction connection, from the worker's dedicated slot pool when one is open (the per-slot path), else the LOOP-scope connection, else none (autonomous consume). An acquire failure raises :class:SlotPoolAcquireError before any span or metric exists: infrastructure, not a job outcome. On the per-slot path the acquired connection also shadows the LOOP-registered asyncpg.Connection for the actor invocation (loop_slot_values), so the actor's own writes join this job's transaction and concurrent slots never share a connection.
2. Create the CONSUMER span with link to the PRODUCER span.
3. Resolve the registered :class:~taskq.obs.ErrorReporter (if any) and validate the payload against actor_ref's payload schema.
4. Build the interim JobContext with the CONSUMER span.
5. Open build_actor_scope to resolve DI kwargs.
6. Hand the resolved kwargs to consume_one_job via a run_actor closure that injects them.
7. TRANSIENT scope is closed on exit (regardless of outcome).
8. Record consumer-path metrics outside the span body.

On the per-slot path the job's enqueuer is a per-job :class:~taskq.client._enqueuer.SubJobEnqueuer bound to the acquired connection, so transactional sub-enqueues and their flush/discard lifecycle are this job's alone.

Source code in src/taskq/worker/dispatch.py

async def dispatch_one_job(
    *,
    backend: Backend,
    deps: WorkerDeps,
    job: JobRow,
    worker_id: UUID,
    registry: ProviderRegistry,
    process_scope: ProcessScope,
    thread_scope: ThreadScope,
    loop_scope: LoopScope,
    actor_ref: ActorRef[BaseModel, BaseModel | None],
    actor_config: ActorConfigLike,
    clock: Clock,
    active_jobs: ActiveJobRegistry | None = None,
    max_retry_backoff: timedelta = DEFAULT_MAX_RETRY_BACKOFF,
    logger_arg: structlog.stdlib.BoundLogger | None = None,
    enqueuer: SubJobEnqueuer,
) -> AttemptOutcome:
    """Dispatch one job through the DI-resolved actor scope.

    1. Acquire the job's transaction connection, from the worker's
       dedicated slot pool when one is open (the per-slot path), else
       the LOOP-scope connection, else none (autonomous consume). An
       acquire failure raises :class:`SlotPoolAcquireError` before any
       span or metric exists: infrastructure, not a job outcome. On the
       per-slot path the acquired connection also shadows the
       LOOP-registered ``asyncpg.Connection`` for the actor invocation
       (``loop_slot_values``), so the actor's own writes join this job's
       transaction and concurrent slots never share a connection.
    2. Create the CONSUMER span with link to the PRODUCER span.
    3. Resolve the registered :class:`~taskq.obs.ErrorReporter` (if any)
       and validate the payload against actor_ref's payload schema.
    4. Build the interim JobContext with the CONSUMER span.
    5. Open build_actor_scope to resolve DI kwargs.
    6. Hand the resolved kwargs to consume_one_job via a run_actor
       closure that injects them.
    7. TRANSIENT scope is closed on exit (regardless of outcome).
    8. Record consumer-path metrics outside the span body.

    On the per-slot path the job's enqueuer is a per-job
    :class:`~taskq.client._enqueuer.SubJobEnqueuer` bound to the
    acquired connection, so transactional sub-enqueues and their
    flush/discard lifecycle are this job's alone.
    """
    link_ctx: trace.SpanContext | None = None
    if job.trace_id and job.span_id:
        try:
            link_ctx = trace.SpanContext(
                trace_id=int(job.trace_id, 16),
                span_id=int(job.span_id, 16),
                is_remote=True,
                trace_flags=trace.TraceFlags(0x01),
            )
        except (ValueError, OverflowError):
            logger.warning(
                "otel-link-skipped",
                reason="malformed_trace_id",
                job_id=str(job.id),
            )
    links = [trace.Link(link_ctx)] if link_ctx else []

    batch_id: str = ""
    if job.metadata:
        raw_bid = job.metadata.get("batch_id")
        if raw_bid is not None:
            batch_id = str(raw_bid)

    # Why the flag check duplicates safe_start_span's own gate: these nine
    # entries (and the str() of the job id among them) are per-job
    # allocations whose only reader is the span. safe_start_span drops the
    # attributes unread when telemetry is off, so build them only when the
    # span can receive them; the exported attribute set is unchanged when
    # it is on.
    consumer_attrs: dict[str, str | int] | None = None
    if otel_enabled():
        consumer_attrs = {
            "messaging.system": "taskq",
            "messaging.destination.name": job.queue,
            "messaging.operation.type": "process",
            "messaging.message.id": str(job.id),
            "messaging.consumer.group.name": deps.settings.worker_group,
            "taskq.actor": job.actor,
            "taskq.attempt": job.attempt,
            "taskq.identity_key": job.identity_key or "",
            "taskq.batch_id": batch_id,
        }

    dispatch_log = logger_arg if logger_arg is not None else logger

    # ── The job's transaction connection ─────────────────────────────
    # When the worker runs a dedicated slot pool (bootstrap opens one
    # whenever a LOOP-scope connection is registered and
    # max_concurrency > 1), each job acquires its own connection for
    # the duration of the dispatch, one transaction per connection, so
    # concurrent slots can never nest savepoints on a shared one. The
    # acquire precedes the span/metrics block on purpose: a job that
    # cannot acquire is infrastructure, not a job outcome, so it must
    # produce no consumer span and no consumed-message record. The
    # slot connection is also what the actor receives: it shadows the
    # LOOP-scope cache for this actor invocation (loop_slot_values
    # below), so an actor's own writes join THIS job's transaction and
    # a registered LOOP-scope connection is never shared across
    # concurrent slots' actors. The registered LOOP-scope
    # connection itself stays untouched, every non-actor reader
    # (bootstrap's activation check, the loop-level enqueuer's
    # provenance inference) still resolves it from the LOOP cache.
    async with AsyncExitStack() as conn_stack:
        job_enqueuer: SubJobEnqueuer = enqueuer
        transaction_conn: ConnLike | None = None
        batch_reissue_owed = False
        actor_loop_slot_values: Mapping[type, object] | None = None
        if deps.slot_pool is not None:
            slot_pool = deps.slot_pool
            acquire_timeout = deps.settings.dispatcher_command_timeout
            try:
                acquired = await slot_pool.acquire(timeout=acquire_timeout)
            except POOL_INFRA_EXCEPTIONS as exc:
                record_slot_pool_acquire_failure()
                logger.error(
                    "slot-pool-acquire-failed",
                    kind="slot_pool_acquire_failed",
                    job_id=str(job.id),
                    pool="slot",
                    acquire_timeout=acquire_timeout,
                    error_class=type(exc).__name__,
                    error_message=str(exc),
                )
                raise SlotPoolAcquireError(acquire_timeout=acquire_timeout) from exc
            transaction_conn = acquired
            # Per-slot LOOP-scope semantics for the actor: the slot
            # connection shadows the LOOP-registered asyncpg.Connection
            # for this invocation (build_actor_scope wraps the LOOP
            # container in a LoopScopeSlotView). The actor's writes then
            # join this job's transaction on this connection, and two
            # concurrent slots' actors can never interleave operations
            # on one connection, asyncpg permits one operation per
            # connection, so the shared shape raised InterfaceError
            # inside healthy actors and burned their retry budget.
            actor_loop_slot_values = {asyncpg.Connection: transaction_conn}

            async def _release_slot_conn(
                # Why default-arg binding for pool/conn: the release must
                # return the connection to the pool that produced it even
                # if a credential reload swaps deps.slot_pool mid-dispatch.
                pool: asyncpg.Pool = slot_pool,
                conn: ConnLike = acquired,
            ) -> None:
                # Why custom release instead of the acquire context's
                # own __aexit__: external cancellation can detach the
                # job's transaction task (the consumer shields an
                # in-flight commit, so the task outlives the dispatch
                # call), and asyncpg's release-reset would then ROLL
                # BACK under that live transaction and hand the
                # connection to a sibling mid-flight, a false
                # `succeeded` row with its sub-jobs gone, and two tasks
                # on one connection. A connection still inside its
                # transaction at release time is terminated instead:
                # the server rolls the transaction back on disconnect,
                # the job row stays `running`, and lock-lease expiry
                # reclaims and retries it, the same loud, retryable
                # outcome a rotation-terminate produces, never a false
                # success.
                #
                # Why the whole body is guarded: this callback runs in
                # the AsyncExitStack unwind, where an exception it
                # raises replaces the job's real outcome, a committed
                # job reported as a dispatch failure. The pool's
                # bounded close (a credential-rotation drain or worker
                # teardown) can terminate the connection or release its
                # proxy underneath the dispatch, and a connection that
                # dead fails even the transaction probe: asyncpg nulls
                # a terminated connection's protocol object (the
                # forwarded probe raises AttributeError), and a proxy
                # the close released refuses the call (InterfaceError).
                # Those states are logged with their cause and
                # swallowed, the bounded close owns a dead
                # connection's disposal, so the unwind never
                # manufactures a failure the job did not have.
                # Anything outside these two families is a programming
                # error and stays loud.
                try:
                    if conn.is_in_transaction():  # pyright: ignore[reportAttributeAccessIssue]  # Why: ConnLike is the object-typed runtime alias for Connection | PoolConnectionProxy; both forward is_in_transaction() to the wrapped protocol.
                        conn.terminate()  # pyright: ignore[reportAttributeAccessIssue]  # Why: same forwarding as is_in_transaction.
                        logger.warning(
                            "slot-conn-terminated-transaction-in-flight",
                            kind="slot_conn_terminated_transaction_in_flight",
                            job_id=str(job.id),
                        )
                        return
                    await pool.release(conn)  # pyright: ignore[reportArgumentType]  # Why: ConnLike is the object-typed runtime alias; release accepts the proxy this acquire produced.
                except asyncpg.InterfaceError as exc:
                    # The pool closed underneath the dispatch: release
                    # is refused, or the close already released the
                    # proxy back, either way the close owns the
                    # connection now.
                    logger.warning(
                        "slot-pool-release-skipped-pool-closed",
                        kind="slot_pool_release_skipped_pool_closed",
                        job_id=str(job.id),
                        error_class=type(exc).__name__,
                        error_message=str(exc),
                    )
                except AttributeError as exc:
                    # The connection was terminated underneath the
                    # dispatch: its protocol object is gone, so even
                    # the probe cannot run. Nothing left to release or
                    # terminate.
                    logger.warning(
                        "slot-pool-release-skipped-conn-dead",
                        kind="slot_pool_release_skipped_conn_dead",
                        job_id=str(job.id),
                        error_class=type(exc).__name__,
                        error_message=str(exc),
                    )

            conn_stack.push_async_callback(_release_slot_conn)
            # The acquired connection must carry the registered connection's
            # declared setup before the actor's DI can receive it. A
            # bootstrap-built pool applied it at connect time (recorded on
            # deps); an injected/foreign pool did not, repair that here,
            # once per physical connection. Armed release first: a failed
            # hook terminates the connection and the unwind must own it.
            await _ensure_registered_init_on_slot_conn(
                acquired,
                deps=deps,
                registry=registry,
                acquire_timeout=acquire_timeout,
                job_id=job.id,
            )
            # Per-job binding: this enqueuer's writes join the slot's
            # transaction and its buffers are this job's alone, so a
            # sibling slot's flush/discard/drain can never reach them.
            # The capacity snapshot is shared with the loop-level
            # enqueuer so per-job construction pays no extra backend
            # read.
            job_enqueuer = SubJobEnqueuer(
                loop_scope_resolved=None,
                worker_pool=deps.worker_pool,
                backend=backend,
                capacity_cache=enqueuer.capacity_cache,
                transaction_conn=transaction_conn,
            )
        else:
            # Why registry-trust instead of isinstance: the LOOP-scope
            # cache is keyed by type and populated through the typed
            # register_value seam, and every other reader of this seam
            # (bootstrap's activation, the enqueuer's provenance
            # inference) trusts that key, an isinstance here would
            # disagree with them and silently drop to the autonomous
            # path for a value the rest of the seam treats as the
            # transaction connection.
            raw_conn = loop_scope.resolved_cache().get(asyncpg.Connection)
            if raw_conn is not None:
                transaction_conn = cast(asyncpg.Connection, raw_conn)  # pyright: ignore[reportUnknownVariableType,reportAssignmentType]  # Why: resolved_cache returns Mapping[type, object]; the DI resolver guarantees the value registered under asyncpg.Connection is one, matching bootstrap's and the enqueuer's trust of the same key.
        # Queue wait from the claimed row's own server-clock stamps, outside
        # the span for sampling independence; a row a test left unstamped
        # records nothing rather than a guess.
        if job.started_at is not None:
            record_queue_wait(
                job.actor,
                job.queue,
                max(0.0, (job.started_at - job.scheduled_at).total_seconds()),
            )
        t0 = time.monotonic()
        outcome: AttemptOutcome = "failed"

        try:
            with safe_start_span(
                f"process {job.actor}",
                kind=SpanKind.CONSUMER,
                attributes=consumer_attrs,
                links=links,
            ) as consumer_span:
                # Resolved first, inside the try: a failure before the
                # actor runs (payload validation, DI resolution) is a
                # terminal failure like any other, and the outer handler
                # below reports it through the same hook. A reporter
                # misregistration surfaces the same way a broken actor
                # dependency does, on the job, through the retry
                # decision, rather than tearing down the consumer loop.
                error_reporter: ErrorReporter | None = None
                try:
                    error_reporter = await _resolve_error_reporter(
                        registry,
                        process_scope=process_scope,
                        thread_scope=thread_scope,
                        loop_scope=loop_scope,
                    )
                    # The row's stored version rides the raise, not the
                    # helper's current-version default, so a row that
                    # predates a payload migration is distinguishable from
                    # a malformed caller payload.
                    validated_payload = validate_actor_payload(
                        actor_ref.payload_type,
                        job.payload,
                        job.actor,
                        payload_schema_ver=str(job.payload_schema_ver),
                    )

                    # Bound once for the whole job: the DI-resolution
                    # context below and the live context consume_one_job
                    # builds share it, off the consumer's own logger so a
                    # job's lines carry one logger name whichever context
                    # emitted them.
                    job_log = bind_job_log(
                        logger_arg if logger_arg is not None else _consumer_log,
                        job,
                        span=consumer_span,
                    )
                    interim_ctx: JobContext[BaseModel] = JobContext(
                        job_id=job.id,
                        actor=job.actor,
                        queue=job.queue,
                        attempt=job.attempt,
                        claim_epoch=job.claim_epoch,
                        snooze_count=job.snooze_count,
                        worker_id=worker_id,
                        payload=validated_payload,
                        jobs=job_enqueuer,
                        log=job_log,
                        span=consumer_span
                        if not isinstance(consumer_span, trace.NonRecordingSpan)
                        else None,
                    )
                    passthrough_kwargs: dict[str, object] = {
                        "payload": validated_payload,
                        "ctx": interim_ctx,
                    }

                    async with build_actor_scope(
                        registry=registry,
                        process_scope=process_scope,
                        thread_scope=thread_scope,
                        loop_scope=loop_scope,
                        actor_func=actor_ref.fn,  # type: ignore[arg-type]  # Why: actor_ref.fn is Callable[..., object] (covers both sync and async); build_actor_scope expects Callable[..., Awaitable[object]] for DI resolution but never calls the function, sync-vs-async dispatch is handled later via actor_ref.is_sync
                        actor_name=actor_ref.name,
                        passthrough_kwargs=passthrough_kwargs,
                        loop_slot_values=actor_loop_slot_values,
                    ) as resolved:

                        async def run_actor_with_di(
                            job_row: JobRow,
                            ctx_arg: JobContext[BaseModel],
                        ) -> object:
                            del job_row
                            actor_kwargs: dict[str, object] = {
                                **resolved.di_kwargs,
                                "payload": ctx_arg.payload,
                            }
                            if actor_ref.wants_ctx:
                                actor_kwargs["ctx"] = ctx_arg
                            if actor_ref.is_sync:
                                return await _run_sync_actor_tracked(
                                    actor_ref.fn,
                                    actor_kwargs,
                                    ctx_arg,
                                )
                            return await actor_ref.fn(**actor_kwargs)  # type: ignore[no-any-return]  # Why: actor_ref.fn is typed Callable[..., object]; runtime result is R.

                        rl_registry: RateLimitRegistry | None = None
                        raw_rl = loop_scope.resolved_cache().get(RateLimitRegistry)
                        if isinstance(raw_rl, RateLimitRegistry):
                            rl_registry = raw_rl

                        redis_client: redis_async.Redis | None = None
                        if _REDIS_CLIENT_TYPE is not None:
                            raw_redis = loop_scope.resolved_cache().get(_REDIS_CLIENT_TYPE)
                            if isinstance(raw_redis, _REDIS_CLIENT_TYPE):
                                redis_client = raw_redis

                        # Fleet-wide per-queue concurrency cap (see
                        # _effective_reservations): O(1) membership test, no
                        # per-job dict/list copies on the no-cap fast path.
                        effective_reservations = _effective_reservations(
                            actor_ref.reservations, job.queue, rl_registry
                        )

                        result = await consume_one_job(
                            backend,
                            job,
                            worker_id,
                            deps=deps,
                            run_actor=run_actor_with_di,
                            actor_config=actor_config,
                            payload_type=actor_ref.payload_type,
                            clock=clock,
                            logger=logger_arg,
                            max_retry_backoff=max_retry_backoff,
                            active_jobs=active_jobs,
                            enqueuer=job_enqueuer,
                            transaction_conn=transaction_conn,
                            validated_payload=validated_payload,
                            rate_limit_registry=rl_registry,
                            rate_limits=actor_ref.rate_limits,
                            reservations=effective_reservations,
                            redis_client=redis_client,
                            worker_pool=deps.worker_pool,
                            settings=deps.settings,
                            error_reporter=error_reporter,
                            fallback_result_ttl=actor_ref.result_ttl,
                            job_log=job_log,
                        )
                        outcome = result

                        # Best-effort: a crash between the terminal write and
                        # the counter increment loses that increment.  See
                        # apply_batch_terminal_outcome docstring for full
                        # safety-net semantics (M7). The hook's reissue-owed
                        # return is captured (batch_reissue_owed below), not
                        # dropped: on the transactional-caller shape a
                        # gated-out completion attempt is unordered with this
                        # dispatch's own commit, and the contract the hook's
                        # docstring states -- the CALLER honors it with one
                        # post-commit re-arbitration -- is honored at the
                        # conn_stack exit below, after every transaction on
                        # the slot connection has ended.
                        try:
                            if await apply_batch_terminal_outcome(
                                backend, job, outcome, transaction_conn=transaction_conn
                            ):
                                batch_reissue_owed = True
                        except Exception:
                            logger.exception("batch-policy-hook-failed", job_id=str(job.id))

                    if outcome == "succeeded":
                        consumer_span.set_status(StatusCode.OK)
                    elif outcome == "failed":
                        consumer_span.set_status(StatusCode.ERROR)

                except asyncio.CancelledError:
                    outcome = "cancelled"
                    consumer_span.set_status(StatusCode.ERROR, "cancelled")
                    # A batch completes on ANY terminal member: the
                    # batch-completion hook runs per-job for all terminal
                    # outcomes (succeeded, failed, cancelled). So a
                    # cancelled last member must complete its batch here,
                    # not a sweep-interval later. Best-effort for the same
                    # reason consume's own mark_cancelled on this path is
                    # best-effort (an infra failure there is logged inside
                    # consume and the row stays running for lock-lease
                    # reclaim): a hook failure here is logged and the M7
                    # stale-batch sweep remains the safety net for batch
                    # status. On the transactional consumer path the
                    # cancel has already aborted the slot's transaction,
                    # so the hook's writes on that connection fail, are
                    # logged, and the sweep recovers, immediate
                    # completion holds on the autonomous path, the normal
                    # case.
                    try:
                        await apply_batch_terminal_outcome(
                            backend, job, "cancelled", transaction_conn=transaction_conn
                        )
                    except Exception:
                        logger.exception("batch-policy-hook-failed", job_id=str(job.id))
                    raise
                except Exception as exc:
                    outcome = "failed"
                    consumer_span.set_status(StatusCode.ERROR)
                    # Unwrap the provider-SystemExit carrier first (the DI
                    # scope's get_or_create seam raises it for a factory's
                    # sys.exit): the failure handler must record the
                    # factory's own exception, never the carrier's name (a
                    # false audit trail).
                    exc = _unwrap_provider_system_exit(exc)
                    handler_log = bind_job_context(
                        dispatch_log,
                        job_id=job.id,
                        actor=job.actor,
                        queue=job.queue,
                        attempt=job.attempt,
                        identity_key=job.identity_key,
                        trace_id="",
                    )
                    try:
                        handler_result: AttemptOutcome = await _handle_generic_exception(
                            backend,
                            job,
                            worker_id,
                            exc,
                            actor_config,
                            max_retry_backoff,
                            consumer_span,
                            handler_log,
                            error_reporter=error_reporter,
                        )
                    except _TERMINAL_WRITE_INFRA_EXCEPTIONS as infra_exc:
                        # An infra-failed terminal write leaves the row
                        # RUNNING, disowned, so lock-lease expiry and the
                        # sweep are the recovery, and no batch counter may
                        # budge on a write that never landed: the same rule
                        # the hook itself applies to the clean-return
                        # path's "noop" (a terminal write that matched
                        # nothing). The hook call therefore lives in the
                        # else below, on a real terminal outcome only.
                        _log_terminal_write_failed(handler_log, job, exc, infra_exc)
                        _disown_job(deps.disowned_jobs, job)
                    else:
                        outcome = handler_result
                        # Best-effort, matching the hook call on consume's
                        # clean return above (M7 sweep semantics documented
                        # on apply_batch_terminal_outcome): non-terminal
                        # handler outcomes ("scheduled") return inside the
                        # hook without touching a counter.
                        try:
                            if await apply_batch_terminal_outcome(
                                backend, job, outcome, transaction_conn=transaction_conn
                            ):
                                batch_reissue_owed = True
                        except Exception:
                            logger.exception("batch-policy-hook-failed", job_id=str(job.id))
        finally:
            elapsed = time.monotonic() - t0
            # A noop means the row moved underneath this dispatch (a
            # reclaim race): nothing was consumed, and the re-dispatch
            # will record the real message and duration, recording here
            # would double-count the message and stretch the histogram
            # with a phantom process.
            if outcome != "noop":
                consumed = _to_consumed_outcome(outcome)
                record_consumed_message(job.actor, job.queue, outcome=consumed)
                record_process_duration(job.actor, job.queue, elapsed, outcome=consumed)

    # The transactional-caller reissue (#589's contract): a completion
    # attempt that gated out while riding this dispatch's transaction was
    # never ordered against its own commit, so the batch is owed one
    # re-arbitration AFTER the transaction ends -- this attempt's snapshot
    # postdates every writer's commit, its probe sees zero open members,
    # and it lands. Without it, a burst-shaped drain leaves the
    # all-terminal batch 'active' with no attempt left to land -- the
    # stuck-batch hole the reissue closes. Best-effort (M7): a failure here
    # is logged and the leader's stale-batch sweep remains the safety net;
    # the CancelledError arm above re-raises past this point, where the
    # sweep covers it. The autonomous shape (no open transaction at hook
    # time) never sets the flag.
    if batch_reissue_owed and batch_id:
        try:
            await backend.complete_batch(UUID(batch_id))
        except Exception:
            logger.exception("batch-policy-reissue-failed", job_id=str(job.id))

    return outcome

Heartbeat and cancellation

heartbeat

Heartbeat loop and cancel-poll seam.

Each tick acquires one connection from heartbeat_pool, opens a single transaction, and atomically extends workers.last_seen_at, jobs lock / heartbeat columns, and reservation_slots leases for jobs locked by this worker, except the jobs the worker has disowned (WorkerDeps.disowned_jobs: finished with, outcome unrecordable), whose leases must lapse so the reclaim sweep can hand them back. The tick's whole command sequence runs under ONE command-timeout budget (see the tick block and :func:_lease_renewal_threshold for the lease-model arithmetic that depends on it), with a bounded rollback-or-close teardown. The jobs-lock renewal is threshold-gated (:func:_lease_renewal_threshold): a row whose lease is still comfortably fresh is left alone so a healthy beat stops paying a non-HOT update per running row per tick, while rows carrying a per-job heartbeat_timeout (the reclaim sweep's heartbeat arm needs their beats fresh) and rows at/under the threshold renew on every beat. After max_heartbeat_failures consecutive failed ticks, transient connection failures and unexpected errors alike, with the count reset only by a fully successful tick, isolate_self proactively transitions running jobs and signals shutdown.

logger module-attribute

logger: BoundLogger = get_logger(__name__)

heartbeat_loop async

heartbeat_loop(
    deps: WorkerDeps,
    worker_id: UUID,
    shutdown: Event,
    *,
    cancel_controller: CancelController | None = None,
    cancel_wake_event: Event | None = None,
) -> None

Source code in src/taskq/worker/heartbeat.py

async def heartbeat_loop(
    deps: WorkerDeps,
    worker_id: UUID,
    shutdown: asyncio.Event,
    *,
    cancel_controller: CancelController | None = None,
    cancel_wake_event: asyncio.Event | None = None,
) -> None:
    interval = deps.settings.heartbeat_interval
    lock_lease = timedelta(seconds=deps.settings.lock_lease)
    schema = deps.settings.schema_name
    # The tick's single command budget: ONE
    # heartbeat_command_timeout for the tick's whole command sequence
    # (BEGIN + writes + probes + hook + COMMIT, and the post-tx drain's
    # share). See the tick block below for why the sequence, not each
    # statement, is the unit the lease model needs bounded.
    tick_command_budget = deps.settings.heartbeat_command_timeout
    # The renewal threshold: a healthy beat only rewrites leases
    # whose remaining time is at or under this. See
    # _lease_renewal_threshold for the sizing derivation (against the
    # tick budget's enforced bound) and the failure-cascade margin
    # proof.
    renewal_threshold = _lease_renewal_threshold(
        lock_lease,
        interval,
        deps.settings.max_heartbeat_failures,
        deps.settings.heartbeat_command_timeout,
    )
    (
        update_worker_liveness_sql,
        update_jobs_lock_sql,
        update_reservation_leases_sql,
    ) = build_heartbeat_sql(schema, renewal_threshold=renewal_threshold)
    select_still_held_sql = _SELECT_STILL_HELD_SQL_TEMPLATE.format(schema=schema)
    refund_lost_claims_sql = _RECONCILE_LOST_CLAIMS_SQL_TEMPLATE.format(schema=schema)

    # Monotonic stamp of the last jobs-lock renewal that landed: the
    # reference the next tick measures its remaining lease against. None
    # until the first renewal, there is nothing to measure before it.
    last_renewal_at: float | None = None
    while not shutdown.is_set():
        _fork_at = take_parent_fork_event()
        if _fork_at is not None:
            # The loop's guaranteed cadence is what makes this the fork
            # report's floor: even a worker idle between jobs learns of a
            # fork within one interval. The parent keeps running - its
            # connections were not written by the child unless the child
            # used them too, and a protocol-dead pooled connection is
            # discarded and rebuilt by the pool - but the operator now has
            # the one line that says a fork happened and what it inherited.
            logger.error(
                "fork-detected-in-worker-process",
                fork_age_secs=round(time.monotonic() - _fork_at, 3),
                detail=(
                    "a fork() happened in this worker process while its "
                    "connections were live; the child inherits every socket "
                    "(pools, LISTEN, redis, the loop's pipes). A job body or "
                    "library must not fork: use subprocess.Popen (default "
                    "close_fds=True drops the inherited descriptors), or "
                    "open fresh resources in the child after the fork."
                ),
            )
        deps.liveness.tick("heartbeat", period=interval)
        _in_tx_failed = False
        _tick_failed = False
        tick_start = time.monotonic()
        try:
            _tick_raised = False
            # The tick's single command budget: None
            # until the acquire lands -- the acquire is bounded by its
            # own timeout (the interval) OUTSIDE this budget -- then the
            # deadline every later phase of the tick shares.
            budget_deadline: float | None = None
            try:
                async with deps.heartbeat_pool.acquire(timeout=interval) as conn:
                    budget_deadline = time.monotonic() + tick_command_budget
                    tx = conn.transaction()
                    try:
                        # ONE command-timeout for the tick's whole command
                        # sequence: BEGIN, the liveness write, the gated
                        # renewal, the reservation-lease write, the
                        # still-held probe, the lost-claim reconcile
                        # probe, the cancel hook's statements,
                        # COMMIT. Each statement's own pool-level
                        # command_timeout stays as the inner backstop; the
                        # budget is what bounds the SEQUENCE, because a
                        # failed tick's cost to the lease model is the
                        # whole tick, not any one statement: a brownout
                        # tick whose first two statements each take just
                        # under one timeout and whose third times out
                        # lasts acquire + ~3 command-timeouts when only
                        # the per-statement bounds apply, which is the
                        # gap bound the tail-less threshold arithmetic
                        # assumed away (the reproduced default-settings
                        # lease-lapse window). A tick whose statements
                        # legitimately need more than one command-timeout
                        # IN TOTAL now fails fast here instead: raise
                        # TASKQ_HEARTBEAT_COMMAND_TIMEOUT for that regime
                        # (the renewal threshold's floor absorbs the
                        # raised value automatically -- see
                        # _lease_renewal_threshold).
                        async with asyncio.timeout(tick_command_budget):
                            await tx.start()
                            # Snapshot: consumers add to the set while this tick
                            # awaits, and an id that arrives mid-tick belongs to
                            # the next tick's exclusion, the prune below must
                            # not read it as "still held" and drop it.
                            disowned = list(deps.disowned_jobs)
                            # The stall-attribution tally rides the liveness write:
                            # one dict merge in the statement the tick already
                            # issues, no extra round trip. An empty tally merges a
                            # no-op, so a quiet process leaves the registered
                            # metadata keys untouched. The idle-fraction window
                            # drains at the same seam (the watchdog thread is its
                            # writer, this loop its reader): the drained aggregate
                            # publishes as ``loop_idle`` beside the stall keys and
                            # feeds the per-window histogram. A failed tick's
                            # drain is metrics-only - the histogram sample stands,
                            # the row's payload dies with the rolled-back
                            # transaction, and the next tick's window starts
                            # fresh (never double-counted).
                            metadata_payload = deps.stall_tally.metadata_value()
                            idle_window = deps.loop_idle.drain()
                            if idle_window is not None:
                                metadata_payload["loop_idle"] = idle_window.as_metadata()
                                _loop_idle_fraction.record(idle_window.idle_fraction)
                            await conn.execute(
                                update_worker_liveness_sql,
                                worker_id,
                                jsonb_param(metadata_payload),
                            )
                            renewal_at = time.monotonic()
                            # The gated renewal: binds the threshold as
                            # $4. Rows with a per-job heartbeat_timeout, rows
                            # with no lease stamp, and rows at/under the
                            # threshold renew; fresh leases are left alone so a
                            # healthy beat stops paying a non-HOT update per
                            # running row per tick. parse_rowcount on the tag
                            # now counts rows RENEWED this tick, not rows held.
                            jobs_tag = await conn.execute(
                                update_jobs_lock_sql,
                                worker_id,
                                lock_lease,
                                disowned,
                                renewal_threshold,
                            )
                            await conn.execute(
                                update_reservation_leases_sql,
                                worker_id,
                                lock_lease,
                                disowned,
                            )
                            if disowned:
                                held_rows = await conn.fetch(
                                    select_still_held_sql, disowned, worker_id
                                )
                                still_held = {row["id"] for row in held_rows}
                                deps.disowned_jobs.difference_update(set(disowned) - still_held)
                            # The claim-loss reconcile: running rows locked
                            # here that nothing holds. The exclusion binds
                            # the THREE maps in one array - registered +
                            # intents (held_ids), queued (parked in the
                            # local queue), and this tick's disowned
                            # snapshot - and the lease-length grace on
                            # started_at spares every in-flight handoff
                            # (see the template's comment). Found ids are
                            # refunded (the claim-time increment of a
                            # never-started execution goes back) and
                            # disowned: this tick's renewal has already
                            # run, the NEXT tick stops renewing them, and
                            # Sweep 1 owns the reclaim from there.
                            lost_excluded = sorted(
                                {
                                    *deps.active_jobs.held_ids(),
                                    *deps.active_jobs.queued_ids(),
                                    *disowned,
                                }
                            )
                            lost_rows = await conn.fetch(
                                refund_lost_claims_sql, worker_id, lost_excluded, lock_lease
                            )
                            if lost_rows:
                                lost_ids: list[UUID] = [row["id"] for row in lost_rows]
                                logger.warning(
                                    "claim-loss-reconciled",
                                    kind="claim_loss_reconciled",
                                    worker_id=str(worker_id),
                                    job_ids=[str(j) for j in lost_ids],
                                )
                            if cancel_controller is not None:
                                try:
                                    await cancel_controller.run_in_tx(conn)  # type: ignore[arg-type]  # Why: asyncpg PoolConnectionProxy is a Connection subclass at runtime; pyright types don't reflect this delegation.
                                except Exception as hook_exc:
                                    _in_tx_failed = True
                                    deps.heartbeat_failures += 1
                                    update_heartbeat_consecutive_failures(
                                        str(worker_id), deps.heartbeat_failures
                                    )
                                    logger.warning(
                                        "heartbeat-hook-failure",
                                        kind="state_change",
                                        cause="heartbeat_hook_failure",
                                        worker_id=str(worker_id),
                                        error=repr(hook_exc),
                                    )
                                    raise OSError(
                                        f"cancel_controller.run_in_tx failed: {hook_exc!r}"
                                    ) from hook_exc
                            # The commit rides the SAME budget: a commit
                            # that would push the sequence past the
                            # command-timeout is cut, the tick fails
                            # without the commit having been confirmed,
                            # and the teardown below closes the
                            # connection (the writes are idempotent
                            # re-stamps, so an unconfirmed commit is
                            # benign: the next tick rewrites them).
                            await tx.commit()
                            # The disown lands AFTER the commit, never
                            # before: the disown is MEMORY state with no
                            # rollback, and the reconcile's exclusion set
                            # folds it - a disown applied before a tick
                            # whose commit then failed (the in-tx cancel
                            # hook's OSError above, a cut commit, a
                            # dropped connection) would survive the
                            # rollback while the REFUND it belongs to
                            # un-applies, and the row would match neither
                            # the renewal nor the reconcile ever again:
                            # the refund is lost permanently, the lease
                            # lapses, and Sweep 1 reclaims at the CHARGED
                            # attempt - issue 458's "crashed, never ran"
                            # record, resurrected by a brownout tick.
                            # Applied post-commit there is no await
                            # between the commit's return and this
                            # synchronous update, so no other task can
                            # observe the gap; a failed commit leaves the
                            # disown unset, the row still matching, and
                            # the next tick re-reconciles it.
                            if lost_rows:
                                deps.disowned_jobs.update(lost_ids)  # pyright: ignore[reportPossiblyUnboundVariable]  # Why: lost_ids is bound exactly when lost_rows is truthy - the same guard guards both; the pair is not reassigned between.
                    except BaseException:
                        # Teardown, bounded by the SAME budget's
                        # remainder, deliberately OUTSIDE the expired
                        # asyncio.timeout scope: an await started after
                        # the scope expired is NOT re-cancelled (measured
                        # , asyncio.timeout fires once), so the teardown
                        # must carry its own bound rather than hide
                        # inside the dead scope. The ROLLBACK and the
                        # CLOSE SHARE that one remainder - the close's
                        # bound is recomputed from the deadline at its
                        # own instant, so a rollback that consumed the
                        # remainder leaves the close an immediate
                        # terminate, never a second full command budget.
                        # This is what holds the failed tick's enforced
                        # cost to acquire + ONE budget (see
                        # _lease_renewal_threshold): a close that took
                        # a FULL second budget after the rollback
                        # had already consumed the remainder would run a
                        # cascade of such ticks
                        # (F+1) * (interval + 2 * command_timeout) past
                        # the last renewal PLUS the last good beat's
                        # tail - measurably past the very cascade floor
                        # the settings validator enforces (the
                        # attack test pins the observed overrun).
                        left = budget_deadline - time.monotonic()
                        rolled_back = False
                        if left > 0:
                            try:
                                async with asyncio.timeout(left):
                                    await tx.rollback()
                                rolled_back = True
                            except TimeoutError:
                                # The rollback did not fit what the tick
                                # had left; fall through to the close.
                                pass
                        if not rolled_back:
                            # Close instead of awaiting a rollback round
                            # trip: the server rolls the transaction back
                            # on disconnect, the close is a local
                            # Terminate write (bounded, terminate() on
                            # timeout, close_conn_bounded never raises),
                            # and the pool discards the closed connection
                            # on release. Bounded by what the budget
                            # still has left (shared with the rollback
                            # above, NOT a second budget): this is the
                            # budget-exhausted path, so it is also the
                            # path that keeps the failed tick inside the
                            # lease model's h + c bound.
                            await close_conn_bounded(
                                conn,  # type: ignore[arg-type]  # Why: PoolConnectionProxy delegates close()/terminate() to the underlying Connection at runtime; pyright's stubs model the proxy as unrelated, the same delegation the run_in_tx call below relies on.
                                "heartbeat-tick-budget",
                                max(0.0, budget_deadline - time.monotonic()),
                                family="mid_run",
                            )
                        # Re-raise the tick's ORIGINAL exception (a bare
                        # raise re-raises it identically): a teardown
                        # timeout or the close's bounds must never
                        # displace what actually failed the tick, least
                        # of all an external cancellation, which must
                        # stay a cancellation.
                        raise
            except BaseException:
                _tick_raised = True
                raise
            finally:
                # Why finally: run_post_tx MUST follow run_in_tx on every tick
                # (see taskq.worker.cancel's module docstring).  Phase-3 jobs
                # queued before an error would otherwise never be drained, and
                # the drain is also where an abandon that could not apply is
                # handed back for a later tick.  Row locks are gone either way
                # by the time this runs, the transaction has committed or
                # rolled back, so mark_abandoned cannot self-deadlock here.
                if cancel_controller is not None:
                    # post_tx shares the tick's ONE command budget: it
                    # runs under whatever the deadline has left, and is
                    # deferred to the next tick when nothing is left (an
                    # expired asyncio.timeout does not re-cancel an
                    # await that starts after expiry, measured, so the
                    # remaining time is enforced with a fresh scope).
                    # The controller's deque persists, so a deferred
                    # drain loses nothing. When the acquire itself
                    # failed, no budget was consumed: post_tx drains the
                    # PREVIOUS ticks' queue under a full budget of its
                    # own (the acquire's own timeout already bounds that
                    # tick).
                    left = (
                        budget_deadline - time.monotonic()
                        if budget_deadline is not None
                        else tick_command_budget
                    )
                    if left <= 0:
                        logger.debug(
                            "heartbeat-post-tx-deferred",
                            worker_id=str(worker_id),
                            reason="tick_command_budget_exhausted",
                        )
                    else:
                        try:
                            async with asyncio.timeout(left):
                                await cancel_controller.run_post_tx()
                        except Exception as post_exc:
                            # A post-tx failure on a healthy tick is the tick's
                            # failure and propagates to the handlers below; on an
                            # already-failing tick it must NOT displace the
                            # original error, which is the root cause worth
                            # reporting. (A mid-drain budget cut lands here as
                            # TimeoutError: the in-flight abandon is shielded
                            # from the tick and re-queued by the drain itself,
                            # so the deferred work rides the next tick.)
                            if not _tick_raised:
                                raise
                            logger.warning(
                                "heartbeat-post-tx-failure",
                                worker_id=str(worker_id),
                                error=repr(post_exc),
                            )
            deps.heartbeat_failures = 0
            update_heartbeat_consecutive_failures(str(worker_id), 0)
            tick_duration_s = time.monotonic() - tick_start
            _tick_duration.record(tick_duration_s)
            # The lease the previous beat's UPDATE stamped had this much
            # left when this one landed: lease minus the gap between the
            # two UPDATEs (both measured at the same point, so network
            # latency cancels). A failed or late tick widens the gap and
            # lowers the sample, which is the signal the lock-expiry
            # alert reads; the config constant would never move. Clamped
            # at 0: a renewal that lands after expiry renews an
            # already-expired lease.
            # Under threshold-gated renewal the sample keeps this
            # beat-cadence meaning deliberately: it is stamped on every
            # successful tick, whether or not the gate renewed any rows,
            # so a late or failing beat lowers it exactly as before and
            # alert thresholds calibrated to the beat cadence
            # keep their semantics. For rows the gate skipped (still
            # above the threshold) the true remaining is anywhere up to
            # the full lease, the sample is the "if this beat renewed
            # everything" floor, never an overstatement of a RENEWED
            # row's remaining, and the threshold's own sizing is pinned
            # by the _lease_renewal_threshold unit and property tests
            # rather than by this histogram.
            if last_renewal_at is not None:
                record_lock_expires_in_seconds(
                    str(worker_id),
                    max(0.0, lock_lease.total_seconds() - (renewal_at - last_renewal_at)),
                )
            last_renewal_at = renewal_at
            logger.debug(
                "heartbeat-tick-success",
                worker_id=str(worker_id),
                tick_duration_ms=int(tick_duration_s * 1000),
                jobs_extended=parse_rowcount(jobs_tag),
                is_leader=deps.is_leader.is_set(),
            )
        except TRANSIENT_PG_ERRORS as e:
            tick_duration_s = time.monotonic() - tick_start
            _tick_duration.record(tick_duration_s)
            _tick_failed = True
            if not _in_tx_failed:
                deps.heartbeat_failures += 1
                update_heartbeat_consecutive_failures(str(worker_id), deps.heartbeat_failures)
            logger.warning(
                "heartbeat-tick-failure",
                worker_id=str(worker_id),
                consecutive_failures=deps.heartbeat_failures,
                error_class=type(e).__name__,
                error=str(e),
            )
            if await _failed_tick_ledger(deps, worker_id, shutdown, e):
                return
        except Exception as e:
            # The unexpected arm counts toward the SAME isolate threshold
            # as the transient arm (see _failed_tick_ledger): a persistent
            # non-transient error must isolate the worker, not loop
            # forever as a functional zombie whose counter never moves.
            tick_duration_s = time.monotonic() - tick_start
            _tick_duration.record(tick_duration_s)
            _tick_failed = True
            deps.heartbeat_failures += 1
            update_heartbeat_consecutive_failures(str(worker_id), deps.heartbeat_failures)
            logger.exception(
                "heartbeat-tick-unexpected-error",
                worker_id=str(worker_id),
                consecutive_failures=deps.heartbeat_failures,
                error_class=type(e).__name__,
            )
            if await _failed_tick_ledger(deps, worker_id, shutdown, e):
                return
        # The wait is anchored to the tick's START, not its end, so the
        # beat cadence is the interval however long a SUCCESSFUL tick
        # took. A fixed post-tick sleep instead makes the cadence
        # tick_duration + interval, and a tick may legitimately run for
        # nearly a whole interval, the pool acquire above is bounded at
        # exactly that. A tick that overruns the interval waits zero and
        # re-enters immediately, which is the correct urgency, it is
        # already late, and cannot become a hot loop, because the next
        # tick's own pool acquire is bounded at the interval and paces it.
        #
        # A FAILED tick is NOT a beat, so it does not get the inter-beat
        # wait either. Sleeping the full remaining interval after a tick
        # that raised made the gap between good beats twice the interval
        # after ONE transient blip (the failed tick stamped nothing, then
        # the loop slept a whole interval anyway): the recovery beat
        # landed at the ``heartbeat_timeout >= 2x interval`` floor the
        # ops guidance calls safe rather than inside it, and a sweep
        # running in the deadline-crossing window reclaimed a live,
        # lease-valid worker's job - a phantom cancel, a duplicate
        # execution of work that was never lost. So a failed tick
        # retries promptly: it waits at most a quarter interval (bounded
        # by what the cadence still owes, so an overrun tick still waits
        # zero), then re-enters. The next tick after a successful
        # recovery beat resumes the normal anchored cadence. The retry
        # cannot hammer: each failure pays the bounded backoff again,
        # and the per-tick pool acquire (bounded at the interval) paces
        # every tick that does run. The ledger counts the blip exactly
        # once - the failed tick incremented the counter, its prompt
        # retry succeeds and resets it - so one transient failure never
        # consumes the isolate budget, and repeated failures still
        # isolate on the (max_heartbeat_failures + 1)-th consecutive
        # one (each failed cycle is now SHORTER, which only moves that
        # decision earlier inside the lease the cascade floor sizes -
        # see _lease_renewal_threshold).
        remaining = max(0.0, interval - (time.monotonic() - tick_start))
        if _tick_failed:
            remaining = min(remaining, _FAILED_TICK_RETRY_FRACTION * interval)
        if cancel_wake_event is not None:
            # Wait out the remainder, but wake immediately on a cancel NOTIFY.
            with contextlib.suppress(asyncio.TimeoutError):
                await asyncio.wait_for(cancel_wake_event.wait(), timeout=remaining)
            cancel_wake_event.clear()
        else:
            await asyncio.sleep(remaining)

isolate_self async

isolate_self(
    deps: WorkerDeps, worker_id: UUID, shutdown: Event
) -> None

Source code in src/taskq/worker/heartbeat.py

async def isolate_self(
    deps: WorkerDeps,
    worker_id: UUID,
    shutdown: asyncio.Event,
) -> None:
    assert deps.settings.pg_dsn_direct is not None
    pg_dsn = str(deps.settings.pg_dsn_direct)
    host = dsn_host(pg_dsn)
    schema = deps.settings.schema_name
    require_schema(schema)
    # Stop the claim path FIRST, before any snapshot or join: this worker
    # is walking away, so the producer must start no new claim rounds, and
    # the producer's own exit pass (drain_local_queue_to_pending) hands
    # back what it claimed but never dispatched. The stop event also
    # closes the residual gap the late exclusion capture inside _inner
    # cannot close (a claim committed to the jobs table but not yet marked
    # in either map, one scheduler step wide): after it, a row a consumer
    # has taken is never dispatched (run.py's stop guard), so a re-pend of
    # such a row cannot double-run.
    deps.producer_stop_event.set()
    select_running_jobs_sql = _SELECT_RUNNING_JOBS_SQL_TEMPLATE.format(schema=schema)
    isolate_job_sql = _ISOLATE_JOB_SQL_TEMPLATE.format(schema=schema)
    insert_attempt_sql = INSERT_ATTEMPT_SQL.format(schema=schema)
    # The reclaim delay's effective ceiling, bound per statement, the
    # same operator knob the sweep binds as its own parameter, so a
    # heartbeat-lost hand-back lands on the same schedule the leader's
    # reclaim would have stamped.
    max_backoff_seconds = deps.settings.max_retry_backoff.total_seconds()
    jobs_pending_count = 0
    jobs_crashed_count = 0
    jobs_cancelled_count = 0
    # Rows whose guarded UPDATE no-oped, transitioned by the leader's
    # sweep between this path's SELECT and UPDATE; their attempt rows
    # belong to the winner, and the count keeps the complete log's
    # arithmetic explainable (selected rows = pending + crashed +
    # cancelled + lost_race).
    jobs_lost_race_count = 0
    # The re-pend's exclusion set is deliberately NOT captured here: a
    # snapshot taken this early re-pends every row claimed
    # during the join window below (the producer keeps claiming for the
    # window's whole bound), handing back rows whose local handlers were
    # live. The capture lives inside _inner, after the join, immediately
    # before the SELECT that binds it.
    # Actors still executing in THIS process: their rows are excluded by
    # that late capture, and they take the same route the shutdown
    # orchestrator's
    # CANCELLING phase gives them: the SHUTDOWN origin stamp plus the cancel event
    # routes each consumer's unwinding through its mark_interrupted arm,
    # which earns the exit-proving hold and releases the row with the
    # spent attempt standing. Where PG is truly unreachable and the
    # interrupt writes fail, the rows stay running and lock-lease expiry
    # reclaims them, the honest residual.
    active_entries = deps.active_jobs.all()
    loop = asyncio.get_running_loop()
    for active in active_entries:
        active.ctx.cancel_event.set()
        if active.cancel_origin is CancelOrigin.NONE:
            active.cancel_origin = CancelOrigin.SHUTDOWN
            active.ctx._set_cancel_origin(CancelOrigin.SHUTDOWN)  # pyright: ignore[reportPrivateUsage]  # Why: the isolate path is the other designated shutdown-side writer of the context's origin stamp, same contract as the orchestrator's CANCELLING phase.
        if active.cancel_phase < CancelPhase.COOPERATIVE:
            active.cancel_phase = CancelPhase.COOPERATIVE
            active.cancel_observed_at = loop.time()
        if not active.task.done():
            active.task.cancel()
    if active_entries:
        # The interrupt writes run inside the consumers' own unwinding,
        # bounded by the terminal-write retry budget; the join here is the
        # outer bound (the grace periods the actor's own unwinding may
        # take, plus close slack). Entries still alive when the join ends
        # keep their rows excluded from the re-pend: the exclusion is
        # captured AFTER this join (inside _inner), so both they and every
        # claim this process took while it ran are in it. Lock-lease
        # expiry is the backstop for rows whose unwinding never finishes.
        join_bound = (
            deps.settings.cancellation_grace_period
            + deps.settings.cleanup_grace_period
            + CLOSE_TIMEOUT_SECS
        )
        live_tasks = [active.task for active in active_entries if not active.task.done()]
        if live_tasks:
            _done, still_running = await asyncio.wait(live_tasks, timeout=join_bound)
            if still_running:
                logger.warning(
                    "isolate-self-actor-join-timeout",
                    kind="isolate_self_actor_join_timeout",
                    worker_id=worker_id,
                    still_running=[
                        str(active.job_id) for active in active_entries if not active.task.done()
                    ],
                )

    try:
        # command_timeout, not just the connect timeout: the connect budget
        # bounds the handshake only, and every statement inside _inner() would
        # otherwise block on the socket read for asyncpg's default 60s each,
        # against the browned-out PG (accepts connections, answers slowly)
        # that a heartbeat-failure cascade produces. The isolate must finish
        # or give up promptly either way: shutdown.set() in the finally is
        # what ends the worker, and an unbounded transaction park would turn
        # the graceful-walk-away path into the watchdog's force exit.
        conn = await asyncpg.connect(
            pg_dsn,
            timeout=5.0,  # pyright: ignore[reportCallIssue]  # Why: asyncpg-stubs does not declare timeout kwarg on connect(); the parameter exists at runtime at 0.31.0.
            command_timeout=deps.settings.dispatcher_command_timeout,
            connection_class=guarded_connection_class(),
        )
        try:

            async def _inner() -> tuple[int, int, int, int]:
                pending = 0
                crashed = 0
                cancelled = 0
                lost_race = 0
                # The reclaim events, batched into ONE insert at the end
                # of the transaction (the sweep's own event writer shape:
                # the microsecond ladder on the row ordinal keeps
                # occurred_at co-monotonic with the bigserial id the
                # watermark protocol reads).
                event_job_ids: list[JobId] = []
                event_details: list[object] = []
                # The re-pend's exclusion set, captured as LATE as the
                # statement boundary allows: AFTER the actor join above,
                # with no await between this capture and the SELECT below
                # that binds it. held_ids() covers BOTH maps, the
                # registered consumers (executing now) AND the claim
                # intents (taken off local_queue, not yet registered). The
                # intents are the easy one to miss: a consumer parked in
                # the take-to-register window has already claimed its row
                # (running, locked, this worker) but owns no registry
                # entry, so a snapshot of ``all()`` alone misses it and the
                # re-pend would hand the row back to the fleet while the
                # local body is about to execute it, a peer claims the
                # re-pended row the moment its scheduled_at arrives and
                # runs it concurrently with the local body, a double run,
                # and the local terminal write then loses the
                # attempt-epoch fence (silently discarded).
                # queued_ids() is included HERE even though the DRAINING
                # exit hand-back excludes it on purpose (held_ids' own
                # docstring): a row parked in local_queue during this
                # window is still this process's to execute - the consumer
                # loops are free to take and run it (only their NEXT
                # iteration honours the stop events, and the
                # shutdown_event-only fall-through runs a job taken on the
                # final turn even after shutdown.set()) - so re-pending it
                # hands it to a peer while the local body runs. If the
                # process dies before the take, lock-lease expiry is the
                # backstop, the same honest residual the interrupt arms
                # carry. The residual one-scheduler-step gap (a claim
                # committed but not yet marked in either map) is closed by
                # the producer_stop_event set at entry, not by this
                # capture: after it, a taken row is never dispatched.
                excluded_ids: list[JobId] = [
                    *deps.active_jobs.held_ids(),
                    *deps.active_jobs.queued_ids(),
                ]
                async with conn.transaction():
                    rows = await conn.fetch(  # pyright: ignore[reportUnknownVariableType]  # Why: conn type suppressed above due to asyncpg-stubs limitation on connect().
                        select_running_jobs_sql,
                        worker_id,
                        # The still-owned rows: their release belongs to the
                        # consumers' interrupt arms (the hold and the spent
                        # attempt), not to this re-pend. An empty array
                        # vacuously matches every row, so the statement
                        # shape never varies.
                        excluded_ids,
                    )
                    for row in rows:  # pyright: ignore[reportUnknownVariableType]  # Why: rows type suppressed above, propagates from conn.fetch() suppression.
                        # The guarded UPDATE is the race arbiter, so its
                        # rowcount, not the SELECT that picked the row ,
                        # decides whether this worker owns the transition.
                        # A row the leader's Sweep 1 reclaimed between the
                        # SELECT and this UPDATE reads UPDATE 0 here, and
                        # the sweep's own attempt row already holds
                        # job_attempts' PRIMARY KEY (job_id, attempt): an
                        # unconditional INSERT on the lost row is a
                        # non-transient constraint error that aborts the
                        # WHOLE transaction and collapses the isolation of
                        # rows that are still this worker's. Only the
                        # winner of the transition writes the attempt row.
                        # fetchrow, not execute+parse_rowcount: the
                        # RETURNING row IS the standing-claim fence's
                        # source of truth (the comment at the attempt
                        # INSERT below) - the row's attempt and stamp as
                        # they stand at the winning transition, which a
                        # reconcile refund committed inside the
                        # SELECT->UPDATE window has already un-stamped.
                        updated = await conn.fetchrow(  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # Why: conn type suppressed above due to asyncpg-stubs limitation on connect(); asyncpg-stubs does not type fetchrow's overload for this sql shape.
                            isolate_job_sql, row["id"], worker_id, max_backoff_seconds
                        )
                        if updated is None:
                            lost_race += 1
                            continue
                        # The classification reads the ARBITER's RETURNING -
                        # the same trustworthy-read doctrine the attempt
                        # row's fence below runs on: a refund that commits
                        # in the SELECT->arbiter window de-charges and
                        # un-stamps the row BEFORE the arbiter takes its
                        # lock, and the arbiter's own CASE then re-pends it
                        # (the budget conjunct re-evaluates on the live
                        # row). A classification from the SELECT snapshot
                        # would call that same row 'crashed' - its snapshot
                        # attempt was already spent - and fabricate a
                        # verdict the row never wrote: the ledger says
                        # pending, the reclaim event's to_state and this
                        # counter say crashed. retry_kind is
                        # enqueue-immutable, so the only snapshot inputs
                        # that can have moved between SELECT and arbiter
                        # are exactly the ones the RETURNING carries.
                        new_status: str = str(updated["status"])
                        if new_status == "pending":
                            pending += 1
                        elif new_status == "cancelled":
                            cancelled += 1
                        else:
                            # The arbiter's CASE has exactly one remaining
                            # arm: crashed - the not-pending,
                            # not-operator-cancelled outcome of the budget
                            # CASE. Total by construction: no unmapped
                            # status can reach here without a CASE
                            # evolution, and the exhaustiveness pin
                            # (tests/test_obs_reclaim_counters.py) fails CI
                            # before that day comes - the same doctrine as
                            # the sweep's _reclaim_disposition.
                            crashed += 1
                        # The reclaim rides the crash-reclaim outbox
                        # channel exactly like the sweep's rows (see the
                        # module comment). reason='lock_expired' is the
                        # CHANNEL key -- the slice poll_reclaim_events
                        # tails and the retention carve-out keeps -- not
                        # a per-row lease claim: this row's lease may
                        # still be live (isolate fires on heartbeat-
                        # connection loss, not lease expiry), exactly as
                        # the sweep's heartbeat_timeout arm stamps the
                        # same channel value for rows whose global lease
                        # never expired. The per-event truth is
                        # cause='isolate_self' (the sweep's rows name the
                        # deadline that fired); worker_id is the
                        # last-known holder, the sweep's detail shape.
                        detail: dict[str, object] = {
                            "from_state": "running",
                            "to_state": new_status,
                            "reason": "lock_expired",
                            "cause": "isolate_self",
                            "worker_id": str(worker_id),
                        }
                        event_job_ids.append(JobId(row["id"]))
                        event_details.append(jsonb_param(detail))
                        # THE STANDING-CLAIM FENCE, read from the ARBITER's
                        # RETURNING - never from the SELECT snapshot above.
                        # The gap this closes (the soak's reconciliation red,
                        # ``attempt counter 1 vs 2 attempt rows``): the
                        # claim-loss reconcile's refund
                        # (_RECONCILE_LOST_CLAIMS_SQL, the same worker's
                        # heartbeat pool, or the shutdown drain's
                        # ATTEMPT_REFUND_SQL arm) can commit in the window
                        # between this path's SELECT and the guarded UPDATE
                        # - the SELECT takes no row lock, the refund matches
                        # exactly the row it sees there (running, locked by
                        # this worker, its charge de-charged and its
                        # started_at un-stamped), and an arbiter judging
                        # only (status + holder, no standing-claim conjunct)
                        # would win the transition anyway, then the INSERT
                        # below would write the SNAPSHOT's (attempt,
                        # started_at) - an epoch minted AFTER the refund
                        # that the live counter never carried. The next
                        # claim re-mints that number
                        # (the refund's whole point) and the ledger closes
                        # the lineage with more attempt rows than the
                        # counter. The arbiter is the trustworthy read: it
                        # takes the row lock, so no refund can interleave
                        # after it, and a refund before it shows up in the
                        # RETURNING as the un-stamped, de-charged row. A
                        # returned NULL stamp is a charge this isolate never
                        # owned: record NO attempt row (the reclaim event
                        # above still lands - the reclaim happened and is
                        # audited; the execution the crashed row would
                        # assert never did), and leave the re-pend standing,
                        # the same fenced shape Sweep 1's reclaim leaves
                        # (its batched INSERT carries the twin fence,
                        # _SWEEP_1_ATTEMPTS_BATCH_SQL's ``WHERE a.started_at
                        # IS NOT NULL``). The fence also keeps
                        # INSERT_ATTEMPT_SQL's raw stamp from hitting
                        # job_attempts.started_at's NOT NULL on exactly this
                        # shape (a NotNullViolation here would abort the
                        # whole isolation transaction and strand every
                        # sibling row with it).
                        if updated["started_at"] is not None:
                            await conn.execute(
                                insert_attempt_sql,
                                row["id"],
                                updated["attempt"],
                                updated["started_at"],
                                "crashed",
                                ERROR_CLASS_HEARTBEAT_LOST,
                                None,
                                None,
                                None,
                                worker_id,
                                "{}",  # metadata, matches the sweep paths' literal
                                # The attempt's due time: the claim-time
                                # scheduled_at from the snapshot SELECT (the
                                # arbiter's re-pend already rescheduled the row,
                                # so its RETURNING cannot carry it; the ONLY
                                # writer between claim and isolate is nobody -
                                # the loop was blocked - so the snapshot value
                                # IS the claim-time due time).
                                row["scheduled_at"],
                            )
                    if event_job_ids:
                        await conn.execute(
                            INSERT_EVENTS_DETAIL_BATCH_SQL.format(schema=schema),
                            event_job_ids,
                            event_details,
                            "state_change",
                        )
                return pending, crashed, cancelled, lost_race

            # shield_with_retrieval, not plain asyncio.shield: on outer
            # cancel the isolation tx keeps running detached on a conn the
            # finally below closes, its late failure must be retrieved and
            # logged, not lost as "Task exception was never retrieved" noise
            # (see taskq._shield).
            (
                jobs_pending_count,
                jobs_crashed_count,
                jobs_cancelled_count,
                jobs_lost_race_count,
            ) = await shield_with_retrieval(_inner())
        finally:
            # Why bounded: isolate_self only runs when PG is already
            # suspected dead (heartbeat failures exceeded), so this close is
            # exactly the dead-PG hang case, unbounded, it would
            # wedge shutdown.set() below. The helper never raises, so a
            # close error cannot mask an in-flight exception or be
            # misreported as an isolate-self failure.
            await close_conn_bounded(conn, "isolate-self", CLOSE_TIMEOUT_SECS, family="mid_run")
    except Exception as exc:
        logger.warning(
            "isolate-self-failure",
            kind="isolate_self_failure",
            worker_id=str(worker_id),
            dsn_host=host,
            error=repr(exc),
        )
    finally:
        shutdown.set()
        logger.warning(
            "isolate-self-complete",
            worker_id=str(worker_id),
            jobs_pending_count=jobs_pending_count,
            jobs_crashed_count=jobs_crashed_count,
            jobs_cancelled_count=jobs_cancelled_count,
            jobs_lost_race_count=jobs_lost_race_count,
        )

cancel

Active-job tracking, cancel-poll hook factory.

This module implements ActiveJobRegistry, the loop-scoped in-process map of running jobs, the _ActiveJob dataclass, and the CancelController class that drives the five-phase cancel-poll loop.

CancelController exposes two methods that heartbeat_loop calls on every tick:

• run_in_tx(conn), runs inside the heartbeat transaction. Phases 1, 2, and the phase-3 eligibility check happen here. Phase-3 jobs are queued into _pending_abandons rather than calling mark_abandoned directly, because mark_abandoned uses a separate pool connection that would deadlock on the row lock that the heartbeat transaction still holds.

• run_post_tx(), called by heartbeat_loop AFTER the transaction block exits (and therefore after the transaction has committed and released its row locks). Drains _pending_abandons, calling mark_abandoned + deregister for each entry.

Key correctness invariants (): - _by_id mutations (register/deregister) are protected by asyncio.Lock. - all() is synchronous: in asyncio's single-threaded model no other coroutine can mutate _by_id between the list-copy and return unless there is an intervening await. all() has no await, so the copy is atomic from the event-loop perspective. The lock is NOT acquired in all(), acquiring an asyncio.Lock requires await and would force a coroutine boundary that breaks the atomicity guarantee. - cancel_observed_at uses asyncio.get_running_loop().time() (monotonic event-loop clock), never time.time() or datetime.now(). - Phase-2 PG write (conn.execute) happens BEFORE task.cancel() in code order with NO intervening await (PG-first invariant). - No try/except inside the phase-2 block, PG-write failures propagate to heartbeat_loop's outer handler. - run_post_tx is always called after run_in_tx on the same tick, even when run_in_tx raises; heartbeat_loop must call it in a finally block (or equivalent) to drain any entries queued before the error.

__all__ module-attribute

__all__ = [
    "ActiveJobRegistry",
    "CancelController",
    "_ActiveJob",
    "make_cancel_controller",
]

CancelController

Bases: Protocol

Structural interface for cancel-poll controllers.

heartbeat_loop calls run_in_tx inside the open heartbeat transaction and run_post_tx immediately after the transaction commits. Implementations must satisfy this two-phase contract.

The concrete production implementation is _CancelController, constructed via make_cancel_controller. Test stubs need only implement these two methods to satisfy the type.

run_in_tx async

run_in_tx(conn: Connection) -> None

Execute cancel-poll phases 1-3 inside the heartbeat transaction.

Source code in src/taskq/worker/cancel.py

async def run_in_tx(self, conn: asyncpg.Connection) -> None:
    """Execute cancel-poll phases 1-3 inside the heartbeat transaction."""
    ...

run_post_tx async

run_post_tx() -> None

Drain phase-3 abandonment queue after the transaction commits.

Source code in src/taskq/worker/cancel.py

async def run_post_tx(self) -> None:
    """Drain phase-3 abandonment queue after the transaction commits."""
    ...

ClaimIntent dataclass

ClaimIntent()

The token a claim take hands back to the taker.

Pure identity: one instance per mark_claimed call, held by the loop iteration that took the row and returned to resolve_claim at exit, so the resolver drops only ITS OWN claim (the issue-461 class: the same key can be re-marked by a later generation's take before the stale generation unwinds). Deliberately empty, the value is the object identity alone.

ActiveJobRegistry

ActiveJobRegistry()

Loop-scoped in-process map of running jobs on this worker.

One instance per worker process, constructed in _main() / WorkerDeps before the TaskGroup is entered. Multiple workers in the same process (test scenarios) each carry their own independent registry.

Thread-safety: not applicable, asyncio workers are single-threaded. The asyncio.Lock on _by_id prevents interleaving between coroutines that await register / await deregister while the heartbeat or consumer is also running.

Public surface per

• register(job_id, task, ctx) -> _ActiveJob (async)
• deregister(job_id, entry) -> None (async, identity-scoped)
• get(job_id) -> _ActiveJob | None (sync)
• all() -> list[_ActiveJob] (sync snapshot copy)
• count() -> int (sync)

Source code in src/taskq/worker/cancel.py

def __init__(self) -> None:
    self._by_id: dict[JobId, _ActiveJob] = {}
    self._lock: asyncio.Lock = asyncio.Lock()
    # Claimed-but-not-yet-registered ids: the window between the
    # consumer's queue take (a bare job.id read, no await) and the
    # register() call is invisible to ``all()`` because the DB row
    # carries no "a consumer took it" mark. Any hand-back pass that
    # ran inside the window would re-pend a row this process is about
    # to execute, and the fleet would run it concurrently. The intent
    # map closes that window: the mark lands with no await after the
    # take, and every hand-back pass excludes both maps. Keyed by id
    # to a CLAIM TOKEN, not a bare set: the token is what makes
    # ``resolve_claim`` identity-scoped (the issue-461 class). A
    # stale attempt's loop iteration can unwind AFTER a same-worker
    # re-claim took the row again (the re-claim's own
    # ``mark_claimed`` overwrote the key); a bare-id discard would
    # erase the LIVE claim's intent and the hand-back passes would
    # re-pend a row this process is about to execute. The resolver
    # drops only the claim it holds.
    self._claim_intents: dict[JobId, ClaimIntent] = {}
    # Claimed-but-not-yet-taken ids: the window between the
    # producer's claim commit and the consumer's queue take. A row
    # parked in local_queue (all consumers busy on long jobs) is
    # running, locked to this worker, and invisible to ``all()`` and
    # the intent map alike; a hand-back pass that excludes only
    # those two maps would re-pend a row this process is about to
    # execute locally. ``mark_claimed`` moves the coverage from this
    # map to the intent map at the take, so every window of the
    # claim-to-register chain is fenced by exactly one map.
    self._queued: set[JobId] = set()

mark_enqueued

mark_enqueued(job_id: JobId) -> None

Record a claim the producer handed to the local queue.

Called by the producer before the queue put (the put is the first await after the claim's return, and the mark must precede any await a hand-back pass could interleave at).

Source code in src/taskq/worker/cancel.py

def mark_enqueued(self, job_id: JobId) -> None:
    """Record a claim the producer handed to the local queue.

    Called by the producer before the queue put (the put is the
    first await after the claim's return, and the mark must precede
    any await a hand-back pass could interleave at).
    """
    self._queued.add(job_id)

queued_ids

queued_ids() -> list[JobId]

Snapshot of the rows parked in local_queue, not yet taken.

Source code in src/taskq/worker/cancel.py

def queued_ids(self) -> list[JobId]:
    """Snapshot of the rows parked in local_queue, not yet taken."""
    return list(self._queued)

mark_claimed

mark_claimed(job_id: JobId) -> ClaimIntent

Record a queue take before any await can let a drain observe the gap.

Must be called with no intervening await after the take: the single-threaded loop makes the record atomic with the take, which is the whole guarantee. The take also moves the row's coverage from the queued map to the intent map, so the claim-to-register chain never has an unfenced window.

Returns the claim token the caller must hand back to resolve_claim at exit: the same key can be re-marked by a later generation's take (a same-worker re-claim) before this caller unwinds, and the token is the only identity that scopes the resolve to this caller's own claim.

Source code in src/taskq/worker/cancel.py

def mark_claimed(self, job_id: JobId) -> ClaimIntent:
    """Record a queue take before any await can let a drain observe the gap.

    Must be called with no intervening await after the take: the
    single-threaded loop makes the record atomic with the take, which
    is the whole guarantee. The take also moves the row's coverage
    from the queued map to the intent map, so the claim-to-register
    chain never has an unfenced window.

    Returns the claim token the caller must hand back to
    ``resolve_claim`` at exit: the same key can be re-marked by a
    later generation's take (a same-worker re-claim) before this
    caller unwinds, and the token is the only identity that scopes
    the resolve to this caller's own claim.
    """
    token = ClaimIntent()
    self._claim_intents[job_id] = token
    self._queued.discard(job_id)
    return token

resolve_claim

resolve_claim(job_id: JobId, token: ClaimIntent) -> None

Drop the claim intent once register covers it or the row is released.

Identity-scoped: the entry is removed only when the map's current token IS the caller's. A stale generation's resolve (its take superseded by a same-worker re-claim's new mark_claimed) removes nothing, and the live claim's intent survives to fence the hand-back passes for the window it exists to cover.

Source code in src/taskq/worker/cancel.py

def resolve_claim(self, job_id: JobId, token: ClaimIntent) -> None:
    """Drop the claim intent once ``register`` covers it or the row is released.

    Identity-scoped: the entry is removed only when the map's current
    token IS the caller's. A stale generation's resolve (its take
    superseded by a same-worker re-claim's new ``mark_claimed``)
    removes nothing, and the live claim's intent survives to fence
    the hand-back passes for the window it exists to cover.
    """
    if self._claim_intents.get(job_id) is token:
        del self._claim_intents[job_id]

held_ids

held_ids() -> list[JobId]

Snapshot of every row this process may still execute: registered and intent.

Deliberately NOT the queued ids (queued_ids()): a row parked in local_queue at a DRAINING exit must be re-pended by the exit hand-back, which excludes this snapshot only — the hand-back passes (drain, isolate) and the heartbeat's lost-claim probe bind different exclusions on purpose, each the narrower of the two its correctness needs.

Source code in src/taskq/worker/cancel.py

def held_ids(self) -> list[JobId]:
    """Snapshot of every row this process may still execute: registered and intent.

    Deliberately NOT the queued ids (``queued_ids()``): a row parked
    in local_queue at a DRAINING exit must be re-pended by the exit
    hand-back, which excludes this snapshot only — the hand-back
    passes (drain, isolate) and the heartbeat's lost-claim probe
    bind different exclusions on purpose, each the narrower of the
    two its correctness needs.
    """
    return list(self._by_id) + list(self._claim_intents)

register async

register(
    job_id: JobId,
    task: Task[object],
    ctx: JobContext[BaseModel],
) -> _ActiveJob

Register a job as in-flight.

The lock ensures no concurrent deregister sees an inconsistent state. The claim intent (if any) is absorbed: the registry now owns the row.

Returns the entry it installed. The caller must hand that entry back to deregister at exit: the same key can be re-registered by a later attempt of the same job (a lapsed lease re-claimed on this worker), and the return value is the only identity that tells the caller's own registration from the live attempt's.

Source code in src/taskq/worker/cancel.py

async def register(
    self,
    job_id: JobId,
    task: asyncio.Task[object],
    ctx: JobContext[BaseModel],
) -> _ActiveJob:
    """Register a job as in-flight.

    The lock ensures no concurrent ``deregister`` sees an inconsistent state.
    The claim intent (if any) is absorbed: the registry now owns the row.

    Returns the entry it installed. The caller must hand that entry back
    to ``deregister`` at exit: the same key can be re-registered by a
    later attempt of the same job (a lapsed lease re-claimed on this
    worker), and the return value is the only identity that tells the
    caller's own registration from the live attempt's.
    """
    entry = _ActiveJob(job_id=job_id, task=task, ctx=ctx)
    async with self._lock:
        # The absorb is deliberately key-scoped, not token-scoped: the
        # registering attempt's chain created the intent at its own
        # take, and the intent's whole purpose is the pre-registration
        # window this call closes. Any registration of the key covers
        # whatever intent stands (a re-claim's registry overwrite
        # makes a stale claim's intent moot the same way).
        self._claim_intents.pop(job_id, None)
        self._by_id[job_id] = entry
    return entry

deregister async

deregister(job_id: JobId, entry: _ActiveJob) -> None

Remove a job from the registry (idempotent, ignores missing keys).

Identity-scoped: the pop happens only when the map's current entry IS the entry this caller registered (or holds from get()). A same-worker re-claim re-registers the same key with a NEW entry (the later attempt's), so a stale attempt's exit removes only its own registration and the live attempt's survives. The bare-id pop this replaces evicted the live attempt's registration behind every defence that reads held_ids() (issue 461). This is the same fence taskq.progress._flush._drop_fenced_out_buffer applies to the progress buffers, the one other same-keyed map whose entries are per-attempt.

Source code in src/taskq/worker/cancel.py

async def deregister(self, job_id: JobId, entry: _ActiveJob) -> None:
    """Remove a job from the registry (idempotent, ignores missing keys).

    Identity-scoped: the pop happens only when the map's current entry
    IS the ``entry`` this caller registered (or holds from ``get()``).
    A same-worker re-claim re-registers the same key with a NEW entry
    (the later attempt's), so a stale attempt's exit removes only its
    own registration and the live attempt's survives. The bare-id pop
    this replaces evicted the live attempt's registration behind every
    defence that reads ``held_ids()`` (issue 461). This is the same
    fence ``taskq.progress._flush._drop_fenced_out_buffer`` applies to
    the progress buffers, the one other same-keyed map whose entries
    are per-attempt.
    """
    async with self._lock:
        if self._by_id.get(job_id) is not entry:
            return
        del self._by_id[job_id]

get

get(job_id: JobId) -> _ActiveJob | None

Return the registry entry for job_id, or None if absent.

Synchronous and lock-free: safe to call from the heartbeat hook or any non-mutating code path between awaits.

Source code in src/taskq/worker/cancel.py

def get(self, job_id: JobId) -> _ActiveJob | None:
    """Return the registry entry for ``job_id``, or ``None`` if absent.

    Synchronous and lock-free: safe to call from the heartbeat hook or any
    non-mutating code path between awaits.
    """
    return self._by_id.get(job_id)

all

all() -> list[_ActiveJob]

Return a snapshot copy of all in-flight entries.

Synchronous: in asyncio's cooperative multitasking, no other coroutine can mutate _by_id while this method runs (there is no await between list(...) and return). The copy ensures the caller's for loop cannot raise RuntimeError: dictionary changed size during iteration even if register/deregister are called in later coroutine steps.

Callers that need a fresh count after iterating must call count() separately.

Source code in src/taskq/worker/cancel.py

def all(self) -> list[_ActiveJob]:
    """Return a snapshot copy of all in-flight entries.

    Synchronous: in asyncio's cooperative multitasking, no other coroutine
    can mutate ``_by_id`` while this method runs (there is no ``await``
    between ``list(...)`` and ``return``).  The copy ensures the caller's
    ``for`` loop cannot raise ``RuntimeError: dictionary changed size during
    iteration`` even if ``register``/``deregister`` are called in later
    coroutine steps.

    Callers that need a fresh count after iterating must call ``count()``
    separately.
    """
    return list(self._by_id.values())

count

count() -> int

Return the number of currently registered in-flight jobs.

Source code in src/taskq/worker/cancel.py

def count(self) -> int:
    """Return the number of currently registered in-flight jobs."""
    return len(self._by_id)

intent_count

intent_count() -> int

Return the number of takes not yet absorbed by a registration.

The claim-to-register window (the consumer's DI resolution, payload validation, and admission acquire all run before register) is invisible to count(), and that window is where a denied admission's bounded local retry now waits, so a producer sizing a claim by registered jobs alone can double-claim every slot its consumers are already holding. The intents are disjoint from both the queued ids (mark_claimed moves the row out of _queued) and the registered entries (register absorbs the intent), so count() + intent_count() is the exact in-process hold.

Source code in src/taskq/worker/cancel.py

def intent_count(self) -> int:
    """Return the number of takes not yet absorbed by a registration.

    The claim-to-register window (the consumer's DI resolution, payload
    validation, and admission acquire all run before ``register``) is
    invisible to ``count()``, and that window is where a denied
    admission's bounded local retry now waits, so a producer sizing a
    claim by registered jobs alone can double-claim every slot its
    consumers are already holding. The intents are disjoint from both
    the queued ids (``mark_claimed`` moves the row out of ``_queued``)
    and the registered entries (``register`` absorbs the intent), so
    ``count() + intent_count()`` is the exact in-process hold.
    """
    return len(self._claim_intents)

__len__

__len__() -> int

Source code in src/taskq/worker/cancel.py

def __len__(self) -> int:
    return len(self._by_id)

make_cancel_controller

make_cancel_controller(
    deps: WorkerDeps, worker_id: UUID, backend: Backend
) -> CancelController

Construct a CancelController for the given worker.

Validates schema_name eagerly (before any ticks run) so misconfiguration is surfaced at startup rather than on the first heartbeat.

Source code in src/taskq/worker/cancel.py

def make_cancel_controller(
    deps: "WorkerDeps",
    worker_id: UUID,
    backend: Backend,
) -> CancelController:
    """Construct a ``CancelController`` for the given worker.

    Validates ``schema_name`` eagerly (before any ticks run) so misconfiguration
    is surfaced at startup rather than on the first heartbeat.
    """
    return _CancelController(deps, worker_id, backend)

Shutdown

shutdown

Two-phase shutdown orchestration.

Consumers (orchestrate_shutdown, health endpoints) MUST observe deps.shutdown_phase set at the START of each phase, BEFORE any per-phase work. Value NONE (0) means the worker is running normally.

Phase ordering invariant: NONE (0) → DRAINING (1) → CANCELLING (2) → FORCING (3) → RELEASING (4).

The early handover: shutdown_start_event fires before phase 1 touches any row, and the leader runtime's election loop obeys it (park, no attempts; a leader resigns at shutdown START, before the drain - leader.py's shutdown-ordering contract). The worker-wide shutdown_event does not fire until the phases complete; nothing in the phases below reads leadership state.

SIGQUIT is not registered; produces a core dump on Linux. Use tini or ulimit -c 0 for containerised deployments.

The second-SIGTERM contract: if the second SIGTERM arrives during FORCING or RELEASING, setting escalate_event is a no-op, the orchestrator is already past CANCELLING.

What the phases owe the work: a deploy is an infrastructure event, so it never terminalises a job. Rows claimed but never started are handed back at DRAINING (attempt refunded; the producer repeats the hand-back on exit so a claim round in flight at the signal is caught too). Rows mid-execution get the cooperative cancel at CANCELLING and the forced cancel at FORCING; an actor that unwinds is interrupted , released back to the fleet, the spent attempt standing; and one still alive past both graces is interrupted with a hold at RELEASING (released only once the process is provably gone). abandoned is never a shutdown verdict: an operator-cancelled row the process dies with is terminalised with the operator's OWN verdict (mark_cancelled, forced) at RELEASING - the same write its unwinding consumer races - never relabelled abandoned.

__all__ module-attribute

__all__ = [
    "ShutdownPhase",
    "drain_local_queue_to_pending",
    "install_signal_handlers",
    "orchestrate_shutdown",
]

ShutdownPhase

Bases: IntEnum

Worker shutdown phase.

NONE , running normally. DRAINING , stop accepting new dispatch, hand back claimed-but-unstarted jobs (attempt refunded). CANCELLING, cooperative cancel of remaining jobs; stamps the shutdown origin on each signalled job. FORCING , force-cancel grace, terminal writes shielded; the cancel escalation write doubles as the origin probe (it lands only on rows already carrying an operator's cancel request). RELEASING , release jobs whose actors never unwound back to the fleet (mark_interrupted: the spent attempt stands, held until this process is provably gone); jobs under an operator cancel are terminalised with the operator's own verdict (mark_cancelled, forced) here.

The value 4 was ABANDONING before the release phase stopped abandoning, the integer is unchanged, so /health JSON and the CLI table keep their numbers; only the label moved to what the phase does.

NONE class-attribute instance-attribute

NONE = 0

DRAINING class-attribute instance-attribute

DRAINING = 1

CANCELLING class-attribute instance-attribute

CANCELLING = 2

FORCING class-attribute instance-attribute

FORCING = 3

RELEASING class-attribute instance-attribute

RELEASING = 4

drain_local_queue_to_pending async

drain_local_queue_to_pending(
    deps: WorkerDeps, worker_id: UUID
) -> int

Re-pend every job this worker claimed but never started.

Issues a single bounded-timeout UPDATE that clears the lock on rows where locked_by_worker = $worker_id AND status = 'running', excluding the jobs with live consumers (deps.active_jobs): CANCELLING owns those, and re-pending one would unlock a row another worker can claim while its consumer still executes it. On pool exhaustion or connection error the helper logs a warning and returns 0 so the recovery sweep acts as the backstop rather than a deadlocked shutdown.

Two callers, one pass each, both on this worker's way out: the orchestrator's DRAINING phase, and the producer loop's own exit (a claim round in flight when the stop event landed commits after the DRAINING pass; the producer's exit pass is the one write that cannot be overtaken by this worker's next claim, because there is none). Both refund the claim's attempt increment through the shared ATTEMPT_REFUND_SQL fragment, a claim that never reached an actor bought nothing, so it spends nothing (the same idiom the snooze arms carry; the refund is floored at 0 and a second pass matches no rows, so the two passes together are exactly-once). mark_interrupted is the deliberate exception: its attempt DID start executing, so it keeps the increment.

Why the started_at IS NOT NULL conjunct: the dispatch claim CTE stamps started_at = clock_timestamp() AT CLAIM (backend/_dispatch_sql.py), so every row the drain exists for is stamped and the conjunct never fences a claimed-but-unstarted row. What it fences is the OTHER refund writer's output: the heartbeat's claim-loss reconcile (worker/heartbeat.py) refunds through the same fragment and UN-STAMPS started_at while leaving the row running and locked (Sweep 1 owns that reclaim). On a running-and-locked row a NULL started_at means exactly "this claim was already refunded", so without the conjunct the drain re-pends the reconciled row and refunds the SAME claim a second time: the counter drops below the epoch a genuine execution's job_attempts row holds, and the row is yanked out of the Sweep-1 reclaim the reconcile preserved. The conjunct is a row-version check, not a memory snapshot, so it holds on the concurrent interleave too: the drain's UPDATE blocking on the reconcile's open transaction re-evaluates the row when the reconcile commits (EvalPlanQual) and the committed NULL stamp drops it from the write set. The refund is exactly-once per claim across both writers.

The disowned exclusion: a row this worker DISOWNED (a terminal write that failed on infrastructure, a slot-pool acquire failure) is not this helper's population. The disown arms cover attempts that may have STARTED executing, and refunding a started attempt re-creates the exact epoch the first execution ran under - the successor's re-claim would re-run the body at the SAME attempt number. The row's recovery stays where the disown arms put it: the lease lapses unrenewed and Sweep 1 reclaims at the charged attempt (the reconcile's documented owner).

The DB row carries no "a consumer took it" mark, so the only honest discriminator for "never started" is this process's own active-jobs registry; the claim-to-register window (a job taken off local_queue but not yet in active_jobs) is invisible to every shutdown arm, CANCELLING iterates the same registry, and stays outside this predicate's guarantee.

Returns: Number of rows updated, or 0 on timeout / connection error.

Source code in src/taskq/worker/shutdown.py

async def drain_local_queue_to_pending(deps: "WorkerDeps", worker_id: UUID) -> int:
    """Re-pend every job this worker claimed but never started.

       Issues a single bounded-timeout UPDATE that clears the lock on rows
       where ``locked_by_worker = $worker_id AND status = 'running'``,
       excluding the jobs with live consumers (``deps.active_jobs``):
       CANCELLING owns those, and re-pending one would unlock a row
       another worker can claim while its consumer still executes it. On
       pool exhaustion or connection error the helper logs a warning and
       returns 0 so the recovery sweep acts as the backstop rather than a
       deadlocked shutdown.

       Two callers, one pass each, both on this worker's way out: the
       orchestrator's DRAINING phase, and the producer loop's own exit (a
       claim round in flight when the stop event landed commits after the
       DRAINING pass; the producer's exit pass is the one write that cannot
       be overtaken by this worker's next claim, because there is none).
       Both refund the claim's attempt increment through the shared
       ``ATTEMPT_REFUND_SQL`` fragment, a claim that never reached an
       actor bought nothing, so it spends nothing (the same idiom the
       snooze arms carry; the refund is floored at 0
       and a second pass matches no rows, so the two passes together are
       exactly-once). ``mark_interrupted`` is the deliberate exception: its
    attempt DID start executing, so it keeps the increment.

       Why the ``started_at IS NOT NULL`` conjunct: the dispatch claim CTE
       stamps ``started_at = clock_timestamp()`` AT CLAIM
       (backend/_dispatch_sql.py), so every row the drain exists for is
       stamped and the conjunct never fences a claimed-but-unstarted row.
       What it fences is the OTHER refund writer's output: the heartbeat's
       claim-loss reconcile (worker/heartbeat.py) refunds through the same
       fragment and UN-STAMPS ``started_at`` while leaving the row
       ``running`` and locked (Sweep 1 owns that reclaim). On a
       running-and-locked row a NULL ``started_at`` means exactly "this
       claim was already refunded", so without the conjunct the drain
       re-pends the reconciled row and refunds the SAME claim a second
       time: the counter drops below the epoch a genuine execution's
       ``job_attempts`` row holds, and the row is yanked out of the
       Sweep-1 reclaim the reconcile preserved. The conjunct is a
       row-version check, not a memory snapshot, so it holds on the
       concurrent interleave too: the drain's UPDATE blocking on the
       reconcile's open transaction re-evaluates the row when the
       reconcile commits (EvalPlanQual) and the committed NULL stamp
       drops it from the write set. The refund is exactly-once per claim
       across both writers.

       The disowned exclusion: a row this worker DISOWNED (a terminal
       write that failed on infrastructure, a slot-pool acquire failure)
       is not this helper's population. The disown arms cover attempts
       that may have STARTED executing, and refunding a started attempt
       re-creates the exact epoch the first execution ran under - the
       successor's re-claim would re-run the body at the SAME attempt
       number. The row's recovery stays where the disown arms put it: the
       lease lapses unrenewed and Sweep 1 reclaims at the charged attempt
       (the reconcile's documented owner).

       The DB row carries no
       "a consumer took it" mark, so the only honest discriminator for
       "never started" is this process's own active-jobs registry; the
       claim-to-register window (a job taken off local_queue but not yet
       in ``active_jobs``) is invisible to every shutdown arm, CANCELLING
       iterates the same registry, and stays outside this predicate's
       guarantee.

       Returns:
           Number of rows updated, or 0 on timeout / connection error.
    """
    schema = deps.settings.schema_name
    require_schema(schema)

    # The attempt refund: the claim stamped attempt + 1 for an execution
    # this hand-back says never happened, so the increment goes back ,
    # the same non-consuming-release idiom the snooze/unavailable and
    # interruption arms carry through ATTEMPT_REFUND_SQL. Without it
    # every rolling deploy spends
    # one retry of every claimed-but-unstarted job's budget. The alias
    # ``j`` is what the shared fragment qualifies on.
    sql = (
        f"UPDATE \"{schema}\".jobs j SET status='pending', locked_by_worker=NULL, "  # noqa: S608  # Why: schema validated against _IDENT_RE before interpolation; asyncpg has no parameter binding for identifiers (same rationale as migrate.py).
        # A drain re-pend hands the row back to the fleet, so it routes
        # by the actor's current assignment from here on (the routing
        # contract in taskq/backend/_dispatch_sql.py) -- the same
        # re-pend class as _SWEEP_1_SQL and the isolate template.
        f"lock_expires_at=NULL, assignment_routed=true, attempt = {ATTEMPT_REFUND_SQL} "
        f"WHERE locked_by_worker=$1 AND status='running' AND j.cancel_phase = 0 "
        "AND NOT EXISTS ("
        f'    SELECT 1 FROM "{schema}".job_attempts a'
        "    WHERE a.job_id = j.id AND a.attempt = j.attempt"
        ")"
        # The exactly-once refund fence across the two refund writers
        # (the docstring's conjunct paragraph): a running-and-locked row
        # with a NULL started_at is the heartbeat reconcile's refunded
        # output, already refunded once, and the drain must leave it for
        # Sweep 1's reclaim.
        " AND j.started_at IS NOT NULL"
    )
    # The ledger guard (the NOT EXISTS above): the refund's premise - "a
    # claim that never reached an actor bought nothing" - is enforced
    # with data, not just the registry's absence. An attempt row already
    # recorded for the row's current attempt number (the isolate-self
    # write, the sweep's reclaim INSERT, a retried terminal write) IS a
    # started attempt: its charge stands, and refunding it would make
    # the next claim revisit the number the PK already carries - the
    # exact collision the ATTEMPT_REFUND_SQL fragment's safety argument
    # excludes. A never-started claim (no registry entry, no ledger row
    # - the drain's whole design) matches nothing and refunds unchanged;
    # a guard-blocked row stays running under its lease, which lapse
    # hands to Sweep 1, the same owner every other started attempt's
    # reclaim has.
    # The cancel fence (``cancel_phase = 0``): a row carrying an operator
    # cancel in flight must NOT re-enter the fleet through the drain: the
    # same fence every other deferral/release arm carries (the snooze and
    # retry-after arms, mark_interrupted's release arm, sweep-1's
    # cancel-first CASE, the isolate template). Without it, a cancel
    # stamped between the operator's write and the consumer's next flag
    # poll would be handed back to dispatch, and the row would execute
    # again under the next holder before any ladder terminalises it. The
    # fenced row stays here: running, owned by this worker, refund never
    # applied; its lease expiry hands it to sweep-1's cancel arm, which
    # terminalises it with the operator's audit intact.
    # The exclusion clauses are only bound when there is something to
    # exclude: an empty registry (the common drained-worker case) keeps
    # the single-parameter statement shape the helper has always issued.
    try:
        async with deps.dispatcher_pool.acquire(timeout=2.0) as conn:
            # The exclusion arrays are captured HERE, after the pool
            # acquire and immediately before the bind - never before it.
            # The capture is the row's classification ("a consumer of
            # this process holds this row, the hand-back must not touch
            # it"), and the statement is the writer that acts on it: any
            # take (a claim intent's mark) or registration landing
            # between an EARLIER capture and this statement's execution
            # would be invisible to the exclusion array, and the UPDATE
            # would re-pend the row UNDER ITS LIVE CONSUMER - the
            # consumer's terminal write then fences out
            # (JOB_FENCE_BOUND_SQL requires status='running' AND
            # locked_by_worker, both cleared by this re-pend) and the
            # row is stranded 'pending' with its body mid-flight. The
            # pool acquire is the one await between the function's entry
            # and the statement, and on a contended pool (the shared-CI
            # shape) it parks this task for as long as the contention
            # lasts - exactly the window a take or a registration lands
            # in. Capturing after it leaves the bind-to-execute round
            # trip as the only gap, the same order every other
            # capture-then-write arm of the shutdown uses.
            #
            # Why the explicit list[UUID] annotation: JobId is
            # NewType(UUID), so the bare comprehension infers
            # list[JobId], and list invariance would refuse the uuid[]
            # bind parameter's declared type below. held_ids() covers
            # BOTH maps: registered consumers (executing now) and claim
            # intents (taken off local_queue, not yet registered, the
            # window the registry's comment documents).
            active_ids: list[JobId] = deps.active_jobs.held_ids()
            # The disowned rows are never this pass's population (the
            # docstring's disowned-exclusion paragraph): their attempts
            # may have started, so the refund is wrong for them, and
            # Sweep 1 owns their reclaim. Why the comprehension instead
            # of sorted(): JobId is NewType(UUID), sorted() erases the
            # NewType, and list invariance would refuse the uuid[] bind
            # parameter's declared type below.
            disowned_ids: list[JobId] = [JobId(job_id) for job_id in sorted(deps.disowned_jobs)]
            params: list[UUID | list[JobId]] = [worker_id]
            if active_ids:
                sql += " AND id <> ALL($2::uuid[])"
                params.append(active_ids)
            if disowned_ids:
                sql += f" AND NOT j.id = ANY(${len(params) + 1}::uuid[])"
                params.append(disowned_ids)
            tag = await conn.execute(sql, *params)
            rowcount = parse_rowcount(tag)
            _log.info(
                "drain-local-queue-completed",
                worker_id=worker_id,
                rows_re_pended=rowcount,
                active_jobs_excluded=len(active_ids),
            )
            return rowcount
    except TRANSIENT_PG_ERRORS as exc:
        _log.warning(
            "drain-local-queue-failed",
            worker_id=worker_id,
            error=str(exc),
        )
        return 0

orchestrate_shutdown async

orchestrate_shutdown(
    deps: WorkerDeps,
    settings: WorkerSettings,
    worker_id: UUID,
    shutdown_event: Event,
    escalate_event: Event | None = None,
    *,
    backend: Backend,
) -> int

Run the four-phase shutdown orchestration.

Phases are DRAINING → CANCELLING → FORCING → RELEASING, followed by TaskQ-owned leader_conn close and shutdown_event.set(). Each phase is assigned to deps.shutdown_phase BEFORE any per-phase work. Returns 0 on clean exit.

Source code in src/taskq/worker/shutdown.py

async def orchestrate_shutdown(
    deps: "WorkerDeps",
    settings: "WorkerSettings",
    worker_id: UUID,
    shutdown_event: asyncio.Event,
    escalate_event: asyncio.Event | None = None,
    *,
    backend: Backend,
) -> int:
    """Run the four-phase shutdown orchestration.

    Phases are DRAINING → CANCELLING → FORCING → RELEASING, followed by
    TaskQ-owned ``leader_conn`` close and ``shutdown_event.set()``.  Each
    phase is assigned to ``deps.shutdown_phase`` BEFORE any per-phase work.
    Returns 0 on clean exit.
    """
    loop = asyncio.get_running_loop()
    t0 = loop.time()

    try:
        # ── Phase 1: DRAINING ──────────────────────────────────────────
        # The early handover (leader.py's shutdown-ordering contract):
        # the stop signal fires BEFORE any phase work, so the leader
        # runtime's election loop parks (no attempts, no renewals) and a
        # leading pod resigns the lease while this pod still drains its
        # jobs - the successor elects into the drain window instead of
        # waiting for it. Everything downstream of this line is written
        # with no leadership premise: the drain's machinery is
        # self-contained (dispatcher pool + backend writes).
        deps.shutdown_start_event.set()
        if deps.shutdown_started_at is None:
            deps.shutdown_started_at = t0
        deps.shutdown_phase = ShutdownPhase.DRAINING
        _log.info(
            "shutdown-phase",
            kind="shutdown_phase",
            phase="DRAINING",
            active_jobs_count=deps.active_jobs.count(),
            elapsed_seconds=0.0,
        )
        deps.producer_stop_event.set()
        # The DRAINING failure containment: the helper contains the
        # transient-PG shapes itself (log + return 0, its backstop
        # doctrine), but a NON-transient one - a data error, a dropped
        # table, the schema-identifier contract - escapes it. Uncontained
        # here, the exception collapses the phase train: CANCELLING never
        # stamps the in-flight jobs' origins, FORCING never delivers
        # task.cancel(), RELEASING never releases, and the jobs the
        # shutdown owes its writes are thrown to the watchdog's os._exit
        # and the lease-expiry sweep while the exception sits unretrieved
        # in the orchestrator task until _main joins it after the whole
        # worker has exited. One pass, no retry (the watchdog owns the
        # budget; the sweep is the backstop - the same no-retry doctrine
        # the release arm carries), the failure loud and the train
        # continuing.
        try:
            await drain_local_queue_to_pending(deps, worker_id)
        except Exception as exc:
            _log.error(
                "drain-phase-failed",
                kind="drain_phase_failed",
                worker_id=str(worker_id),
                error_class=type(exc).__name__,
                error=str(exc),
            )

        # ── Phase 2: CANCELLING ────────────────────────────────────────
        deps.shutdown_phase = ShutdownPhase.CANCELLING
        cancel_grace = settings.cancellation_grace_period
        for active in deps.active_jobs.all():
            active.ctx.cancel_event.set()
            if active.cancel_phase < CancelPhase.COOPERATIVE:
                active.cancel_phase = CancelPhase.COOPERATIVE
                active.cancel_observed_at = loop.time()
            elif active.cancel_observed_at is None:
                active.cancel_observed_at = loop.time()
            # Stamp the shutdown as the cancel's origin, unless an
            # operator's cancel was already observed (the controller's
            # PG-observation arms stamp OPERATOR; the row says who asked).
            # The consumer's terminal routing reads the origin: SHUTDOWN
            # releases the attempt back to the fleet (interrupted), an
            # operator's keeps the cancel ladder. A later operator cancel
            # still wins: the controller overrides the stamp on its next
            # observation, and mark_interrupted's cancel_phase = 0 fence
            # declines the row either way.
            if active.cancel_origin is CancelOrigin.NONE:
                active.cancel_origin = CancelOrigin.SHUTDOWN
                active.ctx._set_cancel_origin(CancelOrigin.SHUTDOWN)  # pyright: ignore[reportPrivateUsage]  # Why: the orchestrator is the designated shutdown-side writer of the context's origin stamp (set alongside cancel_event.set(), per the field's contract).
        _log.info(
            "shutdown-phase",
            kind="shutdown_phase",
            phase="CANCELLING",
            active_jobs_count=deps.active_jobs.count(),
            cancel_origin_counts=_cancel_origin_counts(deps),
            elapsed_seconds=loop.time() - t0,
        )

        deadline = loop.time() + cancel_grace
        while loop.time() < deadline and deps.active_jobs.count() > 0:
            if escalate_event is not None and escalate_event.is_set():
                break
            await asyncio.sleep(0.1)

        # ── Phase 3: FORCING ───────────────────────────────────────────
        deps.shutdown_phase = ShutdownPhase.FORCING
        cleanup_grace = settings.cleanup_grace_period
        _log.info(
            "shutdown-phase",
            kind="shutdown_phase",
            phase="FORCING",
            active_jobs_count=deps.active_jobs.count(),
            elapsed_seconds=loop.time() - t0,
        )
        for active in deps.active_jobs.all():
            try:
                # shield_with_retrieval, not plain asyncio.shield: shutdown
                # races escalating cancellation, so a detached write here can
                # be double-cancelled, its outcome must be retrieved (see
                # taskq._shield).
                #
                # The escalation write doubles as the origin probe: it lands
                # only when the row is at cancel_phase = 1, i.e. carrying an
                # operator's cancel request. A SHUTDOWN-stamped entry whose
                # probe lands was racing an unobserved operator cancel (the
                # heartbeat poll had not seen the row yet), the row says
                # who asked, so the entry re-stamps OPERATOR and stays on
                # the operator ladder (RELEASING abandons it past the
                # graces). A probe that misses on an OPERATOR-stamped entry
                # means the row moved under the ladder (already terminal or
                # reclaimed); that is logged, never silently discarded.
                escalated = await shield_with_retrieval(
                    backend.write_cancel_escalation(active.job_id, worker_id, phase=2)
                )
            except Exception as e:
                _log.warning(
                    "force-cancel-pg-write-failed",
                    job_id=str(active.job_id),
                    error=str(e),
                )
                # The local cancel is delivered even when the row-side
                # write failed: skipping it let a cancellable actor run
                # untouched into RELEASING and be released-with-hold
                # while still alive: the exact overlap the hold exists
                # to prevent. The escalation probe is the row-side
                # half of FORCING; task.cancel() is the process-side
                # half, and only both together advance the entry. The
                # phase stamp mirrors the success path below so the
                # registry's local ladder stays truthful either way.
                active.task.cancel()
                active.cancel_phase = CancelPhase.FORCED
                continue
            if escalated and active.cancel_origin is CancelOrigin.SHUTDOWN:
                active.cancel_origin = CancelOrigin.OPERATOR
                active.ctx._set_cancel_origin(CancelOrigin.OPERATOR)  # pyright: ignore[reportPrivateUsage]  # Why: the orchestrator is the designated shutdown-side writer of the context's origin stamp (per the field's contract).
                _log.info(
                    "force-cancel-escalation-observed-operator-cancel",
                    kind="cancel_origin",
                    job_id=str(active.job_id),
                    worker_id=str(worker_id),
                )
            elif not escalated and active.cancel_origin is CancelOrigin.OPERATOR:
                _log.warning(
                    "force-cancel-escalation-unmatched",
                    job_id=str(active.job_id),
                    worker_id=str(worker_id),
                )
            active.task.cancel()
            active.cancel_phase = CancelPhase.FORCED

        deadline = loop.time() + cleanup_grace
        while loop.time() < deadline and deps.active_jobs.count() > 0:  # noqa: ASYNC110  # Why: poll-for-exit with deadline is the intentional design for shutdown phases; the timed grace period cannot be expressed with Event alone.
            await asyncio.sleep(0.1)

        # ── Phase 4: RELEASING ─────────────────────────────────────────
        # Every entry still registered belongs to an actor that ignored
        # both cancels. The shutdown owes it a release, not a verdict:
        # mark_interrupted hands the row back to the fleet with the spent
        # attempt standing (no refund; the attempt started executing),
        # HELD behind the rest of this
        # process's termination budget: plus the watchdog's exit tail
        # past the deadline itself (the dump-interval lag before the trip
        # is observed and the bounded flush before os._exit), so no other
        # pod can claim the row while this one might still touch it. The
        # release is ordered before the process dies, never after, and the
        # hold closes the overlap where the row is claimable while the
        # dying process could still touch it.
        deps.shutdown_phase = ShutdownPhase.RELEASING
        hold = _release_hold(deps, settings, loop)
        _log.info(
            "shutdown-phase",
            kind="shutdown_phase",
            phase="RELEASING",
            active_jobs_count=deps.active_jobs.count(),
            elapsed_seconds=loop.time() - t0,
        )
        released_count = 0
        noop_count = 0
        cancelled_count = 0
        for active in deps.active_jobs.all():
            # An OPERATOR entry belongs to the cancel ladder, never to a
            # release: one shielded write carrying the operator's OWN
            # verdict. mark_cancelled, fenced to this worker's attempt
            # (the fence mark_interrupted uses below), NOT mark_abandoned:
            # an entry still registered here is a consumer whose unwind
            # overran the cleanup grace - the FORCING probe escalated its
            # row to phase 2 and delivered task.cancel(), and the
            # consumer's own mark_cancelled is in flight (the unwind's
            # commit lags the grace under co-tenant load; the system
            # tier's cancel-racing-SIGTERM scenario). Two writers racing
            # DIFFERENT verdicts on one row is the defect: whoever
            # committed first owned the terminal state, so a row the
            # operator cancelled could land `abandoned`. Writing the SAME
            # verdict from here dissolves the race - one of the two
            # mark_cancelled writes commits, the other's fence no-ops,
            # and the row reads `cancelled` at every timing.
            if active.cancel_origin is CancelOrigin.OPERATOR:
                try:
                    if await shield_with_retrieval(
                        backend.mark_cancelled(
                            active.job_id,
                            worker_id,
                            attempt=active.ctx.attempt,
                            claim_epoch=active.ctx.claim_epoch,
                        )
                    ):
                        cancelled_count += 1
                except asyncio.CancelledError:
                    pass
                except Exception as exc:
                    _log.warning(
                        "cancel-pg-write-failed",
                        job_id=str(active.job_id),
                        error=str(exc),
                    )
                continue
            _buf = deps.progress_buffers.get(active.job_id)
            _seq, _state = _terminal_seq_and_state(_buf)
            try:
                outcome = await shield_with_retrieval(
                    backend.mark_interrupted(
                        active.job_id,
                        worker_id,
                        attempt=active.ctx.attempt,
                        claim_epoch=active.ctx.claim_epoch,
                        hold=hold,
                        progress_seq=_seq,
                        progress_state=_state if _buf is not None and _buf.dirty else None,
                    )
                )
            except asyncio.CancelledError:
                continue
            except Exception as exc:
                # No retry, no disown. The row stays running and locked.
                # Its recovery is the lease-expiry reclaim sweep. What
                # keeps that fallback safe is NOT a settings-load
                # rejection. No such validation exists, by design. Two
                # real protections do the work, and both live in code:
                #
                # 1. The parked consumer's later release write usually
                #    saves the row outright. That park is lease-capped
                #    (WorkerSettings.release_park_lease_cap, applied in
                #    _actor_exit_wait_budget): the cap is exactly the
                #    bound that puts the parked release write ahead of the
                #    earliest lease reclaim for every config that loads.
                #    Do not remove the cap on the belief that startup
                #    validation catches a bad config here. It does not.
                #    With the cap gone, a lease shorter than the
                #    termination budget lets the sweep re-pend this row
                #    while its actor thread still executes.
                # 2. When both writes fail (this one and the consumer's,
                #    whose exhausted retries disown the row), the residue
                #    is the WorkerSettings.release_disown_lease_floor
                #    bound. That one is surfaced as a startup warning by
                #    _emit_startup_warnings, again not a rejection.
                #
                # If you are tempted to add a retry here, first read the
                # mark_cancelled arm's comment in _consumer.py: a retry
                # wait inside this phase races the watchdog's deadline and
                # the teardown that deadline bounds. The sweep is the
                # designed backstop.
                _log.warning(
                    "release-pg-write-failed",
                    job_id=str(active.job_id),
                    error=str(exc),
                )
                continue
            if outcome == "noop":
                # The fence declined: the row is no longer this worker's to
                # release (its own consumer already terminalised it, or an
                # operator cancel landed on it between FORCING and
                # RELEASING). The row is the arbiter of origin: a
                # cancel-carrying row gets the operator's own verdict -
                # mark_cancelled, the same fenced write the consumer's
                # unwind is racing (the OPERATOR arm's comment), so either
                # the consumer's write landed and this fence no-ops, or
                # this one lands and the consumer's does. A row that moved
                # to another owner no-ops the fence either way.
                noop_count += 1
                _log.debug(
                    "release-interrupted-noop",
                    job_id=str(active.job_id),
                    worker_id=str(worker_id),
                    cancel_origin=int(active.cancel_origin),
                )
                try:
                    if await shield_with_retrieval(
                        backend.mark_cancelled(
                            active.job_id,
                            worker_id,
                            attempt=active.ctx.attempt,
                            claim_epoch=active.ctx.claim_epoch,
                        )
                    ):
                        cancelled_count += 1
                except asyncio.CancelledError:
                    pass
                except Exception as exc:
                    _log.warning(
                        "cancel-pg-write-failed",
                        job_id=str(active.job_id),
                        error=str(exc),
                    )
            else:
                released_count += 1
        _log.info(
            "shutdown-phase",
            kind="shutdown_phase",
            phase="RELEASING",
            released=released_count,
            held_seconds=hold.total_seconds(),
            noop=noop_count,
            cancelled=cancelled_count,
            elapsed_seconds=loop.time() - t0,
        )

        # ── leader_conn close ──────────────────────────────────
        # Why the owns_leader_conn guard: the ownership contract ("TaskQ
        # never closes caller-owned resources") forbids closing a
        # caller-provided leader_conn even during shutdown. The reference
        # is also left in place for caller-owned conns.
        #
        # By this point the election loop has long been parked: the
        # shutdown-start event (set before DRAINING) is the signal it
        # obeys, and a leader handed the lease over before the drain via
        # its fenced resign. The leader runtime's TEARDOWN resign
        # (leader.py's run() finally) stays as the backstop for exits
        # that never run this orchestrator (a sibling crash, a bare
        # cancel of _main); on this path it finds no row of ours and
        # no-ops.
        #
        # Why set → null → close, in that order (two races, one ordering):
        # (a) the bounded close can park for seconds, so shutdown_event is
        # set FIRST to stop the loops still reading this event; the
        # election loop itself parked much earlier, on the shutdown-start
        # event, and cannot swap in a fresh (possibly lock-holding) conn
        # mid-park. (b) the early set also releases
        # _main's ``await shutdown_event.wait()`` INSIDE the
        # open_worker_deps context (the orchestrator is awaited only after
        # that context exits), so the deps exit-stack guard unwinds
        # CONCURRENTLY with this parked close, nulling BEFORE the park is
        # what stops the guard entering a second close_conn_bounded on the
        # same conn (one closer's terminate would abort the other's
        # in-flight close and log a spurious conn-teardown-close-error on
        # real asyncpg). A conn swapped in mid-park keeps its reference
        # (nothing nulls after the park); the exit-stack guard closes it.
        conn = deps.leader_conn
        if conn is not None and deps.owns_leader_conn:
            # Why this can be a module-level import from taskq._close:
            # taskq._close imports nothing from taskq.worker, so the
            # deps↔shutdown cycle (deps imports ShutdownPhase from THIS
            # module) is not re-introduced. The bounded helper never raises
            # (timeout → terminate, error → log), so a dead PG cannot wedge
            # shutdown on an unbounded close.
            shutdown_event.set()  # stop the election loop BEFORE the close park
            # BEFORE the park: the deps exit-stack guard unwinds concurrently
            # once shutdown_event is set; it must not enter a second
            # close_conn_bounded on this same conn while the close below is
            # parked.
            deps.leader_conn = None
            await close_conn_bounded(conn, "leader", CLOSE_TIMEOUT_SECS)

        return 0
    finally:
        # ── Signal siblings and exit ───────────────────────────────────
        shutdown_event.set()
        _log.info(
            "shutdown-phase",
            kind="shutdown_phase",
            phase="EXITED",
            active_jobs_count=deps.active_jobs.count(),
            elapsed_seconds=loop.time() - t0,
        )

install_signal_handlers

install_signal_handlers(
    loop: AbstractEventLoop,
    deps: WorkerDeps,
    worker_id: UUID,
    shutdown_event: Event,
    escalate_event: Event,
    backend: Backend,
    orchestrator_holder: list[Task[int]],
) -> None

Register SIGTERM/SIGINT/SIGHUP handlers.

SIGTERM/SIGINT: three-signal escalation counter. First signal schedules orchestrate_shutdown via loop.create_task and appends the created task to orchestrator_holder so that _main can later await it for the exit code. Second signal sets escalate_event to fast-advance CANCELLING → FORCING. Third signal calls sys.exit(1) (Kubernetes SIGKILL is the hard backstop).

SIGHUP: sets deps.reload_event to signal the hot-reload coordinator (:func:~taskq.worker.deps.reload_credentials) that a credential refresh has been requested. The coordinator runs as a sibling task in the worker's TaskGroup and performs the actual pool/connection swap. SIGHUP can be sent multiple times, each one sets the event. The coordinator clears the event before each reload and never after, so a SIGHUP arriving mid-reload (success OR failure) is honored with exactly one follow-up reload: N signals during one reload coalesce into one follow-up, not N. Reload requests arriving while shutdown orchestration is in progress (deps.shutdown_phase is not NONE) are skipped.

The signal counter is closure-scoped, each call to this function creates a fresh, independent counter. The handler callable contains zero await or I/O.

Source code in src/taskq/worker/shutdown.py

def install_signal_handlers(
    loop: asyncio.AbstractEventLoop,
    deps: "WorkerDeps",
    worker_id: UUID,
    shutdown_event: asyncio.Event,
    escalate_event: asyncio.Event,
    backend: Backend,
    orchestrator_holder: list[asyncio.Task[int]],
) -> None:
    """Register SIGTERM/SIGINT/SIGHUP handlers.

    **SIGTERM/SIGINT**: three-signal escalation counter.  First signal
    schedules ``orchestrate_shutdown`` via ``loop.create_task`` and
    appends the created task to ``orchestrator_holder`` so that ``_main``
    can later await it for the exit code.  Second signal sets
    ``escalate_event`` to fast-advance CANCELLING → FORCING.  Third signal
    calls ``sys.exit(1)`` (Kubernetes SIGKILL is the hard backstop).

    **SIGHUP**: sets ``deps.reload_event`` to signal the hot-reload
    coordinator (:func:`~taskq.worker.deps.reload_credentials`) that a
    credential refresh has been requested. The coordinator runs as a
    sibling task in the worker's ``TaskGroup`` and performs the actual
    pool/connection swap. SIGHUP can be sent multiple times, each one
    sets the event. The coordinator clears the event *before* each
    reload and never after, so a SIGHUP arriving mid-reload (success OR
    failure) is honored with exactly one follow-up reload: N signals
    during one reload coalesce into one follow-up, not N. Reload
    requests arriving while shutdown orchestration is in progress
    (``deps.shutdown_phase`` is not NONE) are skipped.

    The signal counter is closure-scoped, each call to this function
    creates a fresh, independent counter.  The handler callable contains
    zero ``await`` or I/O.
    """
    _sig_count = 0

    def _on_shutdown_signal() -> None:
        nonlocal _sig_count
        _sig_count += 1
        if _sig_count == 1:
            # H2 guard: skip if orchestration is already in progress
            # (e.g., drain monitor triggered first). The _sig_count is
            # NOT reset, the next SIGTERM correctly escalates the
            # already-running orchestration rather than starting a new one.
            if _orchestration_in_progress(orchestrator_holder, deps):
                return
            task = loop.create_task(
                orchestrate_shutdown(
                    deps,
                    deps.settings,
                    worker_id,
                    shutdown_event,
                    escalate_event,
                    backend=backend,
                )
            )
            orchestrator_holder.append(task)
        elif _sig_count == 2:
            escalate_event.set()
        else:
            sys.exit(1)

    def _on_reload_signal() -> None:
        deps.reload_event.set()

    for sig in (signal.SIGTERM, signal.SIGINT):
        try:
            loop.add_signal_handler(sig, _on_shutdown_signal)
        except NotImplementedError:
            _log.warning(
                "signal-handlers-unavailable",
                os_name=os.name,
            )
            return
    # SIGHUP, credential hot-reload. Not available on Windows.
    if hasattr(signal, "SIGHUP"):
        try:
            loop.add_signal_handler(signal.SIGHUP, _on_reload_signal)
        except NotImplementedError:
            _log.warning("sighup-handler-unavailable", os_name=os.name)

    # SIGUSR2, on-demand asyncio task-stack dump (names, coros, await
    # sites; no locals or payload values). Live debugging without an
    # image rebuild. Not available on Windows.
    if hasattr(signal, "SIGUSR2"):

        def _on_dump_signal() -> None:
            dump_task_stacks("sigusr2", detector="sigusr2")

        try:
            loop.add_signal_handler(signal.SIGUSR2, _on_dump_signal)
        except NotImplementedError:
            _log.warning("sigusr2-handler-unavailable", os_name=os.name)

Drain (until-idle monitor and exit-code contract)

drain

Drain monitor for until-idle mode.

When spawned as a sibling in the worker's TaskGroup, the drain monitor polls the backend for active jobs in the worker's subscribed queues. When the count stays zero for the settle window (and no jobs are active on this worker), the monitor triggers the normal graceful shutdown via orchestrate_shutdown, the same path SIGTERM takes.

Exit codes

0, all jobs succeeded (drain_failures == 0) 3, some jobs failed (drain_failures > 0) 4, max_runtime exceeded before drain completed

__all__ module-attribute

__all__ = [
    "EXIT_DRAIN_CLEAN",
    "EXIT_DRAIN_TIMEOUT",
    "EXIT_DRAIN_WITH_FAILURES",
    "drain_monitor_loop",
]

EXIT_DRAIN_CLEAN module-attribute

EXIT_DRAIN_CLEAN = 0

EXIT_DRAIN_WITH_FAILURES module-attribute

EXIT_DRAIN_WITH_FAILURES = 3

EXIT_DRAIN_TIMEOUT module-attribute

EXIT_DRAIN_TIMEOUT = 4

drain_monitor_loop async

drain_monitor_loop(
    deps: WorkerDeps,
    settings: WorkerSettings,
    worker_id: UUID,
    shutdown_event: Event,
    escalate_event: Event,
    orchestrator_holder: list[Task[int]],
    backend: Backend,
    *,
    idle_settle_window: float,
    idle_poll_interval: float,
    max_runtime: float | None,
) -> None

Monitor for queue drain and trigger graceful shutdown when idle.

Polls backend.count_active_jobs(queues) and deps.active_jobs.count() every idle_poll_interval. When both are zero, starts the settle timer. If still zero after idle_settle_window, triggers shutdown.

If max_runtime is set and exceeded, triggers shutdown with exit code 4.

Returns after creating the orchestrate_shutdown task. Spawns with may_return=True in the sibling spawner. Does NOT set shutdown_event , orchestrate_shutdown's finally block sets it at the correct point (after all phase work completes), exactly as the SIGTERM signal handler does.

Forgets the liveness registration on every exit path (same contract as the producer loop, run.py): the monitor stops ticking the moment it returns, while the watchdog keeps sweeping until the orchestration's finally sets shutdown_event.

Double-orchestration guard (H2): if orchestrator_holder is already non-empty or deps.shutdown_phase is not ShutdownPhase.NONE, the monitor skips triggering, a SIGTERM-driven orchestration is already in progress.

Source code in src/taskq/worker/drain.py

async def drain_monitor_loop(
    deps: "WorkerDeps",
    settings: "WorkerSettings",
    worker_id: UUID,
    shutdown_event: asyncio.Event,
    escalate_event: asyncio.Event,
    orchestrator_holder: list[asyncio.Task[int]],
    backend: "Backend",
    *,
    idle_settle_window: float,
    idle_poll_interval: float,
    max_runtime: float | None,
) -> None:
    """Monitor for queue drain and trigger graceful shutdown when idle.

     Polls backend.count_active_jobs(queues) and deps.active_jobs.count()
     every idle_poll_interval. When both are zero, starts the settle timer.
     If still zero after idle_settle_window, triggers shutdown.

     If max_runtime is set and exceeded, triggers shutdown with exit code 4.

     Returns after creating the orchestrate_shutdown task. Spawns with
     may_return=True in the sibling spawner. Does NOT set shutdown_event
    , orchestrate_shutdown's finally block sets it at the correct point
     (after all phase work completes), exactly as the SIGTERM signal
     handler does.

     Forgets the liveness registration on every exit path (same contract
     as the producer loop, run.py): the monitor stops ticking the moment
     it returns, while the watchdog keeps sweeping until the
     orchestration's finally sets shutdown_event.

     Double-orchestration guard (H2): if orchestrator_holder is already
     non-empty or deps.shutdown_phase is not ShutdownPhase.NONE, the
     monitor skips triggering, a SIGTERM-driven orchestration is already
     in progress.
    """
    queues = settings.queues
    start_time = time.monotonic()
    idle_since: float | None = None

    _log.info(
        "drain-monitor-start",
        queues=queues,
        settle_window=idle_settle_window,
        poll_interval=idle_poll_interval,
        max_runtime=max_runtime,
        worker_id=str(worker_id),
    )

    # Worst-case seconds ONE iteration can consume: the count bound plus
    # the trailing poll sleep. A slow-but-recovering count_active_jobs
    # that spends the whole bound is a transient this loop rides out
    # (TRANSIENT_PG_ERRORS below), so the liveness tick period, and the
    # staleness budget derived from it, must cover that gap. Registering
    # only the poll interval makes the budget equal the count bound (both
    # max(5*poll, 10)); the gap exceeds it by one poll and detector 2
    # force-exits a healthy worker for a single degraded query. Same
    # timeout + period model as the leader-loop staleness invariant in
    # settings.py.
    count_bound = max(idle_poll_interval * 5, 10.0)
    tick_period = count_bound + idle_poll_interval

    try:
        while not shutdown_event.is_set():
            deps.liveness.tick("drain_monitor", period=tick_period)

            # Check max_runtime
            if max_runtime is not None:
                elapsed = time.monotonic() - start_time
                if elapsed >= max_runtime:
                    _log.info(
                        "drain-monitor-timeout",
                        elapsed=elapsed,
                        max_runtime=max_runtime,
                        worker_id=str(worker_id),
                    )
                    await _trigger_drain_shutdown(
                        deps,
                        settings,
                        worker_id,
                        shutdown_event,
                        escalate_event,
                        orchestrator_holder,
                        backend,
                        exit_code=EXIT_DRAIN_TIMEOUT,
                    )
                    return

            # Check idle condition, catch only recoverable transient errors
            # (F3). Non-recoverable errors propagate and tear down the TaskGroup.
            try:
                queue_count = await asyncio.wait_for(
                    backend.count_active_jobs(queues),
                    timeout=count_bound,
                )
            except TRANSIENT_PG_ERRORS:
                _log.warning("drain-monitor-count-error", worker_id=str(worker_id))
                queue_count = -1  # unknown, don't trigger

            active_count = deps.active_jobs.count()
            is_idle = queue_count == 0 and active_count == 0

            if is_idle:
                if idle_since is None:
                    idle_since = time.monotonic()
                    _log.debug(
                        "drain-monitor-idle-detected",
                        worker_id=str(worker_id),
                        queue_count=queue_count,
                        active_count=active_count,
                    )
                else:
                    idle_elapsed = time.monotonic() - idle_since
                    if idle_elapsed >= idle_settle_window:
                        _log.info(
                            "drain-monitor-drained",
                            worker_id=str(worker_id),
                            idle_elapsed=idle_elapsed,
                            drain_failures=deps.drain_failures,
                        )
                        exit_code = (
                            EXIT_DRAIN_WITH_FAILURES
                            if deps.drain_failures > 0
                            else EXIT_DRAIN_CLEAN
                        )
                        await _trigger_drain_shutdown(
                            deps,
                            settings,
                            worker_id,
                            shutdown_event,
                            escalate_event,
                            orchestrator_holder,
                            backend,
                            exit_code=exit_code,
                        )
                        return
            else:
                if idle_since is not None:
                    _log.debug(
                        "drain-monitor-idle-reset",
                        worker_id=str(worker_id),
                        queue_count=queue_count,
                        active_count=active_count,
                    )
                idle_since = None

            # Wait for poll interval or shutdown
            await _sleep_interruptible(shutdown_event, idle_poll_interval)
    finally:
        # The tick registration must not outlive the loop: this monitor
        # returns (or is cancelled) while orchestrate_shutdown is still
        # running its phases, and shutdown_event is set only in the
        # orchestrator's finally, the watchdog sweeps that whole window.
        # A lingering entry goes stale and detector 2 force-exits the
        # worker mid-grace. Mirrors the producer loop's forget (run.py).
        deps.liveness.forget("drain_monitor")

    _log.info("drain-monitor-exit", reason="shutdown_event", worker_id=str(worker_id))

Workgroup supervisor

workgroup

Workgroup supervisor, spawns and manages multiple TaskQ worker processes.

A workgroup is a single-process orchestrator that manages N child taskq worker subprocesses, each with potentially different queues, poll intervals, concurrency caps, etc. The supervisor restarts children that crash, optionally health-checks them via the database, and cleanly propagates shutdown signals.

Config format (TOML)::

actors = "myapp.actors:registry"

[defaults]
poll_interval = 1.0
max_concurrency = 4

[supervisor]
shutdown_grace = 100.0
backoff_initial = 0.5
backoff_max = 30.0
backoff_factor = 2.0
burst_limit = 10
burst_window = 60.0

[[workers]]
name = "api"
queues = ["default"]
max_concurrency = 8
poll_interval = 0.5
pg_credential_provider = "infra.identity:pg_credentials"

[workers.health]
enabled = true
check_interval = 15
stale_after = 60
startup_grace = 15.0
consecutive_failure_limit = 3

A pg_credential_provider set at the [defaults] level applies to every worker, and TOML has no per-worker null, a worker cannot opt out of a defaults-level provider, so a workgroup mixing provider-backed and provider-less workers must set the field per worker instead of in [defaults].

__all__ module-attribute

__all__ = [
    "DEFAULT_SHUTDOWN_GRACE_SECS",
    "WorkerSpec",
    "WorkgroupConfig",
    "load_workgroup_config",
    "run_forever",
    "worst_case_shutdown_seconds",
]

logger module-attribute

logger: BoundLogger = get_logger(__name__)

DEFAULT_SHUTDOWN_GRACE_SECS module-attribute

DEFAULT_SHUTDOWN_GRACE_SECS: float = 100.0

Default for [supervisor].shutdown_grace, the supervisor's SIGTERM-to-SIGKILL window for its children.

Why 100.0: the window must cover a child's modelled worst-case clean exit, WorkerSettings.worst_case_shutdown_seconds = cancellation grace (30s) + cleanup grace (10s) + the bounded-close teardown tail (8 sequential closes x 5s + 2s publish drain = 42s) = 82s, plus ~22% margin for loop jitter and the shutdown watchdog's non-instantaneous deadline trip (the dump interval and bounded metrics flush land inside the tail, but the tail is a model, not a bound the supervisor enforces). The old 30.0 default sat below even the children's 40s release floor (cancellation + cleanup graces): the supervisor SIGKILLed children before a held release could land, and every interrupted job rode the lease-expiry crash path instead. Operators who want the tighter, release-floor-only window can still set it explicitly; the startup warning names the numbers rather than refusing.

SupervisorConfig dataclass

SupervisorConfig(
    shutdown_grace: float = DEFAULT_SHUTDOWN_GRACE_SECS,
    backoff_initial: float = 0.5,
    backoff_max: float = 30.0,
    backoff_factor: float = 2.0,
    burst_limit: int = 10,
    burst_window: float = 60.0,
    health_pg_dsn: str | None = None,
    health_pg_schema: str | None = None,
)

Global supervisor behaviour tunables.

shutdown_grace class-attribute instance-attribute

shutdown_grace: float = DEFAULT_SHUTDOWN_GRACE_SECS

backoff_initial class-attribute instance-attribute

backoff_initial: float = 0.5

backoff_max class-attribute instance-attribute

backoff_max: float = 30.0

backoff_factor class-attribute instance-attribute

backoff_factor: float = 2.0

burst_limit class-attribute instance-attribute

burst_limit: int = 10

burst_window class-attribute instance-attribute

burst_window: float = 60.0

health_pg_dsn class-attribute instance-attribute

health_pg_dsn: str | None = None

health_pg_schema class-attribute instance-attribute

health_pg_schema: str | None = None

WorkerHealthConfig dataclass

WorkerHealthConfig(
    enabled: bool = False,
    check_interval: float = 15.0,
    stale_after: float = 60.0,
    startup_grace: float = 15.0,
    consecutive_failure_limit: int = 3,
)

Per-worker health-check configuration.

The freshness verdict is server-side and exclusive at the boundary: a worker whose server-measured age equals stale_after exactly is already stale (the predicate is last_seen_at > now() - stale_after, strict >).

enabled class-attribute instance-attribute

enabled: bool = False

check_interval class-attribute instance-attribute

check_interval: float = 15.0

stale_after class-attribute instance-attribute

stale_after: float = 60.0

startup_grace class-attribute instance-attribute

startup_grace: float = 15.0

consecutive_failure_limit class-attribute instance-attribute

consecutive_failure_limit: int = 3

WorkerSpec dataclass

WorkerSpec(
    name: str,
    queues: list[str],
    poll_interval: float = 1.0,
    max_concurrency: int = 8,
    worker_group: str = "default",
    force_update_actor_config: bool = False,
    pg_credential_provider: str | None = None,
    stream_limit: int = _STREAM_LIMIT,
    health: WorkerHealthConfig = WorkerHealthConfig(),
)

Configuration for a single worker process managed by the workgroup.

name instance-attribute

name: str

queues instance-attribute

queues: list[str]

poll_interval class-attribute instance-attribute

poll_interval: float = 1.0

max_concurrency class-attribute instance-attribute

max_concurrency: int = 8

worker_group class-attribute instance-attribute

worker_group: str = 'default'

force_update_actor_config class-attribute instance-attribute

force_update_actor_config: bool = False

pg_credential_provider class-attribute instance-attribute

pg_credential_provider: str | None = None

stream_limit class-attribute instance-attribute

stream_limit: int = _STREAM_LIMIT

health class-attribute instance-attribute

health: WorkerHealthConfig = field(
    default_factory=WorkerHealthConfig
)

cli_args

cli_args() -> list[str]

Build the CLI argument list for this worker.

Source code in src/taskq/worker/workgroup.py

def cli_args(self) -> list[str]:
    """Build the CLI argument list for this worker."""
    args: list[str] = []
    for q in self.queues:
        args.extend(["--queues", q])
    args.extend(["--poll-interval", str(self.poll_interval)])
    args.extend(["--max-concurrency", str(self.max_concurrency)])
    args.extend(["--worker-group", self.worker_group])
    if self.force_update_actor_config:
        args.append("--force-update-actor-config")
    if self.pg_credential_provider is not None:
        args.extend(["--pg-credential-provider", self.pg_credential_provider])
    return args

WorkgroupConfig dataclass

WorkgroupConfig(
    actors: str,
    supervisor: SupervisorConfig = SupervisorConfig(),
    workers: list[WorkerSpec] = list[WorkerSpec](),
    defaults: dict[str, Any] = dict[str, Any](),
)

Top-level workgroup configuration loaded from TOML.

actors instance-attribute

actors: str

supervisor class-attribute instance-attribute

supervisor: SupervisorConfig = field(
    default_factory=SupervisorConfig
)

workers class-attribute instance-attribute

workers: list[WorkerSpec] = field(
    default_factory=list[WorkerSpec]
)

defaults class-attribute instance-attribute

defaults: dict[str, Any] = field(
    default_factory=dict[str, Any]
)

from_toml classmethod

from_toml(path: Path) -> WorkgroupConfig

Load and validate the workgroup TOML configuration.

Source code in src/taskq/worker/workgroup.py

@classmethod
def from_toml(cls, path: Path) -> WorkgroupConfig:
    """Load and validate the workgroup TOML configuration."""
    raw: dict[str, Any] = tomllib.loads(path.read_text())

    actors: Any = raw.get("actors")
    if not actors or not isinstance(actors, str):
        raise ValueError("workgroup config must define 'actors' (str)")
    if ":" not in actors:
        raise ValueError(f"actors must be module:attr syntax, got {actors!r}")

    defaults: dict[str, Any] = dict(raw.get("defaults", {}))

    sup_raw: dict[str, Any] = raw.get("supervisor", {})
    supervisor = SupervisorConfig(
        shutdown_grace=float(sup_raw.get("shutdown_grace", DEFAULT_SHUTDOWN_GRACE_SECS)),
        backoff_initial=float(sup_raw.get("backoff_initial", 0.5)),
        backoff_max=float(sup_raw.get("backoff_max", 30.0)),
        backoff_factor=float(sup_raw.get("backoff_factor", 2.0)),
        burst_limit=int(sup_raw.get("burst_limit", 10)),
        burst_window=float(sup_raw.get("burst_window", 60.0)),
        health_pg_dsn=sup_raw.get("health_pg_dsn"),
        health_pg_schema=sup_raw.get("health_pg_schema"),
    )

    raw_workers: list[dict[str, Any]] = raw.get("workers", [])
    if not raw_workers:
        raise ValueError("workgroup config must define at least one [[workers]] entry")

    names: set[str] = set()
    workers: list[WorkerSpec] = []
    for i, w in enumerate(raw_workers):
        name: Any = w.get("name")
        if not name or not isinstance(name, str):
            raise ValueError(f"workers[{i}] must have a 'name' (str)")
        if name in names:
            raise ValueError(f"duplicate worker name: {name!r}")
        names.add(name)

        health_raw: dict[str, Any] = w.get("health", {})
        health = WorkerHealthConfig(
            enabled=bool(health_raw.get("enabled", False)),
            check_interval=float(health_raw.get("check_interval", 15.0)),
            stale_after=float(health_raw.get("stale_after", 60.0)),
            startup_grace=float(health_raw.get("startup_grace", 15.0)),
            consecutive_failure_limit=int(health_raw.get("consecutive_failure_limit", 3)),
        )

        workers.append(
            WorkerSpec(
                name=name,
                queues=_require_list_str(w, "queues", defaults.get("queues", ["default"])),
                poll_interval=float(w.get("poll_interval", defaults.get("poll_interval", 1.0))),
                max_concurrency=int(
                    w.get("max_concurrency", defaults.get("max_concurrency", 8))
                ),
                worker_group=str(
                    w.get("worker_group", defaults.get("worker_group", "default"))
                ),
                force_update_actor_config=bool(
                    w.get(
                        "force_update_actor_config",
                        defaults.get("force_update_actor_config", False),
                    )
                ),
                pg_credential_provider=_optional_str(
                    w, "pg_credential_provider", defaults.get("pg_credential_provider")
                ),
                stream_limit=int(
                    w.get("stream_limit", defaults.get("stream_limit", _STREAM_LIMIT))
                ),
                health=health,
            )
        )

    cfg = cls(actors=actors, supervisor=supervisor, workers=workers, defaults=defaults)
    _validate_config(cfg)
    return cfg

load_workgroup_config

load_workgroup_config(path: Path) -> WorkgroupConfig

Load a workgroup configuration from a TOML file.

Source code in src/taskq/worker/workgroup.py

def load_workgroup_config(path: Path) -> WorkgroupConfig:
    """Load a workgroup configuration from a TOML file."""
    cfg = WorkgroupConfig.from_toml(path)
    _warn_on_actor_queues_no_child_consumes(cfg, _resolve_actor_registry(cfg.actors))
    return cfg

worst_case_shutdown_seconds

worst_case_shutdown_seconds(
    settings: WorkerSettings,
) -> float

The children's modelled worst-case wall clock from SIGTERM to exit.

Delegates to WorkerSettings.worst_case_shutdown_seconds (cancellation grace + cleanup grace + the bounded-close teardown tail) so the supervisor surface and the worker's own model stay one number, no drift. Exposed here so taskq workgroup validate and operators' sizing scripts can compare a config's shutdown_grace against the floor it must cover, without importing worker bootstrap internals.

Source code in src/taskq/worker/workgroup.py

def worst_case_shutdown_seconds(settings: WorkerSettings) -> float:
    """The children's modelled worst-case wall clock from SIGTERM to exit.

    Delegates to ``WorkerSettings.worst_case_shutdown_seconds`` (cancellation
    grace + cleanup grace + the bounded-close teardown tail) so the
    supervisor surface and the worker's own model stay one number, no
    drift. Exposed here so ``taskq workgroup validate`` and operators'
    sizing scripts can compare a config's ``shutdown_grace`` against the
    floor it must cover, without importing worker bootstrap internals.
    """
    return settings.worst_case_shutdown_seconds

run_forever async

run_forever(config_path: Path) -> None

Load config, spawn children, manage lifecycle until a signal arrives.

Blocks until SIGTERM or SIGINT, then shuts down all children gracefully. SIGHUP is NOT a shutdown signal: it is forwarded to every living child so each child's credential hot-reload runs (see :func:_forward_sighup_to_children); without a handler the default disposition would terminate the supervisor and orphan every child. Returns normally after shutdown, the CLI caller handles the exit code.

Source code in src/taskq/worker/workgroup.py

async def run_forever(config_path: Path) -> None:
    """Load config, spawn children, manage lifecycle until a signal arrives.

    Blocks until SIGTERM or SIGINT, then shuts down all children gracefully.
    SIGHUP is NOT a shutdown signal: it is forwarded to every living child
    so each child's credential hot-reload runs (see
    :func:`_forward_sighup_to_children`); without a handler the default
    disposition would terminate the supervisor and orphan every child.
    Returns normally after shutdown, the CLI caller handles the exit code.
    """
    config = load_workgroup_config(config_path)
    scfg = config.supervisor
    wg_instance = new_uuid()

    logger.info(
        "workgroup.start",
        actors=config.actors,
        instance_id=str(wg_instance),
        worker_count=len(config.workers),
    )

    # ── Resolve health-check PG pool (only when needed) ───────────────
    from taskq.settings import WorkerSettings

    pg_pool: asyncpg.Pool | None = None
    pg_schema: str = "taskq"

    health_workers = [w for w in config.workers if w.health.enabled]
    if health_workers:
        try:
            if scfg.health_pg_dsn:
                pg_dsn = scfg.health_pg_dsn
                pg_schema = scfg.health_pg_schema or "taskq"
                # No WorkerSettings for an explicit health DSN, the module
                # constants (statement_cache_kwargs' fallback) apply.
                stmt_kwargs = statement_cache_kwargs()
            else:
                settings = WorkerSettings.load()
                pg_schema = settings.schema_name
                pg_dsn = str(settings.resolved_pg_dsn_direct)
                stmt_kwargs = statement_cache_kwargs(settings)
            pg_pool = await asyncpg.create_pool(
                pg_dsn,
                min_size=1,
                max_size=len(health_workers) + 1,
                statement_cache_size=stmt_kwargs["statement_cache_size"],
                max_cached_statement_lifetime=stmt_kwargs["max_cached_statement_lifetime"],
                connection_class=guarded_connection_class(),
            )
            logger.info(
                "workgroup.health_pool_ready",
                schema=pg_schema,
                workers=[w.name for w in health_workers],
            )
        except Exception as exc:
            logger.critical("workgroup.health_pool_failed", error=str(exc))
            sys.exit(1)

    # ── Build child state machines ────────────────────────────────────
    children: dict[str, _ChildState] = {}
    for spec in config.workers:
        state = _ChildState(spec=spec)
        state.backoff = scfg.backoff_initial
        children[spec.name] = state

    shutting_down = asyncio.Event()
    # Strong refs to the in-flight SIGHUP forward tasks (asyncio only holds
    # weak refs): a dropped task can be garbage-collected mid-flight.
    _sighup_tasks: set[asyncio.Task[None]] = set()

    # ── Emit warnings for risky config ────────────────────────────────
    for w in config.workers:
        if w.force_update_actor_config:
            logger.warning(
                "workgroup.force_update_actor_config_enabled",
                worker=w.name,
                note="Permanent force-update will silently overwrite actor_config on "
                "every restart. Set to false after the first deploy with config changes.",
            )

    # The children inherit this process's environment, so their timing
    # settings are loadable here: the shutdown-grace window warning needs
    # the graces the children will actually run with. Best-effort: a
    # warning must never keep a workgroup from starting.
    with contextlib.suppress(Exception):
        from taskq.settings import WorkerSettings

        _warn_shutdown_grace_window(scfg, WorkerSettings.load())

    # ── Spawn all children initially ──────────────────────────────────
    for child in children.values():
        try:
            await _spawn_child(child, config.actors, wg_instance)
        except Exception:
            logger.exception(
                "workgroup.spawn_failed",
                worker=child.spec.name,
            )
            # Continue, the liveness monitor retries never-spawned
            # children under the same burst/backoff budget as exited
            # ones.

    # ── Signal handler, just sets the event; real cleanup follows ────
    loop = asyncio.get_running_loop()

    def _on_signal() -> None:
        if not shutting_down.is_set():
            logger.info("workgroup-shutdown-signal")
            shutting_down.set()

    for sig in (signal.SIGTERM, signal.SIGINT):
        loop.add_signal_handler(sig, _on_signal)

    # SIGHUP, credential rotation: forward to the children, never shut down.
    # Without a handler the default disposition kills the supervisor and
    # orphans every child. The single worker registers its own SIGHUP
    # reload handler (taskq.worker.shutdown.install_signal_handlers), so
    # the operator story is one signal for the whole workgroup: each child
    # runs its own hot-reload. Not available on Windows.
    if hasattr(signal, "SIGHUP"):

        def _on_sighup() -> None:
            # Deliberately await-free: signal handlers run between loop
            # callbacks, the forwarding takes each child's restart_lock
            # inside its own task. Ignored mid-shutdown: reloading pools a
            # child is about to tear down churns resources for nothing.
            if shutting_down.is_set():
                return
            logger.info("workgroup-reload-signal", children=len(children))
            # Referenced until done: a task dropped on the floor can be
            # garbage-collected mid-flight, losing the reload forward.
            task = asyncio.create_task(_forward_sighup_to_children(children))
            _sighup_tasks.add(task)
            task.add_done_callback(_sighup_tasks.discard)

        loop.add_signal_handler(signal.SIGHUP, _on_sighup)

    # ── Liveness monitor, detects exited children and restarts them ──
    async def liveness_monitor() -> None:
        while not shutting_down.is_set():
            for child in list(children.values()):
                if child.gave_up:
                    # Burst budget exhausted: one critical, logged when
                    # the budget refused, then this child is never
                    # scheduled again, the monitor must not re-decide
                    # every tick (a 2 Hz critical flood would drown every
                    # other signal on the box).
                    continue
                proc = child.process
                if proc is None:
                    # A child whose spawn failed (initial or restart) has
                    # no process to observe, without this branch it
                    # would stay dead until the whole supervisor
                    # restarts. Retry it under the same burst/backoff
                    # budget as an exited child.
                    async with child.restart_lock:
                        if child.process is not None or shutting_down.is_set():
                            continue
                        delay = _schedule_restart(
                            child,
                            scfg,
                            actors=config.actors,
                            wg_instance=wg_instance,
                            reason="spawn_failed",
                        )
                    if delay is None:
                        continue
                    _start_respawn(child, delay, config.actors, wg_instance, shutting_down)
                    continue
                if proc.returncode is not None:
                    async with child.restart_lock:
                        # Cancel and reap stale stream tasks from the dead
                        # process; the await reaps the cancellation so no
                        # task is left un-awaited.
                        for t in (child.stdout_task, child.stderr_task):
                            if t is not None and not t.done():
                                t.cancel()
                                with contextlib.suppress(asyncio.CancelledError):
                                    await t
                        child.stdout_task = None
                        child.stderr_task = None
                        delay = await _handle_child_exit(
                            child, config.actors, wg_instance, scfg, shutting_down
                        )
                    # Lock released.  The respawn runs detached: the
                    # monitor keeps ticking while this child backs off.
                    if delay is not None:
                        _start_respawn(child, delay, config.actors, wg_instance, shutting_down)
            await asyncio.sleep(0.5)

    # ── Health-check loop, kills hung workers via DB ─────────────────
    async def health_loop() -> None:
        if pg_pool is None:
            return
        while not shutting_down.is_set():
            for child in list(children.values()):
                if not child.spec.health.enabled:
                    continue
                if (
                    child.spawned_at > 0
                    and (time.monotonic() - child.spawned_at) < child.spec.health.startup_grace
                ):
                    continue
                async with child.restart_lock:
                    proc = child.process
                    if proc is None or proc.returncode is not None:
                        continue
                    healthy = await _child_health_check(
                        child, pg_pool, pg_schema, child.spec.health, wg_instance
                    )
                    if not healthy:
                        logger.warning("workgroup-health-kill", worker=child.spec.name)
                        await _kill_child(child)
            with contextlib.suppress(TimeoutError):
                await asyncio.wait_for(
                    shutting_down.wait(),
                    timeout=min(
                        (
                            w.spec.health.check_interval
                            for w in children.values()
                            if w.spec.health.enabled
                        ),
                        default=15.0,
                    ),
                )

    # ── Run foreground + background loops in a TaskGroup ──────────────
    try:
        async with asyncio.TaskGroup() as tg:
            tg.create_task(liveness_monitor())
            tg.create_task(health_loop())
            await shutting_down.wait()
    except* Exception as eg:
        for exc in eg.exceptions:
            logger.error("workgroup-background-task-failed", error=str(exc))
        # Reset child state to a safe baseline before shutdown.
        for child in children.values():
            async with child.restart_lock:
                if child.process is not None and child.process.returncode is not None:
                    child.process = None

    # ── Graceful shutdown ─────────────────────────────────────────────
    logger.info("workgroup-shutdown-begin", grace_s=scfg.shutdown_grace)

    # Forward SIGTERM to all living children (under lock).
    for child in children.values():
        async with child.restart_lock:
            proc = child.process
            if proc is not None and proc.returncode is None:
                with contextlib.suppress(ProcessLookupError):
                    proc.send_signal(signal.SIGTERM)

    # Wait concurrently for all children.
    wait_tasks: list[asyncio.Task[int]] = []
    for child in children.values():
        async with child.restart_lock:
            proc = child.process
            if proc is not None and proc.returncode is None:
                wait_tasks.append(asyncio.create_task(proc.wait()))
    if wait_tasks:
        _done, pending = await asyncio.wait(wait_tasks, timeout=scfg.shutdown_grace)
        for task in pending:
            task.cancel()
        for task in wait_tasks:
            if not task.done():
                with contextlib.suppress(asyncio.CancelledError):
                    await task

    # Force-kill any survivors.
    for child in children.values():
        async with child.restart_lock:
            proc = child.process
            if proc is not None and proc.returncode is None:
                logger.warning(
                    "workgroup.child_shutdown_force_kill",
                    worker=child.spec.name,
                )
                with contextlib.suppress(ProcessLookupError):
                    proc.kill()

    # Cancel stream tasks.
    for child in children.values():
        for t in (child.stdout_task, child.stderr_task):
            if t is not None and not t.done():
                t.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await t

    if pg_pool:
        # Why bounded: the supervisor's health-pool close is the same
        # dead-PG hang class as worker teardown, an unbounded close
        # would wedge the supervisor between workgroup-shutdown-begin and
        # workgroup-shutdown-complete. The helper never raises and terminates the
        # pool on timeout.
        await close_pool_bounded(pg_pool, "workgroup-health", CLOSE_TIMEOUT_SECS)

    logger.info("workgroup-shutdown-complete")