From 0860bc8b44d7e2992b7cd4686b6372ec5d449f71 Mon Sep 17 00:00:00 2001 From: man4ish Date: Sat, 3 Oct 2026 14:58:13 -0500 Subject: [PATCH] fix: atomic reserve-before-work quota check, closing an overrun race Part of M2. The old quota check was check-then-later-decrement: quota_remaining() read the counter before forwarding to RAG, consume_quota() only decremented it after a successful response. Any number of concurrent requests could all observe "quota available" before any of them decremented, and all of them would succeed -- overrunning a near-zero quota by however many were in flight at once. Replaces that pair with reserve_quota() (atomic Redis DECR, compensated back if it would go negative -- only as many concurrent reservations as there are units left can ever succeed) called before forwarding to RAG, and release_quota() (refund) called if the forwarded request then fails upstream, since the reservation assumed success before knowing the outcome. Both are wired into the one shared _handle_billable_literature_call path /literature/answers and /literature/search both already go through, so this closes the race for both at once. Tests: 355 passed, including new coverage for concurrent-reservation exhaustion (exactly one of two requests over a 1-unit quota succeeds, never both) and quota refund on upstream failure. ruff clean. Co-Authored-By: Claude Sonnet 5 --- app/routes/v1.py | 12 +++++++--- app/services/v1_store.py | 47 +++++++++++++++++++++++++++++++++---- tests/test_v1_literature.py | 34 ++++++++++++++++++++++++++- 3 files changed, 85 insertions(+), 8 deletions(-) diff --git a/app/routes/v1.py b/app/routes/v1.py index eaba63c..c6024ac 100644 --- a/app/routes/v1.py +++ b/app/routes/v1.py @@ -148,8 +148,12 @@ async def _handle_billable_literature_call( return _error(409, "idempotency_in_progress", "A request with this Idempotency-Key is still running.", request_id, headers) - remaining = await store.quota_remaining(org_id, resource) - if remaining is not None and remaining <= 0: + # Atomic reserve-before-work: decrements the quota counter now, not + # after the upstream call succeeds, so concurrent requests can never + # all observe "quota available" and all succeed (see + # V1Store.reserve_quota's own docstring for why the old check-then- + # later-decrement pair could overrun a near-zero quota). + if not await store.reserve_quota(org_id, resource): if idempotency_key is not None: await store.idempotency_finish(subject, idempotency_key, fingerprint, 402, None) return _error(402, "quota_exceeded", quota_exceeded_message, request_id, headers) @@ -159,6 +163,9 @@ async def _handle_billable_literature_call( latency_ms = round((time.monotonic() - started) * 1000) if not 200 <= status < 300: + # The reservation above assumed this call would succeed and be + # billed; it didn't, so the unit must be given back. + await store.release_quota(org_id, resource) if idempotency_key is not None: await store.idempotency_finish(subject, idempotency_key, fingerprint, status, response) return _upstream_error(status, response, request_id, headers) @@ -183,7 +190,6 @@ async def _handle_billable_literature_call( "idempotency_key_sha256": request_fingerprint(idempotency_key) if idempotency_key else None, }, ) - await store.consume_quota(org_id, resource) return JSONResponse(public_response, status_code=status, headers={**headers, "X-Request-Id": request_id}) diff --git a/app/services/v1_store.py b/app/services/v1_store.py index b757d00..22fdb18 100644 --- a/app/services/v1_store.py +++ b/app/services/v1_store.py @@ -140,13 +140,52 @@ async def quota_remaining(self, org_id: str, resource: str) -> Optional[int]: except ValueError: return None - async def consume_quota(self, org_id: str, resource: str) -> None: - """Decrement an existing quota after a billed success. Never creates - the key: an org without a quota stays unmetered here.""" + async def reserve_quota(self, org_id: str, resource: str) -> bool: + """Atomically reserve one unit of `resource` for `org_id` before + doing the work it would bill for, and report whether the + reservation succeeded. Replaces the old quota_remaining()-then- + later-consume_quota() pair, which read the counter, did the + (slow) upstream call, and only decremented afterward -- any + number of concurrent requests could all observe "1 remaining" + before any of them decremented, and all of them would then + succeed, overrunning the quota by however many were in flight + at once. Redis's DECR is atomic, so only as many concurrent + reservations as there are units left can ever observe a + non-negative result here; the rest observe negative and + compensate back to zero immediately (never below zero, and + never creating a key that did not already exist -- an org + without a quota key stays unmetered, exactly like before). + + Fails open (reservation succeeds) on a Redis error, the same + posture every other method in this class takes: an outage must + never block a paid request, only leave this particular overrun + protection briefly unenforced until Redis recovers. + """ + key = self._quota_key(org_id, resource) + try: + if await self.redis.get(key) is None: + return True # unmetered: no quota key set for this org/resource + value = await self.redis.decr(key) + except Exception: + return True + if value < 0: + try: + await self.redis.incr(key) + except Exception: + pass + return False + return True + + async def release_quota(self, org_id: str, resource: str) -> None: + """Refund a reservation made by reserve_quota() for a call that was + then not actually billable (the upstream request failed) -- the + reservation already decremented optimistically, before knowing + whether the call would succeed. Never creates the key: mirrors + reserve_quota's own "only adjust an existing counter" rule.""" key = self._quota_key(org_id, resource) try: if await self.redis.get(key) is not None: - await self.redis.decr(key) + await self.redis.incr(key) except Exception: pass diff --git a/tests/test_v1_literature.py b/tests/test_v1_literature.py index 80fa1d9..ee0bfb0 100644 --- a/tests/test_v1_literature.py +++ b/tests/test_v1_literature.py @@ -242,6 +242,38 @@ def test_quota_is_consumed_on_success_and_never_created(client, redis, upstream) assert "gateway:v1:quota:42:literature.answer" not in redis.kv +def test_quota_reservation_is_atomic_under_concurrency(client, redis, upstream): + """Two concurrent requests with only 1 unit of quota left: exactly one + must succeed and one must be quota_exceeded -- not both succeeding + (an overrun) and not both failing (undercounting real capacity). + FakeRedis's incr/decr aren't async-concurrent in the true sense (no + real parallelism in this test process), but this still exercises the + actual reserve-then-compensate sequence reserve_quota runs, not a + mock standing in for it.""" + redis.kv["gateway:v1:quota:42:literature.answer"] = 1 + first = _post(client, idem="race-1") + second = _post(client, idem="race-2") + statuses = sorted([first.status_code, second.status_code]) + assert statuses == [200, 402] + # The quota key never goes negative and ends at exactly zero, not + # some other value a non-atomic check-then-decrement could leave it at. + assert redis.kv["gateway:v1:quota:42:literature.answer"] == 0 + + +def test_quota_is_refunded_when_upstream_call_fails(client, redis, upstream): + """reserve_quota decrements optimistically, before knowing whether the + call will succeed -- a failed upstream call must give the unit back, + or a string of transient RAG failures would silently burn through an + organization's quota for answers it never actually got billed for + (and never received).""" + redis.kv["gateway:v1:quota:42:literature.answer"] = 1 + upstream.return_value = (503, {"detail": "rag down"}) + resp = _post(client) + assert resp.status_code == 503 + assert redis.kv["gateway:v1:quota:42:literature.answer"] == 1 + assert _usage(redis) == [] + + def test_idempotent_retry_replays_without_running_or_billing_again(client, redis, upstream): first = _post(client, idem="retry-1") second = _post(client, idem="retry-1") @@ -362,7 +394,7 @@ def test_search_rejects_missing_question(client, redis, upstream): def test_search_quota_exceeded(client, redis, upstream): - with patch.object(v1.store, "quota_remaining", AsyncMock(return_value=0)): + with patch.object(v1.store, "reserve_quota", AsyncMock(return_value=False)): resp = client.post("/v1/literature/search", json={"question": "q"}, headers={"Authorization": f"Bearer {KEY}"}) assert resp.status_code == 402