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__ = [
    "_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,
    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.

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,
    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.

    ``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:
        return runner.run(
            worker_main_async(
                settings,
                actor_registry=actor_registry,
                di_registry=di_registry,
                cron_registry=cron_registry,
                connections=connections,
                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,
            )
        )

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,
    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,
    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)
    return await _main(
        settings,
        actor_registry=actor_registry,
        _registry=di_registry,
        _cron_registry=schedule_specs,
        connections=connections,
        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,
) -> 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 local_queue.maxsize - local_queue.qsize() pending jobs (pending → running) using FOR UPDATE SKIP LOCKED.
3. Puts each returned :class:JobRow onto local_queue for the consumer tasks.

Exits cleanly when either shutdown_event or producer_stop_event is set.

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,
) -> 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
       ``local_queue.maxsize - local_queue.qsize()`` pending jobs
       (pending → running) using ``FOR UPDATE SKIP LOCKED``.
    3. Puts each returned :class:`JobRow` onto ``local_queue`` for the
       consumer tasks.

    Exits cleanly when either ``shutdown_event`` or ``producer_stop_event``
    is set.
    """
    settings = deps.settings
    queues = settings.queues
    lock_lease_td = timedelta(seconds=settings.lock_lease)
    notify_enabled = getattr(settings, "notify_enabled", False)
    poll_interval = settings.notify_poll_interval if notify_enabled else settings.poll_interval

    _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),
    )

    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(),
                    )
                )
                _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)
            available = local_queue.maxsize - local_queue.qsize()
            if available <= 0:
                await asyncio.sleep(0.1)
                continue

            try:
                jobs = await backend.dispatch_batch(
                    worker_id=worker_id,
                    queues=queues,
                    limit=available,
                    lock_lease=lock_lease_td,
                )
            except Exception:
                _producer_log.exception("dispatch-batch-error", worker_id=str(worker_id))
                with contextlib.suppress(asyncio.CancelledError):
                    await asyncio.sleep(poll_interval)
                continue

            if jobs:
                for job in jobs:
                    await local_queue.put(job)
                if wake_event is not None:
                    wake_event.clear()
                continue

            wake_wait = asyncio.create_task(wake_event.wait()) if wake_event is not None else None
            poll_wait = asyncio.create_task(asyncio.sleep(poll_interval))
            stop_wait = asyncio.create_task(producer_stop_event.wait())
            shutdown_wait = asyncio.create_task(shutdown_event.wait())

            all_waits = [
                w for w in (wake_wait, poll_wait, stop_wait, shutdown_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

            if wake_event is not None:
                wake_event.clear()

    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,
) -> 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,
) -> 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():
        q_get = asyncio.create_task(local_queue.get())
        shut_wait = asyncio.create_task(shutdown_event.wait())
        try:
            _done, pending = await asyncio.wait(
                [q_get, shut_wait],
                return_when=asyncio.FIRST_COMPLETED,
            )
            for task in pending:
                task.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await task
            if shut_wait in _done:
                return
        finally:
            for task in [q_get, shut_wait]:
                if not task.done():
                    task.cancel()

        job: JobRow = q_get.result()

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

        with parent_tags(tuple(job.tags)):
            ctx: JobContext[_StubPayload] = JobContext(
                job_id=job.id,
                actor=job.actor,
                queue=job.queue,
                attempt=job.attempt,
                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 "",
                ),
            )

            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:
                    with contextlib.suppress(asyncio.CancelledError):
                        await asyncio.shield(backend.mark_cancelled(job.id, worker_id))
                    raise
                except TimeoutError:
                    pass

                if ctx.cancellation_requested:
                    await asyncio.shield(backend.mark_cancelled(job.id, worker_id))
                else:
                    await asyncio.shield(backend.mark_succeeded(job.id, worker_id, None))
                # 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 asyncio.shield(backend.mark_cancelled(job.id, worker_id))
                raise

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

            finally:
                await deps.active_jobs.deregister(job.id)

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,
) -> 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,
) -> 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():
        q_get = asyncio.create_task(local_queue.get())
        shut_wait = asyncio.create_task(shutdown_event.wait())
        try:
            _done, pending = await asyncio.wait(
                [q_get, shut_wait],
                return_when=asyncio.FIRST_COMPLETED,
            )
            for task in pending:
                task.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await task
            if shut_wait in _done:
                return
        finally:
            for task in [q_get, shut_wait]:
                if not task.done():
                    task.cancel()

        job: JobRow = q_get.result()

        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.
            try:
                await backend.mark_snoozed(
                    job.id,
                    worker_id,
                    timedelta(seconds=10),
                    metadata_update={"released_reason": "actor-not-found"},
                )
            except Exception:
                _consumer_log.exception(
                    "dispatch-actor-not-found-release-failed",
                    job_id=str(job.id),
                    actor=job.actor,
                )
            continue

        actor_ref = actor_registry[job.actor]
        actor_config = _DispatchActorConfig(
            retry=actor_ref.retry,
            non_retryable_exceptions=actor_ref.non_retryable_exceptions,
            retry_classifier=actor_ref.retry_classifier,
            on_retry_exhausted=actor_ref.on_retry_exhausted,
            on_retry_exhausted_timeout=actor_ref.on_retry_exhausted_timeout,
            on_success=actor_ref.on_success,
            on_success_timeout=actor_ref.on_success_timeout,
        )
        try:
            outcome = await 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_config,
                clock=clock,
                active_jobs=deps.active_jobs,
                enqueuer=enqueuer,
            )
            if outcome == "failed":
                deps.drain_failures += 1
        except Exception:
            _consumer_log.exception("dispatch-failed", job_id=str(job.id))
            deps.drain_failures += 1

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.

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.
    """
    worker_id = new_uuid()
    schema = settings.schema_name
    if not _IDENT_RE.match(schema):
        raise ValueError(f"invalid schema identifier: {schema!r}")
    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)
    metadata: dict[str, object] = {"notify_enabled": notify_enabled}

    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
) -> None

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 returns without raising — the recovery sweep is the backstop and shutdown MUST NOT block on this cleanup.

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

