Rate Limiting

Token bucket, sliding window, concurrency reservation, and the rate-limit registry.

ratelimit

Rate-limiting primitives for TaskQ.

QUEUE_CONCURRENCY_PREFIX module-attribute

QUEUE_CONCURRENCY_PREFIX: Final[str] = "taskq:global:queue:"

Reserved namespace for fleet-wide per-queue concurrency cap reservations.

Primitives whose name starts with this prefix are internal to TaskQ's queue-cap bootstrap path and must be registered via RateLimitRegistry.register_queue_cap_reservation, not via the public RateLimitRegistry.register (which rejects prefixed names to prevent accidental or malicious shadowing of an internal queue cap). Keyed refs must never derive concrete names into this namespace either, see the base_name validators in taskq.ratelimit.refs.

Lives here (rather than in ratelimit.registry) because ratelimit.refs must reference it in its validators, and refs is imported BY registry, defining it in registry would be circular.

__all__ module-attribute

__all__ = [
    "QUEUE_CONCURRENCY_PREFIX",
    "AcquiredResource",
    "ConcurrencyReservation",
    "KeyedRateLimitRef",
    "KeyedReservationRef",
    "RateLimitBackend",
    "RateLimitDecision",
    "RateLimitHandle",
    "RateLimitRef",
    "RateLimitRegistry",
    "RateLimitState",
    "ReservationHandle",
    "ReservationRef",
    "SlidingWindow",
    "SyncResult",
    "TokenBucket",
    "get_redis_pool",
    "queue_concurrency_reservation_name",
    "register_rate_limit_registry",
    "register_redis_pool",
    "registry",
    "sync_rate_limit_buckets",
    "sync_slots",
]

RateLimitBackend

RateLimitBackend = Literal['redis', 'postgres', 'memory']

AcquiredResource

Bases: Protocol

Protocol for a resource handle that can be released.

name property

name: str

release async

release() -> None

Source code in src/taskq/ratelimit/composition.py

async def release(self) -> None: ...

RateLimitHandle dataclass

RateLimitHandle(
    name: str,
    primitive: TokenBucket | SlidingWindow,
    decision: RateLimitDecision,
    redis_client: Redis | None,
    pg_pool: Pool | None,
    clock: Clock | None,
    settings: WorkerSettings | None = None,
    count: float = 1.0,
    refund_on_release: bool = True,
)

Handle for a successfully acquired rate-limit token.

release() is a no-op when refund_on_release is False (post-actor path, token consumption is permanent). When refund_on_release is True (rollback path), release() refunds count tokens via primitive.refund().

name instance-attribute

name: str

primitive instance-attribute

primitive: TokenBucket | SlidingWindow

decision instance-attribute

decision: RateLimitDecision

redis_client instance-attribute

redis_client: Redis | None

pg_pool instance-attribute

pg_pool: Pool | None

clock instance-attribute

clock: Clock | None

settings class-attribute instance-attribute

settings: WorkerSettings | None = field(default=None)

count class-attribute instance-attribute

count: float = 1.0

refund_on_release class-attribute instance-attribute

refund_on_release: bool = True

release async

release() -> None

Source code in src/taskq/ratelimit/composition.py

async def release(self) -> None:
    if not self.refund_on_release:
        return
    await self.primitive.refund(
        self.decision,
        count=self.count,
        redis_client=self.redis_client,
        pg_pool=self.pg_pool,
        clock=self.clock,
        settings=self.settings,
    )

ReservationHandle dataclass

ReservationHandle(
    name: str,
    reservation: ConcurrencyReservation,
    slot_index: int,
    job_id: UUID,
    worker_id: UUID,
    pool: Pool | None,
)

Handle for a successfully acquired reservation slot.

release() calls ConcurrencyReservation.release(slot_index, worker_id, pool) which sets the slot row's job_id to NULL. Idempotent.

slot_index is whatever acquire() returned, in practice a :class:~taskq.ratelimit.reservation.SlotLease, an int subclass that also carries the lease fence. Storing and handing it back unchanged is what makes release() safe against a zombie attempt: a handle from a lease that has since expired and been re-acquired frees nothing, so the bucket cannot silently exceed max_concurrent. Anything that rebuilds a handle must carry the value through as-is rather than re-deriving an int from it.

name instance-attribute

name: str

reservation instance-attribute

reservation: ConcurrencyReservation

slot_index instance-attribute

slot_index: int

job_id instance-attribute

job_id: UUID

worker_id instance-attribute

worker_id: UUID

pool instance-attribute

pool: Pool | None

release async

release() -> None

Source code in src/taskq/ratelimit/composition.py

async def release(self) -> None:
    await self.reservation.release(self.slot_index, self.worker_id, self.pool)

RateLimitDecision dataclass

RateLimitDecision(
    allowed: bool,
    remaining: float,
    retry_after: timedelta | None,
    bucket_name: str,
    backend: RateLimitBackend,
    request_id: str | None = None,
    previous_state: dict[str, object] | None = None,
)

allowed instance-attribute

allowed: bool

remaining instance-attribute

remaining: float

retry_after instance-attribute

retry_after: timedelta | None

bucket_name instance-attribute

bucket_name: str

backend instance-attribute

backend: RateLimitBackend

request_id class-attribute instance-attribute

request_id: str | None = None

previous_state class-attribute instance-attribute

previous_state: dict[str, object] | None = None

RateLimitState dataclass

RateLimitState(
    bucket_name: str,
    backend: RateLimitBackend,
    is_exhausted: bool,
    tokens_remaining: float = 0.0,
    remaining: float = 0.0,
    retry_after: timedelta | None = None,
    capacity: float | None = None,
    limit: int | None = None,
    window: timedelta | None = None,
    style: str | None = None,
    refill_per_second: float | None = None,
)

Read-only snapshot of a rate-limit bucket's current state.

Returned by TokenBucket.peek(), SlidingWindow.peek(), and ConcurrencyReservation.peek(). Fields are backend-agnostic: TB backends populate tokens_remaining and capacity; SW backends populate remaining, limit, window, and style.

bucket_name instance-attribute

bucket_name: str

backend instance-attribute

backend: RateLimitBackend

is_exhausted instance-attribute

is_exhausted: bool

tokens_remaining class-attribute instance-attribute

tokens_remaining: float = 0.0

remaining class-attribute instance-attribute

remaining: float = 0.0

retry_after class-attribute instance-attribute

retry_after: timedelta | None = None

capacity class-attribute instance-attribute

capacity: float | None = None

limit class-attribute instance-attribute

limit: int | None = None

window class-attribute instance-attribute

window: timedelta | None = None

style class-attribute instance-attribute

style: str | None = None

refill_per_second class-attribute instance-attribute

refill_per_second: float | None = None

KeyedRateLimitRef

Bases: BaseModel

Reference to a per-key token bucket, derived from the payload.

Mirrors :class:KeyedReservationRef but for rate limits: base_name namespaces the derived buckets (concrete name is f"{base_name}:{key}"), key_fn derives the key from the actor's validated payload, and capacity / refill_per_second configure every bucket derived from this ref identically (all keys share the same per-key budget).

payload_type is the :class:~pydantic.BaseModel subclass that the payload will be validated against. Use :meth:typed for type-safe construction that binds key_fn to the same payload_type: KeyedRateLimitRef.typed(MyPayload, base_name="api-per-tenant", key_fn=lambda p: p.tenant_id, capacity=10, refill_per_second=1.0).

A consumer calling an external API with per-tenant rate limits would declare rate_limits=[KeyedRateLimitRef.typed(MyPayload, base_name="api-per-tenant", key_fn=lambda p: p.tenant_id, capacity=10, refill_per_second=1.0)] to give each tenant its own independent token budget, with each tenant's bucket materializing on first use.

Concurrency caps vs. rate limits. A concurrency limiter (how many jobs at once, e.g. :class:KeyedReservationRef / :class:~taskq.ratelimit.reservation.ConcurrencyReservation) and a rate limiter (how many per unit time, e.g. :class:~taskq.ratelimit.token_bucket.TokenBucket) solve different problems: N concurrent slots with fast responses can still burst well past a per-time-unit budget, so both may be needed together on the same actor.

Backend selection. The backend field (default "redis") controls which storage backend the materialized :class:~taskq.ratelimit.token_bucket.TokenBucket uses, identical to the backend constructor parameter on a static TokenBucket. In a deployment without Redis configured, set backend="postgres" or backend="memory" to avoid the Redis-required failure mode, a keyed bucket with backend="redis" but no redis_client raises RuntimeError on acquire (not caught by with_pg_fallback, which only handles ConnectionError/TimeoutError).

PG fallback inheritance. _resolve_rate_limit_name constructs a plain :class:~taskq.ratelimit.token_bucket.TokenBucket (with the backend from this ref, default "redis") and calls its normal .acquire(). The existing with_pg_fallback path in token_bucket._acquire_redis_wrapped is therefore inherited automatically, on Redis ConnectionError/TimeoutError, the acquire falls back to the PG rate_limit_buckets table governed by settings.rate_limit_pg_fallback_enabled. No second fallback mechanism is built or needed.

Dual growth bounds. Per-key Redis memory is self-bounding because the token-bucket Lua script sets an EXPIRE TTL on each bucket's Redis hash (computed from capacity/refill_per_second). The Python-process-local dict/registry growth is bounded separately by :meth:~taskq.ratelimit.registry.RateLimitRegistry.evict_idle_keyed_rate_limits, which evicts idle entries from the in-memory registry. These are two independent bounds, Redis TTL bounds Redis memory; registry eviction bounds Python memory.

Concrete per-key :class:~taskq.ratelimit.token_bucket.TokenBucket instances are registered lazily on first acquisition and are not automatically removed, see :meth:~taskq.ratelimit.registry.RateLimitRegistry.evict_idle_keyed_rate_limits for bounding registry growth under high key cardinality.

.. note:: The declared type Callable[[BaseModel], str] is deliberately unsound at the field level, :meth:typed stores a Callable[[P], str] (contravariance prevents direct assignment). Runtime safety is enforced by the registry's isinstance check against payload_type before calling key_fn. Direct invocation of ref.key_fn(model) is unchecked, prefer :meth:typed for compile-time safety.

model_config class-attribute instance-attribute

model_config = ConfigDict(arbitrary_types_allowed=True)

base_name instance-attribute

base_name: str

key_fn instance-attribute

key_fn: Callable[[BaseModel], str]

payload_type instance-attribute

payload_type: type[BaseModel]

capacity instance-attribute

capacity: float

refill_per_second instance-attribute

refill_per_second: float

backend class-attribute instance-attribute

backend: RateLimitBackend = 'redis'

typed classmethod

typed(
    payload_type: type[P],
    *,
    base_name: str,
    key_fn: Callable[[P], str],
    capacity: float,
    refill_per_second: float,
    backend: RateLimitBackend = "redis",
) -> KeyedRateLimitRef

Type-safe constructor that binds key_fn to payload_type.

The key_fn parameter is typed as Callable[[P], str] where P is the provided payload_type, at static-analysis time the caller's lambda or function is checked against the concrete model's attributes (e.g. lambda p: p.tenant_id is verified against payload_type.tenant_id).

At runtime the registry passes the validated payload model to key_fn, so the callable always receives an instance of payload_type (verified by isinstance in the registry , same-type payloads pass through directly; different-type or dict payloads are re-validated via model_validate).

Source code in src/taskq/ratelimit/refs.py

@classmethod
def typed[P: BaseModel](
    cls,
    payload_type: type[P],
    *,
    base_name: str,
    key_fn: Callable[[P], str],
    capacity: float,
    refill_per_second: float,
    backend: RateLimitBackend = "redis",
) -> "KeyedRateLimitRef":
    """Type-safe constructor that binds ``key_fn`` to ``payload_type``.

    The ``key_fn`` parameter is typed as ``Callable[[P], str]`` where
    ``P`` is the provided ``payload_type``, at static-analysis time
    the caller's lambda or function is checked against the concrete
    model's attributes (e.g. ``lambda p: p.tenant_id`` is verified
    against ``payload_type.tenant_id``).

    At runtime the registry passes the validated payload model to
    ``key_fn``, so the callable always receives an instance of
    ``payload_type`` (verified by ``isinstance`` in the registry ,
    same-type payloads pass through directly; different-type or dict
    payloads are re-validated via ``model_validate``).
    """
    return cls(
        base_name=base_name,
        key_fn=key_fn,  # type: ignore[arg-type]  # Why: Callable[[P], str] is not assignable to Callable[[BaseModel], str] due to contravariance, but at runtime the registry only passes the validated model P (verified by isinstance check against ref.payload_type).
        payload_type=payload_type,
        capacity=capacity,
        refill_per_second=refill_per_second,
        backend=backend,
    )

KeyedReservationRef

Bases: BaseModel

Reference to a per-key concurrency reservation, derived from the payload.

base_name namespaces the derived reservations (the concrete name registered for a given key is f"{base_name}:{key}") so distinct KeyedReservationRef declarations never collide. key_fn receives the actor's validated payload (as a :class:~pydantic.BaseModel instance, validated from the job row's raw JSON payload) and must return a non-empty string, typically a tenant, session, or account identifier already present on the payload.

payload_type is the :class:~pydantic.BaseModel subclass that the payload will be validated against. Use :meth:typed for type-safe construction that binds key_fn to the same payload_type: KeyedReservationRef.typed(MyPayload, base_name="geocode-session", key_fn=lambda p: p.session_id, slots=3, lease=timedelta(minutes=5)).

slots and lease configure every reservation derived from this ref identically (all keys share the same per-key cap and lease duration); use a separate KeyedReservationRef if different keys need different caps.

Concrete per-key reservations are registered lazily on first acquisition and are not automatically removed, see :meth:~taskq.ratelimit.registry.RateLimitRegistry.evict_idle_keyed_reservations for bounding registry growth under high key cardinality.

.. note:: The declared type Callable[[BaseModel], str] is deliberately unsound at the field level, :meth:typed stores a Callable[[P], str] (contravariance prevents direct assignment). Runtime safety is enforced by the registry's isinstance check against payload_type before calling key_fn. Direct invocation of ref.key_fn(model) is unchecked, prefer :meth:typed for compile-time safety.

model_config class-attribute instance-attribute

model_config = ConfigDict(arbitrary_types_allowed=True)

base_name instance-attribute

base_name: str

key_fn instance-attribute

key_fn: Callable[[BaseModel], str]

payload_type instance-attribute

payload_type: type[BaseModel]

slots instance-attribute

slots: int

lease instance-attribute

lease: timedelta

typed classmethod

typed(
    payload_type: type[P],
    *,
    base_name: str,
    key_fn: Callable[[P], str],
    slots: int,
    lease: timedelta,
) -> KeyedReservationRef

Type-safe constructor that binds key_fn to payload_type.

The key_fn parameter is typed as Callable[[P], str] where P is the provided payload_type, at static-analysis time the caller's lambda or function is checked against the concrete model's attributes (e.g. lambda p: p.session_id is verified against payload_type.session_id).

At runtime the registry passes the validated payload model to key_fn, so the callable always receives an instance of payload_type (verified by isinstance in the registry , same-type payloads pass through directly; different-type or dict payloads are re-validated via model_validate).

Source code in src/taskq/ratelimit/refs.py

@classmethod
def typed[P: BaseModel](
    cls,
    payload_type: type[P],
    *,
    base_name: str,
    key_fn: Callable[[P], str],
    slots: int,
    lease: timedelta,
) -> "KeyedReservationRef":
    """Type-safe constructor that binds ``key_fn`` to ``payload_type``.

    The ``key_fn`` parameter is typed as ``Callable[[P], str]`` where
    ``P`` is the provided ``payload_type``, at static-analysis time
    the caller's lambda or function is checked against the concrete
    model's attributes (e.g. ``lambda p: p.session_id`` is verified
    against ``payload_type.session_id``).

    At runtime the registry passes the validated payload model to
    ``key_fn``, so the callable always receives an instance of
    ``payload_type`` (verified by ``isinstance`` in the registry ,
    same-type payloads pass through directly; different-type or dict
    payloads are re-validated via ``model_validate``).
    """
    return cls(
        base_name=base_name,
        key_fn=key_fn,  # type: ignore[arg-type]  # Why: Callable[[P], str] is not assignable to Callable[[BaseModel], str] due to contravariance, but at runtime the registry only passes the validated model P (verified by isinstance check against ref.payload_type).
        payload_type=payload_type,
        slots=slots,
        lease=lease,
    )

RateLimitRef

Bases: BaseModel

Typed reference to a rate-limit primitive by name.

name instance-attribute

name: str

count class-attribute instance-attribute

count: float = 1.0

ReservationRef

