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

worker_main(
    settings: WorkerSettings,
    *,
    actor_registry: Mapping[str, ActorRef[Any, Any]]
    | None = None,
    di_registry: ProviderRegistry | None = None,
    cron_registry: list[CronScheduleSpec] | None = None,
) -> int

Worker process entry point.

Runs _main under an asyncio.Runner and returns its int result. Uses Runner (not asyncio.run) for finer control over teardown.

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.

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.

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,
) -> int:
    """Worker process entry point.

    Runs ``_main`` under an ``asyncio.Runner`` and returns its int result.
    Uses ``Runner`` (not ``asyncio.run``) for finer control over teardown.

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

    ``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.
    """
    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)
    with asyncio.Runner() as runner:
        return runner.run(
            _main(
                settings,
                actor_registry=actor_registry,
                _registry=di_registry,
                _cron_registry=schedule_specs,
            )
        )

__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"):
        from taskq.worker._bootstrap import _emit_sub_enqueue_startup_warnings, _main

        return {
            "_main": _main,
            "_emit_sub_enqueue_startup_warnings": _emit_sub_enqueue_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()):
            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)

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.

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.

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

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

        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=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=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=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:
            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,
            )
        except Exception:
            _consumer_log.exception("dispatch-failed", job_id=job.id)

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 (TimeoutError, asyncpg.PostgresConnectionError, OSError) 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 (TimeoutError, asyncpg.PostgresConnectionError, OSError) 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",
    "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,
)

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
) -> 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)
    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:
    _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:
        await self._close_leader_owned_conns()
        _active_leaders.discard(self)

archive_expiry_sweep async

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

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

async def archive_expiry_sweep(
    conn: ConnLike,
    *,
    batch_size: int = 10000,
    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] = {}
    expire_before = datetime.now(UTC)
    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

prune_terminal_jobs async

prune_terminal_jobs(
    conn: ConnLike,
    *,
    retention_per_status: dict[str, timedelta],
    archive_retention: timedelta,
    batch_size: int = 10000,
    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 = 10000,
    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

    for status in _TERMINAL_STATUSES:
        retention = retention_per_status.get(status, timedelta(days=30))
        cutoff = datetime.now(UTC) - retention
        cutoffs[status] = cutoff
        sql = _ARCHIVE_CTE_SQL.format(schema=schema)

        while True:
            rows = await conn.fetch(sql, status, cutoff, 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():
            actor_cutoff = datetime.now(UTC) - actor_retention
            sql = _ARCHIVE_CTE_ACTOR_SQL.format(schema=schema)
            for status in _TERMINAL_STATUSES:
                if actor_retention >= retention_per_status.get(status, timedelta(days=30)):
                    continue
                while True:
                    rows = await conn.fetch(
                        sql,
                        status,
                        actor_cutoff,
                        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 = timedelta(hours=24),
    logger_arg: BoundLogger | None = None,
    enqueuer: SubJobEnqueuer,
) -> None

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 = timedelta(hours=24),
    logger_arg: structlog.stdlib.BoundLogger | None = None,
    enqueuer: SubJobEnqueuer,
) -> None:
    """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=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: str = "failed"

    try:
        with safe_start_span(
            f"process {job.actor}",
            kind=SpanKind.CONSUMER,
            attributes=consumer_attrs,
            links=links,
        ) as consumer_span:
            try:
                validated_payload = actor_ref.payload_type.model_validate(job.payload)

                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

                    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=actor_ref.reservations,
                        redis_client=redis_client,
                        worker_pool=deps.worker_pool,
                        settings=deps.settings,
                    )
                    outcome = result

                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:
                    await _handle_generic_exception(
                        backend,
                        job,
                        worker_id,
                        exc,
                        actor_config,
                        clock,
                        max_retry_backoff,
                        consumer_span,
                        handler_log,
                    )
                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)

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():
        _in_tx_failed = False
        tick_start = time.monotonic()
        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)
            # Transaction committed: row locks released.  Run post-tx work
            # (phase-3 mark_abandoned calls) now that deadlock is impossible.
            if cancel_controller is not None:
                await cancel_controller.run_post_tx()
            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 (
            TimeoutError,
            asyncpg.PostgresConnectionError,
            asyncpg.QueryCanceledError,
            OSError,
        ) 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

    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]:
                pending = 0
                crashed = 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() return.
                        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
                        else:
                            crashed += 1
                        await conn.execute(isolate_job_sql, row["id"], worker_id)
                        metadata: dict[str, object] = {}
                        await conn.execute(
                            insert_attempt_sql,
                            row["id"],
                            row["attempt"],
                            row["started_at"],
                            "crashed",
                            "HeartbeatLost",
                            None,
                            None,
                            None,
                            worker_id,
                            dumps_str(metadata),
                        )
                return pending, crashed

            jobs_pending_count, jobs_crashed_count = await asyncio.shield(_inner())
        finally:
            await conn.close()
    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,
        )

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 (TimeoutError, asyncpg.PostgresConnectionError, OSError) 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 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
    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 ──────────────────────────────────────────
        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=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=active.job_id,
                    error=str(exc),
                )

        # ── leader_conn close ──────────────────────────────────
        if deps.leader_conn is not None:
            with contextlib.suppress(asyncpg.PostgresConnectionError, OSError):
                await deps.leader_conn.close()
            deps.leader_conn = None

        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 handlers with a 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).

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 handlers with a 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).

    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_signal() -> None:
        nonlocal _sig_count
        _sig_count += 1
        if _sig_count == 1:
            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)

    for sig in (signal.SIGTERM, signal.SIGINT):
        try:
            loop.add_signal_handler(sig, _on_signal)
        except NotImplementedError:
            _log.warning(
                "signal-handlers-unavailable",
                os_name=os.name,
            )
            return