Skip to content

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 
41
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 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
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 
68
69
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 
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
@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 
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
@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 
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
def __init__(self) -> None:
    self._rate_limits: dict[str, TokenBucket | SlidingWindow] = {}
    self._reservations: dict[str, ConcurrencyReservation] = {}
    # 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] = {}
    # 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_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 
266
267
268
269
270
271
272
273
274
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 
276
277
278
279
280
281
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 
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
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 
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
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 
370
371
372
373
374
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 
376
377
378
379
380
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 
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
@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 
 838
 839
 840
 841
 842
 843
 844
 845
 846
 847
 848
 849
 850
 851
 852
 853
 854
 855
 856
 857
 858
 859
 860
 861
 862
 863
 864
 865
 866
 867
 868
 869
 870
 871
 872
 873
 874
 875
 876
 877
 878
 879
 880
 881
 882
 883
 884
 885
 886
 887
 888
 889
 890
 891
 892
 893
 894
 895
 896
 897
 898
 899
 900
 901
 902
 903
 904
 905
 906
 907
 908
 909
 910
 911
 912
 913
 914
 915
 916
 917
 918
 919
 920
 921
 922
 923
 924
 925
 926
 927
 928
 929
 930
 931
 932
 933
 934
 935
 936
 937
 938
 939
 940
 941
 942
 943
 944
 945
 946
 947
 948
 949
 950
 951
 952
 953
 954
 955
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
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 reservation.acquire(
                job_id,
                worker_id,
                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 
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
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,
) -> dict[str, RateLimitState]

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

Source code in src/taskq/ratelimit/registry.py 
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
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,
) -> dict[str, RateLimitState]:
    """Peek all registered rate limits. Returns {name: RateLimitState}."""
    results: dict[str, RateLimitState] = {}
    for name, prim in list(self._rate_limits.items()):
        try:
            if isinstance(prim, TokenBucket):
                results[name] = await prim.peek(
                    redis_client=redis_client,
                    pg_pool=pg_pool,
                    clock=clock,
                    settings=settings,
                )
            else:
                results[name] = await prim.peek(
                    redis_client=redis_client,
                    pg_pool=pg_pool,
                    clock=clock,
                    settings=settings,
                )
        except Exception as exc:
            logger.warning(
                "ratelimit-peek-failed",
                bucket_name=name,
                error=str(exc),
            )
    return results

reset async 

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

Reset a rate-limit bucket to full capacity.

Source code in src/taskq/ratelimit/registry.py 
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
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,
) -> None:
    """Reset a rate-limit bucket to full capacity."""
    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 isinstance(primitive, TokenBucket):
        await primitive.reset(
            redis_client=redis_client,
            pg_pool=pg_pool,
            clock=clock,
            settings=settings,
        )
    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 
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
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()
        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)

evict_idle_keyed_reservations 

evict_idle_keyed_reservations(idle_for: timedelta) -> 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.

