Professional Code Generator

1. Solution

"""In-memory token-bucket rate limiting for FastAPI/Starlette.

File: rate_limiter.py
Run the demo with: uvicorn rate_limiter:app
"""

from __future__ import annotations

import asyncio
import hashlib
import math
import threading
import time
from collections import OrderedDict
from collections.abc import Callable, Iterable
from dataclasses import dataclass

from fastapi import FastAPI
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import JSONResponse, Response
from starlette.types import ASGIApp

Clock = Callable[[], float]

DEFAULT_CAPACITY = 100
DEFAULT_REFILL_RATE = DEFAULT_CAPACITY / 60  # 100 requests per 60 s
DEFAULT_MAX_CLIENTS = 10_000
API_KEY_HEADER = "X-API-Key"

# Float arithmetic like 12 * (5/60) gives 0.9999999999999999. Without a
# tolerance, a client that waited exactly long enough would still be refused.
_FLOAT_TOLERANCE = 1e-9


# --------------------------------------------------------------------------- #
# Validation helpers
# --------------------------------------------------------------------------- #
def _require_positive_int(name: str, value: object) -> None:
    # bool is a subclass of int. Reject it explicitly so True is not read as 1.
    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
        raise ValueError(f"{name} must be a positive int, got {value!r}")


def _require_positive_finite(name: str, value: object) -> None:
    if (
        isinstance(value, bool)
        or not isinstance(value, (int, float))
        or not math.isfinite(value)
        or value <= 0
    ):
        raise ValueError(f"{name} must be a positive finite number, got {value!r}")


def _validate_tokens(tokens: int, capacity: int) -> None:
    _require_positive_int("tokens", tokens)
    if tokens > capacity:
        # This request could never succeed. Fail loudly instead of returning 429 forever.
        raise ValueError(f"tokens ({tokens}) exceeds bucket capacity ({capacity})")


def _whole_seconds(seconds: float) -> int:
    """Round a duration up to whole seconds, ignoring float noise."""
    return max(0, math.ceil(seconds - _FLOAT_TOLERANCE))


# --------------------------------------------------------------------------- #
# Token bucket
# --------------------------------------------------------------------------- #
@dataclass(frozen=True, slots=True)
class ConsumeResult:
    """Outcome of a consume attempt, with everything needed for response headers."""

    allowed: bool
    limit: int
    remaining: int
    retry_after: float  # seconds until the request could succeed (0 if allowed)
    reset_after: float  # seconds until the bucket is full again


class TokenBucket:
    """A token bucket that refills continuously at `refill_rate` tokens per second.

    Refill is computed lazily on each call, so an idle bucket costs no CPU.
    Each bucket has its own asyncio.Lock, so different clients never wait on each other.
    """

    __slots__ = ("capacity", "refill_rate", "_clock", "_tokens", "_updated_at", "_lock")

    def __init__(
        self,
        capacity: int = DEFAULT_CAPACITY,
        refill_rate: float = DEFAULT_REFILL_RATE,
        clock: Clock = time.monotonic,
    ) -> None:
        _require_positive_int("capacity", capacity)
        _require_positive_finite("refill_rate", refill_rate)
        self.capacity = capacity
        self.refill_rate = float(refill_rate)
        self._clock = clock
        self._tokens = float(capacity)  # new clients start with a full bucket
        self._updated_at = clock()
        self._lock = asyncio.Lock()

    async def try_consume(self, tokens: int = 1) -> bool:
        """Take `tokens` if available. Return True if the request is allowed."""
        return (await self.consume(tokens)).allowed

    async def consume(self, tokens: int = 1) -> ConsumeResult:
        """Try to take `tokens` and report the resulting bucket state."""
        _validate_tokens(tokens, self.capacity)
        async with self._lock:
            self._refill()
            allowed = self._tokens + _FLOAT_TOLERANCE >= tokens
            if allowed:
                self._tokens = max(0.0, self._tokens - tokens)
                retry_after = 0.0
            else:
                retry_after = (tokens - self._tokens) / self.refill_rate
            return ConsumeResult(
                allowed=allowed,
                limit=self.capacity,
                remaining=math.floor(self._tokens + _FLOAT_TOLERANCE),
                retry_after=retry_after,
                reset_after=(self.capacity - self._tokens) / self.refill_rate,
            )

    def _refill(self) -> None:
        now = self._clock()
        # Clamp the elapsed time at 0, in case an injected clock ever goes backwards.
        elapsed = max(0.0, now - self._updated_at)
        self._tokens = min(float(self.capacity), self._tokens + elapsed * self.refill_rate)
        self._updated_at = now


