WIP checkpoint: Redis rate-limiting dependency

This commit is contained in:
2026-08-18 15:40:10 +00:00
parent d4abe65451
commit 41b1c3959e
+131 -66
View File
@@ -2,21 +2,36 @@
This module provides a reusable FastAPI dependency that rate-limits requests
per endpoint scope (optionally per-user) using an atomic Lua script on Redis.
If Redis is unavailable the dependency can either allow the request through
(fail-open) or reject it with a ``503`` (fail-closed), depending on the policy
for the limit.
Rate limits are intentionally dependency-only: they do not wrap a whole router
as middleware because per-scope capacity/refill and per-user bucketing make a
dependency (with explicit ``Request`` access) the most predictable place to
implement them.
If Redis is unavailable the limiter can either allow the request through
(fail-open) or reject it with a ``503`` (fail-closed), depending on the
policy configured for that limit. Chat send uses fail-closed by default
because flooding/abuse is worse than a temporary outage.
Example
-------
from app.rate_limit import rate_limit, named_limit
chat_send_limit = rate_limit(
named_limit("chat_send", capacity=5, refill_rate=0.2, fail_open=False),
user_id_dependency=Depends(get_current_user_id),
)
@router.post("/messages")
async def send_message(
payload: MessageCreate,
current_user: CurrentUser,
db: DbDep,
_: None = Depends(chat_send_limit),
):
...
"""
from __future__ import annotations
import time
from dataclasses import dataclass
from typing import Any, Callable
from typing import Any, Awaitable, Callable
from fastapi import Depends, HTTPException, Request, status
from redis.asyncio import Redis
@@ -66,8 +81,8 @@ else
retry_after = math.ceil((cost - tokens) / refill_rate)
end
-- Keep bucket state small/expire-able. Retain the stamp as a float so
-- partial tokens accumulate smoothly between requests.
-- Keep the bucket state small and expire-able. Retaining the timestamp as a
-- float lets partial tokens accumulate smoothly between requests.
redis.call("HSET", key, "tokens", tostring(tokens), "ts", tostring(now))
redis.call("PEXPIRE", key, 600000)
@@ -76,7 +91,7 @@ return {allowed, tostring(tokens), retry_after}
class RateLimitExceeded(HTTPException):
"""Raised when a rate limit bucket has no tokens left."""
"""Raised when a rate-limit bucket has no tokens left."""
def __init__(self, retry_after: int) -> None:
super().__init__(
@@ -92,8 +107,8 @@ class RateLimit:
``capacity`` is the burst size and ``refill_rate`` is the steady-state
token accumulation rate (tokens/second). For example,
``RateLimit(10, 1.0)`` allows a burst of 10 requests and then one request
per second on average.
``RateLimit(name="chat_send", capacity=10, refill_rate=1.0)`` allows a
burst of 10 requests and then one request per second on average.
"""
name: str
@@ -104,36 +119,68 @@ class RateLimit:
fail_open: bool = True
def __post_init__(self) -> None:
if not self.name:
raise ValueError("name must be non-empty")
if self.capacity <= 0:
raise ValueError("capacity must be positive")
if self.refill_rate <= 0:
raise ValueError("refill_rate must be positive")
def default_user_scope(request: Request, user_id: int | None) -> str:
"""Scope a bucket per authenticated user, or per client IP if anonymous."""
if user_id is not None:
return f"user:{user_id}"
client = request.client
return f"ip:{client.host if client else 'unknown'}"
@dataclass(frozen=True)
class RateLimitConfig:
"""A named rate limit plus how it is applied to a request."""
limit: RateLimit
# Builds the Redis key prefix for the request. The authenticated user id
# is passed as the first argument (or ``None`` for anonymous limits).
scope: Callable[[Request, int | None], str]
# Builds the Redis key fragment for the request. The resolved user id is
# passed as the second argument (or ``None`` for anonymous limits).
scope: Callable[[Request, int | None], str] = default_user_scope
@dataclass(frozen=True)
class RateLimitResult:
"""Result of a rate-limit check (used by tests)."""
"""Outcome of a rate-limit check (used by tests)."""
allowed: bool
remaining: int
retry_after: int
_redis_client: Redis | None = None
def get_redis() -> Redis:
"""Return a shared async Redis client for rate limiting.
The client is created lazily (``Redis.from_url`` does not connect until a
command is issued), so this dependency never blocks or fails just because
Redis is down.
"""
global _redis_client
if _redis_client is None:
_redis_client = Redis.from_url(
settings.REDIS_URL,
encoding="utf-8",
decode_responses=True,
socket_connect_timeout=settings.REDIS_TIMEOUT_SECONDS,
socket_timeout=settings.REDIS_TIMEOUT_SECONDS,
)
return _redis_client
class RateLimiter:
"""Stateless dependency factory that enforces token buckets in Redis.
"""Enforce a token bucket in Redis.
Buckets are keyed as ``rate:{limit.name}:{scope_key}`` so multiple
endpoints can share a named limit while still being bucketed separately.
endpoints can share a named limit while remaining independently bucketed.
"""
def __init__(self, config: RateLimitConfig, redis: Redis | None = None) -> None:
@@ -141,32 +188,25 @@ class RateLimiter:
self._redis = redis
def _redis_client(self) -> Redis:
return self._redis or Redis.from_url(
settings.REDIS_URL,
encoding="utf-8",
decode_responses=True,
)
return self._redis or get_redis()
def _scope_key(self, request: Request, user_id: int | None) -> str:
return self._config.scope(request, user_id)
def _bucket_key(self, request: Request, user_id: int | None) -> str:
limit = self._config.limit
scope_key = self._config.scope(request, user_id)
return f"rate:{limit.name}:{scope_key}"
async def __call__(
self,
request: Request,
user_id: int | None = None,
) -> None:
"""Enforce the rate limit, raising 429 or 503 as appropriate.
async def check(self, request: Request, user_id: int | None = None) -> RateLimitResult:
"""Atomically refill/take a token and return the outcome.
``user_id`` is normally injected by another dependency (see
``RateLimitDependency`` below); the limiter itself only needs the
resolved value.
Redis failures are mapped to ``allowed=True`` (fail-open) or
``allowed=False`` (fail-closed, represented with ``retry_after=-1``)
according to the limit's policy, without raising.
"""
limit = self._config.limit
key = f"rate:{limit.name}:{self._scope_key(request, user_id)}"
key = self._bucket_key(request, user_id)
try:
redis = self._redis_client()
result = await redis.eval(
result = await self._redis_client().eval(
_TOKEN_BUCKET_LUA,
1,
key,
@@ -177,52 +217,77 @@ class RateLimiter:
)
except RedisError:
if limit.fail_open:
return
return RateLimitResult(allowed=True, remaining=0, retry_after=0)
# Fail-closed: no reliable retry estimate is available.
return RateLimitResult(allowed=False, remaining=0, retry_after=-1)
# Redis may return ints or strings depending on decode_responses.
allowed = bool(int(result[0]))
remaining = int(result[1])
retry_after = int(result[2])
return RateLimitResult(allowed=allowed, remaining=remaining, retry_after=retry_after)
async def __call__(self, request: Request, user_id: int | None = None) -> None:
"""Enforce the rate limit, raising 429 or 503 as appropriate."""
result = await self.check(request, user_id)
if result.allowed:
return
if result.retry_after < 0:
# Redis was unavailable and this limit is fail-closed.
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Rate limiting service unavailable",
) from None
# Redis may return ints or bytes depending on decode_responses.
allowed = bool(int(result[0]))
retry_after = int(result[2])
if not allowed:
raise RateLimitExceeded(retry_after=retry_after)
)
raise RateLimitExceeded(retry_after=result.retry_after)
def named_limit(name: str, **kwargs: Any) -> RateLimitConfig:
"""Build a named rate-limit config with a sensible default scope.
def named_limit(
name: str,
*,
capacity: int,
refill_rate: float,
fail_open: bool = True,
) -> RateLimitConfig:
"""Create a named rate-limit config with the default per-user/IP scope.
The default scope buckets per authenticated user (falling back to the
client IP for unauthenticated requests), which matches the chat-send and
report-message use cases.
``capacity`` is the burst size and ``refill_rate`` is tokens/second.
"""
return RateLimitConfig(
limit=RateLimit(name=name, **kwargs),
scope=default_user_scope,
limit=RateLimit(
name=name,
capacity=capacity,
refill_rate=refill_rate,
fail_open=fail_open,
)
)
def default_user_scope(request: Request, user_id: int | None) -> str:
"""Return the authenticated user id, or client IP for anonymous callers."""
if user_id is not None:
return f"user:{user_id}"
return f"ip:{request.client.host if request.client else 'unknown'}"
def rate_limit(
config: RateLimitConfig,
*,
user_id_dependency: Any = None,
) -> Callable[[Request, Any], Awaitable[None]]:
"""Build a FastAPI dependency enforcing ``config``.
``user_id_dependency`` optionally injects the authenticated user id, e.g.
``Depends(get_current_user_id)``. If omitted the limiter buckets by client
IP (the ``default_user_scope`` fallback), which is suitable for anonymous
endpoints but should be paired with an auth resolver for per-user limits.
def rate_limit(config: RateLimitConfig) -> Callable[..., Any]:
"""FastAPI dependency that resolves the current user and enforces a limit.
This keeps call sites declarative, e.g.::
ChatSendDep = Depends(rate_limit(named_limit("chat_send", ...)))
The returned callable is intended for ``Depends(...)``.
"""
async def _anonymous_user_id() -> int | None:
return None
resolver = user_id_dependency if user_id_dependency is not None else Depends(_anonymous_user_id)
async def dependency(
request: Request,
user_id: int | None = None,
limiter: RateLimiter = Depends(lambda: RateLimiter(config)),
user_id: int | None = resolver,
) -> None:
limiter = RateLimiter(config)
await limiter(request, user_id)
return dependency