Per-Client Rate Limiting in FastAPI with a Token Bucket Dependency
import asyncio
import time
from dataclasses import dataclass, field
@dataclass
class TokenBucket:
capacity: float
refill_rate: float
tokens: float = field(default=None)
last_refill: float = field(default_factory=time.monotonic)
def __post_init__(self):
if self.tokens is None:
self.tokens = self.capacity
def _refill(self):
now = time.monotonic()
elapsed = now - self.last_refill
if elapsed > 0:
self.tokens = min(self.capacity, self.tokens + elapsed * self.refill_rate)
self.last_refill = now
def try_consume(self, amount=1.0):
self._refill()
if self.tokens >= amount:
self.tokens -= amount
return True
return False
def retry_after(self, amount=1.0):
self._refill()
deficit = amount - self.tokens
if deficit <= 0:
return 0.0
return deficit / self.refill_rate
class BucketStore:
def __init__(self):
self._buckets = {}
self._lock = asyncio.Lock()
async def get(self, key, capacity, refill_rate):
async with self._lock:
bucket = self._buckets.get(key)
if bucket is None:
bucket = TokenBucket(capacity=capacity, refill_rate=refill_rate)
self._buckets[key] = bucket
return bucket
from fastapi import Request, HTTPException, status
from token_bucket import BucketStore
_store = BucketStore()
class RateLimiter:
def __init__(self, capacity, refill_rate, scope="default"):
self.capacity = float(capacity)
self.refill_rate = float(refill_rate)
self.scope = scope
def _client_key(self, request):
user = getattr(request.state, "user_id", None)
if user:
return f"{self.scope}:user:{user}"
forwarded = request.headers.get("x-forwarded-for")
if forwarded:
ip = forwarded.split(",")[0].strip()
else:
ip = request.client.host if request.client else "unknown"
return f"{self.scope}:ip:{ip}"
async def __call__(self, request: Request):
key = self._client_key(request)
bucket = await _store.get(key, self.capacity, self.refill_rate)
if bucket.try_consume():
return
wait = bucket.retry_after()
raise HTTPException(
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
detail="Rate limit exceeded",
headers={"Retry-After": str(int(wait) + 1)},
)
from fastapi import FastAPI, Depends
from rate_limit import RateLimiter
app = FastAPI()
# Strict: 5 token bucket refilling at 1/sec for auth attempts.
login_limit = RateLimiter(capacity=5, refill_rate=1, scope="login")
# Looser: 60 burst, 20/sec sustained for general reads.
read_limit = RateLimiter(capacity=60, refill_rate=20, scope="read")
@app.post("/login", dependencies=[Depends(login_limit)])
async def login(payload: dict):
return {"status": "authenticated", "user": payload.get("username")}
@app.get("/items/{item_id}", dependencies=[Depends(read_limit)])
async def get_item(item_id: int):
return {"id": item_id, "name": f"item-{item_id}"}
@app.get("/health")
async def health():
return {"ok": True}
This snippet shows how to enforce per-client rate limits in a FastAPI application using the token bucket algorithm, wired in as a reusable route dependency. The token bucket is a good fit for HTTP rate limiting because it allows short bursts up to the bucket capacity while still bounding the sustained rate to refill_rate tokens per second. Compared to a fixed-window counter, it avoids the boundary problem where a client can send 2 * limit requests around a window edge, and it degrades gracefully: a client that has been idle accumulates tokens and can burst, then throttles back to the steady rate.
In token_bucket.py, the TokenBucket dataclass holds capacity, refill_rate, the current tokens, and a last_refill timestamp. The core is try_consume, which lazily refills based on elapsed wall-clock time rather than running a background timer — this keeps memory small and avoids a ticker per client. It computes elapsed * refill_rate, clamps tokens at capacity, and only deducts when at least one token is available. retry_after returns how many seconds a rejected caller should wait for one token, which is used to populate the Retry-After header. The BucketStore keeps one bucket per client key in a plain dict guarded by an asyncio.Lock, so that concurrent requests for the same client mutate the bucket atomically. Because everything lives in process memory, this design targets a single-process or sticky-session deployment; a multi-worker setup would need a shared store like Redis, and that trade-off is deliberate for simplicity and speed.
In rate_limit.py, the RateLimiter class is a callable dependency, which is the idiomatic FastAPI way to make a dependency configurable. It is instantiated with the desired capacity and refill_rate, and its __call__ receives the Request. The _client_key helper prefers an authenticated identity, falling back to the X-Forwarded-For header and finally the socket peer address, so limits attach to a real caller rather than the reverse proxy. When try_consume fails, it raises HTTPException with status 429 and a Retry-After header, which FastAPI turns into a proper JSON error response. Different limits can coexist by creating multiple RateLimiter instances with different budgets.
In main.py, two limiters are declared — a strict one for the login route and a looser one for a general read endpoint — showing that limits are per-route policy, not global. Each is attached via Depends, so the dependency runs before the handler and short-circuits with a 429 when the bucket is empty. A key pitfall to remember is clock skew and process restarts wiping buckets; for stronger guarantees the same interface can be backed by an external store while keeping the dependency signature unchanged.
Share this code
Here's the card — post it anywhere.