# --------------------------------------------------------------------------- #
# Per-client registry with LRU eviction
# --------------------------------------------------------------------------- #
class RateLimiter:
    """Keeps one TokenBucket per client key, capped at `max_clients` (LRU eviction)."""

    def __init__(
        self,
        capacity: int = DEFAULT_CAPACITY,
        refill_rate: float = DEFAULT_REFILL_RATE,
        max_clients: int = DEFAULT_MAX_CLIENTS,
        clock: Clock = time.monotonic,
    ) -> None:
        _require_positive_int("capacity", capacity)
        _require_positive_finite("refill_rate", refill_rate)
        _require_positive_int("max_clients", max_clients)
        self.capacity = capacity
        self.refill_rate = float(refill_rate)
        self.max_clients = max_clients
        self._clock = clock
        self._buckets: OrderedDict[str, TokenBucket] = OrderedDict()
        # This lock only guards the O(1) dict lookup/insert/evict. It is never
        # held across an await, so it cannot stall the event loop. It also keeps
        # the registry safe if it is touched from worker threads.
        self._registry_lock = threading.Lock()

    def __len__(self) -> int:
        return len(self._buckets)

    def __contains__(self, key: object) -> bool:
        return key in self._buckets

    async def consume(self, key: str, tokens: int = 1) -> ConsumeResult:
        """Charge `tokens` to the client identified by `key`."""
        _validate_tokens(tokens, self.capacity)
        return await self._get_bucket(key).consume(tokens)

    async def try_consume(self, key: str, tokens: int = 1) -> bool:
        return (await self.consume(key, tokens)).allowed

    def _get_bucket(self, key: str) -> TokenBucket:
        if not isinstance(key, str) or not key:
            raise ValueError(f"client key must be a non-empty str, got {key!r}")
        with self._registry_lock:
            bucket = self._buckets.get(key)
            if bucket is not None:
                self._buckets.move_to_end(key)  # mark as most recently used
                return bucket
            bucket = TokenBucket(self.capacity, self.refill_rate, self._clock)
            self._buckets[key] = bucket
            if len(self._buckets) > self.max_clients:
                self._buckets.popitem(last=False)  # evict least recently used
            return bucket


# --------------------------------------------------------------------------- #
# Middleware
# --------------------------------------------------------------------------- #
def client_key(request: Request) -> str:
    """Identify the caller: API key if present, otherwise the client IP.

    API keys are hashed. This bounds memory for oversized header values and
    keeps raw secrets out of process memory and debug dumps. The "key:" and
    "ip:" prefixes stop an API key from colliding with an IP address.
    """
    api_key = request.headers.get(API_KEY_HEADER, "").strip()
    if api_key:
        return "key:" + hashlib.sha256(api_key.encode()).hexdigest()
    host = request.client.host if request.client else "unknown"
    return f"ip:{host}"


def _rate_limit_headers(result: ConsumeResult) -> dict[str, str]:
    return {
        "X-RateLimit-Limit": str(result.limit),
        "X-RateLimit-Remaining": str(result.remaining),
        "X-RateLimit-Reset": str(_whole_seconds(result.reset_after)),
    }


class RateLimitMiddleware(BaseHTTPMiddleware):
    """Throttle requests per client. Requests over the limit get a 429 JSON response."""

    def __init__(
        self,
        app: ASGIApp,
        limiter: RateLimiter | None = None,
        exempt_paths: Iterable[str] = ("/health",),
    ) -> None:
        super().__init__(app)
        self._limiter = limiter if limiter is not None else RateLimiter()
        self._exempt_paths = frozenset(exempt_paths)

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        if request.url.path in self._exempt_paths:
            return await call_next(request)

        result = await self._limiter.consume(client_key(request))
        headers = _rate_limit_headers(result)

        if not result.allowed:
            retry_after = max(1, _whole_seconds(result.retry_after))
            return JSONResponse(
                status_code=429,
                content={"error": "rate_limit_exceeded", "retry_after": retry_after},
                headers={**headers, "Retry-After": str(retry_after)},
            )

        response = await call_next(request)
        response.headers.update(headers)
        return response