Bases: BaseModel

Typed reference to a concurrency reservation primitive by name.

name instance-attribute

name: str

RateLimitRegistry

RateLimitRegistry()

Unified registry for rate-limit and reservation primitives.

Stores two separate dicts: _rate_limits for TokenBucket / SlidingWindow and _reservations for ConcurrencyReservation. Cross-dict name collision is allowed, they live in separate namespaces.

Ownership. A registry is an ordinary ownable object: construct one per process and pass it to worker_main(..., rate_limit_registry=...) / create_router(..., rate_limit_registry=...), or rely on the module-level registry singleton (the default at every entry point). Actor-declared primitive instances are registered by the worker bootstrap's collection pass; use explicit .register() for primitives shared outside actor dispatch. :meth:clear resets all state and is a test aid only, NOT safe while a worker is running.

Source code in src/taskq/ratelimit/registry.py

def __init__(self) -> None:
    self._rate_limits: dict[str, TokenBucket | SlidingWindow] = {}
    self._reservations: dict[str, ConcurrencyReservation] = {}
    # Per-reservation release signals, created on demand: a consumer
    # whose admission acquire was denied parks its bounded local retry
    # on its bucket's event instead of a bare sleep, so a LIVE holder's
    # completion (a release this process performed) wakes the retry
    # within one loop turn - the freed slot is then contested against
    # fresh claimants' acquires, which the retry's head start wins.
    # The event is a WAKE SIGNAL, never a token: any number of waiters
    # wake on one release and re-probe, each denial re-arming its own
    # wait. Cross-process releases (a peer worker's) are not seen here;
    # the retry's timeout floor covers them.
    # Reclaimed with the keyed lifecycle: evict_idle_keyed_reservations()
    # pops the entry alongside the bucket itself (same idle criterion),
    # so high-cardinality keyed reservations under sustained denials do
    # not grow this dict without bound. A waiter holding a popped event
    # object is never stranded: its wait is bounded (the retry's backoff
    # floor) and its return - wake or timeout - falls through to the
    # availability re-probe, which is the correctness anchor; a
    # re-materialized bucket gets a fresh event, so no stale set()
    # crosses the eviction boundary.
    self._reservation_release_events: dict[str, asyncio.Event] = {}
    # Names of reservations materialized from a KeyedReservationRef
    # (as opposed to a static @actor(reservations=["name"]) entry),
    # and the monotonic time each was last acquired, used only by
    # evict_idle_keyed_reservations() to bound registry growth under
    # high key cardinality. Never consulted by acquire_for_actor.
    self._keyed_reservation_last_used: dict[str, float] = {}
    # Names of rate limits materialized from a KeyedRateLimitRef
    # (as opposed to a static @actor(rate_limits=["name"]) entry),
    # and the monotonic time each was last acquired, used only by
    # evict_idle_keyed_rate_limits() to bound registry growth under
    # high key cardinality. Never consulted by acquire_for_actor.
    self._keyed_rate_limit_last_used: dict[str, float] = {}
    # Schema of the rate_limit_buckets row each keyed rate limit
    # published, captured at publish time, a TokenBucket carries no
    # schema of its own (unlike ConcurrencyReservation), so this is
    # the only record of where the row landed. Read by
    # evict_idle_keyed_rate_limits() to record the bucket's row for
    # reclamation; rides the keyed lifecycle like the heal stamps
    # (dropped on eviction, re-captured on re-materialization's
    # publish).
    self._keyed_rate_limit_row_schemas: dict[str, str] = {}
    # Evicted keyed-reservation bucket names whose reservation_slots
    # rows are still to be deleted, keyed by the schema those rows
    # live in. Each schema's names live in a dict[str, None] used as
    # an insertion-ordered set: the drain takes the FRONT batch and
    # re-appends that tick's survivors (names whose rows were all
    # still held) to the BACK, so a held bucket waits at most one
    # full rotation and never starves the buckets behind it. Bounded:
    # eviction refuses to record past the pending cap (the
    # settings-derived max_keyed_reservations at both production call
    # sites, the constant default for direct callers), and the drain
    # removes names whose rows are gone, drops names that
    # re-registered, and pops schema keys whose set is empty.
    # evict_idle_keyed_reservations() records; the sweep loop's
    # drain_pending_reservation_reclaims() deletes.
    self._pending_reservation_reclaims: dict[str, dict[str, None]] = {}
    # Evicted keyed rate-limit bucket names whose published
    # rate_limit_buckets rows are still to be deleted, the
    # rate-limit twin of _pending_reservation_reclaims: same
    # insertion-ordered per-schema sets, same cap-vetoed recording,
    # same drain (which deletes the front slice per schema; a bucket
    # row has no holder, so nothing survives the DELETE and no
    # survivor rotation is needed). evict_idle_keyed_rate_limits()
    # records; drain_pending_reservation_reclaims() deletes.
    self._pending_rate_limit_reclaims: dict[str, dict[str, None]] = {}
    # Monotonic time of the last acquire-path heal attempt per keyed
    # reservation name, gates the existence probe in the
    # ReservationUnavailable heal to one attempt per
    # _KEYED_RECLAIM_HEAL_WINDOW per bucket. Rides the keyed
    # registration's lifecycle: discarded on re-registration, pruned
    # by the same eviction pass that drops the tracking entry.
    self._keyed_reservation_heal_attempted: dict[str, float] = {}
    # Monotonic time of the last heal-failure WARNING per keyed
    # reservation name. The heal itself retries on every denial (a
    # failed attempt rolls its own stamp back), so the log needs a
    # separate gate: one warning per bucket per heal window, while
    # the failure counter records every attempt (metrics aggregate;
    # a per-denial warning line is what floods). Rides the same
    # lifecycle as _keyed_reservation_heal_attempted.
    self._keyed_reservation_heal_failure_logged: dict[str, float] = {}
    # Monotonic timestamps of the last opportunistic eviction scan on
    # each acquisition path, used to amortize the O(n) scan to at most
    # once per _OPPORTUNISTIC_EVICT_MIN_INTERVAL under sustained cap-hit
    # denials (see the constant's docstring). -inf so the first cap-hit
    # after startup always scans. evict_idle_keyed_*() are synchronous
    # with no await points, so check-and-stamp is atomic within the
    # event loop, concurrent scanners cannot pile up.
    self._keyed_reservation_last_eviction_scan: float = float("-inf")
    self._keyed_rate_limit_last_eviction_scan: float = float("-inf")

rate_limits property

rate_limits: dict[str, TokenBucket | SlidingWindow]

reservations property

reservations: dict[str, ConcurrencyReservation]

has_keyed_reservations property

has_keyed_reservations: bool

has_keyed_rate_limits property

has_keyed_rate_limits: bool

has_pending_reservation_reclaims property

has_pending_reservation_reclaims: bool

Whether any evicted keyed bucket still awaits its row deletion.

Covers both pending sets, reservation slot rows and published rate-limit bucket rows. The sweep loop uses this single gate to skip the drain entirely (no connection acquired) when there is nothing to reclaim, so a pending set the gate cannot see would never be drained.

is_keyed_rate_limit

is_keyed_rate_limit(name: str) -> bool

Whether name is a keyed-MATERIALISED rate-limit entry.

A keyed ref's first acquire derives base_name:key and lazily registers a primitive tracked in _keyed_rate_limit_last_used; a STATICALLY pre-registered entry is deliberately never stamped there (the idle-eviction sweep must not evict it). This predicate is that same discrimination, published for the metrics sampler: named series are for statically-registered buckets only (bounded by the code the user ships), keyed materialisations — whose names embed caller-controlled keys — are aggregated so they can never become label cardinality.

Source code in src/taskq/ratelimit/registry.py

def is_keyed_rate_limit(self, name: str) -> bool:
    """Whether *name* is a keyed-MATERIALISED rate-limit entry.

    A keyed ref's first acquire derives ``base_name:key`` and lazily
    registers a primitive tracked in ``_keyed_rate_limit_last_used``;
    a STATICALLY pre-registered entry is deliberately never stamped
    there (the idle-eviction sweep must not evict it). This
    predicate is that same discrimination, published for the
    metrics sampler: named series are for statically-registered
    buckets only (bounded by the code the user ships), keyed
    materialisations — whose names embed caller-controlled keys —
    are aggregated so they can never become label cardinality.
    """
    return name in self._keyed_rate_limit_last_used

has_reservation

has_reservation(name: str) -> bool

O(1) membership test against the live reservations dict.

Unlike the :attr:reservations property this does NOT defensively copy the dict, use it on per-job hot paths (e.g. the dispatch queue-cap check), where copying the whole registry per call is prohibitive at high keyed-entry cardinality.

Source code in src/taskq/ratelimit/registry.py

def has_reservation(self, name: str) -> bool:
    """O(1) membership test against the live reservations dict.

    Unlike the :attr:`reservations` property this does NOT defensively
    copy the dict, use it on per-job hot paths (e.g. the dispatch
    queue-cap check), where copying the whole registry per call is
    prohibitive at high keyed-entry cardinality.
    """
    return name in self._reservations

has_rate_limit

has_rate_limit(name: str) -> bool

O(1) membership test against the live rate-limits dict.

See :meth:has_reservation, the same no-copy guarantee applies.

Source code in src/taskq/ratelimit/registry.py

def has_rate_limit(self, name: str) -> bool:
    """O(1) membership test against the live rate-limits dict.

    See :meth:`has_reservation`, the same no-copy guarantee applies.
    """
    return name in self._rate_limits

register

register(
    primitive: TokenBucket
    | SlidingWindow
    | ConcurrencyReservation,
) -> None

Source code in src/taskq/ratelimit/registry.py

def register(
    self,
    primitive: TokenBucket | SlidingWindow | ConcurrencyReservation,
) -> None:
    if primitive.name.startswith(QUEUE_CONCURRENCY_PREFIX):
        raise ValueError(
            f"name {primitive.name!r} starts with the reserved prefix "
            f"{QUEUE_CONCURRENCY_PREFIX!r}, internal queue-cap reservations "
            f"must be registered via register_queue_cap_reservation()"
        )
    if isinstance(primitive, ConcurrencyReservation):
        self._register_reservation_unchecked(primitive)
        return

    name = primitive.name
    existing = self._rate_limits.get(name)
    if existing is not None:
        if _same_config(existing, primitive):
            logger.debug(
                "registry-register-idempotent-noop",
                kind="rate_limit",
                name=name,
            )
            return
        raise ValueError(
            f"rate-limit name already registered with a different config: "
            f"{name!r}, existing={existing!r}, new={primitive!r}"
        )
    self._rate_limits[name] = primitive
    logger.debug(
        "registry-registered",
        kind="rate_limit",
        name=name,
    )

register_queue_cap_reservation

register_queue_cap_reservation(
    reservation: ConcurrencyReservation,
) -> None

Register a fleet-wide queue-cap reservation in the reserved namespace.

This is the ONLY way to register a reservation whose name starts with :data:QUEUE_CONCURRENCY_PREFIX. The public :meth:register rejects such names to prevent users from accidentally shadowing internal queue caps. Idempotency and conflict detection are identical to :meth:register (duplicate-name-with-different-config → ValueError; duplicate-name-with-same-config → idempotent no-op).

Source code in src/taskq/ratelimit/registry.py

def register_queue_cap_reservation(
    self,
    reservation: ConcurrencyReservation,
) -> None:
    """Register a fleet-wide queue-cap reservation in the reserved namespace.

    This is the ONLY way to register a reservation whose name starts with
    :data:`QUEUE_CONCURRENCY_PREFIX`. The public :meth:`register` rejects
    such names to prevent users from accidentally shadowing internal
    queue caps. Idempotency and conflict detection are identical to
    :meth:`register` (duplicate-name-with-different-config →
    ``ValueError``; duplicate-name-with-same-config → idempotent no-op).
    """
    if not reservation.name.startswith(QUEUE_CONCURRENCY_PREFIX):
        raise ValueError(
            f"register_queue_cap_reservation() requires a name starting with "
            f"{QUEUE_CONCURRENCY_PREFIX!r}, got {reservation.name!r}"
        )
    self._register_reservation_unchecked(reservation)

get_rate_limit

get_rate_limit(name: str) -> TokenBucket | SlidingWindow

Source code in src/taskq/ratelimit/registry.py

def get_rate_limit(self, name: str) -> TokenBucket | SlidingWindow:
    try:
        return self._rate_limits[name]
    except KeyError:
        raise KeyError(name) from None

get_reservation

get_reservation(name: str) -> ConcurrencyReservation

Source code in src/taskq/ratelimit/registry.py

def get_reservation(self, name: str) -> ConcurrencyReservation:
    try:
        return self._reservations[name]
    except KeyError:
        raise KeyError(name) from None

acquire async

acquire(
    name: str,
    count: float = 1.0,
    *,
    redis_client: Redis | None = None,
    pg_pool: Pool | None = None,
    clock: Clock | None = None,
    settings: WorkerSettings | None = None,
) -> AsyncGenerator[RateLimitDecision, None]

Source code in src/taskq/ratelimit/registry.py

@asynccontextmanager
async def acquire(
    self,
    name: str,
    count: float = 1.0,
    *,
    redis_client: "redis_async.Redis | None" = None,
    pg_pool: "asyncpg.Pool | None" = None,
    clock: "Clock | None" = None,
    settings: "WorkerSettings | None" = None,
) -> AsyncGenerator[RateLimitDecision, None]:
    if name in self._reservations:
        raise TypeError(
            f"name {name!r} is a ConcurrencyReservation, "
            f"registry.acquire() is only for rate limits; "
            f"reservation acquisition requires a job_id"
        )
    if name not in self._rate_limits:
        raise KeyError(name)

    primitive = self._rate_limits[name]
    if isinstance(primitive, TokenBucket):
        decision = await primitive.acquire(
            count,
            redis_client=redis_client,
            pg_pool=pg_pool,
            clock=clock,
            settings=settings,
        )
    else:
        decision = await primitive.acquire(
            redis_client=redis_client,
            pg_pool=pg_pool,
            clock=clock,
            settings=settings,
        )
    yield decision

acquire_for_actor async

acquire_for_actor(
    rate_limits: Sequence[
        str
        | KeyedRateLimitRef
        | TokenBucket
        | SlidingWindow
    ],
    reservations: Sequence[
        str | KeyedReservationRef | ConcurrencyReservation
    ],
    *,
    job_id: UUID,
    worker_id: UUID,
    payload: dict[str, object] | BaseModel | None = None,
    redis_client: Redis | None = None,
    pg_pool: Pool | None = None,
    clock: Clock | None = None,
    settings: WorkerSettings | None = None,
) -> list[AcquiredResource]

AND-composition: acquire reservations first, then rate limits.

