Rate Limiting

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

ratelimit

Rate-limiting primitives for TaskQ.

__all__ module-attribute

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

RateLimitBackend

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

AcquiredResource

Bases: Protocol

Protocol for a resource handle that can be released.

name property

name: str

release async

release() -> None

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

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

RateLimitHandle dataclass

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

Handle for a successfully acquired rate-limit token.

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

name instance-attribute

name: str

primitive instance-attribute

primitive: TokenBucket | SlidingWindow

decision instance-attribute

decision: RateLimitDecision

redis_client instance-attribute

redis_client: Redis | None

pg_pool instance-attribute

pg_pool: Pool | None

clock instance-attribute

clock: Clock | None

settings class-attribute instance-attribute

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

count class-attribute instance-attribute

count: float = 1.0

refund_on_release class-attribute instance-attribute

refund_on_release: bool = True

release async

release() -> None

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

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

ReservationHandle dataclass

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

Handle for a successfully acquired reservation slot.

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

name instance-attribute

name: str

reservation instance-attribute

reservation: ConcurrencyReservation

slot_index instance-attribute

slot_index: int

job_id instance-attribute

job_id: UUID

worker_id instance-attribute

worker_id: UUID

pool instance-attribute

pool: Pool | None

release async

release() -> None

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

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

RateLimitDecision dataclass

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

allowed instance-attribute

allowed: bool

remaining instance-attribute

remaining: float

retry_after instance-attribute

retry_after: timedelta | None

bucket_name instance-attribute

bucket_name: str

backend instance-attribute

backend: RateLimitBackend

request_id class-attribute instance-attribute

request_id: str | None = None

previous_state class-attribute instance-attribute

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

RateLimitState dataclass

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

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

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

bucket_name instance-attribute

bucket_name: str

backend instance-attribute

backend: RateLimitBackend

is_exhausted instance-attribute

is_exhausted: bool

tokens_remaining class-attribute instance-attribute

tokens_remaining: float = 0.0

remaining class-attribute instance-attribute

remaining: float = 0.0

retry_after class-attribute instance-attribute

retry_after: timedelta | None = None

capacity class-attribute instance-attribute

capacity: float | None = None

limit class-attribute instance-attribute

limit: int | None = None

window class-attribute instance-attribute

window: timedelta | None = None

style class-attribute instance-attribute

style: str | None = None

refill_per_second class-attribute instance-attribute

refill_per_second: float | None = None

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 dict[str, object], the same shape stored on the job row) and must return a non-empty string — typically a tenant, session, or account identifier already present on the payload.

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.

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[[dict[str, object]], str]

slots instance-attribute

slots: int

lease instance-attribute

lease: timedelta

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.

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

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] = {}

rate_limits property

rate_limits: dict[str, TokenBucket | SlidingWindow]

reservations property

reservations: dict[str, ConcurrencyReservation]

has_keyed_reservations property

has_keyed_reservations: bool

register

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

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

def register(
    self,
    primitive: TokenBucket | SlidingWindow | ConcurrencyReservation,
) -> None:
    if isinstance(primitive, ConcurrencyReservation):
        name = primitive.name
        existing_reservation = self._reservations.get(name)
        if existing_reservation is not None:
            if _same_config(existing_reservation, primitive):
                logger.debug(
                    "registry-register-idempotent-noop",
                    kind="reservation",
                    name=name,
                )
                return
            raise ValueError(
                f"reservation name already registered with a different config: "
                f"{name!r} — existing={existing_reservation!r}, new={primitive!r}"
            )
        self._reservations[name] = primitive
        logger.debug(
            "registry-registered",
            kind="reservation",
            name=name,
        )
        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,
    )

get_rate_limit

get_rate_limit(name: str) -> TokenBucket | SlidingWindow

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

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

get_reservation

get_reservation(name: str) -> ConcurrencyReservation

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

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

acquire async

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

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

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

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

acquire_for_actor async

acquire_for_actor(
    rate_limits: list[str],
    reservations: Sequence[str | KeyedReservationRef],
    *,
    job_id: UUID,
    worker_id: UUID,
    payload: dict[str, object] | 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) or :class:KeyedReservationRef instances (resolved dynamically per job from payload — see :meth:_resolve_reservation_name). payload is required if any entry is a KeyedReservationRef.

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

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

