fix(gateway): move the budget counter at charge time, reconcile as a floor - #28
Open
marcorivm wants to merge 1 commit into
Open
fix(gateway): move the budget counter at charge time, reconcile as a floor#28marcorivm wants to merge 1 commit into
marcorivm wants to merge 1 commit into
Conversation
…floor `pre_forward` reads the hot spend counter to decide whether to deny with 402, but the counter only moved once the flush loop's PostgreSQL upsert completed. Requests arriving between a charge and its flush all read the same stale total and were all admitted. Scale correction worth recording: the window is NOT the 5s flush interval. `collect_batch` awaits `rx.recv()` with a 5s TIMEOUT and returns as soon as one event arrives, then `fill` drains the rest with `try_recv`. So under load the lag is one loop turn; the 5s only applies when the channel is idle, i.e. when there is no spend to flush anyway. The race was real but millisecond-scale, not seconds-scale. Three changes: - `CacheStore::incr_by(key, delta, ttl)` — a signed add returning the new total, distinct from `incr`'s +1 u64 rate-limit counter. The in-memory implementation does the read-modify-write under DashMap's entry lock, so concurrent charges on one budget accumulate. A get_raw/set_raw pair would keep only the last, which is how spend slipped past a cap. - The flush loop now applies charges to the counter BEFORE the database round-trip, since the counter is what enforcement reads. - `record_spend` reconciles instead of overwriting. The durable row is a FLOOR: raise the counter to it when the cache is behind (cold start, eviction, a lost increment), never lower it. A blind `set_raw(total)` would roll back charges already applied for the next batch and let that spend through twice. Deliberately NOT Redis. The cache is a single-instance in-process DashMap and is already atomic per key, so a single gateway gains nothing here from Redis — what Redis buys is durability across restarts, which the PostgreSQL floor already covers on rehydrate. Adopting it stays an option; it is not a prerequisite for this fix. Fail-open throughout, unchanged: a cache miss or error still admits the request. 3 new tests, including 64 concurrent charges accumulating exactly. 618 passing.
marcorivm
force-pushed
the
feat/budget-metering-accuracy
branch
from
August 8, 2026 19:35
48ddc75 to
cd0a0ac
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Addresses the first of the three design comments from #8's review. 2 files, ~90 lines. Stacked on #27.
First, a correction to the review comment
The comment said overshoot is bounded by "how much can arrive in ≤5s across however many parallel agents share this secret." That overstates it, and I'd rather say so than ship a fix premised on a wrong number.
collect_batchawaitsrx.recv()with a 5-second timeout — it returns the instant one event arrives, thenfilldrains the rest with non-blockingtry_recv. So the 5s interval only applies when the channel is idle, i.e. when there's no spend to flush anyway. Under load the counter lagged by one loop turn, not five seconds.The race was real, just millisecond-scale rather than seconds-scale.
What changed
CacheStore::incr_by(key, delta, ttl) -> Option<i64>— a signed add returning the new total, distinct from the existingincr(+1,u64, for rate limiting). The in-memory implementation does the read-modify-write under DashMap's entry lock, so concurrent charges accumulate. Aget_raw/set_rawpair keeps only the last write — which is precisely how spend slipped past a cap.Charges apply to the counter before the database round-trip, since the counter is what
pre_forwardreads to decide on the 402.record_spendreconciles instead of overwriting. This is the subtle one:The durable row is a floor — raise the counter to it when the cache is behind (cold start, eviction, a lost increment), never lower it. The old blind
set_raw(total)would have rolled back charges already applied for the next batch and let that spend through a second time. Keeping the overwrite alongside charge-time increments would have been worse than either alone.Deliberately not Redis
The cache is a single-instance in-process
DashMapand is already atomic per key, so a single gateway gains nothing here from Redis. What Redis buys is durability across restarts — which the PostgreSQL floor already covers on rehydrate — and shared state across replicas, which isn't the deployment shape. Adopting it stays an option; it is not a prerequisite for this fix, and theincr_byprimitive works identically behind a Redis backend later.Fail-open is unchanged throughout: a cache miss or error still admits the request.
Tests
3 new, the important one being 64 concurrent
tokio::spawned charges that must sum to exactly 64,000 — it fails against aget_raw/set_rawimplementation. Plus the durable-floor interleaving, and an unparseable value being treated as 0 rather than poisoning spend accounting.Still open from #8's review
The other two comments — the hand-maintained pricing table silently metering unknown models as free, and the cross-language constant duplication between
budget.rsandbudget-service.ts— are not addressed here. They're independent and I'd rather they land as their own change than pad this one.