reservations entries may be plain names (resolved against statically pre-registered primitives), :class:KeyedReservationRef instances (resolved dynamically per job from payload, see :meth:_resolve_reservation_name), or :class:ConcurrencyReservation instances (normalized to their .name up front, the instance must already be registered, e.g. by the worker bootstrap's actor-declaration collection pass; an unregistered instance raises KeyError exactly like an unknown name). rate_limits entries may likewise be plain names, :class:KeyedRateLimitRef instances, or :class:TokenBucket / :class:SlidingWindow instances. payload is required if any entry is a KeyedReservationRef or KeyedRateLimitRef. It may be a dict (validated via ref.payload_type.model_validate) or a BaseModel (used directly if it matches ref.payload_type, otherwise re-validated via model_dump() → model_validate()).

Returns the list of AcquiredResource handles on full success. Raises ReservationUnavailable on any denial, rollback is performed internally before re-raising (already-acquired resources released in reverse order, each failure logged at ERROR).

Source code in src/taskq/ratelimit/registry.py

async def acquire_for_actor(
    self,
    rate_limits: Sequence["str | KeyedRateLimitRef | TokenBucket | SlidingWindow"],
    reservations: Sequence["str | KeyedReservationRef | ConcurrencyReservation"],
    *,
    job_id: "UUID",
    worker_id: "UUID",
    payload: dict[str, object] | BaseModel | None = None,
    redis_client: "redis_async.Redis | None" = None,
    pg_pool: "asyncpg.Pool | None" = None,
    clock: "Clock | None" = None,
    settings: "WorkerSettings | None" = None,
) -> list[AcquiredResource]:
    """AND-composition: acquire reservations first, then rate limits.

    ``reservations`` entries may be plain names (resolved against
    statically pre-registered primitives), :class:`KeyedReservationRef`
    instances (resolved dynamically per job from ``payload``, see
    :meth:`_resolve_reservation_name`), or
    :class:`ConcurrencyReservation` instances (normalized to their
    ``.name`` up front, the instance must already be registered,
    e.g. by the worker bootstrap's actor-declaration collection
    pass; an unregistered instance raises ``KeyError`` exactly like
    an unknown name). ``rate_limits`` entries may likewise be plain
    names, :class:`KeyedRateLimitRef` instances, or
    :class:`TokenBucket` / :class:`SlidingWindow` instances.
    ``payload`` is required if any entry is a ``KeyedReservationRef``
    or ``KeyedRateLimitRef``. It may be a ``dict`` (validated via
    ``ref.payload_type.model_validate``) or a ``BaseModel`` (used
    directly if it matches ``ref.payload_type``, otherwise re-validated
    via ``model_dump()`` → ``model_validate()``).

    Returns the list of ``AcquiredResource`` handles on full success.
    Raises ``ReservationUnavailable`` on any denial, rollback is performed
    internally before re-raising (already-acquired resources released in
    reverse order, each failure logged at ERROR).
    """
    # Normalize primitive instances to their names BEFORE any use ,
    # _ref_display (below) only handles str | keyed refs and would
    # AttributeError on a primitive instance, and every dict lookup
    # and handle construction sees names only. No acquisition-time
    # auto-registration: bootstrap is the fail-fast point; an
    # unregistered instance raises KeyError from the dict lookups below.
    rl_seq: list[str | KeyedRateLimitRef] = [
        rl.name if isinstance(rl, TokenBucket | SlidingWindow) else rl for rl in rate_limits
    ]
    res_seq: list[str | KeyedReservationRef] = [
        res.name if isinstance(res, ConcurrencyReservation) else res for res in reservations
    ]
    acquired: list[AcquiredResource] = []
    try:
        for res_ref in res_seq:
            res_name = await self._resolve_reservation_name(
                res_ref, payload, pg_pool=pg_pool, settings=settings
            )
            reservation = self._reservations[res_name]
            slot_index = await self._acquire_reservation_slot_with_heal(
                reservation,
                job_id=job_id,
                worker_id=worker_id,
                pg_pool=pg_pool,
            )
            acquired.append(
                ReservationHandle(
                    name=res_name,
                    reservation=reservation,
                    slot_index=slot_index,
                    job_id=job_id,
                    worker_id=worker_id,
                    pool=pg_pool,
                )
            )

        for rl_ref in rl_seq:
            rl_name = await self._resolve_rate_limit_name(
                rl_ref, payload, settings=settings, pg_pool=pg_pool
            )
            rl = self._rate_limits[rl_name]
            if isinstance(rl, TokenBucket):
                result = await rl.acquire(
                    1.0,
                    redis_client=redis_client,
                    pg_pool=pg_pool,
                    clock=clock,
                    settings=settings,
                )
            else:
                result = await rl.acquire(
                    redis_client=redis_client,
                    pg_pool=pg_pool,
                    clock=clock,
                    settings=settings,
                )
            if not result.allowed:
                retry_td = (
                    result.retry_after
                    if result.retry_after is not None
                    else DEFAULT_RESERVATION_BACKOFF
                )
                logger.info(
                    "composition-denied",
                    job_id=str(job_id),
                    rate_limits=[_ref_display(r) for r in rl_seq],
                    reservations=[_ref_display(r) for r in res_seq],
                    allowed=False,
                    retry_after_seconds=retry_td.total_seconds(),
                    failed_bucket=rl_name,
                )
                raise ReservationUnavailable(
                    bucket_name=rl_name,
                    retry_after=retry_td,
                    source="rate_limit",
                )
            acquired.append(
                RateLimitHandle(
                    name=rl_name,
                    primitive=rl,
                    decision=result,
                    redis_client=redis_client,
                    pg_pool=pg_pool,
                    clock=clock,
                    settings=settings,
                    count=1.0,
                    refund_on_release=True,
                )
            )

        logger.debug(
            "composition-acquired",
            job_id=str(job_id),
            rate_limits=[_ref_display(r) for r in rl_seq],
            reservations=[_ref_display(r) for r in res_seq],
            allowed=True,
            retry_after=None,
            handle_count=len(acquired),
        )
        return acquired
    except Exception:
        # CancelledError deliberately bypasses this rollback:
        # asyncio.CancelledError derives from BaseException, not
        # Exception, so a cancellation landing mid-composition leaves
        # any already-acquired handles in place. That is an accepted,
        # bounded, self-healing leak, NOT an oversight: leaked
        # reservation slots are reclaimed by lease expiry (the
        # lock-expiry sweep, within ~30s), and consumed rate-limit
        # tokens are bounded by the bucket's Redis EXPIRE TTL. Rolling
        # back here would mean network I/O (handle.release()) while the
        # task is being torn down, delaying cancellation, with a
        # second cancel able to interrupt the release itself, which is
        # worse than a leak with an existing reclaim path.
        for handle in reversed(acquired):
            try:
                await handle.release()
            except Exception as exc:
                backend = (
                    handle.decision.backend
                    if isinstance(handle, RateLimitHandle)
                    else "postgres"
                )
                logger.error(
                    "ratelimit-rollback-failure",
                    handle_name=handle.name,
                    operation="release",
                    error=str(exc),
                    acquired_count=len(acquired),
                )
                record_ratelimit_refund_failure(handle.name, backend)
        raise

peek async

peek(
    name: str,
    *,
    redis_client: Redis | None = None,
    pg_pool: Pool | None = None,
    clock: Clock | None = None,
    settings: WorkerSettings | None = None,
) -> RateLimitState

Look up a rate-limit primitive by name and return its current state.

Source code in src/taskq/ratelimit/registry.py

async def peek(
    self,
    name: str,
    *,
    redis_client: "redis_async.Redis | None" = None,
    pg_pool: "asyncpg.Pool | None" = None,
    clock: "Clock | None" = None,
    settings: "WorkerSettings | None" = None,
) -> RateLimitState:
    """Look up a rate-limit primitive by name and return its current state."""
    if name in self._reservations:
        raise TypeError(
            f"name {name!r} is a ConcurrencyReservation, "
            f"peek() on reservations is not supported via this method"
        )
    if name not in self._rate_limits:
        raise KeyError(name)

    primitive = self._rate_limits[name]
    if isinstance(primitive, TokenBucket):
        return await primitive.peek(
            redis_client=redis_client,
            pg_pool=pg_pool,
            clock=clock,
            settings=settings,
        )
    else:
        return await primitive.peek(
            redis_client=redis_client,
            pg_pool=pg_pool,
            clock=clock,
            settings=settings,
        )

peek_all async

peek_all(
    *,
    redis_client: Redis | None = None,
    pg_pool: Pool | None = None,
    clock: Clock | None = None,
    settings: WorkerSettings | None = None,
    timeout: float | None = None,
) -> dict[str, RateLimitState]

Peek all registered rate limits. Returns {name: RateLimitState}.

Each bucket's read is a separate Redis/PG round trip, so a call costs O(buckets) round trips, and the registry can hold up to max_keyed_rate_limits keyed-materialised buckets. timeout bounds the WHOLE pass: a bucket whose store hangs (a black-holed broker answers no read) must not park the caller past it. Raises :class:TimeoutError when the bound fires, per-bucket failures are still caught and logged per bucket, but a read that never RETURNS is indistinguishable from a dead registry at page-render time, so the caller learns of the bound instead of rendering a half-empty live-state map as if it were current. None (the default) keeps the unbounded shape for callers that manage their own deadline.

Source code in src/taskq/ratelimit/registry.py

async def peek_all(
    self,
    *,
    redis_client: "redis_async.Redis | None" = None,
    pg_pool: "asyncpg.Pool | None" = None,
    clock: "Clock | None" = None,
    settings: "WorkerSettings | None" = None,
    timeout: "float | None" = None,  # noqa: ASYNC109  # Why: the bound is a per-call deadline the caller passes, not an enclosing asyncio.timeout scope; the registry's other bounded methods (drain_pending_reservation_reclaims) take the same shape.
) -> dict[str, RateLimitState]:
    """Peek all registered rate limits. Returns {name: RateLimitState}.

    Each bucket's read is a separate Redis/PG round trip, so a call
    costs O(buckets) round trips, and the registry can hold up to
    ``max_keyed_rate_limits`` keyed-materialised buckets. *timeout*
    bounds the WHOLE pass: a bucket whose store hangs (a black-holed
    broker answers no read) must not park the caller past it. Raises
    :class:`TimeoutError` when the bound fires, per-bucket failures
    are still caught and logged per bucket, but a read that never
    RETURNS is indistinguishable from a dead registry at page-render
    time, so the caller learns of the bound instead of rendering a
    half-empty live-state map as if it were current. ``None`` (the
    default) keeps the unbounded shape for callers that manage their
    own deadline.
    """
    if timeout is not None:
        return await asyncio.wait_for(
            self._peek_all(redis_client, pg_pool, clock, settings), timeout=timeout
        )
    return await self._peek_all(redis_client, pg_pool, clock, settings)

reset async

reset(
    name: str,
    *,
    redis_client: Redis | None = None,
    pg_pool: Pool | None = None,
    clock: Clock | None = None,
    settings: WorkerSettings | None = None,
    timeout: float | None = None,
) -> None

Reset a rate-limit bucket to full capacity.

timeout bounds the reset's backend round trip (a Redis DEL or a PG upsert against a dead store can hang the caller forever); raises :class:TimeoutError when it fires. None (the default) keeps the unbounded shape for callers that manage their own deadline.

Source code in src/taskq/ratelimit/registry.py

async def reset(
    self,
    name: str,
    *,
    redis_client: "redis_async.Redis | None" = None,
    pg_pool: "asyncpg.Pool | None" = None,
    clock: "Clock | None" = None,
    settings: "WorkerSettings | None" = None,
    timeout: "float | None" = None,  # noqa: ASYNC109  # Why: the bound is a per-call deadline the caller passes, not an enclosing asyncio.timeout scope.
) -> None:
    """Reset a rate-limit bucket to full capacity.

    *timeout* bounds the reset's backend round trip (a Redis DEL or a
    PG upsert against a dead store can hang the caller forever);
    raises :class:`TimeoutError` when it fires. ``None`` (the default)
    keeps the unbounded shape for callers that manage their own
    deadline.
    """
    if name in self._reservations:
        raise TypeError(
            f"name {name!r} is a ConcurrencyReservation, "
            f"reset() on reservations is not supported"
        )
    if name not in self._rate_limits:
        raise KeyError(name)

    primitive = self._rate_limits[name]
    if timeout is not None:
        await asyncio.wait_for(
            primitive.reset(
                redis_client=redis_client,
                pg_pool=pg_pool,
                clock=clock,
                settings=settings,
            ),
            timeout=timeout,
        )
    else:
        await primitive.reset(
            redis_client=redis_client,
            pg_pool=pg_pool,
            clock=clock,
            settings=settings,
        )

release_for_actor async

release_for_actor(
    acquired: list[AcquiredResource],
    *,
    pg_pool: Pool | None = None,
) -> None

Release acquired resources after actor completion.

Sets refund_on_release=False on all RateLimitHandle instances before iterating (token consumption is permanent after actor ran). Releases in reverse acquisition order. Each release failure is caught, logged at ERROR, and loop continues (same pattern as rollback).

Why pg_pool is unused: each handle captured the pool it needs at acquisition time, so handle.release() is self-contained. The parameter mirrors :meth:acquire_for_actor, which does need it, so the acquire/release pair reads as one symmetric surface at the call site (taskq.worker._consumer passes the same worker_pool to both). It is public API, so it is kept rather than removed.

Source code in src/taskq/ratelimit/registry.py

async def release_for_actor(
    self,
    acquired: list[AcquiredResource],
    *,
    pg_pool: "asyncpg.Pool | None" = None,
) -> None:
    """Release acquired resources after actor completion.

    Sets ``refund_on_release=False`` on all ``RateLimitHandle`` instances
    before iterating (token consumption is permanent after actor ran).
    Releases in reverse acquisition order.  Each release failure is caught,
    logged at ERROR, and loop continues (same pattern as rollback).

    Why *pg_pool* is unused: each handle captured the pool it needs at
    acquisition time, so ``handle.release()`` is self-contained. The
    parameter mirrors :meth:`acquire_for_actor`, which does need it, so
    the acquire/release pair reads as one symmetric surface at the call
    site (``taskq.worker._consumer`` passes the same ``worker_pool`` to
    both). It is public API, so it is kept rather than removed.
    """
    for handle in acquired:
        if isinstance(handle, RateLimitHandle):
            handle.refund_on_release = False

    for handle in reversed(acquired):
        try:
            await handle.release()
            if isinstance(handle, ReservationHandle):
                self._notify_bucket_release(handle.name)
        except Exception as exc:
            backend = (
                handle.decision.backend if isinstance(handle, RateLimitHandle) else "postgres"
            )
            logger.error(
                "ratelimit-rollback-failure",
                handle_name=handle.name,
                operation="release",
                error=str(exc),
                acquired_count=len(acquired),
            )
            record_ratelimit_refund_failure(handle.name, backend)

bucket_release_event

bucket_release_event(bucket_name: str) -> asyncio.Event

Return (creating on first use) the release-wake event for bucket_name. The retry path in taskq.worker._consumer waits on this between acquire attempts; see _notify_bucket_release.

Source code in src/taskq/ratelimit/registry.py

def bucket_release_event(self, bucket_name: str) -> "asyncio.Event":
    """Return (creating on first use) the release-wake event for
    *bucket_name*. The retry path in ``taskq.worker._consumer`` waits
    on this between acquire attempts; see ``_notify_bucket_release``.
    """
    event = self._reservation_release_events.get(bucket_name)
    if event is None:
        event = asyncio.Event()
        self._reservation_release_events[bucket_name] = event
    return event

evict_idle_keyed_reservations

evict_idle_keyed_reservations(
    idle_for: timedelta,
    *,
    max_pending_reclaims: int | None = None,
) -> int

Drop registry entries for keyed reservations idle at least idle_for.

Reservations derived from a :class:KeyedReservationRef are registered lazily and never removed automatically, under high key cardinality (e.g. one reservation per import session over a long worker lifetime) this dict grows without bound. Each worker's 30-second sweep calls this automatically against its own registry (not leader-gated) with a 1-hour idle threshold; call directly for custom eviction windows.

Removes the in-memory registry entry, its acquire-recency tracking, and its release-wake event (the denial-retry's wake signal, popped with the same idle criterion so sustained denials under high key cardinality cannot grow the event dict without bound), and RECORDS the bucket in the pending-reclaim set so its reservation_slots rows are deleted by :meth:drain_pending_reservation_reclaims on the same sweep cadence. The lock-expiry sweep is an UPDATE ... SET job_id = NULL, it clears a row's holder but never deletes the row, and sync_slots iterates only currently-registered reservations, so without the pending-reclaim drain an evicted bucket's rows would be orphaned permanently (steady-state cardinality: slots x every key ever seen, unbounded in the caller-controlled key space). A slot still held by a live lease survives the drain's idle-guarded DELETE and stays pending until its lease expires.

max_pending_reclaims caps the pending-reclaim set (default: :data:taskq.constants.DEFAULT_MAX_KEYED_RESERVATIONS; both production call sites, the per-worker sweep and the acquisition path's opportunistic eviction, pass the settings-derived WorkerSettings.max_keyed_reservations, so the pending set is bounded by the same ceiling that bounds the tracked entries; direct callers without a settings object get the constant). At the cap the eviction is VETOED, the entry stays registered and re-scanned on the next sweep, so no structure grows unbounded and no rows are orphaned by an eviction that could not be recorded. The visible signals are the pending-depth gauge (taskq.ratelimit.reclaim_pending, the steady signal), the registry-keyed-reclaim-pending-cap-veto warning (one aggregated line per eviction call that shed evictions, pending at cap means reclamation is falling behind), and, if the veto persists up to the registry's own entry cap, the existing registry-keyed-reservation-limit-exceeded soft-cap warning.

A key that is acquired again after eviction is re-registered on next use (idempotent, see :meth:_resolve_reservation_name), and the drain drops its name from the pending set without deleting anything (a re-activated key owns its rows again), so eviction is always safe to call, including concurrently with in-flight acquisitions for other keys.

Returns the number of entries evicted.

Source code in src/taskq/ratelimit/registry.py

def evict_idle_keyed_reservations(
    self,
    idle_for: "timedelta",
    *,
    max_pending_reclaims: "int | None" = None,
) -> int:
    """Drop registry entries for keyed reservations idle at least ``idle_for``.

    Reservations derived from a :class:`KeyedReservationRef` are
    registered lazily and never removed automatically, under high key
    cardinality (e.g. one reservation per import session over a long
    worker lifetime) this dict grows without bound. Each worker's
    30-second sweep calls this automatically against its own registry
    (not leader-gated) with a 1-hour idle threshold; call directly for
    custom eviction windows.

    Removes the in-memory registry entry, its acquire-recency
    tracking, and its release-wake event (the denial-retry's wake
    signal, popped with the same idle criterion so sustained denials
    under high key cardinality cannot grow the event dict without
    bound), and RECORDS the bucket in the pending-reclaim set so
    its ``reservation_slots`` rows are deleted by
    :meth:`drain_pending_reservation_reclaims` on the same sweep
    cadence. The lock-expiry sweep is an ``UPDATE ... SET job_id =
    NULL``, it clears a row's holder but never deletes the row, and
    ``sync_slots`` iterates only currently-registered reservations,
    so without the pending-reclaim drain an evicted bucket's rows
    would be orphaned permanently (steady-state cardinality: slots x
    every key ever seen, unbounded in the caller-controlled key
    space). A slot still held by a live lease survives the drain's
    idle-guarded DELETE and stays pending until its lease expires.

    *max_pending_reclaims* caps the pending-reclaim set (default:
    :data:`taskq.constants.DEFAULT_MAX_KEYED_RESERVATIONS`; both
    production call sites, the per-worker sweep and the acquisition
    path's opportunistic eviction, pass the settings-derived
    ``WorkerSettings.max_keyed_reservations``, so the pending set is
    bounded by the same ceiling that bounds the tracked entries;
    direct callers without a settings object get the constant). At
    the cap the eviction is VETOED, the entry stays registered and
    re-scanned on the next sweep, so no structure grows unbounded
    and no rows are orphaned by an eviction that could not be
    recorded. The visible signals are the pending-depth gauge
    (``taskq.ratelimit.reclaim_pending``, the steady signal), the
    ``registry-keyed-reclaim-pending-cap-veto`` warning (one
    aggregated line per eviction call that shed evictions, pending
    at cap means reclamation is falling behind), and, if the veto
    persists up to the registry's own entry cap, the existing
    ``registry-keyed-reservation-limit-exceeded`` soft-cap warning.

    A key that is acquired again after eviction is
    re-registered on next use (idempotent, see
    :meth:`_resolve_reservation_name`), and the drain drops its name
    from the pending set without deleting anything (a re-activated
    key owns its rows again), so eviction is always safe to call,
    including concurrently with in-flight acquisitions for other
    keys.

    Returns the number of entries evicted.
    """
    cap = (
        DEFAULT_MAX_KEYED_RESERVATIONS if max_pending_reclaims is None else max_pending_reclaims
    )
    vetoed = 0

    def _admit(name: str, prim: ConcurrencyReservation) -> bool:
        nonlocal vetoed
        if self._record_pending_reclaim(
            self._pending_reservation_reclaims, prim.schema, name, cap=cap
        ):
            return True
        vetoed += 1
        return False

    evicted = self._evict_idle_keyed(
        self._keyed_reservation_last_used,
        self._reservations,
        idle_for,
        "registry-evicted-idle-keyed-reservations",
        admit=_admit,
    )
    # The heal stamps ride the registration's lifecycle: an evicted
    # bucket's window must not survive into a future re-registration
    # of the same concrete name. The release-wake event rides the
    # same lifecycle (see the dict's comment above): a keyed bucket
    # that denies waiters creates an event on demand, and without
    # this pop every distinct denied key would leave one behind
    # forever - unbounded in the caller-controlled key space under
    # sustained denials. A waiter holding the popped object is never
    # stranded: its wait is bounded and falls through to the
    # availability re-probe, and a re-materialized bucket gets a
    # fresh event (no stale set() crosses the eviction boundary).
    for name in evicted:
        self._keyed_reservation_heal_attempted.pop(name, None)
        self._keyed_reservation_heal_failure_logged.pop(name, None)
        self._reservation_release_events.pop(name, None)
    update_keyed_reclaim_pending(self._pending_reclaim_total())
    if vetoed:
        logger.warning(
            "registry-keyed-reclaim-pending-cap-veto",
            vetoed=vetoed,
            cap=cap,
            pending=self._pending_reclaim_total(),
        )
    return len(evicted)

evict_idle_keyed_rate_limits

evict_idle_keyed_rate_limits(
    idle_for: timedelta,
    *,
    max_pending_reclaims: int | None = None,
) -> int

Drop registry entries for keyed rate limits idle at least idle_for.

Rate limits derived from a :class:KeyedRateLimitRef are registered lazily and never removed automatically, under high key cardinality (e.g. one token bucket per tenant over a long worker lifetime) this dict grows without bound. Each worker's 30-second sweep calls this automatically against its own registry (not leader-gated) with a 1-hour idle threshold; call directly for custom eviction windows.

Removes the in-memory registry entry and its acquire-recency tracking. The underlying Redis hash is deliberately NOT touched: per-key Redis memory is already self-bounding via the Lua script's EXPIRE TTL on the bucket's hash (see :meth:_resolve_rate_limit_name), that TTL governs Redis, not the PG rate_limit_buckets row the materialization path publishes, so a bucket resolved with a PG pool has its schema captured at resolution time (publish outcome irrelevant, the acquire path can preseed the row) and is recorded here for row reclamation: :meth:drain_pending_reservation_reclaims deletes the published row on the same sweep cadence, the exact shape of the reservation-side reclamation (steady state without it: one rate_limit_buckets row per key ever seen, unbounded in the caller-controlled key space). A key that is acquired again after eviction is re-registered on next use (idempotent, see :meth:_resolve_rate_limit_name), and the drain drops its name from the pending set without deleting anything (a re-activated key owns its row again), so eviction is always safe to call, including concurrently with in-flight acquisitions for other keys.

max_pending_reclaims caps the rate-limit pending-reclaim set (default: :data:taskq.constants.DEFAULT_MAX_KEYED_RESERVATIONS, the shared keyed-pending ceiling; the acquisition path's opportunistic eviction passes the settings-derived WorkerSettings.max_keyed_rate_limits, so the pending set is bounded by the same ceiling that bounds the tracked entries). At the cap the eviction is VETOED, the entry stays registered and re-scanned on the next sweep, so no structure grows unbounded and no published row is orphaned by an eviction that could not be recorded (the same fail-closed bound as the reservation side; the pending-depth gauge and the registry-keyed-reclaim-pending-cap-veto warning are the visible signals).

Exemption: memory fixed-quota buckets. A backend="memory" bucket with refill_per_second == 0 that has consumed part of its quota is NOT evicted (see :meth:TokenBucket.holds_consumed_quota): its token state lives only on the in-process instance, so eviction would silently reset the drained quota to full, and the next acquire would over-admit against a budget the tenant already spent. The exemption applies to both callers of this method (the per-worker sweep and the cap-pressure opportunistic eviction). Trade-off, deliberately chosen: an exempt bucket counts against settings.max_keyed_rate_limits until its quota returns to full (refund/reset) or the process restarts, so under sustained high-cardinality memory fixed-quota keys the cardinality cap can fill and deny NEW keys, the cap fails CLOSED with a warning rather than silently resetting quotas, which is the correct failure direction for a limiter.

PG fixed-quota buckets are NOT exempt: their quota state lives in the rate_limit_buckets row, and no reclamation path here can lose it, the drain's DELETE below and the maintenance leader's fleet sweep share the consumed-quota veto (_no_consumed_quota_sql) that keeps a partly-spent row, and a re-materialized bucket resumes from that row (the acquire path preseeds ON CONFLICT DO NOTHING and reads the surviving state). Evicting the registry entry is pure bookkeeping recycling, which is what keeps the max_keyed_rate_limits cap from filling with never-again-used fixed-quota keys and refusing every new key.

Returns the number of entries evicted.

Source code in src/taskq/ratelimit/registry.py

def evict_idle_keyed_rate_limits(
    self,
    idle_for: "timedelta",
    *,
    max_pending_reclaims: "int | None" = None,
) -> int:
    """Drop registry entries for keyed rate limits idle at least ``idle_for``.

    Rate limits derived from a :class:`KeyedRateLimitRef` are registered
    lazily and never removed automatically, under high key cardinality
    (e.g. one token bucket per tenant over a long worker lifetime) this
    dict grows without bound. Each worker's 30-second sweep calls this
    automatically against its own registry (not leader-gated) with a
    1-hour idle threshold; call directly for custom eviction windows.

    Removes the in-memory registry entry and its acquire-recency
    tracking. The underlying Redis hash is deliberately NOT touched:
    per-key Redis memory is already self-bounding via the Lua script's
    ``EXPIRE`` TTL on the bucket's hash (see
    :meth:`_resolve_rate_limit_name`), that TTL governs Redis, not
    the PG ``rate_limit_buckets`` row the materialization path
    publishes, so a bucket resolved with a PG pool has its schema
    captured at resolution time (publish outcome irrelevant, the
    acquire path can preseed the row) and is recorded here for row
    reclamation: :meth:`drain_pending_reservation_reclaims` deletes
    the published row on the same sweep cadence, the exact shape of
    the reservation-side reclamation (steady state without it: one
    ``rate_limit_buckets`` row per key ever seen, unbounded in the
    caller-controlled key space). A key that is acquired again after
    eviction is re-registered on next use (idempotent, see
    :meth:`_resolve_rate_limit_name`), and the drain drops its name
    from the pending set without deleting anything (a re-activated
    key owns its row again), so eviction is always safe to call,
    including concurrently with in-flight acquisitions for other
    keys.

    *max_pending_reclaims* caps the rate-limit pending-reclaim set
    (default: :data:`taskq.constants.DEFAULT_MAX_KEYED_RESERVATIONS`,
    the shared keyed-pending ceiling; the acquisition path's
    opportunistic eviction passes the settings-derived
    ``WorkerSettings.max_keyed_rate_limits``, so the pending set is
    bounded by the same ceiling that bounds the tracked entries). At
    the cap the eviction is VETOED, the entry stays registered and
    re-scanned on the next sweep, so no structure grows unbounded
    and no published row is orphaned by an eviction that could not
    be recorded (the same fail-closed bound as the reservation
    side; the pending-depth gauge and the
    ``registry-keyed-reclaim-pending-cap-veto`` warning are the
    visible signals).

    **Exemption: memory fixed-quota buckets.** A ``backend="memory"``
    bucket with ``refill_per_second == 0`` that has consumed part of its
    quota is NOT evicted (see :meth:`TokenBucket.holds_consumed_quota`):
    its token state lives only on the in-process instance, so eviction
    would silently reset the drained quota to full, and the next acquire
    would over-admit against a budget the tenant already spent. The
    exemption applies to both callers of this method (the per-worker
    sweep and the cap-pressure opportunistic eviction). Trade-off,
    deliberately chosen: an exempt bucket counts against
    ``settings.max_keyed_rate_limits`` until its quota returns to full
    (refund/reset) or the process restarts, so under sustained
    high-cardinality memory fixed-quota keys the cardinality cap can
    fill and deny NEW keys, the cap fails CLOSED with a warning rather
    than silently resetting quotas, which is the correct failure
    direction for a limiter.

    PG fixed-quota buckets are NOT exempt: their quota state lives in
    the ``rate_limit_buckets`` row, and no reclamation path here can
    lose it, the drain's DELETE below and the maintenance leader's
    fleet sweep share the consumed-quota veto
    (``_no_consumed_quota_sql``) that keeps a partly-spent row, and a
    re-materialized bucket resumes from that row (the acquire path
    preseeds ``ON CONFLICT DO NOTHING`` and reads the surviving state).
    Evicting the registry entry is pure bookkeeping recycling, which is
    what keeps the ``max_keyed_rate_limits`` cap from filling with
    never-again-used fixed-quota keys and refusing every new key.

    Returns the number of entries evicted.
    """
    cap = (
        DEFAULT_MAX_KEYED_RESERVATIONS if max_pending_reclaims is None else max_pending_reclaims
    )
    vetoed = 0

    def _admit(name: str, prim: TokenBucket | SlidingWindow) -> bool:
        nonlocal vetoed
        schema = self._keyed_rate_limit_row_schemas.get(name)
        if schema is None:
            # Never resolved with a PG pool on this worker: no publish,
            # no preseed (a pool-less acquire cannot touch PG), so this
            # registry created no row, nothing to reclaim from it. A
            # pool-bearing sibling that resolved the same key holds its
            # own capture for the concrete name and reclaims the row
            # through its own eviction+drain.
            return True
        if self._record_pending_reclaim(
            self._pending_rate_limit_reclaims, schema, name, cap=cap
        ):
            return True
        vetoed += 1
        return False

    evicted = self._evict_idle_keyed(
        self._keyed_rate_limit_last_used,
        self._rate_limits,
        idle_for,
        "registry-evicted-idle-keyed-rate-limits",
        preserve=_preserves_fixed_quota_state,
        admit=_admit,
    )
    # The publish-schema capture rides the registration's lifecycle:
    # an evicted bucket's schema must not survive into a future
    # re-registration of the same concrete name (a vetoed entry keeps
    # its capture, it is still registered and still owns its row).
    for name in evicted:
        self._keyed_rate_limit_row_schemas.pop(name, None)
    update_keyed_reclaim_pending(self._pending_reclaim_total())
    if vetoed:
        logger.warning(
            "registry-keyed-reclaim-pending-cap-veto",
            vetoed=vetoed,
            cap=cap,
            pending=self._pending_reclaim_total(),
        )
    return len(evicted)

drain_pending_reservation_reclaims async

drain_pending_reservation_reclaims(
    pool: Pool,
    *,
    batch_names: int = _DEFAULT_RECLAIM_BATCH_NAMES,
    acquire_timeout: float | None = None,
) -> int

Delete the PG rows of evicted keyed buckets.

Two row kinds, one bounded pass per schema per call: reservation_slots rows of evicted keyed reservations, and the published rate_limit_buckets rows of evicted keyed rate limits. For each kind and schema, at most batch_names pending bucket names and ONE batched DELETE. The slice is the FRONT of the per-schema insertion-ordered pending set. Driven by the per-worker sweep loop on the sweep cadence, immediately after the keyed evictions that feed it; a no-op (no connection acquired) when nothing is pending on either side.

Semantics, what stays pending and why:

• A name that re-registered since eviction is dropped from its pending set WITHOUT a statement: a re-activated key owns its rows again.
• The reservation DELETE removes only free or lease-expired rows , a slot still held by a live lease survives (the over-admission invariant). A HELD slot is therefore the only reason a reservation name stays pending after its slice ran: the existence probe names the buckets whose rows survived the DELETE, and those names rotate to the BACK of the queue, the lease expires, the holder's release or the lock-expiry sweep frees the row, and a later drain (at most one full rotation away) deletes it. A name with no rows left, fully deleted this tick, or never materialized at all, leaves the pending set.
• The rate-limit DELETE needs no HOLDER guard and no survivor probe, a bucket row has no holder or lease, so nothing survives the statement on that axis and every sliced name leaves the pending set in one pass, the queue drains FIFO with no survivors to rotate. It DOES carry a guard, and that guard is critical: the consumed-quota veto (_no_consumed_quota_sql, the same predicate the fleet sweep applies) refuses to delete a row whose fixed quota is partly spent, so an evicted spent bucket's row survives the drain and a re-materialized key resumes its spent state instead of resetting to full capacity (pinned by tests/test_keyed_fixed_quota_eviction.py). A vetoed name still leaves the pending set, the row stays, and the fleet sweep is the backstop once the quota is no longer consumed. A live bucket never reaches the statement at all: only evicted keyed names are ever recorded, and the re-registered check above drops a re-activated key first.
• A schema key whose pending set empties (every name re-registered, or every row reclaimed) is popped after the pass, so a later drain with nothing pending acquires no connection at all.

Instrumentation, on the failure path as much as the success path: duration in a finally (a timed-out drain still leaves a duration sample), deleted rows only as RETURNING-confirmed (a timed-out statement cannot masquerade as an empty drain), the drain-failure counter on exception (a STORAGE signal, a failing drain strands rows; the pending-depth gauge shows the backlog forming), and the pending-depth gauge after each drain. Raises on failure, callers guard (the sweep loop warns and continues on the next tick).

acquire_timeout bounds the pool wait, for callers on a loop that must never block indefinitely on an exhausted pool (the sweep loop passes its dispatcher command timeout); the default matches the registry's other pool-bearing methods (unbounded).

Returns the number of rows deleted.

Source code in src/taskq/ratelimit/registry.py

async def drain_pending_reservation_reclaims(
    self,
    pool: "asyncpg.Pool",
    *,
    batch_names: int = _DEFAULT_RECLAIM_BATCH_NAMES,
    acquire_timeout: "float | None" = None,
) -> int:
    """Delete the PG rows of evicted keyed buckets.

    Two row kinds, one bounded pass per schema per call:
    ``reservation_slots`` rows of evicted keyed reservations, and the
    published ``rate_limit_buckets`` rows of evicted keyed rate
    limits. For each kind and schema, at most *batch_names* pending
    bucket names and ONE batched DELETE. The slice is the FRONT of
    the per-schema insertion-ordered pending set. Driven by the
    per-worker sweep loop on the sweep cadence, immediately after
    the keyed evictions that feed it; a no-op (no connection
    acquired) when nothing is pending on either side.

    Semantics, what stays pending and why:

    - A name that re-registered since eviction is dropped from its
      pending set WITHOUT a statement: a re-activated key owns its
      rows again.
    - The reservation DELETE removes only free or lease-expired rows
     , a slot still held by a live lease survives (the
      over-admission invariant). A HELD slot is therefore the only
      reason a reservation name stays pending after its slice ran:
      the existence probe names the buckets whose rows survived the
      DELETE, and those names rotate to the BACK of the queue, the
      lease expires, the holder's release or the lock-expiry sweep
      frees the row, and a later drain (at most one full rotation
      away) deletes it. A name with no rows left, fully deleted
      this tick, or never materialized at all, leaves the pending
      set.
    - The rate-limit DELETE needs no HOLDER guard and no survivor
      probe, a bucket row has no holder or lease, so nothing
      survives the statement on that axis and every sliced name
      leaves the pending set in one pass, the queue drains FIFO
      with no survivors to rotate. It DOES carry a guard, and that
      guard is critical: the consumed-quota veto
      (``_no_consumed_quota_sql``, the same predicate the fleet
      sweep applies) refuses to delete a row whose fixed quota is
      partly spent, so an evicted spent bucket's row survives the
      drain and a re-materialized key resumes its spent state
      instead of resetting to full capacity (pinned by
      tests/test_keyed_fixed_quota_eviction.py). A vetoed name
      still leaves the pending set, the row stays, and the
      fleet sweep is the backstop once the quota is no longer
      consumed. A live bucket never reaches the statement at all:
      only evicted keyed names are ever recorded, and the
      re-registered check above drops a re-activated key first.
    - A schema key whose pending set empties (every name
      re-registered, or every row reclaimed) is popped after the
      pass, so a later drain with nothing pending acquires no
      connection at all.

    Instrumentation, on the failure path as much as the success
    path: duration in a ``finally`` (a timed-out drain still leaves
    a duration sample), deleted rows only as RETURNING-confirmed
    (a timed-out statement cannot masquerade as an empty drain),
    the drain-failure counter on exception (a STORAGE signal, a
    failing drain strands rows; the pending-depth gauge shows the
    backlog forming), and the pending-depth gauge after each drain.
    Raises on failure, callers guard (the sweep loop warns and
    continues on the next tick).

    *acquire_timeout* bounds the pool wait, for callers on a loop
    that must never block indefinitely on an exhausted pool (the
    sweep loop passes its dispatcher command timeout); the default
    matches the registry's other pool-bearing methods (unbounded).

    Returns the number of rows deleted.
    """
    if not self._pending_reservation_reclaims and not self._pending_rate_limit_reclaims:
        return 0
    start = monotonic()
    total_deleted = 0
    try:
        async with pool.acquire(timeout=acquire_timeout) as conn:
            for schema in list(self._pending_reservation_reclaims):
                require_schema(schema)
                pending = self._pending_reservation_reclaims.get(schema)
                if not pending:
                    self._drop_pending_schema_if_empty(
                        self._pending_reservation_reclaims, schema
                    )
                    continue
                # Round-robin: the front batch (insertion order) is
                # served first; survivors re-enter at the back below.
                candidates = list(pending)[:batch_names]
                slice_names: list[str] = []
                for name in candidates:
                    if name in self._reservations or name in self._keyed_reservation_last_used:
                        # Re-registered since eviction: the bucket is
                        # live again and owns its rows, drop the
                        # pending entry, touch nothing in PG.
                        pending.pop(name, None)
                    else:
                        slice_names.append(name)
                if not slice_names:
                    self._drop_pending_schema_if_empty(
                        self._pending_reservation_reclaims, schema
                    )
                    continue
                deleted_rows = await conn.fetch(
                    _RECLAIM_SLICE_DELETE_SQL_TEMPLATE.format(schema=schema),
                    slice_names,
                )
                total_deleted += len(deleted_rows)
                surviving: set[str] = {
                    row["bucket_name"]
                    for row in await conn.fetch(
                        _RECLAIM_SLICE_EXISTING_SQL_TEMPLATE.format(schema=schema),
                        slice_names,
                    )
                }
                # Tolerant pops: the DELETE/probe awaits can interleave
                # with another drain pass on the same registry, a
                # missing key means the name was already removed, and
                # re-removing it must not turn into a spurious drain
                # failure.
                for name in slice_names:
                    pending.pop(name, None)
                for name in surviving:
                    # A survivor (rows still held) re-enters at the
                    # BACK: its next chance comes after every other
                    # pending bucket's, never before them.
                    pending[name] = None
                self._drop_pending_schema_if_empty(self._pending_reservation_reclaims, schema)
            for schema in list(self._pending_rate_limit_reclaims):
                require_schema(schema)
                pending = self._pending_rate_limit_reclaims.get(schema)
                if not pending:
                    self._drop_pending_schema_if_empty(
                        self._pending_rate_limit_reclaims, schema
                    )
                    continue
                # Front batch, same as the reservation pass; with no
                # survivors there is nothing to rotate, the names
                # that fit the slice are served in insertion order.
                candidates = list(pending)[:batch_names]
                slice_names: list[str] = []
                for name in candidates:
                    if name in self._rate_limits or name in self._keyed_rate_limit_last_used:
                        # Re-registered since eviction: the bucket is
                        # live again and owns its published row, drop
                        # the pending entry, touch nothing in PG.
                        pending.pop(name, None)
                    else:
                        slice_names.append(name)
                if not slice_names:
                    self._drop_pending_schema_if_empty(
                        self._pending_rate_limit_reclaims, schema
                    )
                    continue
                deleted_rows = await conn.fetch(
                    _RECLAIM_RATE_LIMIT_SLICE_DELETE_SQL_TEMPLATE.format(schema=schema),
                    slice_names,
                )
                total_deleted += len(deleted_rows)
                for name in slice_names:
                    pending.pop(name, None)
                self._drop_pending_schema_if_empty(self._pending_rate_limit_reclaims, schema)
    except Exception as exc:
        record_reservation_reclaim_drain_failure(type(exc).__name__)
        raise
    finally:
        record_reservation_reclaim_drain_duration(monotonic() - start)
        record_reservation_reclaim_drain_rows(total_deleted)
        update_keyed_reclaim_pending(self._pending_reclaim_total())
    return total_deleted

clear

clear() -> None

Reset ALL mutable registry state, a test aid, NOT safe while running.

Clears the nine dicts (_rate_limits, _reservations, _keyed_reservation_last_used, _keyed_rate_limit_last_used, _keyed_rate_limit_row_schemas, _pending_reservation_reclaims, _pending_rate_limit_reclaims, _keyed_reservation_heal_attempted, _keyed_reservation_heal_failure_logged) AND resets the two opportunistic-eviction scan timestamps (_keyed_reservation_last_eviction_scan / _keyed_rate_limit_last_eviction_scan) to float("-inf"). Omitting the timestamps would leave the opportunistic-eviction throttle stamped, silently suppressing scans for up to _OPPORTUNISTIC_EVICT_MIN_INTERVAL (30 s) in the next test; omitting a pending-reclaim set, a publish-schema capture, or a heal stamp would leak one test's evictions into the next.

Not safe to call while a worker is running, concurrent dispatch / sweep iteration over the dicts would observe inconsistent state. Use for per-test isolation only.

Like the eviction methods, this resets IN-PROCESS bookkeeping only, it does NOT touch Redis bucket hashes or Postgres reservation_slots / rate_limit_buckets rows; backend state persists and will be observed on next acquire.

Source code in src/taskq/ratelimit/registry.py

def clear(self) -> None:
    """Reset ALL mutable registry state, a test aid, NOT safe while running.

    Clears the nine dicts (``_rate_limits``, ``_reservations``,
    ``_keyed_reservation_last_used``, ``_keyed_rate_limit_last_used``,
    ``_keyed_rate_limit_row_schemas``,
    ``_pending_reservation_reclaims``,
    ``_pending_rate_limit_reclaims``,
    ``_keyed_reservation_heal_attempted``,
    ``_keyed_reservation_heal_failure_logged``)
    AND resets the two opportunistic-eviction scan timestamps
    (``_keyed_reservation_last_eviction_scan`` /
    ``_keyed_rate_limit_last_eviction_scan``) to ``float("-inf")``.
    Omitting the timestamps would leave the opportunistic-eviction
    throttle stamped, silently suppressing scans for up to
    ``_OPPORTUNISTIC_EVICT_MIN_INTERVAL`` (30 s) in the next test;
    omitting a pending-reclaim set, a publish-schema capture, or a
    heal stamp would leak one test's evictions into the next.

    **Not safe to call while a worker is running**, concurrent
    dispatch / sweep iteration over the dicts would observe
    inconsistent state. Use for per-test isolation only.

    Like the eviction methods, this resets IN-PROCESS bookkeeping
    only, it does NOT touch Redis bucket hashes or Postgres
    ``reservation_slots`` / ``rate_limit_buckets`` rows; backend
    state persists and will be observed on next acquire.
    """
    self._rate_limits.clear()
    self._reservations.clear()
    self._keyed_reservation_last_used.clear()
    self._keyed_rate_limit_last_used.clear()
    self._keyed_rate_limit_row_schemas.clear()
    self._pending_reservation_reclaims.clear()
    self._pending_rate_limit_reclaims.clear()
    self._keyed_reservation_heal_attempted.clear()
    self._keyed_reservation_heal_failure_logged.clear()
    self._keyed_reservation_last_eviction_scan = float("-inf")
    self._keyed_rate_limit_last_eviction_scan = float("-inf")
    update_keyed_reclaim_pending(0)

ConcurrencyReservation

ConcurrencyReservation(
    name: str,
    slots: int,
    lease: timedelta | float,
    lock_lease: timedelta | None = None,
    *,
    clock: Clock | None = None,
    schema: str = "taskq",
    keyed: bool = False,
)

Concurrency reservation using pre-allocated slot rows.

Raises :class:ValueError if slots < 1 or lease <= 0. Raises :class:ReservationUnavailable when no slot is available.

keyed (default False) marks a reservation materialised from a :class:~taskq.ratelimit.refs.KeyedReservationRef: its slot rows carry the fleet-reclaimable keyed flag and a last_used_at stamp (refreshed by this class's acquire/release/ensure statements), so the maintenance leader's sweep_idle_keyed_rows can reclaim them after the OWNING PROCESS dies, the in-process registry bookkeeping that would otherwise name them dies with it. A statically declared reservation keeps the default: its rows are never fleet-reclaimable (there is no keyed lifecycle, no acquire-path heal, and no re-materialisation, a sweep that deleted them would leave a permanently denying limiter).

Source code in src/taskq/ratelimit/reservation.py

def __init__(
    self,
    name: str,
    slots: int,
    lease: timedelta | float,
    lock_lease: timedelta | None = None,
    *,
    clock: Clock | None = None,
    schema: str = "taskq",
    keyed: bool = False,
) -> None:
    if slots < 1:
        raise ValueError(f"slots must be >= 1, got {slots}")

    if isinstance(lease, timedelta):
        if lease <= timedelta(0):
            raise ValueError(f"lease must be > 0, got {lease!r}")
        lease_td = lease
    else:
        if lease <= 0:
            raise ValueError(f"lease must be > 0, got {lease}")
        lease_td = timedelta(seconds=lease)

    self._name = name
    self._slots = slots
    self._lease = lease_td
    self._lock_lease = lock_lease
    self._schema = schema
    self._keyed = keyed

    if lock_lease is not None and lease_td < lock_lease:
        logger.warning(
            "reservation-lease-shorter-than-lock-lease",
            bucket_name=name,
            lease_seconds=lease_td.total_seconds(),
            lock_lease_seconds=lock_lease.total_seconds(),
        )

    _validate_schema(schema)
    self._ensure_sql = _ENSURE_SLOTS_SQL_TEMPLATE.format(schema=schema)
    self._acquire_sql = _ACQUIRE_SQL_TEMPLATE.format(schema=schema)
    self._release_sql = _RELEASE_SQL_TEMPLATE.format(schema=schema)
    self._release_fenced_sql = _RELEASE_FENCED_SQL_TEMPLATE.format(schema=schema)

    if clock is not None:
        self._table: _InMemorySlotTable | None = _InMemorySlotTable(clock)
    else:
        self._table = None

__slots__ class-attribute instance-attribute

__slots__ = (
    "_acquire_sql",
    "_ensure_sql",
    "_keyed",
    "_lease",
    "_lock_lease",
    "_name",
    "_release_fenced_sql",
    "_release_sql",
    "_schema",
    "_slots",
    "_table",
)

schema property

schema: str

The PG schema this reservation's slot table lives in.

Workers filter registry-global reservations by their own schema at startup (a process-global registry may carry reservations declared for other schemas/databases, touching those would write into the wrong schema or fail noisily).

keyed property

keyed: bool

Whether this reservation was keyed-materialised.

Carried on the instance so the row-writing statements can stamp the fleet-reclaimable mark without every call site re-deriving it, the constructor is the single point that knows the reservation's lifecycle origin.

name property

name: str

slots property

slots: int

lease property

lease: timedelta

bucket_name property

bucket_name: str

table property

table: _InMemorySlotTable

The in-memory slot table (requires clock at construction).

ensure_slots async

ensure_slots(
    pool: Pool, *, timeout: float | None = None
) -> None

Idempotent pre-allocation of slot rows.

timeout bounds the pool acquire (forwarded to asyncpg.Pool.acquire): None, the historical default, waits unboundedly under pool starvation, which is how a bootstrap startup could hang forever on a contended dispatcher pool. Callers in bounded contexts (the bootstrap's ensure-slots loop passes its dispatcher_command_timeout budget) must pass it.

Inserts the bucket's full slot row set with this reservation's fleet-reclaimable mark and a fresh last_used_at; the conflict arm (rows already present, re-materialisation over survivors, or another process's earlier materialisation) refreshes last_used_at on every conflicting row and converges the keyed mark toward immortality without ever touching the holder/lease columns (held state is untouched):

• a keyed materialisation (or its heal) re-marks rows that are already keyed and refreshes their staleness, re-materialised survivors restart the horizon;
• a keyed materialisation over rows marked keyed=false leaves them false: that mark is a STATIC claim (born false, possibly declared in a process whose registry this worker's process-local collision guard cannot see), and a static reservation has no acquire-path heal, swept rows would deny forever;
• a later static declaration of the same name (the bootstrap's startup ensure) retires keyed rows to keyed=false , never-sweep again, the immortality direction a live static declaration owns.

Source code in src/taskq/ratelimit/reservation.py

async def ensure_slots(
    self,
    pool: "asyncpg.Pool",
    *,
    timeout: float | None = None,  # noqa: ASYNC109  # Why: mirrors asyncpg.Pool.acquire(timeout=...); the bound IS the forwarded driver kwarg, not an enclosing asyncio.timeout scope.
) -> None:
    """Idempotent pre-allocation of slot rows.

    *timeout* bounds the pool acquire (forwarded to
    ``asyncpg.Pool.acquire``): ``None``, the historical default, waits
    unboundedly under pool starvation, which is how a bootstrap startup
    could hang forever on a contended dispatcher pool. Callers in
    bounded contexts (the bootstrap's ensure-slots loop passes its
    ``dispatcher_command_timeout`` budget) must pass it.

    Inserts the bucket's full slot row set with this reservation's
    fleet-reclaimable mark and a fresh ``last_used_at``; the conflict
    arm (rows already present, re-materialisation over survivors,
    or another process's earlier materialisation) refreshes
    ``last_used_at`` on every conflicting row and converges the
    ``keyed`` mark toward immortality without ever touching the
    holder/lease columns (held state is untouched):

    - a keyed materialisation (or its heal) re-marks rows that are
      already keyed and refreshes their staleness, re-materialised
      survivors restart the horizon;
    - a keyed materialisation over rows marked ``keyed=false``
      leaves them false: that mark is a STATIC claim (born false,
      possibly declared in a process whose registry this worker's
      process-local collision guard cannot see), and a static
      reservation has no acquire-path heal, swept rows would deny
      forever;
    - a later static declaration of the same name (the bootstrap's
      startup ensure) retires keyed rows to ``keyed=false`` ,
      never-sweep again, the immortality direction a live static
      declaration owns.
    """
    async with pool.acquire(timeout=timeout) as conn:
        await conn.execute(self._ensure_sql, self._name, self._slots, self._keyed)

slot_rows_exist async

slot_rows_exist(
    pool: Pool, *, timeout: float | None = None
) -> bool

Whether any reservation_slots row exists for this bucket.

The acquire-path heal for keyed reservations uses this to distinguish a bucket whose rows were deleted out from under it (zero rows, re-materialise via :meth:ensure_slots) from ordinary contention (rows present, all held, deny). A read-only existence probe: it never writes. timeout bounds the pool acquire exactly as:meth:ensure_slots' does.

Source code in src/taskq/ratelimit/reservation.py

async def slot_rows_exist(
    self,
    pool: "asyncpg.Pool",
    *,
    timeout: float | None = None,  # noqa: ASYNC109  # Why: mirrors asyncpg.Pool.acquire(timeout=...); same forward contract as ensure_slots.
) -> bool:
    """Whether any ``reservation_slots`` row exists for this bucket.

    The acquire-path heal for keyed reservations uses this to
    distinguish a bucket whose rows were deleted out from under it
    (zero rows, re-materialise via :meth:`ensure_slots`) from
    ordinary contention (rows present, all held, deny). A read-only
    existence probe: it never writes. *timeout* bounds the pool
    acquire exactly as:meth:`ensure_slots`' does.
    """
    async with pool.acquire(timeout=timeout) as conn:
        return (
            await conn.fetchval(
                _SYNC_EXISTING_SQL_TEMPLATE.format(schema=self._schema),
                self._name,
            )
            is not None
        )

acquire async

acquire(
    job_id: UUID,
    worker_id: UUID,
    pool: Pool | None = None,
    *,
    timeout: float | None = None,
) -> SlotLease

Acquire a slot. Returns the acquired slot_index.

The return value is a :class:SlotLease, an int that also carries the fence :meth:release needs to tell this lease apart from an earlier, expired one on the same slot. Pass it back to :meth:release (which is what ReservationHandle does) to get that protection.

When pool is None, the in-memory table (clock= at construction) is used. Raises :class:ReservationUnavailable when no slot is available.

timeout bounds the pool acquire (forwarded to asyncpg.Pool.acquire; None, the default, waits unboundedly under pool starvation, the shape). Callers that already run under their own budget wrapper can leave it unset; callers with a budget in hand should pass it.

Source code in src/taskq/ratelimit/reservation.py

async def acquire(
    self,
    job_id: UUID,
    worker_id: UUID,
    pool: "asyncpg.Pool | None" = None,
    *,
    timeout: float | None = None,  # noqa: ASYNC109  # Why: mirrors asyncpg.Pool.acquire(timeout=...); same forward contract as ensure_slots.
) -> SlotLease:
    """Acquire a slot. Returns the acquired ``slot_index``.

    The return value is a :class:`SlotLease`, an ``int`` that also
    carries the fence :meth:`release` needs to tell this lease apart from
    an earlier, expired one on the same slot.  Pass it back to
    :meth:`release` (which is what ``ReservationHandle`` does) to get that
    protection.

    When *pool* is ``None``, the in-memory table (``clock=`` at
    construction) is used.  Raises :class:`ReservationUnavailable` when
    no slot is available.

    *timeout* bounds the pool acquire (forwarded to
    ``asyncpg.Pool.acquire``; ``None``, the default, waits
    unboundedly under pool starvation, the shape). Callers that
    already run under their own budget wrapper can leave it unset;
    callers with a budget in hand should pass it.
    """
    if pool is None:
        if self._table is None:
            raise RuntimeError(
                "pool=None but no in-memory table, pass clock= at "
                "construction for in-memory acquire, or supply a PG pool"
            )
        # The keyed mark rides the materialisation exactly as the PG
        # ensure_slots statement stamps it (the fleet-reclaimable mark
        # is born on the rows the keyed lifecycle creates).
        self._table.ensure_slots(self._name, self._slots, keyed=self._keyed)
        slot_index = self._table.acquire(
            self._name,
            job_id,
            worker_id,
            self._lease,
        )
        logger.debug(
            "reservation-acquired",
            bucket_name=self._name,
            slot_index=slot_index,
            job_id=str(job_id),
            worker_id=worker_id,
            backend="memory",
        )
        return slot_index

    # No explicit transaction: the acquire is one data-modifying-CTE
    # statement, atomic on its own, its FOR UPDATE SKIP LOCKED row
    # lock lives exactly as long as the statement, so BEGIN/COMMIT
    # would be two extra round trips per reserved job.
    async with pool.acquire(timeout=timeout) as conn:
        row = await conn.fetchrow(
            self._acquire_sql,
            self._name,
            job_id,
            worker_id,
            self._lease.total_seconds(),
        )

    if row is None or row["slot_index"] is None:
        # Denial branch of the acquire statement: the acquired fields
        # are NULL and the hint column carries the earliest held
        # lease's remaining seconds (NULL when no live-held row was
        # readable). The margin is added client-side so both arms ,
        # this and the in-memory twin, share one definition of it.
        hint_seconds = row["retry_after_seconds"] if row is not None else None
        retry_after = (
            timedelta(seconds=float(hint_seconds)) + RESERVATION_RETRY_HINT_MARGIN
            if hint_seconds is not None
            else DEFAULT_RESERVATION_BACKOFF
        )
        logger.info(
            "reservation-unavailable",
            bucket_name=self._name,
            retry_after_seconds=retry_after.total_seconds(),
        )
        raise ReservationUnavailable(self._name, retry_after)

    slot_lease = SlotLease(row["slot_index"], row["acquired_at"])
    logger.debug(
        "reservation-acquired",
        bucket_name=self._name,
        slot_index=int(slot_lease),
        job_id=str(job_id),
        worker_id=worker_id,
    )
    return slot_lease

release async

release(
    slot_index: int,
    worker_id: UUID,
    pool: Pool | None = None,
) -> None

Release slot. No-op if worker_id mismatch.

Also a no-op when slot_index is a :class:SlotLease whose fence no longer matches the slot row, i.e. the caller is releasing a lease that has since expired and been re-acquired. Without that gate a zombie attempt (heartbeats stalled, job reclaimed and redispatched, original coroutine never cancelled) frees the slot its own LIVE successor holds, and the bucket silently exceeds max_concurrent: the job id cannot discriminate, because a retry reuses the same job row, and neither can the worker id, because the redispatch commonly lands on the same worker. A plain int releases unfenced, as before.

When pool is None, the in-memory table is used.

Source code in src/taskq/ratelimit/reservation.py

async def release(
    self,
    slot_index: int,
    worker_id: UUID,
    pool: "asyncpg.Pool | None" = None,
) -> None:
    """Release slot. No-op if ``worker_id`` mismatch.

    Also a no-op when *slot_index* is a :class:`SlotLease` whose fence no
    longer matches the slot row, i.e. the caller is releasing a lease that
    has since expired and been re-acquired.  Without that gate a zombie
    attempt (heartbeats stalled, job reclaimed and redispatched, original
    coroutine never cancelled) frees the slot its own LIVE successor
    holds, and the bucket silently exceeds ``max_concurrent``: the job id
    cannot discriminate, because a retry reuses the same job row, and
    neither can the worker id, because the redispatch commonly lands on
    the same worker.  A plain ``int`` releases unfenced, as before.

    When *pool* is ``None``, the in-memory table is used.
    """
    if pool is None:
        if self._table is None:
            raise RuntimeError(
                "pool=None but no in-memory table, pass clock= at "
                "construction for in-memory release, or supply a PG pool"
            )
        self._table.release(self._name, slot_index, worker_id)
        logger.debug(
            "reservation-released",
            bucket_name=self._name,
            slot_index=slot_index,
            worker_id=worker_id,
            backend="memory",
        )
        return

    fence = _fence_of(slot_index)
    async with pool.acquire() as conn:
        if fence is None:
            await conn.execute(
                self._release_sql,
                self._name,
                slot_index,
                worker_id,
            )
        else:
            await conn.execute(
                self._release_fenced_sql,
                self._name,
                int(slot_index),
                worker_id,
                fence,
            )
    logger.debug(
        "reservation-released",
        bucket_name=self._name,
        slot_index=slot_index,
        worker_id=worker_id,
    )

peek async

peek(pool: Pool | None = None) -> dict[str, object]

Return {"free_count": int, "total_slots": int} for the bucket.

When pool is None, the in-memory table is used.

Source code in src/taskq/ratelimit/reservation.py

async def peek(self, pool: "asyncpg.Pool | None" = None) -> dict[str, object]:
    """Return ``{"free_count": int, "total_slots": int}`` for the bucket.

    When *pool* is ``None``, the in-memory table is used.
    """
    if pool is None:
        if self._table is None:
            raise RuntimeError(
                "pool=None but no in-memory table, pass clock= at "
                "construction for in-memory peek, or supply a PG pool"
            )
        free, held = self._table.peek_slots(self._name)
        return {"free_count": free, "total_slots": self._slots, "held_count": held}

    require_schema(self._schema)
    schema = self._schema

    # Schema-name interpolation ; schema_name is
    # pre-validated against _IDENT_RE at WorkerSettings load time.
    peek_sql = (
        f"SELECT count(*) FILTER (WHERE job_id IS NULL OR lease_expires_at < clock_timestamp()) AS free_count, "  # noqa: S608
        f"count(*) AS total_slots, "
        f"count(*) FILTER (WHERE job_id IS NOT NULL AND lease_expires_at >= clock_timestamp()) AS held_count "
        f'FROM "{schema}".reservation_slots WHERE bucket_name = $1'
    )
    async with pool.acquire() as conn:
        row = await conn.fetchrow(peek_sql, self._name)

    if row is None:
        return {"free_count": self._slots, "total_slots": self._slots, "held_count": 0}

    return {
        "free_count": int(row["free_count"]),
        "total_slots": int(row["total_slots"]),
        "held_count": int(row["held_count"]),
    }

SyncResult dataclass

SyncResult(
    inserted: list[tuple[str, int]],
    deleted: list[tuple[str, int]],
    skipped_held: list[tuple[str, int]],
)

Result of a sync_slots call.

inserted instance-attribute

inserted: list[tuple[str, int]]

deleted instance-attribute

deleted: list[tuple[str, int]]

skipped_held instance-attribute

skipped_held: list[tuple[str, int]]

SlidingWindow

SlidingWindow(
    name: str,
    limit: int,
    window: timedelta,
    backend: Literal[
        "redis", "postgres", "memory"
    ] = "redis",
    style: SlidingWindowStyle = "log",
    ttl: timedelta | None = None,
)

Sliding-window rate limiter with pluggable backends.

Raises :class:ValueError if limit < 1, window <= timedelta(0), or style is not "log" or "gcra".

Source code in src/taskq/ratelimit/sliding_window.py

def __init__(
    self,
    name: str,
    limit: int,
    window: timedelta,
    backend: Literal["redis", "postgres", "memory"] = "redis",
    style: SlidingWindowStyle = "log",
    ttl: timedelta | None = None,
) -> None:
    if limit < 1:
        raise ValueError(f"limit must be >= 1, got {limit}")
    if window <= timedelta(0):
        raise ValueError(f"window must be > timedelta(0), got {window}")
    if style not in _VALID_STYLES:
        raise ValueError(f"style must be 'log' or 'gcra', got {style!r}")

    self._name = name
    self._limit = limit
    self._window = window
    self._backend: RateLimitBackend = backend
    self._style: SlidingWindowStyle = style

    if ttl is not None:
        self._ttl = ttl
    elif style == "gcra":
        self._ttl = window + _TTL_SAFETY_MARGIN
    else:
        self._ttl = 2 * window + _TTL_SAFETY_MARGIN

    self._mem_log: _InMemorySlidingWindowLog | None = None
    self._mem_gcra: _InMemorySlidingWindowGCRA | None = None
    if backend == "memory":
        window_ms = int(window.total_seconds() * 1000)
        if style == "log":
            self._mem_log = _InMemorySlidingWindowLog(name, limit, window_ms)
        else:
            self._mem_gcra = _InMemorySlidingWindowGCRA(name, limit, window_ms)

    self._redis_log_script: AsyncScript | None = None
    self._redis_log_script_client: redis_async.Redis | None = None
    self._redis_gcra_script: AsyncScript | None = None
    self._redis_gcra_script_client: redis_async.Redis | None = None
    self._redis_gcra_refund_script: AsyncScript | None = None
    self._redis_gcra_refund_script_client: redis_async.Redis | None = None
    self._script_lock: asyncio.Lock = asyncio.Lock()

__slots__ class-attribute instance-attribute

__slots__ = (
    "_backend",
    "_limit",
    "_mem_gcra",
    "_mem_log",
    "_name",
    "_redis_gcra_refund_script",
    "_redis_gcra_refund_script_client",
    "_redis_gcra_script",
    "_redis_gcra_script_client",
    "_redis_log_script",
    "_redis_log_script_client",
    "_script_lock",
    "_style",
    "_ttl",
    "_window",
)

name property

name: str

limit property

limit: int

window property

window: timedelta

backend property

backend: Literal['redis', 'postgres', 'memory']

style property

style: SlidingWindowStyle

ttl property

ttl: timedelta | None

acquire async

acquire(
    *,
    redis_client: Redis | None = None,
    pg_pool: Pool | None = None,
    clock: Clock | None = None,
    settings: WorkerSettings | None = None,
) -> RateLimitDecision

Acquire one admission slot.

The shared admission state is measured in the data store's own clock domain: PG paths use clock_timestamp() and Redis paths use Redis TIME inside their scripts, so callers on nodes with divergent Python clocks are all measured against the same window. The injected clock drives the memory backend only (its single domain) and remains part of the public call shape, the unified TokenBucket contract.

Source code in src/taskq/ratelimit/sliding_window.py

async def acquire(
    self,
    *,
    redis_client: "redis_async.Redis | None" = None,
    pg_pool: "asyncpg.Pool | None" = None,
    clock: Clock | None = None,
    settings: "WorkerSettings | None" = None,
) -> RateLimitDecision:
    """Acquire one admission slot.

    The shared admission state is measured in the data store's own
    clock domain: PG paths use ``clock_timestamp()`` and Redis paths
    use Redis ``TIME`` inside their scripts, so callers on nodes with
    divergent Python clocks are all measured against the same window.
    The injected *clock* drives the memory backend only (its single
    domain) and remains part of the public call shape, the unified
    TokenBucket contract.
    """
    request_id: UUID | None = new_uuid() if self._style == "log" else None

    match (self._backend, self._style):
        case ("memory", "log"):
            if clock is None:
                raise RuntimeError("clock not injected for memory backend")
            now_ms = int(clock.now().timestamp() * 1000)
            return await self._acquire_memory_log(now_ms, request_id)
        case ("memory", "gcra"):
            if clock is None:
                raise RuntimeError("clock not injected for memory backend")
            now_ms = int(clock.now().timestamp() * 1000)
            return await self._acquire_memory_gcra(now_ms)
        case ("redis", "log"):
            return await _acquire_redis_log_wrapped(
                self, request_id, redis_client, pg_pool, settings
            )
        case ("redis", "gcra"):
            return await _acquire_redis_gcra_wrapped(self, redis_client, pg_pool, settings)
        case ("postgres", "log"):
            return await _acquire_pg_log(self, pg_pool, settings, request_id)
        case ("postgres", "gcra"):
            return await _acquire_pg_gcra(self, pg_pool, settings)
        case _:
            assert_never((self._backend, self._style))

refund async

refund(
    decision: RateLimitDecision,
    *,
    count: float = 1.0,
    redis_client: Redis | None = None,
    pg_pool: Pool | None = None,
    clock: Clock | None = None,
    settings: WorkerSettings | None = None,
) -> None

Source code in src/taskq/ratelimit/sliding_window.py

async def refund(
    self,
    decision: RateLimitDecision,
    *,
    count: float = 1.0,
    redis_client: "redis_async.Redis | None" = None,
    pg_pool: "asyncpg.Pool | None" = None,
    clock: Clock | None = None,
    settings: "WorkerSettings | None" = None,
) -> None:
    # Why decision.backend and not self._backend: with backend="redis" and
    # rate_limit_pg_fallback_enabled (the default), an acquire during a
    # Redis outage falls through to Postgres and records the admission
    # THERE. The decision says which store actually holds it; the
    # primitive's configuration only says where it prefers to go.
    # Dispatching on the latter left the Postgres window permanently
    # holding an admission that was released, and, for GCRA, whose
    # previous_state differs per backend, raised KeyError out of the
    # release path outright. Kept identical to TokenBucket.refund, which
    # dispatches the same way for the same reason.
    #
    # Why clock and count are unused: a sliding-window refund removes
    # the one logged entry named by *decision*'s request id, so there
    # is nothing to re-add and no timestamp to read (TokenBucket.refund
    # does use count). Both stay in the signature because
    # RateLimitRegistry dispatches refund/peek/reset polymorphically
    # over both primitives with one fixed keyword block, dropping
    # either here would raise TypeError there.
    match (decision.backend, self._style):
        case ("redis", "log"):
            await _refund_redis_log(self, decision, redis_client, settings)
        case ("redis", "gcra"):
            await _refund_redis_gcra(self, decision, redis_client, settings)
        case ("memory", "log"):
            await self._refund_memory_log(decision)
        case ("memory", "gcra"):
            await self._refund_memory_gcra(decision)
        case ("postgres", "log"):
            await _refund_pg_log(self, decision, pg_pool, settings)
        case ("postgres", "gcra"):
            await _refund_pg_gcra(self, decision, pg_pool, settings)
        case _:
            assert_never((decision.backend, self._style))

peek async

peek(
    *,
    redis_client: Redis | None = None,
    pg_pool: Pool | None = None,
    clock: Clock | None = None,
    settings: WorkerSettings | None = None,
) -> RateLimitState

Read-only state snapshot.

PG and Redis peeks measure against the store's own clock (the same domain their admission state is stamped in); the memory backend uses the injected clock (its single domain), the unified TokenBucket contract: clock is required only on memory.

Source code in src/taskq/ratelimit/sliding_window.py

async def peek(
    self,
    *,
    redis_client: "redis_async.Redis | None" = None,
    pg_pool: "asyncpg.Pool | None" = None,
    clock: Clock | None = None,
    settings: "WorkerSettings | None" = None,
) -> RateLimitState:
    """Read-only state snapshot.

    PG and Redis peeks measure against the store's own clock (the same
    domain their admission state is stamped in); the memory backend
    uses the injected *clock* (its single domain), the unified
    TokenBucket contract: *clock* is required only on memory.
    """
    match (self._backend, self._style):
        case ("memory", "log"):
            if clock is None:
                raise RuntimeError("clock not injected for memory backend")
            now_ms = int(clock.now().timestamp() * 1000)
            return await self._peek_memory_log(now_ms)
        case ("memory", "gcra"):
            if clock is None:
                raise RuntimeError("clock not injected for memory backend")
            now_ms = int(clock.now().timestamp() * 1000)
            return await self._peek_memory_gcra(now_ms)
        case ("redis", "log"):
            return await _peek_redis_log(self, redis_client, settings)
        case ("redis", "gcra"):
            return await _peek_redis_gcra(self, redis_client, settings)
        case ("postgres", "log"):
            return await _peek_pg_log(self, pg_pool, settings)
        case ("postgres", "gcra"):
            return await _peek_pg_gcra(self, pg_pool, settings)
        case _:
            assert_never((self._backend, self._style))

reset async

reset(
    *,
    redis_client: Redis | None = None,
    pg_pool: Pool | None = None,
    clock: Clock | None = None,
    settings: WorkerSettings | None = None,
) -> None

Source code in src/taskq/ratelimit/sliding_window.py

async def reset(
    self,
    *,
    redis_client: "redis_async.Redis | None" = None,
    pg_pool: "asyncpg.Pool | None" = None,
    clock: Clock | None = None,
    settings: "WorkerSettings | None" = None,
) -> None:
    # Why clock is unused: resetting clears the window outright, with
    # no timestamp to stamp (TokenBucket.reset does re-stamp, so it
    # uses its clock). Retained for the polymorphic keyword block
    # described on refund above.
    match (self._backend, self._style):
        case ("memory", "log"):
            await self._reset_memory_log()
        case ("memory", "gcra"):
            await self._reset_memory_gcra()
        case ("redis", "log"):
            await _reset_redis_log(self, redis_client, settings)
        case ("redis", "gcra"):
            await _reset_redis_gcra(self, redis_client, settings)
        case ("postgres", "log"):
            await _reset_pg_log(self, pg_pool, settings)
        case ("postgres", "gcra"):
            await _reset_pg_gcra(self, pg_pool, settings)
        case _:
            assert_never((self._backend, self._style))

    logger.warning(
        "ratelimit-reset",
        bucket_name=self._name,
        backend=self._backend,
        style=self._style,
    )

TokenBucket

TokenBucket(
    name: str,
    capacity: float,
    refill_per_second: float,
    backend: RateLimitBackend = "redis",
    ttl: timedelta | None = None,
    *,
    keyed: bool = False,
)

Token-bucket rate limiter with pluggable backends.

Raises :class:ValueError if capacity <= 0 or refill_per_second < 0.

keyed (default False) marks a bucket materialised from a :class:~taskq.ratelimit.refs.KeyedRateLimitRef whose rate_limit_buckets row is FLEET-reclaimable: keyed-materialised AND PG-state-backed. The restriction to backend="postgres" is the point, not an accident: only a PG-state-backed bucket's acquire path touches its PG row, so only there does the row's last_used_at stamp (refreshed by the preseed/upsert/refund statements below) truthfully track use, a redis-backend keyed bucket's healthy acquire never touches PG, and its PG row is outage-fallback state plus admin metadata that the fleet sweep must never delete (the stamp would be a false staleness signal for an actively-used bucket). The registry's materialisation arm is the single call site that knows both facts and passes the flag accordingly; every other constructor call keeps the static default.

Source code in src/taskq/ratelimit/token_bucket.py

def __init__(
    self,
    name: str,
    capacity: float,
    refill_per_second: float,
    backend: RateLimitBackend = "redis",
    ttl: timedelta | None = None,
    *,
    keyed: bool = False,
) -> None:
    if capacity <= 0:
        raise ValueError(f"capacity must be > 0, got {capacity}")
    if refill_per_second < 0:
        raise ValueError(f"refill_per_second must be >= 0, got {refill_per_second}")

    self._name = name
    self._capacity = capacity
    self._refill = refill_per_second
    self._backend: RateLimitBackend = backend
    self._keyed = keyed

    self._ttl = ttl if ttl is not None else _default_ttl(capacity, refill_per_second)

    self._mem_bucket: _InMemoryBucket | None = None
    if backend == "memory":
        self._mem_bucket = _InMemoryBucket(name, capacity, refill_per_second)

    self._redis_script: AsyncScript | None = None
    self._redis_script_client: redis_async.Redis | None = None
    self._redis_refund_script: AsyncScript | None = None
    self._redis_refund_script_client: redis_async.Redis | None = None
    self._script_lock: asyncio.Lock = asyncio.Lock()

__slots__ class-attribute instance-attribute

__slots__ = (
    "_backend",
    "_capacity",
    "_keyed",
    "_mem_bucket",
    "_name",
    "_redis_refund_script",
    "_redis_refund_script_client",
    "_redis_script",
    "_redis_script_client",
    "_refill",
    "_script_lock",
    "_ttl",
)

name property

name: str

capacity property

capacity: float

refill_per_second property

refill_per_second: float

backend property

backend: RateLimitBackend

ttl property

ttl: timedelta

keyed property

keyed: bool

Whether this bucket's rate_limit_buckets row is fleet-reclaimable (see the class docstring for the exact marking rule).

holds_consumed_quota

holds_consumed_quota() -> bool

True if idle-evicting this bucket would silently reset a consumed fixed quota.

Idle eviction exists to bound registry growth, not to hand back capacity a tenant already spent. A fixed quota (refill_per_second == 0) is drained forever by design, so nothing about it recovers on its own and idleness is never evidence that it is safe to discard. Refilling buckets are not exempt: their state converges back toward full anyway, so eviction loses at most one refill window's worth of tokens, an accepted, bounded divergence.

What eviction destroys differs by where the state lives, and only ONE backend actually loses the budget:

• memory: token state lives on the instance and nowhere else, so eviction discards it and the next acquire materializes at FULL capacity. The instance is right here, so the exemption is exact: only a bucket that has actually spent some quota is held.
• postgres, the rate_limit_buckets row IS the state, and evicting the REGISTRY entry does not delete it: both row-delete paths (the maintenance leader's fleet sweep and the per-worker pending-reclaim drain, sharing _no_consumed_quota_sql) veto deleting a row whose fixed quota is partly spent, and re-materialization resumes from the surviving row (the acquire preseeds ON CONFLICT DO NOTHING and reads the existing state under the row lock). The registry entry is pure bookkeeping, dropping it loses nothing, so no PG fixed-quota bucket is held. Holding them anyway is worse than a wasted entry: every PG fixed-quota key ever seen would count against max_keyed_rate_limits forever, so once the cap filled, every NEW key would be refused with ReservationUnavailable and its jobs snooze-looped until process restart.
• redis: the backend keeps fixed-quota state for 24 h of its own accord (see _compute_ttl_seconds), so a re-materialized bucket resumes prior state there and eviction is state-safe.

Reads _tokens without the bucket's async lock; safe because the only caller (the registry's idle-eviction sweep) runs synchronously in the event loop with no await between this read and the dict pop, so the value is consistent at the sweep instant.

Source code in src/taskq/ratelimit/token_bucket.py

def holds_consumed_quota(self) -> bool:
    """True if idle-evicting this bucket would silently reset a consumed fixed quota.

    Idle eviction exists to bound registry growth, not to hand back
    capacity a tenant already spent. A fixed quota
    (``refill_per_second == 0``) is drained forever by design, so
    nothing about it recovers on its own and idleness is never
    evidence that it is safe to discard. Refilling buckets are not
    exempt: their state converges back toward full anyway, so eviction
    loses at most one refill window's worth of tokens, an accepted,
    bounded divergence.

    What eviction destroys differs by where the state lives, and only
    ONE backend actually loses the budget:

    * **memory**: token state lives on the instance and nowhere else,
      so eviction discards it and the next acquire materializes at
      FULL capacity. The instance is right here, so the exemption is
      exact: only a bucket that has actually spent some quota is held.
    * **postgres**, the ``rate_limit_buckets`` row IS the state, and
      evicting the REGISTRY entry does not delete it: both row-delete
      paths (the maintenance leader's fleet sweep and the per-worker
      pending-reclaim drain, sharing ``_no_consumed_quota_sql``) veto
      deleting a row whose fixed quota is partly spent, and
      re-materialization resumes from the surviving row (the acquire
      preseeds ``ON CONFLICT DO NOTHING`` and reads the existing state
      under the row lock). The registry entry is pure bookkeeping,
      dropping it loses nothing, so no PG fixed-quota bucket is held.
      Holding them anyway is worse than a
      wasted entry: every PG fixed-quota key ever seen would count against
      ``max_keyed_rate_limits`` forever, so once the cap filled, every
      NEW key would be refused with ``ReservationUnavailable`` and its jobs
      snooze-looped until process restart.
    * **redis**: the backend keeps fixed-quota state for 24 h of its
      own accord (see ``_compute_ttl_seconds``), so a re-materialized
      bucket resumes prior state there and eviction is state-safe.

    Reads ``_tokens`` without the bucket's async lock; safe because the
    only caller (the registry's idle-eviction sweep) runs synchronously
    in the event loop with no await between this read and the dict pop,
    so the value is consistent at the sweep instant.
    """
    if self._refill != 0.0:
        return False
    if self._backend != "memory":
        # postgres: the row-delete vetoes carry the state-safety
        # guarantee (see the docstring); redis: the 24 h TTL does.
        # Neither's quota can be reset by a REGISTRY eviction.
        return False
    # Why the protected read: _InMemoryBucket._tokens is this module's
    # own accumulator, and the registry's idle-eviction sweep (the
    # only caller) runs synchronously with no await between this read
    # and the dict pop, the docstring above documents the
    # consistency argument; a public accessor would widen the surface
    # for one internal read.
    tokens: float | None = (
        self._mem_bucket._tokens  # pyright: ignore[reportPrivateUsage]  # Why: same-module internal accumulator; see the comment above.
        if self._mem_bucket is not None
        else None
    )
    return tokens is not None and tokens < self._capacity

acquire async

acquire(
    count: float = 1.0,
    *,
    redis_client: Redis | None = None,
    pg_pool: Pool | None = None,
    clock: Clock | None = None,
    settings: WorkerSettings | None = None,
) -> RateLimitDecision

Spend count tokens, or return the denial with its retry hint.

The shared admission state is measured in the data store's own clock domain: PG paths use statement_timestamp() and Redis paths use Redis TIME inside their scripts, so callers on nodes with divergent Python clocks are all measured against the same bucket. The injected clock drives the memory backend only (its single domain) and remains part of the public call shape, the unified limiter contract shared with :class:~taskq.ratelimit.sliding_window.SlidingWindow.

With rate_limit_pg_fallback_enabled (the default) the redis backend fails over to Postgres on connection-family errors, and the returned decision's backend names the store that actually paid, the field refund dispatches on. A PG row-lock timeout fails closed (a denial, never an exception, never an admission); a store reply no honest write could have produced raises :class:RateLimitStoreCorrupt.

Source code in src/taskq/ratelimit/token_bucket.py

async def acquire(
    self,
    count: float = 1.0,
    *,
    redis_client: "redis_async.Redis | None" = None,
    pg_pool: "asyncpg.Pool | None" = None,
    clock: Clock | None = None,
    settings: "WorkerSettings | None" = None,
) -> RateLimitDecision:
    """Spend *count* tokens, or return the denial with its retry hint.

    The shared admission state is measured in the data store's own
    clock domain: PG paths use ``statement_timestamp()`` and Redis
    paths use Redis ``TIME`` inside their scripts, so callers on
    nodes with divergent Python clocks are all measured against the
    same bucket. The injected *clock* drives the memory backend only
    (its single domain) and remains part of the public call shape,
    the unified limiter contract shared with
    :class:`~taskq.ratelimit.sliding_window.SlidingWindow`.

    With ``rate_limit_pg_fallback_enabled`` (the default) the redis
    backend fails over to Postgres on connection-family errors, and
    the returned decision's ``backend`` names the store that
    actually paid, the field ``refund`` dispatches on. A PG
    row-lock timeout fails closed (a denial, never an exception,
    never an admission); a store reply no honest write could have
    produced raises :class:`RateLimitStoreCorrupt`.
    """
    if self._backend == "memory":
        return await self._acquire_memory(count, clock)
    if self._backend == "redis":
        return await self._acquire_redis_wrapped(count, redis_client, pg_pool, settings)
    if self._backend == "postgres":
        return await self._acquire_pg(count, pg_pool, settings)

    raise RuntimeError(f"unknown backend: {self._backend!r}")

refund async

refund(
    decision: RateLimitDecision,
    *,
    count: float = 1.0,
    redis_client: Redis | None = None,
    pg_pool: Pool | None = None,
    clock: Clock | None = None,
    settings: WorkerSettings | None = None,
) -> None

Return count tokens to the bucket that actually paid them.

Dispatches on decision's backend (see the comment in the body for why not the bucket's own: a redis acquire during a Redis outage may have spent its token in Postgres), so the refund lands in the store whose quota was consumed. The memory path re-adds under the bucket lock capped at capacity; the Redis path runs the refund Lua script (elapsed refill plus count, capped); the PG path rewrites the row under the same bounded FOR UPDATE lock the acquire takes, and raises on lock-budget exhaustion rather than silently reporting a lost refund (for a fixed-quota bucket that loss is permanent). clock is accepted and unused, a token-bucket refund needs no timestamp; it stays in the signature because RateLimitRegistry dispatches refund/peek/reset polymorphically over both primitives with one fixed keyword block.

Source code in src/taskq/ratelimit/token_bucket.py

async def refund(
    self,
    decision: RateLimitDecision,
    *,
    count: float = 1.0,
    redis_client: "redis_async.Redis | None" = None,
    pg_pool: "asyncpg.Pool | None" = None,
    clock: Clock | None = None,
    settings: "WorkerSettings | None" = None,
) -> None:
    """Return *count* tokens to the bucket that actually paid them.

    Dispatches on *decision*'s ``backend`` (see the comment in the
    body for why not the bucket's own: a redis acquire during a
    Redis outage may have spent its token in Postgres), so the
    refund lands in the store whose quota was consumed. The memory
    path re-adds under the bucket lock capped at capacity; the
    Redis path runs the refund Lua script (elapsed refill plus
    *count*, capped); the PG path rewrites the row under the same
    bounded ``FOR UPDATE`` lock the acquire takes, and raises on
    lock-budget exhaustion rather than silently reporting a lost
    refund (for a fixed-quota bucket that loss is permanent).
    *clock* is accepted and unused, a token-bucket refund needs no
    timestamp; it stays in the signature because
    RateLimitRegistry dispatches refund/peek/reset polymorphically
    over both primitives with one fixed keyword block.
    """
    # Why decision.backend and not self._backend: with backend="redis" and
    # rate_limit_pg_fallback_enabled (the default), an acquire during a
    # Redis outage falls through to Postgres and consumes the token THERE.
    # The decision records which store actually paid; the primitive's own
    # configuration only records where it prefers to go. Dispatching on the
    # latter refunded Redis for a token Postgres spent, inflating one
    # store's quota and destroying the other's, and for a fixed-quota
    # bucket (refill_per_second == 0) nothing ever puts the Postgres token
    # back, so that loss is permanent.
    #
    # Why clock is unused: a token-bucket refund returns *count* tokens to
    # the bucket and needs no timestamp. It stays in the signature because
    # RateLimitRegistry dispatches refund/peek/reset polymorphically over
    # TokenBucket and SlidingWindow with one fixed keyword block
    # (redis_client, pg_pool, clock, settings), see registry.reset_limit's
    # call sites. Dropping it would raise TypeError there, not merely break
    # symmetry.
    if decision.backend == "memory":
        await self._refund_memory(count)
    elif decision.backend == "redis":
        await self._refund_redis(count, redis_client, settings)
    elif decision.backend == "postgres":
        await self._refund_pg(count, pg_pool, settings)

peek async

peek(
    *,
    redis_client: Redis | None = None,
    pg_pool: Pool | None = None,
    clock: Clock | None = None,
    settings: WorkerSettings | None = None,
) -> RateLimitState

Source code in src/taskq/ratelimit/token_bucket.py

async def peek(
    self,
    *,
    redis_client: "redis_async.Redis | None" = None,
    pg_pool: "asyncpg.Pool | None" = None,
    clock: Clock | None = None,
    settings: "WorkerSettings | None" = None,
) -> RateLimitState:
    if self._backend == "memory":
        return await self._peek_memory(clock)
    if self._backend == "redis":
        return await self._peek_redis(redis_client, settings)
    if self._backend == "postgres":
        return await self._peek_pg(pg_pool, settings)

    raise RuntimeError(f"unknown backend: {self._backend!r}")

reset async

reset(
    *,
    redis_client: Redis | None = None,
    pg_pool: Pool | None = None,
    clock: Clock | None = None,
    settings: WorkerSettings | None = None,
) -> None

Source code in src/taskq/ratelimit/token_bucket.py

async def reset(
    self,
    *,
    redis_client: "redis_async.Redis | None" = None,
    pg_pool: "asyncpg.Pool | None" = None,
    clock: Clock | None = None,
    settings: "WorkerSettings | None" = None,
) -> None:
    if self._backend == "memory":
        await self._reset_memory(clock)
    elif self._backend == "redis":
        await self._reset_redis(redis_client, settings)
    elif self._backend == "postgres":
        await self._reset_pg(pg_pool, settings)
    else:
        raise RuntimeError(f"unknown backend: {self._backend!r}")

    logger.warning(
        "ratelimit-reset",
        bucket_name=self._name,
        backend=self._backend,
    )

get_redis_pool async

get_redis_pool(
    settings: WorkerSettings,
) -> AsyncIterator[Any]

Yield a Redis client for the worker loop lifetime.

The return type is AsyncIterator[Any] rather than AsyncIterator[redis.asyncio.Redis] so the DI system can introspect the annotation without requiring the [redis] extra at runtime. The actual yielded value is a redis.asyncio.Redis instance.

Raises :class:RuntimeError when settings.redis_url is None (Redis not configured but a Redis-backed rate limiter was registered). Raises :class:ImportError when the [redis] extra is not installed.

Source code in src/taskq/ratelimit/_provider.py

async def get_redis_pool(
    settings: WorkerSettings,
) -> AsyncIterator[Any]:
    """Yield a Redis client for the worker loop lifetime.

    The return type is ``AsyncIterator[Any]`` rather than
    ``AsyncIterator[redis.asyncio.Redis]`` so the DI system can introspect
    the annotation without requiring the ``[redis]`` extra at runtime.
    The actual yielded value is a ``redis.asyncio.Redis`` instance.

    Raises :class:`RuntimeError` when ``settings.redis_url`` is ``None``
    (Redis not configured but a Redis-backed rate limiter was registered).
    Raises :class:`ImportError` when the ``[redis]`` extra is not installed.
    """
    if settings.redis_url is None:
        raise RuntimeError(
            "Redis not configured but a Redis-backed rate limiter "
            "(TokenBucket/SlidingWindow) was registered"
        )
    from taskq._redis_client import build_redis_client

    client = build_redis_client(str(settings.redis_url))
    try:
        yield client
    finally:
        # Why bounded: this close runs on the worker exit stack, which unwinds
        # each teardown bare, an unbounded close here is a teardown tail
        # outside the accounted budget. One close bound, one mental model:
        # same CLOSE_TIMEOUT_SECS module-global seam as every other
        # TaskQ-initiated redis close (taskq._close; read at call time so
        # tests can shrink it).
        await close_redis_bounded(client, "ratelimit", CLOSE_TIMEOUT_SECS)

register_rate_limit_registry

register_rate_limit_registry(
    di_registry: ProviderRegistry,
    rl_registry: RateLimitRegistry,
) -> None

Idempotent registration of the resolved RateLimitRegistry instance.

Registers the given :class:RateLimitRegistry as a Scope.LOOP value so it is available at dispatch time via DI resolution. Skips when a provider is already registered, a user pre-registered LOOP-scope value provider wins (bootstrap resolution rule 2); after the bootstrap kind/scope checks the DI-cached instance and the bootstrap instance are provably the same object. When the module singleton is resolved (the default), the DI-registered instance and :data:taskq.ratelimit.registry.registry are the same object, both paths observe identical state.

Source code in src/taskq/ratelimit/_provider.py

def register_rate_limit_registry(
    di_registry: ProviderRegistry,
    rl_registry: RateLimitRegistry,
) -> None:
    """Idempotent registration of the resolved RateLimitRegistry instance.

    Registers the given :class:`RateLimitRegistry` as a ``Scope.LOOP`` value
    so it is available at dispatch time via DI resolution.  Skips when a
    provider is already registered, a user pre-registered LOOP-scope value
    provider wins (bootstrap resolution rule 2); after the bootstrap
    kind/scope checks the DI-cached instance and the bootstrap
    instance are provably the same object.  When the module singleton is
    resolved (the default), the DI-registered instance and
    :data:`taskq.ratelimit.registry.registry` are the same object, both
    paths observe identical state.
    """
    if di_registry.has_provider(RateLimitRegistry):
        return
    di_registry.register_value(RateLimitRegistry, Scope.LOOP, rl_registry)

register_redis_pool

register_redis_pool(registry: ProviderRegistry) -> None

Idempotent registration of the LOOP-scoped Redis pool factory.

Calls registry.register_factory(redis.asyncio.Redis, Scope.LOOP, get_redis_pool) only when registry.has_provider(redis.asyncio.Redis) is False, so user-supplied registrations take precedence.

Silently skips registration when the [redis] extra is not installed.

Source code in src/taskq/ratelimit/_provider.py

def register_redis_pool(registry: ProviderRegistry) -> None:
    """Idempotent registration of the LOOP-scoped Redis pool factory.

    Calls ``registry.register_factory(redis.asyncio.Redis, Scope.LOOP,
    get_redis_pool)`` only when ``registry.has_provider(redis.asyncio.Redis)``
    is ``False``, so user-supplied registrations take precedence.

    Silently skips registration when the ``[redis]`` extra is not installed.
    """
    try:
        import redis.asyncio as redis_async
    except ImportError:
        return

    if registry.has_provider(redis_async.Redis):
        return
    registry.register_factory(redis_async.Redis, Scope.LOOP, get_redis_pool)

queue_concurrency_reservation_name

queue_concurrency_reservation_name(queue: str) -> str

Return the registry name for the fleet-wide concurrency cap of queue.

The taskq:global:queue: prefix namespaces these internally-generated reservations apart from user-declared ones. A reservation only exists in the registry for a queue if that queue's max_concurrent column was set in the queues table (read from Postgres at worker startup); there is no other config source. All workers sharing the schema register and acquire against the same PG reservation_slots rows, giving a true fleet-wide cap per queue.

Source code in src/taskq/ratelimit/registry.py

def queue_concurrency_reservation_name(queue: str) -> str:
    """Return the registry name for the fleet-wide concurrency cap of *queue*.

    The ``taskq:global:queue:`` prefix namespaces these internally-generated
    reservations apart from user-declared ones.  A reservation only exists
    in the registry for a queue if that queue's ``max_concurrent`` column
    was set in the ``queues`` table (read from Postgres at worker startup);
    there is no other config source. All workers sharing the schema
    register and acquire against the same PG ``reservation_slots`` rows,
    giving a true fleet-wide cap per queue.
    """
    return f"{QUEUE_CONCURRENCY_PREFIX}{queue}"

sync_rate_limit_buckets async

sync_rate_limit_buckets(
    rl_registry: RateLimitRegistry,
    pool: Pool,
    *,
    schema: str = "taskq",
) -> None

Publish every registered rate limit to rate_limit_buckets.

Each worker calls this at startup so the admin UI can discover configured buckets from PG without depending on the in-memory singleton being populated in the admin process. Keyed buckets materialized lazily AFTER startup are published individually by the acquisition path, see :meth:RateLimitRegistry._resolve_rate_limit_name.

Uses ON CONFLICT DO NOTHING so concurrent workers and restarts are idempotent. Only PG-backed primitives are written; memory-only and log-style sliding windows (which have no PG backend) are skipped.

Source code in src/taskq/ratelimit/registry.py

async def sync_rate_limit_buckets(
    rl_registry: RateLimitRegistry,
    pool: "asyncpg.Pool",
    *,
    schema: str = "taskq",
) -> None:
    """Publish every registered rate limit to ``rate_limit_buckets``.

    Each worker calls this at startup so the admin UI can discover
    configured buckets from PG without depending on the in-memory
    singleton being populated in the admin process.  Keyed buckets
    materialized lazily AFTER startup are published individually by the
    acquisition path, see
    :meth:`RateLimitRegistry._resolve_rate_limit_name`.

    Uses ``ON CONFLICT DO NOTHING`` so concurrent workers and restarts
    are idempotent.  Only PG-backed primitives are written; memory-only
    and log-style sliding windows (which have no PG backend) are skipped.
    """
    require_schema(schema)

    for name, prim in rl_registry.rate_limits.items():
        if isinstance(prim, TokenBucket):
            kind = "token_bucket"
        else:
            if prim.style == "gcra":
                kind = "gcra"
            else:
                continue

        await _upsert_rate_limit_bucket_row(pool, schema, name, kind)

        logger.debug(
            "rl-bucket-synced",
            bucket_name=name,
            kind=kind,
        )

sync_slots async

sync_slots(
    reservations: list[ConcurrencyReservation],
    pool: Pool,
    *,
    schema: str = "taskq",
    timeout: float | None = None,
) -> SyncResult

Synchronise slot rows to match the registered reservation config.

For each reservation: insert missing slots (filling gaps from prior held-slot-preserving shrinks), delete excess free slots, and report held slots that could not be deleted.

timeout bounds the WHOLE pass, one connection acquire plus a transaction of statements per reservation, so a call costs O(reservations) round trips and a wedged store must not park the caller past the bound. Raises :class:TimeoutError when it fires (a reservation whose transaction already committed stays synced; the pass is safe to re-run). None (the default) keeps the unbounded shape for callers that manage their own deadline.

"Free" / "held" use the same definition as the acquire CTE: a slot row with an EXPIRED lease is acquirable, hence deletable; only rows with a live (unexpired) lease are reported as skipped_held. The anomalous job_id NOT NULL / lease_expires_at NULL state (no code path produces it, but the nullable schema permits it) is conservatively reported as held rather than deleted, matching the in-memory table. Treating expired-lease rows as held would let a dead worker's leaked slots defeat a shrink indefinitely, the rows stay acquirable, so the old larger cap would keep being honored. Deleting an expired-lease row is safe: lease expiry is the design's abandonment signal (the acquire CTE would hand the same slot to a new job anyway), and a late release by the old holder is a harmless no-op (worker_id mismatch or missing row). A row acquired concurrently with the delete is protected by its row lock: the delete blocks, then re-evaluates its predicate against the post-acquire row (job_id now set, lease in the future) and skips it.

Source code in src/taskq/ratelimit/reservation.py

async def sync_slots(
    reservations: list[ConcurrencyReservation],
    pool: "asyncpg.Pool",
    *,
    schema: str = "taskq",
    timeout: "float | None" = None,  # noqa: ASYNC109  # Why: the bound is a per-call deadline the caller passes, not an enclosing asyncio.timeout scope.
) -> SyncResult:
    """Synchronise slot rows to match the registered reservation config.

    For each reservation: insert missing slots (filling gaps from prior
    held-slot-preserving shrinks), delete excess free slots, and report
    held slots that could not be deleted.

    *timeout* bounds the WHOLE pass, one connection acquire plus a
    transaction of statements per reservation, so a call costs
    O(reservations) round trips and a wedged store must not park the
    caller past the bound. Raises :class:`TimeoutError` when it fires
    (a reservation whose transaction already committed stays synced;
    the pass is safe to re-run). ``None`` (the default) keeps the
    unbounded shape for callers that manage their own deadline.

    "Free" / "held" use the same definition as the acquire CTE: a slot row
    with an EXPIRED lease is acquirable, hence deletable; only rows with a
    live (unexpired) lease are reported as ``skipped_held``. The anomalous
    ``job_id NOT NULL / lease_expires_at NULL`` state (no code path
    produces it, but the nullable schema permits it) is conservatively
    reported as held rather than deleted, matching the in-memory table.
    Treating expired-lease rows as held would let a dead worker's leaked slots
    defeat a shrink indefinitely, the rows stay acquirable, so the old
    larger cap would keep being honored. Deleting an expired-lease row is
    safe: lease expiry is the design's abandonment signal (the acquire CTE
    would hand the same slot to a new job anyway), and a late release by
    the old holder is a harmless no-op (``worker_id`` mismatch or missing
    row). A row acquired concurrently with the delete is protected by its
    row lock: the delete blocks, then re-evaluates its predicate against
    the post-acquire row (``job_id`` now set, lease in the future) and
    skips it.
    """
    _validate_schema(schema)

    async def _sync_all() -> SyncResult:
        all_inserted: list[tuple[str, int]] = []
        all_deleted: list[tuple[str, int]] = []
        all_skipped: list[tuple[str, int]] = []

        for res in reservations:
            n_inserted = 0
            n_deleted = 0
            n_skipped = 0

            async with pool.acquire() as conn, conn.transaction():
                existing_sql = _SYNC_EXISTING_SQL_TEMPLATE.format(schema=schema)
                existing_rows = await conn.fetch(existing_sql, res.name)
                existing_indices: set[int] = {row["slot_index"] for row in existing_rows}

                desired_set = set(range(res.slots))
                missing_indices = sorted(desired_set - existing_indices)
                excess_indices = sorted(existing_indices - desired_set)

                if missing_indices:
                    insert_sql = _SYNC_INSERT_SQL_TEMPLATE.format(schema=schema)
                    rows = await conn.fetch(
                        insert_sql,
                        res.name,
                        missing_indices,
                    )
                    for row in rows:
                        all_inserted.append((res.name, row["slot_index"]))
                    n_inserted = len(rows)

                if excess_indices:
                    held_sql = _SYNC_HELD_SQL_TEMPLATE.format(schema=schema)
                    held_rows = await conn.fetch(
                        held_sql,
                        res.name,
                        excess_indices,
                    )
                    for row in held_rows:
                        all_skipped.append((res.name, row["slot_index"]))
                    n_skipped = len(held_rows)

                    delete_sql = _SYNC_DELETE_SQL_TEMPLATE.format(schema=schema)
                    deleted_rows = await conn.fetch(
                        delete_sql,
                        res.name,
                        excess_indices,
                    )
                    for row in deleted_rows:
                        all_deleted.append((res.name, row["slot_index"]))
                    n_deleted = len(deleted_rows)

            logger.debug(
                "reservation-sync-slots",
                bucket_name=res.name,
                inserted=n_inserted,
                deleted=n_deleted,
                skipped=n_skipped,
            )

        return SyncResult(
            inserted=all_inserted,
            deleted=all_deleted,
            skipped_held=all_skipped,
        )

    if timeout is not None:
        return await asyncio.wait_for(_sync_all(), timeout=timeout)
    return await _sync_all()