async def acquire_for_actor(
    self,
    rate_limits: list[str],
    reservations: Sequence["str | KeyedReservationRef"],
    *,
    job_id: "UUID",
    worker_id: "UUID",
    payload: dict[str, object] | 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) or :class:`KeyedReservationRef`
    instances (resolved dynamically per job from ``payload`` — see
    :meth:`_resolve_reservation_name`). ``payload`` is required if any
    entry is a ``KeyedReservationRef``.

    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).
    """
    acquired: list[AcquiredResource] = []
    try:
        for res_ref in reservations:
            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_name in rate_limits:
            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=job_id,
                    rate_limits=rate_limits,
                    reservations=reservations,
                    allowed=False,
                    retry_after=retry_td,
                    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=job_id,
            rate_limits=rate_limits,
            reservations=reservations,
            allowed=True,
            retry_after=None,
            handle_count=len(acquired),
        )
        return acquired
    except ReservationUnavailable:
        for handle in reversed(acquired):
            try:
                await handle.release()
            except Exception as exc:
                backend = (
                    handle.decision.backend
                    if isinstance(handle, RateLimitHandle)
                    else "postgres"
                )
                logger.error(
                    "ratelimit-rollback-failure",
                    handle_name=handle.name,
                    operation="release",
                    error=str(exc),
                    acquired_count=len(acquired),
                )
                record_ratelimit_refund_failure(handle.name, backend)
        raise

peek async

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

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

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

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

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

peek_all async

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

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

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

async def peek_all(
    self,
    *,
    redis_client: "redis_async.Redis | None" = None,
    pg_pool: "asyncpg.Pool | None" = None,
    clock: "Clock | None" = None,
    settings: "WorkerSettings | None" = None,
) -> 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

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

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

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

    Sets ``refund_on_release=False`` on all ``RateLimitHandle`` instances
    before iterating (token consumption is permanent after actor ran).
    Releases in reverse acquisition order.  Each release failure is caught,
    logged at ERROR, and loop continues (same pattern as rollback).
    """
    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. The leader sweep calls this automatically 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

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. The leader sweep
    calls this automatically 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.
    """
    cutoff = monotonic() - idle_for.total_seconds()
    stale = [
        name
        for name, last_used in self._keyed_reservation_last_used.items()
        if last_used < cutoff
    ]
    for name in stale:
        self._reservations.pop(name, None)
        del self._keyed_reservation_last_used[name]
    if stale:
        logger.debug("registry-evicted-idle-keyed-reservations", count=len(stale))
    return len(stale)

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

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=lease_td,
            lock_lease=lock_lease,
        )

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

    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_sql",
    "_schema",
    "_slots",
    "_table",
)

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

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

Acquire a slot. Returns slot_index.

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

async def acquire(
    self,
    job_id: UUID,
    worker_id: UUID,
    pool: "asyncpg.Pool | None" = None,
) -> int:
    """Acquire a slot. Returns ``slot_index``.

    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=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_index: int = row["slot_index"]
    logger.debug(
        "reservation-acquired",
        bucket_name=self._name,
        slot_index=slot_index,
        job_id=job_id,
        worker_id=worker_id,
    )
    return slot_index

release async

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

Release slot. No-op if worker_id mismatch.

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

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

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

    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

    async with pool.acquire() as conn:
        await conn.execute(
            self._release_sql,
            self._name,
            slot_index,
            worker_id,
        )
    logger.debug(
        "reservation-released",
        bucket_name=self._name,
        slot_index=slot_index,
        worker_id=worker_id,
    )

peek async

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

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

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

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

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

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

    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 < now()) AS free_count, "  # noqa: S608
        f"count(*) AS total_slots, "
        f"count(*) FILTER (WHERE job_id IS NOT NULL AND lease_expires_at >= now()) AS held_count "
        f'FROM "{schema}".reservation_slots WHERE bucket_name = $1'
    )
    async with pool.acquire() as conn:
        row = await conn.fetchrow(peek_sql, self._name)

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

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

SyncResult dataclass

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

Result of a sync_slots call.

inserted instance-attribute

inserted: list[tuple[str, int]]

deleted instance-attribute

deleted: list[tuple[str, int]]

skipped_held instance-attribute

skipped_held: list[tuple[str, int]]

SlidingWindow

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

Sliding-window rate limiter with pluggable backends.

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

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

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

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

    if ttl is not None:
        self._ttl = ttl
    elif style == "gcra":
        self._ttl = window + timedelta(milliseconds=60_000)
    else:
        self._ttl = 2 * window + timedelta(milliseconds=60_000)

    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

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

async def acquire(
    self,
    *,
    redis_client: "redis_async.Redis | None" = None,
    pg_pool: "asyncpg.Pool | None" = None,
    clock: Clock | None = None,
    settings: "WorkerSettings | None" = None,
) -> RateLimitDecision:
    if clock is None:
        raise RuntimeError("clock not injected for sliding window acquire")

    now_ms = int(clock.now().timestamp() * 1000)
    request_id: UUID | None = new_uuid() if self._style == "log" else None

    match (self._backend, self._style):
        case ("memory", "log"):
            return await self._acquire_memory_log(now_ms, request_id)
        case ("memory", "gcra"):
            return await self._acquire_memory_gcra(now_ms)
        case ("redis", "log"):
            return await _acquire_redis_log_wrapped(
                self, now_ms, request_id, redis_client, pg_pool, clock, settings
            )
        case ("redis", "gcra"):
            return await _acquire_redis_gcra_wrapped(
                self, now_ms, redis_client, pg_pool, clock, settings
            )
        case ("postgres", "log"):
            return await _acquire_pg_log(self, pg_pool, clock, settings, request_id)
        case ("postgres", "gcra"):
            return await _acquire_pg_gcra(self, pg_pool, clock, settings)
        case _:
            assert_never((self._backend, self._style))