Only removes the in-memory registry entry and its acquire-recency tracking — it does NOT touch the underlying Postgres reservation_slots rows for that name; those are already reclaimed independently by the existing lock-expiry sweep. A key that is acquired again after eviction is simply re-registered on next use (idempotent — see :meth:_resolve_reservation_name), 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 
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
def evict_idle_keyed_reservations(self, idle_for: "timedelta") -> 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.

    Only removes the in-memory registry entry and its
    acquire-recency tracking — it does NOT touch the underlying
    Postgres ``reservation_slots`` rows for that name; those are
    already reclaimed independently by the existing lock-expiry sweep.
    A key that is acquired again after eviction is simply
    re-registered on next use (idempotent — see
    :meth:`_resolve_reservation_name`), so eviction is always safe to
    call, including concurrently with in-flight acquisitions for
    other keys.

    Returns the number of entries evicted.
    """
    return self._evict_idle_keyed(
        self._keyed_reservation_last_used,
        self._reservations,
        idle_for,
        "registry-evicted-idle-keyed-reservations",
    )

evict_idle_keyed_rate_limits 

evict_idle_keyed_rate_limits(idle_for: timedelta) -> 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.

Only removes the in-memory registry entry and its acquire-recency tracking — it does NOT touch the underlying Redis hash for that bucket; 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). A key that is acquired again after eviction is simply re-registered on next use (idempotent — see :meth:_resolve_rate_limit_name), so eviction is always safe to call, including concurrently with in-flight acquisitions for other keys. Token buckets are not automatically removed from Redis; their independent TTL handles that side.

Exemption — memory fixed-quota buckets. A backend="memory" bucket with refill_per_second == 0 that has consumed any of its quota is NOT evicted (see :meth:TokenBucket.holds_consumed_memory_quota): its token state lives on the instance, so eviction would silently reset the drained quota to full on next acquire — whereas Redis deliberately retains that same state for 24 h. 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 fixed-quota keys the cardinality cap can permanently 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.

Returns the number of entries evicted.

Source code in src/taskq/ratelimit/registry.py 
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
def evict_idle_keyed_rate_limits(self, idle_for: "timedelta") -> 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.

    Only removes the in-memory registry entry and its acquire-recency
    tracking — it does NOT touch the underlying Redis hash for that
    bucket; 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`). A key that is acquired again
    after eviction is simply re-registered on next use (idempotent — see
    :meth:`_resolve_rate_limit_name`), so eviction is always safe to
    call, including concurrently with in-flight acquisitions for other
    keys. Token buckets are not automatically removed from Redis; their
    independent TTL handles that side.

    **Exemption — memory fixed-quota buckets.** A ``backend="memory"``
    bucket with ``refill_per_second == 0`` that has consumed any of its
    quota is NOT evicted (see
    :meth:`TokenBucket.holds_consumed_memory_quota`): its token state
    lives on the instance, so eviction would silently reset the drained
    quota to full on next acquire — whereas Redis deliberately retains
    that same state for 24 h. 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 fixed-quota keys the cardinality cap can
    permanently 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.

    Returns the number of entries evicted.
    """
    return self._evict_idle_keyed(
        self._keyed_rate_limit_last_used,
        self._rate_limits,
        idle_for,
        "registry-evicted-idle-keyed-rate-limits",
        preserve=_preserves_memory_fixed_quota_state,
    )

clear 

clear() -> None

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

Clears the four dicts (_rate_limits, _reservations, _keyed_reservation_last_used, _keyed_rate_limit_last_used) 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.

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 rows; backend state persists and will be observed on next acquire.

Source code in src/taskq/ratelimit/registry.py 
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
def clear(self) -> None:
    """Reset ALL mutable registry state — a test aid, NOT safe while running.

    Clears the four dicts (``_rate_limits``, ``_reservations``,
    ``_keyed_reservation_last_used``, ``_keyed_rate_limit_last_used``)
    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.

    **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`` 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_reservation_last_eviction_scan = float("-inf")
    self._keyed_rate_limit_last_eviction_scan = float("-inf")

ConcurrencyReservation 

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

Concurrency reservation using pre-allocated slot rows.

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

Source code in src/taskq/ratelimit/reservation.py 
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
def __init__(
    self,
    name: str,
    slots: int,
    lease: timedelta | float,
    lock_lease: timedelta | None = None,
    *,
    clock: Clock | None = None,
    schema: str = "taskq",
) -> 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

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

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

Idempotent pre-allocation of slot rows.

Source code in src/taskq/ratelimit/reservation.py 
456
457
458
459
async def ensure_slots(self, pool: "asyncpg.Pool") -> None:
    """Idempotent pre-allocation of slot rows."""
    async with pool.acquire() as conn:
        await conn.execute(self._ensure_sql, self._name, self._slots)

acquire async 

acquire(
    job_id: UUID, worker_id: UUID, pool: Pool | 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.

Source code in src/taskq/ratelimit/reservation.py 
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
async def acquire(
    self,
    job_id: UUID,
    worker_id: UUID,
    pool: "asyncpg.Pool | 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.
    """
    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"
            )
        self._table.ensure_slots(self._name, self._slots)
        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

    async with pool.acquire() as conn, conn.transaction():
        row = await conn.fetchrow(
            self._acquire_sql,
            self._name,
            job_id,
            worker_id,
            self._lease.total_seconds(),
        )

    if row is None:
        logger.info(
            "reservation-unavailable",
            bucket_name=self._name,
        )
        raise ReservationUnavailable(self._name, DEFAULT_RESERVATION_BACKOFF)

    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 
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
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 
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
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}

    if not _IDENT_RE.match(self._schema):
        raise ValueError(f"invalid schema identifier: {self._schema!r}")
    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 
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
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_gcra_script: AsyncScript | None = None
    self._redis_gcra_refund_script: AsyncScript | 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_script",
    "_redis_log_script",
    "_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 
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
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 
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
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 
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
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 
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
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,
)

