PR-2841: Token-bucket rate limiter for ingestion API
Replace the naive per-second cap with a smoothed bucket; reduce 429 storms
Be conservative in what you do, be liberal in what you accept from others.
What this PR does
The current ingestion API caps requests at 100/second per tenant via a hard counter that resets at the top of each second. That produces thundering herd on the second boundary: tenants whose clients respect rate limits batch their work to fire just-after-reset, which spikes our load right when rate-limit memory is empty. Result: occasional storms of 429 responses when a too-clever client retries within the same second.
This PR replaces that mechanism with a token-bucket smoother:
- 100 tokens per bucket per tenant
- Tokens replenish at 100/second linearly
- A request consumes 1 token; if the bucket is empty, request fails 429
- Burst capacity = bucket size; smoothing happens via continuous replenishment
This is Thariq’s “Code Review & Understanding” use case rendered through stoichos. The diff is rendered with per-line classification, the rationale is in margin asides, and the test output sits inline as output blocks rather than competing with the source.
The change
The implementation is a single new module plus a callsite swap. Most of the diff is in the new file:
+ # apps/ingestion/rate_limit.py+ import time+ from dataclasses import dataclass, field+ from typing import Dict++ @dataclass+ class TokenBucket:+ capacity: int+ refill_per_sec: float+ tokens: float = field(init=False)+ last_refill: float = field(init=False)++ def __post_init__(self) -> None:+ self.tokens = float(self.capacity)+ self.last_refill = time.monotonic()++ def consume(self, n: int = 1) -> bool:+ now = time.monotonic()+ elapsed = now - self.last_refill+ self.tokens = min(self.capacity, self.tokens + elapsed * self.refill_per_sec)+ self.last_refill = now+ if self.tokens >= n:+ self.tokens -= n+ return True+ return False++ class TenantBuckets:+ def __init__(self, capacity: int = 100, refill_per_sec: float = 100.0) -> None:+ self._buckets: Dict[str, TokenBucket] = {}+ self._capacity = capacity+ self._refill = refill_per_sec++ def consume(self, tenant_id: str) -> bool:+ bucket = self._buckets.get(tenant_id)+ if bucket is None:+ bucket = TokenBucket(self._capacity, self._refill)+ self._buckets[tenant_id] = bucket+ return bucket.consume()
Note time.monotonic() not time.time() — wall-clock can move backward (NTP drift, leap seconds). Buckets that depend on backward time movement will lock up.
The callsite swap removes the old per-second counter:
@@ -8,12 +8,7 @@ from .response import json_error-_request_counts: dict[str, list[int]] = {}--def _check_rate_limit_legacy(tenant: str) -> bool:- now_sec = int(time.time())- counts = _request_counts.setdefault(tenant, [now_sec, 0])- if counts[0] != now_sec:- counts[0] = now_sec- counts[1] = 0- counts[1] += 1- return counts[1] <= 100+from .rate_limit import TenantBuckets++_buckets = TenantBuckets(capacity=100, refill_per_sec=100.0)
@@ -41,7 +36,7 @@ async def ingest_document(request: Request) -> Response: tenant = request.tenant_id- if not _check_rate_limit_legacy(tenant):+ if not _buckets.consume(tenant): return json_error(429, "rate_limited")
What I want reviewed
TenantBuckets._buckets grows unboundedly as new tenants appear. We need eviction. The natural shape is a TTL cache (e.g. cachetools.TTLCache) that drops buckets unused for >1h. Should I add it in this PR or follow up? I’d lean follow-up; this PR is already 187 lines.
consume mutates self.tokens and self.last_refill without a lock. The ingestion API is single-threaded async, so this is safe per-event-loop. If we ever go multi-process for the API, we need to swap the in-process bucket for a Redis-backed equivalent. This is documented in the module docstring; flagging here for visibility.
Test results
Unit tests for the bucket math:
$ pytest apps/ingestion/test_rate_limit.py -v test_rate_limit.py::test_bucket_starts_full PASSED test_rate_limit.py::test_bucket_drains_under_load PASSED test_rate_limit.py::test_bucket_refills_linearly PASSED test_rate_limit.py::test_bucket_does_not_overshoot_capacity PASSED test_rate_limit.py::test_consume_returns_false_when_empty PASSED test_rate_limit.py::test_separate_tenants_have_separate_buckets PASSED ============================== 6 passed in 0.41s ==============================
Load test against staging with the new limiter:
$ ./scripts/loadtest.sh --tenant=t-load-1 --rps=120 --duration=60s[load] starting: target_rps=120 duration=60s tenant=t-load-1[load] 60s elapsed; 7,200 requests sent[load] success: 6,012 (83.5%)[load] rate_limited: 1,188 (16.5%)[load] mean_latency_ms: 14.2[load] p99_latency_ms: 31.7[load] no 5xx errors
Compare to the same load against the old limiter (run yesterday, same configuration):
[load] success: 5,801 (80.6%) [load] rate_limited: 1,372 (19.0%) [load] errors: 27 (0.4%) ← 5xx during second-boundary spikes [load] mean_latency_ms: 18.6 [load] p99_latency_ms: 87.3
The token bucket reduces P99 latency by 64% (87ms → 32ms) and eliminates the second-boundary 5xx errors entirely.
Migration
Zero-downtime. The callsite swap is a single file. The new module is independent. Roll forward; if the bucket misbehaves, revert is a single-line change.
@ops for the memory-growth question. @sre to confirm load test parameters match production traffic shape.