async def deregister_worker(pool: asyncpg.Pool, settings: WorkerSettings, worker_id: UUID) -> None:
    """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
    returns without raising — the recovery sweep is the backstop and
    shutdown MUST NOT block on this cleanup.
    """
    schema = settings.schema_name
    if not _IDENT_RE.match(schema):
        raise ValueError(f"invalid schema identifier: {schema!r}")

    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:
            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),
        )

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, and reservation slots. Non-leader pods retry election periodically and skip the gated work. Failover SLA: Worker killed ≤ heartbeat_interval + 1 s Partition detect ≤ watchdog_interval + heartbeat_interval + 2 s PG failover ≤ heartbeat_interval Watchdog detect ≤ watchdog_interval + heartbeat_interval

ARCHIVE_EXPIRY_LOCK_NAME module-attribute

ARCHIVE_EXPIRY_LOCK_NAME: str = 'taskq:archive_expiry'

PRUNE_LOCK_NAME module-attribute

PRUNE_LOCK_NAME: str = 'taskq:prune'

__all__ module-attribute

__all__ = [
    "ARCHIVE_EXPIRY_LOCK_NAME",
    "MAINTENANCE_LEADER_LOCK_NAME",
    "PRUNE_LOCK_NAME",
    "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__)

MAINTENANCE_LEADER_LOCK_NAME module-attribute

MAINTENANCE_LEADER_LOCK_NAME: str = (
    "taskq:maintenance_leader"
)

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,
)

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,
) -> None:
    self._deps = deps
    self._worker_id = worker_id
    self._backend = backend
    self._clock = clock
    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

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 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._reservation_slots_loop(shutdown), name="leader.reservation_slots"
            )
            tg.create_task(self._stranded_jobs_loop(shutdown), name="leader.stranded_jobs")
            await shutdown.wait()
    finally:
        # 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)

archive_expiry_sweep async