refund async

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

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

async def refund(
    self,
    decision: RateLimitDecision,
    *,
    count: float = 1.0,
    redis_client: "redis_async.Redis | None" = None,
    pg_pool: "asyncpg.Pool | None" = None,
    clock: Clock | None = None,
    settings: "WorkerSettings | None" = None,
) -> None:
    match (self._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((self._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

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

async def peek(
    self,
    *,
    redis_client: "redis_async.Redis | None" = None,
    pg_pool: "asyncpg.Pool | None" = None,
    clock: Clock | None = None,
    settings: "WorkerSettings | None" = None,
) -> RateLimitState:
    if clock is None:
        raise RuntimeError("clock not injected for sliding window peek")

    now_ms = int(clock.now().timestamp() * 1000)

    match (self._backend, self._style):
        case ("memory", "log"):
            return await self._peek_memory_log(now_ms)
        case ("memory", "gcra"):
            return await self._peek_memory_gcra(now_ms)
        case ("redis", "log"):
            return await _peek_redis_log(self, now_ms, redis_client, settings)
        case ("redis", "gcra"):
            return await _peek_redis_gcra(self, now_ms, redis_client, settings)
        case ("postgres", "log"):
            return await _peek_pg_log(self, now_ms, pg_pool, clock, settings)
        case ("postgres", "gcra"):
            return await _peek_pg_gcra(self, now_ms, pg_pool, clock, settings)
        case _:
            assert_never((self._backend, self._style))

reset async

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

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

async def reset(
    self,
    *,
    redis_client: "redis_async.Redis | None" = None,
    pg_pool: "asyncpg.Pool | None" = None,
    clock: Clock | None = None,
    settings: "WorkerSettings | None" = None,
) -> None:
    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

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

    if ttl is not None:
        self._ttl = ttl
    elif refill_per_second == 0.0:
        self._ttl = _DEFAULT_FIXED_QUOTA_TTL
    else:
        self._ttl = timedelta(seconds=math.ceil(capacity / refill_per_second * 2) + 60)

    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

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

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, clock, settings)
    if self._backend == "postgres":
        return await self._acquire_pg(count, pg_pool, clock, 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

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:
    if self._backend == "memory":
        await self._refund_memory(count)
    elif self._backend == "redis":
        await self._refund_redis(decision, count, redis_client, clock, settings)
    elif self._backend == "postgres":
        await self._refund_pg(count, pg_pool, clock, settings)

peek async

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

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

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

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

reset async

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

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

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

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

get_redis_pool async

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

Yield a Redis client for the worker loop lifetime.

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

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

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

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

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

    Raises :class:`RuntimeError` when ``settings.redis_url`` is ``None``
    (Redis not configured but a Redis-backed rate limiter was registered).
    Raises :class:`ImportError` when the ``[redis]`` extra is not installed.
    """
    if settings.redis_url is None:
        raise RuntimeError(
            "Redis not configured but a Redis-backed rate limiter "
            "(TokenBucket/SlidingWindow) was registered"
        )
    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:
        await client.aclose()

register_rate_limit_registry

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

Idempotent registration of the LOOP-scope RateLimitRegistry singleton.

Registers the given :class:RateLimitRegistry as a Scope.LOOP value so it is available at dispatch time via DI resolution. The same object is also importable as :data:taskq.ratelimit.registry.registry — both paths observe identical state.

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

def register_rate_limit_registry(
    di_registry: ProviderRegistry,
    rl_registry: RateLimitRegistry,
) -> None:
    """Idempotent registration of the LOOP-scope RateLimitRegistry singleton.

    Registers the given :class:`RateLimitRegistry` as a ``Scope.LOOP`` value
    so it is available at dispatch time via DI resolution.  The same object
    is also importable as :data:`taskq.ratelimit.registry.registry` — both
    paths observe identical state.
    """
    if di_registry.has_provider(RateLimitRegistry):
        return
    di_registry.register_value(RateLimitRegistry, Scope.LOOP, rl_registry)

register_redis_pool

register_redis_pool(registry: ProviderRegistry) -> None

Idempotent registration of the LOOP-scoped Redis pool factory.

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

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

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

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

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

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

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

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.

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

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

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

    Each worker calls this at startup so the admin UI can discover
    configured buckets from PG without depending on the in-memory
    singleton being populated in the admin process.

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

    upsert_sql = (
        f'INSERT INTO "{schema}".rate_limit_buckets (bucket_name, kind, state, updated_at) '  # noqa: S608
        f"VALUES ($1, $2, '{{}}'::jsonb, now()) "
        f"ON CONFLICT (bucket_name) DO NOTHING"
    )

    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

        async with pool.acquire() as conn:
            await conn.execute(upsert_sql, 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.

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

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