# --------------------------------------------------------------------------- #
# Demo app
# --------------------------------------------------------------------------- #
def create_app(limiter: RateLimiter | None = None) -> FastAPI:
    """Build the demo app. Pass a limiter to customise the policy or clock (tests)."""
    app = FastAPI(title="Rate limiter demo")
    app.add_middleware(
        RateLimitMiddleware,
        limiter=limiter if limiter is not None else RateLimiter(),
        exempt_paths={"/health"},
    )

    @app.get("/health")
    async def health() -> dict[str, str]:
        return {"status": "ok"}

    @app.get("/data")
    async def data() -> dict[str, list[int]]:
        return {"items": [1, 2, 3]}

    return app


app = create_app()

2. How it works

• Lazy refill. Each bucket stores (tokens, last_updated) and adds elapsed * refill_rate only when it is accessed. There are no timers and no background tasks, and idle clients use no CPU. time.monotonic is the default clock, so wall-clock changes don't affect it. The clock is injectable, so tests can use a fake one.
• Per-bucket asyncio.Lock. Each client's check-and-decrement runs under its own lock, so one client's traffic never waits on another's. The registry's threading.Lock covers only an O(1) OrderedDict get/insert/evict. It is never held across an await.
• LRU cap. OrderedDict.move_to_end and popitem(last=False) give O(1) LRU eviction at 10,000 clients. __slots__ buckets plus a 64-char hashed key come to roughly a few hundred bytes per client. That puts 10k clients at a few MB, well under the 50 MB budget.
• Header semantics (my assumptions):
  • X-RateLimit-Reset is the number of seconds until the bucket is full again, not an epoch timestamp.
  • retry_after is rounded up to a whole number of seconds, with a minimum of 1. It is also sent as a standard Retry-After header.
  • Exempt paths (/health) get no rate-limit headers, because they aren't rate-limited.
• Float tolerance. A tolerance of 1e-9 absorbs results like 12 * (5/60) = 0.9999…. Without it, a client that waited exactly long enough would still be refused.
• Fail loudly. Invalid inputs raise ValueError instead of being silently clamped. This covers a zero or negative capacity, a NaN rate, True used as an int, an empty key, and tokens > capacity.

3. Usage

uvicorn rate_limiter:app --port 8000
curl -i -H "X-API-Key: alice" localhost:8000/data

HTTP/1.1 200 OK
x-ratelimit-limit: 100
x-ratelimit-remaining: 99
x-ratelimit-reset: 1
{"items":[1,2,3]}

The 101st request within the window returns:

HTTP/1.1 429 Too Many Requests
retry-after: 1
x-ratelimit-limit: 100
x-ratelimit-remaining: 0
x-ratelimit-reset: 60
{"error":"rate_limit_exceeded","retry_after":1}

A custom policy, e.g. 10 requests per second with bursts up to 20:

app = create_app(RateLimiter(capacity=20, refill_rate=10))

4. Tests

"""File: test_rate_limiter.py. Run with: pytest -q"""

import asyncio

import pytest
from fastapi.testclient import TestClient

from rate_limiter import RateLimiter, create_app

LIMIT = 5
REFILL_RATE = LIMIT / 60  # 5 requests per minute -> 1 token every 12 s
ALICE = {"X-API-Key": "alice"}


class FakeClock:
    def __init__(self, start: float = 1_000.0) -> None:
        self.now = start

    def __call__(self) -> float:
        return self.now

    def advance(self, seconds: float) -> None:
        self.now += seconds


@pytest.fixture
def clock() -> FakeClock:
    return FakeClock()


@pytest.fixture
def client(clock: FakeClock):
    limiter = RateLimiter(capacity=LIMIT, refill_rate=REFILL_RATE, clock=clock)
    with TestClient(create_app(limiter)) as test_client:
        yield test_client