archive_expiry_sweep(
    conn: ConnLike,
    *,
    batch_size: int = DEFAULT_PRUNE_BATCH_SIZE,
    schema: str = "taskq",
) -> ArchiveExpiryResult

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",
) -> ArchiveExpiryResult:
    if not _IDENT_RE.match(schema):
        raise ValueError(f"invalid schema identifier: {schema!r}")
    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 < clock_timestamp()` in _EXPIRY_CTE_SQL).
    # Read from the database anyway so the reported instant cannot disagree
    # with the predicate that actually ran — 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)
    sql = _EXPIRY_CTE_SQL.format(schema=schema)

    while True:
        rows = await conn.fetch(sql, batch_size)
        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 < batch_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",
) -> int

Delete worker rows whose last_seen_at exceeds staleness.

The caller's worker_id is never deleted. Returns the number of rows removed. Worker-level cascade (maintenance_leader, job_attempts) is handled by the DDL ON DELETE clauses — no extra sweeping needed.

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

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

    The caller's *worker_id* is never deleted. Returns the number of rows
    removed. Worker-level cascade (``maintenance_leader``, ``job_attempts``)
    is handled by the DDL ``ON DELETE`` clauses — no extra sweeping needed.
    """
    if not _IDENT_RE.match(schema):
        raise ValueError(f"invalid schema identifier: {schema!r}")
    sql = _CLEANUP_STALE_WORKERS_SQL.format(schema=schema)
    tag = await conn.execute(sql, staleness, worker_id)
    return int(tag.rsplit(" ", 1)[-1]) if tag else 0

complete_stale_batches async

complete_stale_batches(
    conn: ConnLike, *, schema: str = "taskq"
) -> 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). Returns the number of batches completed.

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

async def complete_stale_batches(
    conn: ConnLike,
    *,
    schema: str = "taskq",
) -> 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).
    Returns the number of batches completed.
    """
    if not _IDENT_RE.match(schema):
        raise ValueError(f"invalid schema identifier: {schema!r}")
    rows = await conn.fetch(
        _COMPLETE_STALE_BATCHES_SQL.format(schema=schema, terminal_not_in=_TERMINAL_NOT_IN)
    )
    return len(rows)

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,
) -> PruneResult

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,
) -> PruneResult:
    if not _IDENT_RE.match(schema):
        raise ValueError(f"invalid schema identifier: {schema!r}")
    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 (clock_timestamp() - $2::interval), 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
        sql = _ARCHIVE_CTE_SQL.format(schema=schema)

        while True:
            rows = await conn.fetch(sql, status, retention, batch_size, archive_interval)
            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 < batch_size:
                break

    if actor_overrides:
        for actor_name, actor_retention in actor_overrides.items():
            sql = _ARCHIVE_CTE_ACTOR_SQL.format(schema=schema)
            for status in TERMINAL_STATUSES:
                if actor_retention >= retention_per_status.get(status, DEFAULT_PRUNE_RETENTION):
                    continue
                while True:
                    rows = await conn.fetch(
                        sql,
                        status,
                        actor_retention,
                        batch_size,
                        archive_interval,
                        actor_name,
                    )
                    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 < batch_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,
    )

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__)

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. Create the CONSUMER span with link to the PRODUCER span.
2. Validate the payload against actor_ref's payload schema.
3. Build the interim JobContext with the CONSUMER span.
4. Open build_actor_scope to resolve DI kwargs.
5. Hand the resolved kwargs to consume_one_job via a run_actor closure that injects them.
6. TRANSIENT scope is closed on exit (regardless of outcome).
7. Record consumer-path metrics outside the span body.

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. Create the CONSUMER span with link to the PRODUCER span.
    2. Validate the payload against actor_ref's payload schema.
    3. Build the interim JobContext with the CONSUMER span.
    4. Open build_actor_scope to resolve DI kwargs.
    5. Hand the resolved kwargs to consume_one_job via a run_actor
       closure that injects them.
    6. TRANSIENT scope is closed on exit (regardless of outcome).
    7. Record consumer-path metrics outside the span body.
    """
    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)

    consumer_attrs: dict[str, str | int] = {
        "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

    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:
            try:
                validated_payload = validate_actor_payload(
                    actor_ref.payload_type,
                    job.payload,
                    job.actor,
                )

                span_ctx = consumer_span.get_span_context()
                dispatch_trace_id: str = ""
                if span_ctx.is_valid:
                    dispatch_trace_id = format(span_ctx.trace_id, "032x")

                interim_ctx: JobContext[BaseModel] = JobContext(
                    job_id=job.id,
                    actor=job.actor,
                    queue=job.queue,
                    attempt=job.attempt,
                    worker_id=worker_id,
                    payload=validated_payload,
                    jobs=enqueuer,
                    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=dispatch_trace_id,
                        batch_id=batch_id or None,
                    ),
                    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,
                ) 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 asyncio.to_thread(actor_ref.fn, **actor_kwargs)
                        return await actor_ref.fn(**actor_kwargs)  # type: ignore[no-any-return]  # Why: actor_ref.fn is typed Callable[..., object]; runtime result is R.

                    loop_conn: asyncpg.Connection | None = None
                    raw_conn = loop_scope.resolved_cache().get(asyncpg.Connection)
                    if isinstance(raw_conn, asyncpg.Connection):
                        loop_conn = raw_conn  # pyright: ignore[reportUnknownVariableType]  # Why: asyncpg.Connection is generic (Connection[Record]); resolved_cache returns Mapping[type, object] so isinstance narrows to Connection[Unknown] — the record type is irrelevant for the transaction lifecycle.

                    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
                    try:
                        import redis.asyncio as _redis_mod  # type: ignore[no-redef]  # Why: runtime import for DI lookup; TYPE_CHECKING import is for annotations only

                        raw_redis = loop_scope.resolved_cache().get(_redis_mod.Redis)
                        if isinstance(raw_redis, _redis_mod.Redis):
                            redis_client = raw_redis
                    except ImportError:
                        pass

                    # 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=enqueuer,
                        loop_conn=loop_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,
                        fallback_result_ttl=actor_ref.result_ttl,
                    )
                    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).
                    try:
                        await apply_batch_terminal_outcome(
                            backend, job, outcome, loop_conn=loop_conn
                        )
                    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")
                raise
            except Exception as exc:
                outcome = "failed"
                consumer_span.set_status(StatusCode.ERROR)
                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,
                    )
                    outcome = handler_result
                except _TERMINAL_WRITE_INFRA_EXCEPTIONS as infra_exc:
                    _log_terminal_write_failed(handler_log, job, exc, infra_exc)
    finally:
        elapsed = time.monotonic() - t0
        record_consumed_message(job.actor, job.queue, outcome=_to_consumed_outcome(outcome))
        record_process_duration(job.actor, job.queue, elapsed)

    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. After max_heartbeat_failures consecutive connection failures, 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
    (
        update_worker_liveness_sql,
        update_jobs_lock_sql,
        update_reservation_leases_sql,
        update_leader_ping_sql,
    ) = build_heartbeat_sql(schema)

    while not shutdown.is_set():
        deps.liveness.tick("heartbeat", period=interval)
        _in_tx_failed = False
        tick_start = time.monotonic()
        try:
            _tick_raised = False
            try:
                async with (
                    deps.heartbeat_pool.acquire(timeout=interval) as conn,
                    conn.transaction(),
                ):
                    await conn.execute(update_worker_liveness_sql, worker_id)
                    jobs_tag = await conn.execute(update_jobs_lock_sql, worker_id, lock_lease)
                    await conn.execute(update_reservation_leases_sql, worker_id, lock_lease)
                    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
                    if deps.is_leader.is_set():
                        await conn.execute(update_leader_ping_sql, worker_id)
            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:
                    try:
                        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.
                        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)
            record_lock_expires_in_seconds(str(worker_id), lock_lease.total_seconds())
            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)
            if not _in_tx_failed:
                deps.heartbeat_failures += 1
                update_heartbeat_consecutive_failures(str(worker_id), deps.heartbeat_failures)
            record_heartbeat_miss(str(worker_id))
            logger.warning(
                "heartbeat-tick-failure",
                worker_id=str(worker_id),
                consecutive_failures=deps.heartbeat_failures,
                error_class=type(e).__name__,
                error=str(e),
            )
            early_warn_threshold = deps.settings.max_heartbeat_failures // 2
            if early_warn_threshold > 0 and deps.heartbeat_failures == early_warn_threshold:
                logger.warning(
                    "heartbeat-failures-approaching-limit",
                    worker_id=str(worker_id),
                    consecutive_failures=deps.heartbeat_failures,
                    max_heartbeat_failures=deps.settings.max_heartbeat_failures,
                    error_class=type(e).__name__,
                )
            if deps.heartbeat_failures > deps.settings.max_heartbeat_failures:
                await isolate_self(deps, worker_id, shutdown)
                return
        except Exception:
            tick_duration_s = time.monotonic() - tick_start
            _tick_duration.record(tick_duration_s)
            logger.exception(
                "heartbeat-tick-unexpected-error",
                worker_id=str(worker_id),
            )
        if cancel_wake_event is not None:
            # Wait up to interval, but wake immediately on a cancel NOTIFY.
            with contextlib.suppress(asyncio.TimeoutError):
                await asyncio.wait_for(cancel_wake_event.wait(), timeout=interval)
            cancel_wake_event.clear()
        else:
            await asyncio.sleep(interval)

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
    if not _IDENT_RE.match(schema):
        raise ValueError(f"invalid schema identifier: {schema!r}")
    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)
    jobs_pending_count = 0
    jobs_crashed_count = 0
    jobs_cancelled_count = 0

    try:
        conn = await asyncpg.connect(pg_dsn, timeout=5.0)  # pyright: ignore[reportCallIssue, reportUnknownVariableType]  # Why: asyncpg-stubs does not declare timeout kwarg on connect(); the parameter exists at runtime at 0.31.0.  asyncpg default is 60s — far too long when PG is already problematic.
        try:

            async def _inner() -> tuple[int, int, int]:
                pending = 0
                crashed = 0
                cancelled = 0
                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
                    )
                    for row in rows:  # pyright: ignore[reportUnknownVariableType]  # Why: rows type suppressed above — propagates from conn.fetch() suppression.
                        is_pending = (  # pyright: ignore[reportUnknownVariableType]  # Why: row column accessor types unknown — propagates from conn.fetch() suppression.
                            row["attempt"] < row["max_attempts"]
                            and row["retry_kind"] != "non_retryable"
                        )
                        if is_pending:
                            pending += 1
                        elif row["cancel_phase"] != 0:
                            # Mirrors _ISOLATE_JOB_SQL_TEMPLATE's CASE arm:
                            # exhausted + cancel in-flight → 'cancelled'.
                            cancelled += 1
                        else:
                            crashed += 1
                        await conn.execute(isolate_job_sql, row["id"], worker_id)
                        await conn.execute(
                            insert_attempt_sql,
                            row["id"],
                            row["attempt"],
                            row["started_at"],
                            "crashed",
                            "HeartbeatLost",
                            None,
                            None,
                            None,
                            worker_id,
                            "{}",  # metadata — matches the sweep paths' literal
                        )
                return pending, crashed, cancelled

            jobs_pending_count, jobs_crashed_count, jobs_cancelled_count = await asyncio.shield(
                _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 (#38) — unbounded, it would
            # wedge shutdown.set() below. The helper never raises, so a
            # close error can no longer mask an in-flight exception or be
            # misreported as an isolate-self failure.
            await close_conn_bounded(conn, "isolate-self", CLOSE_TIMEOUT_SECS, mid_run=True)
    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,
        )

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."""
    ...

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) -> None (async)
• deregister(job_id) -> None (async)
• 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()

register async

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

Register a job as in-flight.

The lock ensures no concurrent deregister sees an inconsistent state.

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

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

    The lock ensures no concurrent ``deregister`` sees an inconsistent state.
    """
    entry = _ActiveJob(job_id=job_id, task=task, ctx=ctx)
    async with self._lock:
        self._by_id[job_id] = entry

deregister async

deregister(job_id: JobId) -> None

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

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

async def deregister(self, job_id: JobId) -> None:
    """Remove a job from the registry (idempotent — ignores missing keys)."""
    async with self._lock:
        self._by_id.pop(job_id, None)

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)

__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) → ABANDONING (4).

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 ABANDONING, setting escalate_event is a no-op — the orchestrator is already past CANCELLING.

__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, finish in-flight jobs. CANCELLING — cooperative cancel of remaining jobs. FORCING — force-cancel grace, terminal writes shielded. ABANDONING — pod must be replaced to reclaim slots.

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

ABANDONING class-attribute instance-attribute

ABANDONING = 4

drain_local_queue_to_pending async

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

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

Issues a single bounded-timeout UPDATE that clears the lock on rows where locked_by_worker = $worker_id AND status = 'running' AND started_at IS NULL. 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.

Returns:

Type	Description
int	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 locked but never started.

    Issues a single bounded-timeout UPDATE that clears the lock on rows
    where ``locked_by_worker = $worker_id AND status = 'running' AND
    started_at IS NULL``.  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.

    Returns:
        Number of rows updated, or 0 on timeout / connection error.
    """
    schema = deps.settings.schema_name
    if not _IDENT_RE.match(schema):
        raise ValueError(f"invalid schema identifier: {schema!r}")

    sql = (
        f"UPDATE \"{schema}\".jobs 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).
        f"lock_expires_at=NULL "
        f"WHERE locked_by_worker=$1 AND status='running' AND started_at IS NULL"
    )

    try:
        async with deps.dispatcher_pool.acquire(timeout=2.0) as conn:
            tag = await conn.execute(sql, worker_id)
            rowcount = parse_rowcount(tag)
            _log.info(
                "drain-local-queue-completed",
                worker_id=worker_id,
                rows_re_pended=rowcount,
            )
            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 → ABANDONING, 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 → ABANDONING, 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 ──────────────────────────────────────────
        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()
        await drain_local_queue_to_pending(deps, worker_id)

        # ── Phase 2: CANCELLING ────────────────────────────────────────
        deps.shutdown_phase = ShutdownPhase.CANCELLING
        cancel_grace = settings.cancellation_grace_period
        _log.info(
            "shutdown-phase",
            kind="shutdown_phase",
            phase="CANCELLING",
            active_jobs_count=deps.active_jobs.count(),
            elapsed_seconds=loop.time() - t0,
        )
        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()

        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:
                await asyncio.shield(
                    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),
                )
                continue
            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: ABANDONING ────────────────────────────────────────
        deps.shutdown_phase = ShutdownPhase.ABANDONING
        _log.info(
            "shutdown-phase",
            kind="shutdown_phase",
            phase="ABANDONING",
            active_jobs_count=deps.active_jobs.count(),
            elapsed_seconds=loop.time() - t0,
        )
        for active in deps.active_jobs.all():
            try:
                await asyncio.shield(backend.mark_abandoned(active.job_id))
            except asyncio.CancelledError:
                pass
            except Exception as exc:
                _log.warning(
                    "abandon-pg-write-failed",
                    job_id=str(active.job_id),
                    error=str(exc),
                )

        # ── 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: the leader
        # election loop keeps running until shutdown_event fires (finally
        # block below), and a None leader_conn would make it open a
        # *fresh* conn and possibly re-acquire the advisory lock
        # mid-shutdown. For TaskQ-owned conns, close+null releases the
        # advisory lock early so a replacement pod can take over before
        # the SIGTERM budget expires.
        #
        # 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 election loop — a still-live loop could
        # otherwise drop the closing conn and swap in a fresh (possibly
        # lock-holding) one 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)