Token-bucket rate limiter with pluggable backends.

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

Source code in src/taskq/ratelimit/token_bucket.py 
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
def __init__(
    self,
    name: str,
    capacity: float,
    refill_per_second: float,
    backend: RateLimitBackend = "redis",
    ttl: timedelta | None = None,
) -> 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._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_refund_script: AsyncScript | None = None
    self._script_lock: asyncio.Lock = asyncio.Lock()

__slots__ class-attribute instance-attribute 

__slots__ = (
    "_backend",
    "_capacity",
    "_mem_bucket",
    "_name",
    "_redis_refund_script",
    "_redis_script",
    "_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

holds_consumed_memory_quota 

holds_consumed_memory_quota() -> bool

True if idle-evicting this bucket's registry entry would silently reset a consumed fixed quota.

Only the memory backend stores token state on the instance (_mem_bucket); Redis and PG keep state in the store — Redis deliberately for 24 h for fixed-quota buckets (see _compute_ttl_seconds) — so a re-materialized bucket seamlessly resumes prior state there, and eviction of those backends' registry entries is always state-safe.

For a memory fixed-quota (refill_per_second == 0) bucket that has consumed any of its quota, eviction destroys that state permanently: the next acquire materializes a fresh instance at FULL capacity, silently resetting a quota designed to never refill. (Refilling buckets are not exempt: their state converges back toward full on its own, so eviction loses at most one refill window's worth of tokens — an accepted, bounded divergence.)

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 
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
def holds_consumed_memory_quota(self) -> bool:
    """True if idle-evicting this bucket's registry entry would silently reset a consumed fixed quota.

    Only the memory backend stores token state on the instance
    (``_mem_bucket``); Redis and PG keep state in the store — Redis
    deliberately for 24 h for fixed-quota buckets (see
    ``_compute_ttl_seconds``) — so a re-materialized bucket seamlessly
    resumes prior state there, and eviction of those backends' registry
    entries is always state-safe.

    For a memory fixed-quota (``refill_per_second == 0``) bucket that
    has consumed any of its quota, eviction destroys that state
    permanently: the next acquire materializes a fresh instance at FULL
    capacity, silently resetting a quota designed to never refill.
    (Refilling buckets are not exempt: their state converges back
    toward full on its own, so eviction loses at most one refill
    window's worth of tokens — an accepted, bounded divergence.)

    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.
    """
    return (
        self._backend == "memory"
        and self._refill == 0.0
        and self._mem_bucket is not None
        and self._mem_bucket._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
Source code in src/taskq/ratelimit/token_bucket.py 
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
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:
    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
Source code in src/taskq/ratelimit/token_bucket.py 
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
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 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 
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
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 
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
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 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
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"
        )
    import redis.asyncio as redis_async

    client = redis_async.from_url(
        str(settings.redis_url),
        decode_responses=False,  # Why: raw bytes are safer for binary payloads and cluster-safety across shards
    )
    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 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
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 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
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 
132
133
134
135
136
137
138
139
140
141
142
143
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 
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
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.
    """
    if not _IDENT_RE.match(schema):
        raise ValueError(f"invalid schema identifier: {schema!r}")

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

"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 
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
async def sync_slots(
    reservations: list[ConcurrencyReservation],
    pool: "asyncpg.Pool",
    *,
    schema: str = "taskq",
) -> 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.

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

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