def test_allows_n_requests_then_rejects_n_plus_1(client: TestClient) -> None:
    for expected_remaining in range(LIMIT - 1, -1, -1):
        response = client.get("/data", headers=ALICE)
        assert response.status_code == 200
        assert response.headers["X-RateLimit-Limit"] == str(LIMIT)
        assert response.headers["X-RateLimit-Remaining"] == str(expected_remaining)

    rejected = client.get("/data", headers=ALICE)
    assert rejected.status_code == 429
    assert rejected.json() == {"error": "rate_limit_exceeded", "retry_after": 12}
    assert rejected.headers["X-RateLimit-Remaining"] == "0"
    assert rejected.headers["X-RateLimit-Reset"] == "60"

    # Other clients (API key or IP fallback) and exempt routes are unaffected.
    assert client.get("/data", headers={"X-API-Key": "bob"}).status_code == 200
    assert client.get("/data").status_code == 200
    assert client.get("/health").status_code == 200


def test_client_recovers_after_refill(client: TestClient, clock: FakeClock) -> None:
    for _ in range(LIMIT):
        client.get("/data", headers=ALICE)
    assert client.get("/data", headers=ALICE).status_code == 429

    clock.advance(12)  # exactly one token's worth, which exercises the float tolerance
    assert client.get("/data", headers=ALICE).status_code == 200
    assert client.get("/data", headers=ALICE).status_code == 429

    clock.advance(60)  # a full window refills to capacity, never beyond it
    statuses = [client.get("/data", headers=ALICE).status_code for _ in range(LIMIT + 1)]
    assert statuses == [200] * LIMIT + [429]


def test_lru_evicts_least_recently_used_client(clock: FakeClock) -> None:
    limiter = RateLimiter(capacity=1, refill_rate=0.001, max_clients=2, clock=clock)

    async def scenario() -> None:
        assert await limiter.try_consume("a")
        assert await limiter.try_consume("b")
        assert not await limiter.try_consume("a")  # touching "a" makes "b" the LRU
        assert await limiter.try_consume("c")      # evicts "b", not "a"

    asyncio.run(scenario())
    assert len(limiter) == 2
    assert "a" in limiter and "c" in limiter and "b" not in limiter


@pytest.mark.parametrize(
    "kwargs",
    [{"capacity": 0}, {"capacity": 1.5}, {"capacity": True},
     {"refill_rate": 0}, {"refill_rate": float("nan")}, {"max_clients": -1}],
)
def test_invalid_configuration_is_rejected(kwargs: dict) -> None:
    with pytest.raises(ValueError):
        RateLimiter(**kwargs)

The tests use TestClient (sync) and asyncio.run, so pytest-asyncio is not needed. The fake clock means no test sleeps.

5. Notes

Assumptions

• try_consume is async. The required per-bucket asyncio.Lock can only be acquired from a coroutine.
• The concurrency guarantee has two parts:
  • Buckets are safe across concurrent tasks on one event loop.
  • The registry is also thread-safe.
  • A single bucket shared across multiple event loops or threads would need a threading.Lock instead.

Known limitations

• Behind nginx, IP fallback sees nginx's address. Every keyless client then shares one bucket. Run uvicorn --proxy-headers --forwarded-allow-ips=<nginx-ip> so that request.client.host comes from X-Forwarded-For. Set the allowed IPs precisely, or clients can spoof the header.
• Limits are per process. With --workers N, each client effectively gets N× the limit. That is fine for the single-process setup described.
• Eviction resets a client's bucket. An evicted client comes back with full tokens. This only matters with more than 10k active clients, which won't happen at your scale.
• BaseHTTPMiddleware adds some per-request overhead. It is negligible at 20 req/s.

Improvements for your internal analytics API

1. Don't trust unknown API keys. Right now a misbehaving client can rotate X-API-Key values to get a fresh bucket every time. Check the key against a known set (e.g. loaded from env/config) before keying on it, and fall back to IP for unknown keys. The same config could hold per-client policies, e.g. a higher capacity for the dashboard service.
2. Make throttling visible. Log each 429 with the client key hash and path, and expose a counter, e.g. on the /health endpoint or a /metrics endpoint. This lets you tell a noisy client apart from a policy that is too strict before users report it.