diff --git a/admin_test.go b/admin_test.go index 9b2de92..e888b9d 100644 --- a/admin_test.go +++ b/admin_test.go @@ -1407,6 +1407,10 @@ func (s *countingMultiStore) getMulti(keys []string) ([]int64, error) { func (s *countingMultiStore) incrMulti(entries []counterIncr) ([]int64, error) { return make([]int64, len(entries)), nil } +func (s *countingMultiStore) incrAndGetMulti(entries []counterIncr, reads []string) ([]int64, []int64, error) { + readVals, err := s.getMulti(reads) + return make([]int64, len(entries)), readVals, err +} func TestLimiterCurrentUsage_BatchesOneGetMultiCallForAllScopes(t *testing.T) { t.Parallel() diff --git a/limits.go b/limits.go index 4ae90e4..becdc2a 100644 --- a/limits.go +++ b/limits.go @@ -213,13 +213,38 @@ type counterStore interface { // call; memoryStore's in-process map needs no such optimization but // implements the same contract for interface conformance), returning // each entry's new counter value in the same order (perf review, - // 2026-08-21: checkAndCount/account previously paid one round trip per - // counter — up to 9 and 27 respectively for a 3-scope request — this - // collapses each down to one). An error fails the whole batch, never a - // partial result, mirrored by the limiter's storeIncrMulti as a single - // fail-open/fail-closed decision, the same shape storeGetMulti already - // applies to a batch read. + // 2026-08-21: account previously paid one round trip per counter — up + // to 27 for a 3-scope request with all three metrics nonzero — this + // collapses it down to one; checkAndCount's own increments moved on to + // incrAndGetMulti below, perf review round 3, 2026-08-22). An error + // fails the whole batch, never a partial result, mirrored by the + // limiter's storeIncrMulti as a single fail-open/fail-closed decision, + // the same shape storeGetMulti already applies to a batch read. incrMulti(entries []counterIncr) ([]int64, error) + // incrAndGetMulti applies incrMulti's own per-entry contract to every + // entry in entries AND reads every key in reads, in ONE round trip + // where the backend supports it (redisStore pipelines every entry's + // INCRBY+EXPIRE pair together with every read key's GET into a single + // pipeline call; memoryStore's in-process map needs no such + // optimization but implements the same contract for interface + // conformance) — fusing checkAndCount's request-counter increments + // with its own token/cost budget reads into the single round trip + // that used to cost up to 9 separate ones for a 2-scope request, both + // fully limited (perf review round 3, 2026-08-22: reads carries the + // tokin/tokout/cost keys checkAndCount's own evaluateScope helper used + // to read via its own separate storeGetMulti calls — see + // buildBudgetProbes, below). This is safe because every read key is + // DISJOINT from every entry's key (req:* vs tokin:*/tokout:*/cost:*), + // and tokens/cost are only ever WRITTEN by account() after the + // upstream response completes — never by this call — so there is no + // read-after-write ordering hazard in reading them alongside the + // increments. Returns incrMulti's own per-entry post-increment values + // and getMulti's own per-key values (0 for a missing/expired key), + // each in entries'/reads' own order. An error fails the whole batch, + // mirrored by the limiter's storeIncrAndGetMulti as one + // fail-open/fail-closed decision, the same shape storeIncrMulti/ + // storeGetMulti already apply to their own separate batches. + incrAndGetMulti(entries []counterIncr, reads []string) (incrVals, readVals []int64, err error) } // memoryEntry is one counter's value and expiry in memoryStore. @@ -414,6 +439,19 @@ func (m *memoryStore) incrMulti(entries []counterIncr) ([]int64, error) { return out, nil } +// incrAndGetMulti implements counterStore by composing incrMulti and +// getMulti: memoryStore's in-process map access is already effectively +// free per key, so there is no round-trip cost to fuse away here — this +// exists purely so memoryStore satisfies counterStore's incrAndGetMulti +// contract (perf review round 3, 2026-08-22) for the limiter's nil-store +// fallback path and its own failOpen fallback, mirroring incrMulti/ +// getMulti's own "implements the contract, buys nothing locally" role. +func (m *memoryStore) incrAndGetMulti(entries []counterIncr, reads []string) ([]int64, []int64, error) { + incrVals, _ := m.incrMulti(entries) // memoryStore.incrMulti never errors + readVals, _ := m.getMulti(reads) // memoryStore.getMulti never errors + return incrVals, readVals, nil +} + // sweepLocked deletes every expired entry. Callers must hold m.mu. func (m *memoryStore) sweepLocked(now time.Time) { for k, e := range m.data { @@ -783,9 +821,11 @@ func (l *limiter) failPolicyGetMulti(keys []string) ([]int64, bool) { // trip against the configured store (or the in-process fallback) for the // whole slice, applying the identical fail-open/fail-closed/latched // policy storeIncrBy applies per key. ok is false only in the fail-closed -// case — a caller must refuse the request (checkAndCount) or drop the -// sample (account) for that, rather than treating a nil/short slice as -// real counter readings. +// case — a caller must drop the sample (account) or skip the write +// (countTargetRequests, recordProviderAttempt) for that, rather than +// treating a nil/short slice as real counter readings. checkAndCount's own +// increments moved to storeIncrAndGetMulti (perf review round 3, +// 2026-08-22), so it is no longer a caller of this method. func (l *limiter) storeIncrMulti(entries []counterIncr) (v []int64, ok bool) { if l.store == nil { v, _ = l.fallback.incrMulti(entries) // fallback never errors @@ -811,6 +851,40 @@ func (l *limiter) failPolicyIncrMulti(entries []counterIncr) ([]int64, bool) { return v, true } +// storeIncrAndGetMulti mirrors storeIncrMulti for checkAndCount's fused +// admission round trip (perf review round 3, 2026-08-22): one call against +// the configured store (or the in-process fallback) for BOTH entries' own +// increments and reads' own current values, applying the identical +// fail-open/fail-closed/latched policy storeIncrMulti/storeGetMulti already +// apply to their own separate batches. ok is false only in the fail-closed +// case — checkAndCount must refuse the request for that, rather than +// treating a nil/short slice as real counter readings. +func (l *limiter) storeIncrAndGetMulti(entries []counterIncr, reads []string) (incrVals, readVals []int64, ok bool) { + if l.store == nil { + incrVals, readVals, _ = l.fallback.incrAndGetMulti(entries, reads) // fallback never errors + return incrVals, readVals, true + } + if l.storeLatched() { + return l.failPolicyIncrAndGetMulti(entries, reads) + } + incrVals, readVals, err := l.store.incrAndGetMulti(entries, reads) + if err == nil { + return incrVals, readVals, true + } + l.recordStoreFailure(err) + return l.failPolicyIncrAndGetMulti(entries, reads) +} + +// failPolicyIncrAndGetMulti mirrors failPolicyIncrMulti for the fused +// admission round trip. +func (l *limiter) failPolicyIncrAndGetMulti(entries []counterIncr, reads []string) ([]int64, []int64, bool) { + if !l.failOpen { + return nil, nil, false + } + incrVals, readVals, _ := l.fallback.incrAndGetMulti(entries, reads) + return incrVals, readVals, true +} + // storeDownViolation is the violation checkAndCount/budgetViolation return // when the configured store is unreachable and failOpen is false: the // request is refused instead of silently enforcing limits against a @@ -904,30 +978,107 @@ func (l *limiter) getCounter(kind, id, metric, window string, t time.Time) (int6 // counts by. const checkAndCountKeysPerScope = 3 +// budgetProbe names one token/cost budget checkAndCount's fused admission +// round trip must evaluate — the same values the pre-fusion implementation +// checked via its own evaluateScope/budgetViolation/tokenBudgetViolation +// helpers (removed by perf review round 3, 2026-08-22), now inlined into +// checkAndCount itself so their GET reads can ride the same round trip as +// checkAndCount's own request-counter increments (see incrAndGetMulti, +// counterStore's own doc comment above, for why that fusion is safe). +type budgetProbe struct { + name string // e.g. "tokens-per-day" — echoed into the violation message verbatim + window string // windowDay or windowMonth; retryAfterSeconds is computed against this + scopeIdx int // index into the scopes slice checkAndCount was given + limit int64 // already micro-USD for a cost probe (usdToMicros); always > 0 + keyStart int // index into buildBudgetProbes' own reads slice + tokens bool // true: keyStart/keyStart+1 are tokin/tokout, summed; false: keyStart alone is cost +} + +// buildBudgetProbes returns, for every scope in scopes that has a +// TokensPerDay/TokensPerMonth/CostPerDayUSD/CostPerMonthUSD limit set +// (limit<=0 means unlimited and needs no read at all — the same +// "limit<=0 returns immediately, before touching the store" short-circuit +// the pre-fusion budgetViolation/tokenBudgetViolation applied), the probe +// describing that budget and the flat list of counterStore keys +// checkAndCount's fused round trip must read to evaluate it. Probes are +// appended in the SAME order the pre-fusion evaluateScope checked them, +// scope by scope: tokens-per-day, tokens-per-month, cost-per-day, +// cost-per-month within each scope, scopes themselves in their own given +// order — so a caller that walks probes in order after the round trip +// completes reproduces the identical violation precedence. +func buildBudgetProbes(scopes []limitScope, now time.Time) (probes []budgetProbe, reads []string) { + for i, sc := range scopes { + if sc.limits == nil { + continue + } + lim := sc.limits + add := func(name, window string, limit int64, tokens bool) { + if limit <= 0 { + return + } + probes = append(probes, budgetProbe{scopeIdx: i, name: name, window: window, limit: limit, tokens: tokens, keyStart: len(reads)}) + if tokens { + reads = append(reads, + windowKey(sc.kind, sc.id, metricTokIn, window, now), + windowKey(sc.kind, sc.id, metricTokOut, window, now)) + } else { + reads = append(reads, windowKey(sc.kind, sc.id, metricCost, window, now)) + } + } + add("tokens-per-day", windowDay, lim.TokensPerDay, true) + add("tokens-per-month", windowMonth, lim.TokensPerMonth, true) + add("cost-per-day", windowDay, usdToMicros(lim.CostPerDayUSD), false) + add("cost-per-month", windowMonth, usdToMicros(lim.CostPerMonthUSD), false) + } + return probes, reads +} + // checkAndCount increments every scope's req:min, req:day, and req:hour // counters — unconditionally, before any evaluation, so a request that // ultimately gets refused by one scope's limit still counts toward every -// other scope's rate tracking — then evaluates each scope's set limits in -// order and returns the first violation found, or nil if the request may -// proceed. If the batch fails closed (a configured store errored and +// other scope's rate tracking — then evaluates each scope's set limits, in +// order, and returns the first violation found, or nil if the request may +// proceed. If the round trip fails closed (a configured store errored and // failOpen is false), it returns a storeDownViolation immediately — the // request is refused rather than evaluated against partial or // fallback-only counters. // -// Every scope's three counters are built into ONE storeIncrMulti call -// (perf review, 2026-08-21): a 3-scope request (user, group, total — -// withTotalScope, routes_unified.go) previously paid up to 9 separate -// round trips here, one per counter. incrMulti's own reply already -// carries each entry's post-increment value in order, so no extra read -// is needed to recover req:min/req:day for evaluation below — only -// req:hour's slot goes unread, since windowHour is stats-only (its own -// doc comment) and never evaluated. +// Every scope's three increments AND every scope's token/cost budget reads +// ride ONE storeIncrAndGetMulti round trip (perf review round 3, +// 2026-08-22, fusing on top of round 1's own "one round trip for the +// increments" fix): a 2-scope request (user, group — buildLimitScopes, +// routes_unified.go) with both scopes fully limited previously paid that +// one round trip for the increments, PLUS up to 4 more SEPARATE, SERIAL +// round trips per scope for its token/cost budgets (evaluateScope calling +// tokenBudgetViolation twice and budgetViolation twice) — 9 total before +// this fix. incrAndGetMulti's own reply already carries each entry's +// post-increment value AND each budget key's current value, in order, so +// evaluation below needs no further store call at all. This is safe +// because every budget key (tokin/tokout/cost:*) is disjoint from every +// increment's key (req:*), and tokens/cost are only ever written by +// account() AFTER the upstream response completes — never by this call — +// so reading them alongside the increments carries no read-after-write +// hazard; see incrAndGetMulti's own doc comment (counterStore, above) for +// the full argument. +// +// Violation precedence is unchanged from the pre-fusion implementation: +// scopes are evaluated in their own given order (buildLimitScopes puts a +// user before their group; withTotalScope appends the synthetic total +// scope last), and within one scope: requests-per-minute, requests-per-day, +// tokens-per-day, tokens-per-month, cost-per-day, cost-per-month — the +// first violation found anywhere in that walk is returned immediately, and +// evaluation of any later scope never runs. evaluateScope/budgetViolation/ +// tokenBudgetViolation (the pre-fusion, separate-round-trip +// implementation) are removed by this change — buildBudgetProbes plus the +// loop below inline their exact logic instead, which is what lets their +// store reads ride this one round trip. // // Requests limits compare against the value just incremented in this call. // Token and cost limits compare against the value already accumulated by -// account (via get): a request that would start over budget is refused, -// but because usage is only known after the upstream call completes, one -// request may still push a scope over its budget — an accepted trade-off. +// account (via this call's own reads): a request that would start over +// budget is refused, but because usage is only known after the upstream +// call completes, one request may still push a scope over its budget — an +// accepted trade-off, unchanged from before this round. func (l *limiter) checkAndCount(scopes []limitScope) *limitViolation { now := l.now() @@ -939,58 +1090,43 @@ func (l *limiter) checkAndCount(scopes []limitScope) *limitViolation { newCounterIncr(sc.kind, sc.id, metricReq, windowHour, now, 1, hourWindowTTL), ) } + probes, reads := buildBudgetProbes(scopes, now) - vals, ok := l.storeIncrMulti(entries) - if !ok || len(vals) != len(entries) { + incrVals, readVals, ok := l.storeIncrAndGetMulti(entries, reads) + if !ok || len(incrVals) != len(entries) || len(readVals) != len(reads) { return storeDownViolation() } for i, sc := range scopes { - minCount := vals[i*checkAndCountKeysPerScope] - dayCount := vals[i*checkAndCountKeysPerScope+1] - // vals[i*checkAndCountKeysPerScope+2] is the req:hour count — + if sc.limits == nil { + continue + } + minCount := incrVals[i*checkAndCountKeysPerScope] + dayCount := incrVals[i*checkAndCountKeysPerScope+1] + // incrVals[i*checkAndCountKeysPerScope+2] is the req:hour count — // stats-only, deliberately never read here. - if v := l.evaluateScope(sc, minCount, dayCount, now); v != nil { + if v := requestLimitViolation(sc, "requests-per-minute", sc.limits.RequestsPerMinute, minCount, windowMin, now); v != nil { return v } - } - return nil -} - -// evaluateScope checks one scope's already-set limits against its -// just-incremented request counts and its accumulated token/cost counts. -// A nil limits field means nothing to evaluate for this scope — the real, -// commonly-exercised path since the v0.21 accounting fix: buildLimitScopes -// (routes_unified.go) now ALWAYS builds a user and a group scope, so a -// limit-less user or group reaches here on every request, with its req -// counters already incremented by checkAndCount (that increment runs before -// any evaluation) and nothing for this function to enforce against them. -// This is exactly the mechanism that decouples accounting from -// enforcement: the scope is counted unconditionally, evaluated only when it -// actually has a limit configured. -func (l *limiter) evaluateScope(sc limitScope, minCount, dayCount int64, now time.Time) *limitViolation { - if sc.limits == nil { - return nil - } - lim := sc.limits - - if v := requestLimitViolation(sc, "requests-per-minute", lim.RequestsPerMinute, minCount, windowMin, now); v != nil { - return v - } - if v := requestLimitViolation(sc, "requests-per-day", lim.RequestsPerDay, dayCount, windowDay, now); v != nil { - return v - } - if v := l.tokenBudgetViolation(sc, "tokens-per-day", lim.TokensPerDay, windowDay, now); v != nil { - return v - } - if v := l.tokenBudgetViolation(sc, "tokens-per-month", lim.TokensPerMonth, windowMonth, now); v != nil { - return v - } - if v := l.budgetViolation(sc, metricCost, "cost-per-day", usdToMicros(lim.CostPerDayUSD), windowDay, now); v != nil { - return v - } - if v := l.budgetViolation(sc, metricCost, "cost-per-month", usdToMicros(lim.CostPerMonthUSD), windowMonth, now); v != nil { - return v + if v := requestLimitViolation(sc, "requests-per-day", sc.limits.RequestsPerDay, dayCount, windowDay, now); v != nil { + return v + } + for _, p := range probes { + if p.scopeIdx != i { + continue + } + used := readVals[p.keyStart] + if p.tokens { + used += readVals[p.keyStart+1] + } + if used < p.limit { + continue + } + return &limitViolation{ + message: fmt.Sprintf("%s %q exceeded %s budget", sc.kind, sc.id, p.name), + retryAfter: retryAfterSeconds(now, p.window), + } + } } return nil } @@ -1007,60 +1143,6 @@ func requestLimitViolation(sc limitScope, name string, limit, count int64, windo } } -// budgetViolation reports a violation when the scope's accumulated -// counter for metric/window is already at or above limit. limit<=0 means -// unlimited. A failed-closed read (configured store errored, failOpen -// false) returns storeDownViolation rather than silently passing the -// check — a store outage must never look the same as staying under -// budget. -func (l *limiter) budgetViolation(sc limitScope, metric, name string, limit int64, window string, now time.Time) *limitViolation { - if limit <= 0 { - return nil - } - used, ok := l.getCounter(sc.kind, sc.id, metric, window, now) - if !ok { - return storeDownViolation() - } - if used < limit { - return nil - } - return &limitViolation{ - message: fmt.Sprintf("%s %q exceeded %s budget", sc.kind, sc.id, name), - retryAfter: retryAfterSeconds(now, window), - } -} - -// tokenBudgetViolation mirrors budgetViolation for a token limit, which — -// unlike a request or cost limit — enforces a TOTAL budget over two split -// counters (metricTokIn, metricTokOut; v0.2 data-layer task): it reads -// both of the scope's accumulated tok-in/tok-out counters for window in -// ONE storeGetMulti round trip and compares their sum against limit. -// limit<=0 means unlimited. A fail-closed read returns -// storeDownViolation, matching budgetViolation's own contract; name is -// still e.g. "tokens-per-day", so the violation message still says -// "tokens" regardless of which direction pushed the total over. -func (l *limiter) tokenBudgetViolation(sc limitScope, name string, limit int64, window string, now time.Time) *limitViolation { - if limit <= 0 { - return nil - } - keys := []string{ - windowKey(sc.kind, sc.id, metricTokIn, window, now), - windowKey(sc.kind, sc.id, metricTokOut, window, now), - } - vals, ok := l.storeGetMulti(keys) - if !ok || len(vals) != 2 { - return storeDownViolation() - } - used := vals[0] + vals[1] - if used < limit { - return nil - } - return &limitViolation{ - message: fmt.Sprintf("%s %q exceeded %s budget", sc.kind, sc.id, name), - retryAfter: retryAfterSeconds(now, window), - } -} - // account records u's prompt/completion tokens — split into metricTokIn/ // metricTokOut (v0.2 data-layer task) — and costMicros against every // scope's hour, day, and month counters. Each metric/direction is skipped diff --git a/limits_test.go b/limits_test.go index b1b56ed..d436967 100644 --- a/limits_test.go +++ b/limits_test.go @@ -366,6 +366,10 @@ func (s *historyCountingStore) getMulti(keys []string) ([]int64, error) { func (s *historyCountingStore) incrMulti(entries []counterIncr) ([]int64, error) { return make([]int64, len(entries)), nil } +func (s *historyCountingStore) incrAndGetMulti(entries []counterIncr, reads []string) ([]int64, []int64, error) { + readVals, err := s.getMulti(reads) + return make([]int64, len(entries)), readVals, err +} // TestLimiterHistory_StoreDown asserts limiter.history reports ok=false // when the configured store errors and failOpen is false — the caller @@ -404,7 +408,7 @@ func TestWithTotalScope_AppendsUnlimitedTotalScope(t *testing.T) { // TestWithTotalScope_CountedButNeverEvaluated drives checkAndCount many // times against a scopes slice built via withTotalScope: the total // scope's own limits are always nil, so no volume of traffic can ever -// make it violate — evaluateScope's nil-limits check skips it — while its +// make it violate — checkAndCount's own nil-limits check skips it — while its // req counter still climbs by exactly one per call, proving // checkAndCount's counting loop does not skip a nil-limits scope the way // evaluation does. @@ -788,6 +792,9 @@ func (s *erroringStore) incrBy(string, int64, time.Duration) (int64, error) { re func (s *erroringStore) get(string) (int64, error) { return 0, s.err } func (s *erroringStore) getMulti([]string) ([]int64, error) { return nil, s.err } func (s *erroringStore) incrMulti([]counterIncr) ([]int64, error) { return nil, s.err } +func (s *erroringStore) incrAndGetMulti([]counterIncr, []string) ([]int64, []int64, error) { + return nil, nil, s.err +} // TestLimiter_FailOpen_StoreErrorUsesFallback is carried-item (c): a store // error with failOpen=true must not block the request — the operation @@ -832,13 +839,14 @@ func TestLimiter_FailClosed_StoreErrorReturnsStoreDownViolation(t *testing.T) { } // succeedIncrFailGetStore is a counterStore stub whose incrBy/incrMulti -// always succeed and whose get/getMulti always error. It isolates -// budgetViolation's own storeDown branch (limits.go) from checkAndCount's -// earlier incrMulti one: an erroringStore that fails every method makes -// checkAndCount's batched req:min/req:day/req:hour increment fail closed -// and return before evaluateScope — and therefore budgetViolation — ever -// runs, so a test built on it cannot actually prove budgetViolation's own -// branch works. +// always succeed and whose get/getMulti/incrAndGetMulti always error. Its +// incrMulti/get/getMulti stay independently overridable (used by any +// OTHER caller — account, currentUsage, and so on — that still calls them +// separately) even though checkAndCount itself, since the round-3 fusion +// (2026-08-22), calls incrAndGetMulti exclusively: a store that cannot +// serve a read reliably fails that ONE combined round trip as a whole, +// there being no longer a separate "increment succeeded, read failed" +// split to isolate at the store level for checkAndCount specifically. type succeedIncrFailGetStore struct { getErr error } @@ -857,16 +865,18 @@ func (s *succeedIncrFailGetStore) incrMulti(entries []counterIncr) ([]int64, err } return out, nil } +func (s *succeedIncrFailGetStore) incrAndGetMulti([]counterIncr, []string) ([]int64, []int64, error) { + return nil, nil, s.getErr +} // TestLimiter_FailClosed_BudgetReadReturnsStoreDownViolation covers the -// budgetViolation fail-closed path specifically (as opposed to the -// request-counter incrMulti path TestLimiter_FailClosed_ -// StoreErrorReturnsStoreDownViolation covers above): a store error on a -// token/cost budget read also refuses the request when failOpen is false. -// The store's incrMulti succeeds so checkAndCount's batched -// req:min/req:day/req:hour increment passes and evaluateScope actually -// reaches budgetViolation; a TokensPerDay limit makes evaluateScope call -// it. +// budget-read fail-closed path: a store that cannot be trusted to serve a +// token/cost budget read refuses the request when failOpen is false, via +// checkAndCount's own storeDownViolation short-circuit (limits.go) — the +// same contract the pre-fusion budgetViolation helper enforced before +// checkAndCount's admission round trip was fused (perf review round 3, +// 2026-08-22). A TokensPerDay limit is what makes checkAndCount build a +// budget probe/read for this scope at all. func TestLimiter_FailClosed_BudgetReadReturnsStoreDownViolation(t *testing.T) { store := &succeedIncrFailGetStore{getErr: errors.New("boom")} l := newLimiter(store, false) @@ -973,6 +983,16 @@ func (s *countingErrorStore) incrMulti([]counterIncr) ([]int64, error) { return nil, s.err } +// incrAndGetMulti counts against the same incrCalls field incrMulti uses: +// checkAndCount's own admission round trip calls incrAndGetMulti +// exclusively since the round-3 fusion (2026-08-22, perf review), so +// TestLimiter_StoreDownLatch_SkipsStoreCallsWithinWindow's "the store must +// not be called" assertion needs this call counted the same way. +func (s *countingErrorStore) incrAndGetMulti([]counterIncr, []string) ([]int64, []int64, error) { + s.incrCalls++ + return nil, nil, s.err +} + // TestLimiter_StoreDownLatch_SkipsStoreCallsWithinWindow is review-round-3 // item 2: after a store operation fails, every operation for the next // storeDownLatchFor must short-circuit straight to the fail-open policy @@ -1081,17 +1101,21 @@ func TestLimiter_FailPolicyGet(t *testing.T) { // --- perf review, 2026-08-21: single-batch counter writes --- -// countingIncrStore is a counterStore stub whose incrMulti/get/getMulti -// all succeed against a real backing map (so repeated calls actually -// accumulate, proving checkAndCount/account read back correct values -// after batching, not just that batching happened at all), counting how -// many times incrMulti and getMulti were each called. incrBy is never -// exercised by checkAndCount/account after the perf-review batching and -// just returns a zero value. +// countingIncrStore is a counterStore stub whose incrMulti/getMulti/ +// incrAndGetMulti all succeed against a real backing map (so repeated +// calls actually accumulate, proving checkAndCount/account read back +// correct values after batching, not just that batching happened at all), +// counting how many times each was called. account still calls incrMulti +// directly; checkAndCount, since the round-3 fusion (2026-08-22, perf +// review), calls incrAndGetMulti exclusively — so the two counters below +// track genuinely different callers, not the same round trip under two +// names. incrBy is never exercised by either caller and just returns a +// zero value. type countingIncrStore struct { - values map[string]int64 - incrMultiCalls int - getMultiCalls int + values map[string]int64 + incrMultiCalls int + incrAndGetMultiCalls int + getMultiCalls int } func (s *countingIncrStore) incrBy(string, int64, time.Duration) (int64, error) { return 0, nil } @@ -1127,13 +1151,35 @@ func (s *countingIncrStore) incrMulti(entries []counterIncr) ([]int64, error) { return out, nil } -// TestCheckAndCount_OneIncrMultiCallRegardlessOfScopeCount is the -// perf-review (2026-08-21) case for checkAndCount: a 3-scope request -// (user, group, total — 9 counters, min/day/hour x 3 scopes) must cost -// exactly ONE incrMulti round trip, not one per counter, and the batched -// reply must still drive correct evaluation — a real requestsPerMinute -// violation on the 3rd call, with the exact right counter value. -func TestCheckAndCount_OneIncrMultiCallRegardlessOfScopeCount(t *testing.T) { +func (s *countingIncrStore) incrAndGetMulti(entries []counterIncr, reads []string) ([]int64, []int64, error) { + s.incrAndGetMultiCalls++ + if s.values == nil { + s.values = make(map[string]int64) + } + incrVals := make([]int64, len(entries)) + for i, e := range entries { + s.values[e.key] += e.delta + incrVals[i] = s.values[e.key] + } + readVals := make([]int64, len(reads)) + for i, k := range reads { + readVals[i] = s.values[k] + } + return incrVals, readVals, nil +} + +// TestCheckAndCount_OneIncrAndGetMultiCallRegardlessOfScopeCount is the +// perf-review round-3 (2026-08-22) case for checkAndCount: a 3-scope +// request (user, group, total — 9 req counters, min/day/hour x 3 scopes) +// must cost exactly ONE incrAndGetMulti round trip, not one per counter +// and not a separate round trip for token/cost budget reads, and the +// batched reply must still drive correct evaluation — a real +// requestsPerMinute violation on the 3rd call, with the exact right +// counter value. Supersedes the pre-fusion incrMulti-only version of this +// test (perf review, 2026-08-21): checkAndCount no longer calls incrMulti +// at all, only incrAndGetMulti (account still does, covered separately +// below by TestAccount_OneIncrMultiCallRegardlessOfScopeOrMetricCount). +func TestCheckAndCount_OneIncrAndGetMultiCallRegardlessOfScopeCount(t *testing.T) { store := &countingIncrStore{} l := newLimiter(store, true) now := time.Date(2026, 8, 21, 9, 0, 0, 0, time.UTC) @@ -1148,15 +1194,15 @@ func TestCheckAndCount_OneIncrMultiCallRegardlessOfScopeCount(t *testing.T) { if v := l.checkAndCount(scopes); v != nil { t.Fatalf("1st call should pass, got %+v", v) } - if store.incrMultiCalls != 1 { - t.Errorf("incrMultiCalls after 1st checkAndCount = %d, want 1", store.incrMultiCalls) + if store.incrAndGetMultiCalls != 1 { + t.Errorf("incrAndGetMultiCalls after 1st checkAndCount = %d, want 1", store.incrAndGetMultiCalls) } if v := l.checkAndCount(scopes); v != nil { t.Fatalf("2nd call should pass, got %+v", v) } - if store.incrMultiCalls != 2 { - t.Errorf("incrMultiCalls after 2nd checkAndCount = %d, want 2 (one call per checkAndCount, not per scope or per counter)", store.incrMultiCalls) + if store.incrAndGetMultiCalls != 2 { + t.Errorf("incrAndGetMultiCalls after 2nd checkAndCount = %d, want 2 (one call per checkAndCount, not per scope or per counter)", store.incrAndGetMultiCalls) } v := l.checkAndCount(scopes) @@ -1166,8 +1212,11 @@ func TestCheckAndCount_OneIncrMultiCallRegardlessOfScopeCount(t *testing.T) { if v.storeDown { t.Error("v.storeDown must be false — this is a real limit breach, not a store failure") } - if store.incrMultiCalls != 3 { - t.Errorf("incrMultiCalls after 3rd checkAndCount = %d, want 3", store.incrMultiCalls) + if store.incrAndGetMultiCalls != 3 { + t.Errorf("incrAndGetMultiCalls after 3rd checkAndCount = %d, want 3", store.incrAndGetMultiCalls) + } + if store.incrMultiCalls != 0 { + t.Errorf("incrMultiCalls = %d, want 0 (checkAndCount must never call the separate incrMulti after fusion)", store.incrMultiCalls) } // Values correct: batching three scopes into one call must not @@ -1259,6 +1308,168 @@ func TestTokenBudget_ExactBoundaryViolates(t *testing.T) { } } +// --- perf review round 3, 2026-08-22: fused admission round trip --- + +// TestCheckAndCount_FusedViolationPrecedence is the fusion-equivalence +// case: checkAndCount's admission round trip now fuses the request-counter +// increments with every scope's token/cost budget reads into ONE store +// call (buildBudgetProbes, incrAndGetMulti), replacing the pre-fusion +// evaluateScope/budgetViolation/tokenBudgetViolation helpers that made +// those reads as separate, serial round trips. This table pins that the +// VIOLATION PRECEDENCE those helpers enforced is byte-for-byte unchanged: +// within one scope, requests-per-minute, then requests-per-day, then +// tokens-per-day, then tokens-per-month, then cost-per-day, then +// cost-per-month; across scopes, the scope earlier in the given slice +// (buildLimitScopes puts a user before their group) wins. Every subtest +// engineers TWO simultaneously-true violation conditions and asserts only +// the higher-precedence one is ever reported. +func TestCheckAndCount_FusedViolationPrecedence(t *testing.T) { + t.Run("requests-per-minute beats requests-per-day in the same scope", func(t *testing.T) { + l := newLimiter(nil, true) + scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{RequestsPerMinute: 1, RequestsPerDay: 1}}} + if v := l.checkAndCount(scopes); v != nil { + t.Fatalf("1st call must pass (both counters land exactly on their limit), got %+v", v) + } + v := l.checkAndCount(scopes) + if v == nil { + t.Fatal("2nd call must violate: both requests-per-minute and requests-per-day are now over limit") + } + assert.Contains(t, v.message, "requests-per-minute") + assert.NotContains(t, v.message, "requests-per-day") + }) + + t.Run("requests-per-day beats tokens-per-day in the same scope", func(t *testing.T) { + l := newLimiter(nil, true) + scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{RequestsPerDay: 1, TokensPerDay: 10}}} + if v := l.checkAndCount(scopes); v != nil { + t.Fatalf("1st call must pass (req:day lands exactly on its limit, no tokens seeded yet), got %+v", v) + } + l.account(scopes, usage{prompt: 20}, 0) // tokin:day now 20, over the 10 budget + v := l.checkAndCount(scopes) + if v == nil { + t.Fatal("2nd call must violate: both requests-per-day and tokens-per-day are now over limit") + } + assert.Contains(t, v.message, "requests-per-day") + assert.NotContains(t, v.message, "tokens-per-day") + }) + + t.Run("tokens-per-day beats tokens-per-month in the same scope", func(t *testing.T) { + l := newLimiter(nil, true) + scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{TokensPerDay: 10, TokensPerMonth: 10}}} + l.account(scopes, usage{prompt: 20}, 0) // tokin:day and tokin:month both now 20, both over budget + v := l.checkAndCount(scopes) + if v == nil { + t.Fatal("must violate: both tokens-per-day and tokens-per-month are over limit") + } + assert.Contains(t, v.message, "tokens-per-day") + assert.NotContains(t, v.message, "tokens-per-month") + }) + + t.Run("tokens-per-month beats cost-per-day in the same scope", func(t *testing.T) { + l := newLimiter(nil, true) + scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{TokensPerMonth: 10, CostPerDayUSD: 0.0001}}} + l.account(scopes, usage{prompt: 20}, 200) // tokin:month = 20 (>10); cost:day = 200 micros (>100 micros) + v := l.checkAndCount(scopes) + if v == nil { + t.Fatal("must violate: both tokens-per-month and cost-per-day are over limit") + } + assert.Contains(t, v.message, "tokens-per-month") + assert.NotContains(t, v.message, "cost-per-day") + }) + + t.Run("cost-per-day beats cost-per-month in the same scope", func(t *testing.T) { + l := newLimiter(nil, true) + scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{CostPerDayUSD: 0.0001, CostPerMonthUSD: 0.0001}}} + l.account(scopes, usage{}, 200) // cost:day and cost:month both now 200 micros, both over the 100-micro budget + v := l.checkAndCount(scopes) + if v == nil { + t.Fatal("must violate: both cost-per-day and cost-per-month are over limit") + } + assert.Contains(t, v.message, "cost-per-day") + assert.NotContains(t, v.message, "cost-per-month") + }) + + t.Run("user scope beats group scope when both violate", func(t *testing.T) { + l := newLimiter(nil, true) + scopes := []limitScope{ + {kind: "user", id: "u", limits: &LimitsConfig{RequestsPerMinute: 1}}, + {kind: "group", id: "g", limits: &LimitsConfig{RequestsPerMinute: 1}}, + } + if v := l.checkAndCount(scopes); v != nil { + t.Fatalf("1st call must pass, got %+v", v) + } + v := l.checkAndCount(scopes) + if v == nil { + t.Fatal("2nd call must violate: both user and group are now over their requests-per-minute limit") + } + assert.Contains(t, v.message, `user "u"`) + }) + + t.Run("group scope evaluated when the earlier user scope has no violation", func(t *testing.T) { + l := newLimiter(nil, true) + scopes := []limitScope{ + {kind: "user", id: "u", limits: nil}, + {kind: "group", id: "g", limits: &LimitsConfig{RequestsPerMinute: 1}}, + } + if v := l.checkAndCount(scopes); v != nil { + t.Fatalf("1st call must pass, got %+v", v) + } + v := l.checkAndCount(scopes) + if v == nil { + t.Fatal("2nd call must violate: group is now over its requests-per-minute limit") + } + assert.Contains(t, v.message, `group "g"`) + }) +} + +// TestCheckAndCount_FusedRoundTrip_MemoryFallback proves the fused +// admission round trip works correctly end to end against memoryStore — +// the limiter's ALWAYS-available in-process fallback (newLimiter(nil, _) +// uses it directly; a configured store also falls back to it on error or +// while store-latched) — covering both an increment-only scope and a +// scope with every budget type configured at once, in one round trip. +func TestCheckAndCount_FusedRoundTrip_MemoryFallback(t *testing.T) { + l := newLimiter(nil, true) + now := time.Date(2026, 8, 22, 9, 0, 0, 0, time.UTC) + l.nowFn = func() time.Time { return now } + + scopes := []limitScope{ + {kind: "user", id: "u", limits: &LimitsConfig{ + RequestsPerMinute: 100, RequestsPerDay: 100, + TokensPerDay: 1000, TokensPerMonth: 1000, + CostPerDayUSD: 1, CostPerMonthUSD: 1, + }}, + {kind: totalScopeKind, id: totalScopeID, limits: nil}, + } + + // Well under every budget: must pass and must have actually + // incremented every request counter (proving the fused write side + // still lands correctly, not just that no violation was reported). + if v := l.checkAndCount(scopes); v != nil { + t.Fatalf("want no violation comfortably under every limit, got %+v", v) + } + uMin, ok := l.getCounter("user", "u", metricReq, windowMin, now) + if !ok || uMin != 1 { + t.Errorf("user req:min = %d, ok=%v, want 1", uMin, ok) + } + totalMin, ok := l.getCounter(totalScopeKind, totalScopeID, metricReq, windowMin, now) + if !ok || totalMin != 1 { + t.Errorf("total req:min = %d, ok=%v, want 1 (the unlimited total scope is still counted)", totalMin, ok) + } + + // Push tokens over budget via account, then confirm the SAME fused + // round trip catches it on the very next call. + l.account(scopes, usage{prompt: 2000}, 0) + v := l.checkAndCount(scopes) + if v == nil { + t.Fatal("want a tokens-per-day violation now that usage exceeds the 1000 budget") + } + if v.storeDown { + t.Error("v.storeDown must be false — this is a real limit breach, not a store failure") + } + assert.Contains(t, v.message, "tokens-per-day") +} + // --- review round 2, 2026-08-21: enforcement-aware fallback TTL clamp --- // TestClampTTL_ExactMathTable pins clampTTL's ceiling-vs-floor precedence diff --git a/redis_store.go b/redis_store.go index d6563ce..3ab9887 100644 --- a/redis_store.go +++ b/redis_store.go @@ -120,6 +120,75 @@ func (s *redisStore) incrMulti(entries []counterIncr) ([]int64, error) { return out, nil } +// incrAndGetMulti implements counterStore: every entry's INCRBY+EXPIRE pair +// AND every read key's GET, sent as ONE pipeline (perf review round 3, +// 2026-08-22) — 2*len(entries)+len(reads) commands, one round trip +// regardless of how many of each there are, fusing checkAndCount's own +// request-counter increments with its token/cost budget reads (limits.go's +// counterStore doc comment explains why that fusion is safe). Reply layout +// mirrors incrMulti's own: entry i's INCRBY reply sits at reply index 2*i +// (EXPIRE's reply at 2*i+1 is not otherwise inspected, same reasoning as +// incrMulti's own doc comment), followed by one GET reply per read key, in +// order — a missing read key (RESP null bulk) maps to 0, matching +// getBatch's own convention. The same at-least-once semantics as incrMulti +// (see its own doc comment) apply to the INCRBY half of this call; the GET +// half is read-only and carries no such caveat. +func (s *redisStore) incrAndGetMulti(entries []counterIncr, reads []string) ([]int64, []int64, error) { + cmds := make([][]string, 0, len(entries)*2+len(reads)) + for _, e := range entries { + cmds = append(cmds, + []string{"INCRBY", e.key, strconv.FormatInt(e.delta, 10)}, + []string{"EXPIRE", e.key, strconv.FormatInt(ttlToSeconds(e.ttl), 10)}, + ) + } + for _, k := range reads { + cmds = append(cmds, []string{"GET", k}) + } + + replies, err := s.client.pipeline(cmds) + if err != nil { + return nil, nil, fmt.Errorf("redisStore: incrAndGetMulti: %w", err) + } + if len(replies) != len(cmds) { + return nil, nil, fmt.Errorf("redisStore: incrAndGetMulti: expected %d replies, got %d", len(cmds), len(replies)) + } + + incrVals := make([]int64, len(entries)) + for i, e := range entries { + reply := replies[i*2] + if re, ok := reply.(error); ok { + return nil, nil, fmt.Errorf("redisStore: incrAndGetMulti %q: INCRBY failed: %w", e.key, re) + } + v, ok := reply.(int64) + if !ok { + return nil, nil, fmt.Errorf("redisStore: incrAndGetMulti %q: unexpected INCRBY reply type %T", e.key, reply) + } + incrVals[i] = v + } + + base := len(entries) * 2 + readVals := make([]int64, len(reads)) + for i, k := range reads { + reply := replies[base+i] + if reply == nil { + continue // missing key -> 0, matching getBatch's own convention + } + if re, ok := reply.(error); ok { + return nil, nil, fmt.Errorf("redisStore: incrAndGetMulti %q: GET failed: %w", k, re) + } + b, ok := reply.([]byte) + if !ok { + return nil, nil, fmt.Errorf("redisStore: incrAndGetMulti %q: unexpected GET reply type %T", k, reply) + } + n, perr := strconv.ParseInt(string(b), 10, 64) + if perr != nil { + return nil, nil, fmt.Errorf("redisStore: incrAndGetMulti %q: non-integer value %q: %w", k, b, perr) + } + readVals[i] = n + } + return incrVals, readVals, nil +} + // getMulti implements counterStore: one pipelined GET per key // (respClient.getBatch), a single round trip regardless of len(keys) — // used by limiter.currentUsage (admin dashboard, spec §4) to read a diff --git a/redis_store_test.go b/redis_store_test.go index a7188c5..083f796 100644 --- a/redis_store_test.go +++ b/redis_store_test.go @@ -70,6 +70,69 @@ func TestRedisStore_IncrMulti_PipelinesAllPairsInOneCall(t *testing.T) { } } +// TestRedisStore_IncrAndGetMulti_PipelinesEverythingInOneCall is the +// perf-review round-3 (2026-08-22) case: incrAndGetMulti sends every +// entry's INCRBY+EXPIRE pair AND every read key's GET as ONE pipeline — +// no SELECT (or any other command) interleaved, which would only happen +// if it dialled a fresh connection or split the call into more than one +// round trip — and returns entries' INCRBY replies and reads' GET replies +// each in their own given order, not the reply stream's flat position. +func TestRedisStore_IncrAndGetMulti_PipelinesEverythingInOneCall(t *testing.T) { + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"INCRBY", "llmgw:user:a:req:min:202608220900", "1"}, reply: []byte(":3\r\n")}, + {wantArgs: []string{"EXPIRE", "llmgw:user:a:req:min:202608220900", "120"}, reply: []byte(":1\r\n")}, + {wantArgs: []string{"INCRBY", "llmgw:user:a:req:day:20260822", "1"}, reply: []byte(":7\r\n")}, + {wantArgs: []string{"EXPIRE", "llmgw:user:a:req:day:20260822", "3024000"}, reply: []byte(":1\r\n")}, + {wantArgs: []string{"GET", "llmgw:user:a:tokin:day:20260822"}, reply: []byte("$3\r\n150\r\n")}, + {wantArgs: []string{"GET", "llmgw:user:a:tokout:day:20260822"}, reply: []byte("$-1\r\n")}, // missing -> 0 + {wantArgs: []string{"GET", "llmgw:user:a:cost:day:20260822"}, reply: []byte("$2\r\n50\r\n")}, + }) + + store := newRedisStore(newRESPClient(ln.Addr().String(), "", 0)) + entries := []counterIncr{ + {key: "llmgw:user:a:req:min:202608220900", delta: 1, ttl: minWindowTTL}, + {key: "llmgw:user:a:req:day:20260822", delta: 1, ttl: dayWindowTTL}, + } + reads := []string{ + "llmgw:user:a:tokin:day:20260822", + "llmgw:user:a:tokout:day:20260822", + "llmgw:user:a:cost:day:20260822", + } + incrVals, readVals, err := store.incrAndGetMulti(entries, reads) + if err != nil { + t.Fatalf("incrAndGetMulti: %v", err) + } + assert.Equal(t, []int64{3, 7}, incrVals) + assert.Equal(t, []int64{150, 0, 50}, readVals) +} + +// TestRedisStore_IncrAndGetMulti_EmptyReads asserts a nil/empty reads +// slice (no scope in the request has a token/cost budget configured — the +// common case) still sends and returns the increment half correctly, with +// no GET commands at all. +func TestRedisStore_IncrAndGetMulti_EmptyReads(t *testing.T) { + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"INCRBY", "k", "1"}, reply: []byte(":1\r\n")}, + {wantArgs: []string{"EXPIRE", "k", "60"}, reply: []byte(":1\r\n")}, + }) + + store := newRedisStore(newRESPClient(ln.Addr().String(), "", 0)) + incrVals, readVals, err := store.incrAndGetMulti([]counterIncr{{key: "k", delta: 1, ttl: time.Minute}}, nil) + if err != nil { + t.Fatalf("incrAndGetMulti: %v", err) + } + if len(incrVals) != 1 || incrVals[0] != 1 { + t.Errorf("incrVals = %v, want [1]", incrVals) + } + if len(readVals) != 0 { + t.Errorf("readVals = %v, want empty", readVals) + } +} + // TestRedisStore_IncrBy_RoundsSubSecondTTLUp asserts a ttl under one // second is never sent to EXPIRE as 0 (which would delete the key // immediately) — it is rounded up to a 1s floor. diff --git a/routes_unified.go b/routes_unified.go index 32a9e43..d8523ed 100644 --- a/routes_unified.go +++ b/routes_unified.go @@ -234,7 +234,7 @@ func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, gr // same request. // // u.limits or grp.limits may be nil — that scope simply carries no limit for -// evaluateScope (limits.go) to enforce, and evaluateScope's own nil-limits +// checkAndCount (limits.go) to enforce, and checkAndCount's own nil-limits // check skips it during evaluation. It is NOT omitted from the slice: usage // accounting (checkAndCount's own counting half, and account) is // unconditional and must never be coupled to whether a limit happens to be diff --git a/routes_unified_test.go b/routes_unified_test.go index cb703ea..936bd9f 100644 --- a/routes_unified_test.go +++ b/routes_unified_test.go @@ -87,6 +87,10 @@ func (alwaysErrStore) incrMulti([]counterIncr) ([]int64, error) { return nil, errStoreDownStub } +func (alwaysErrStore) incrAndGetMulti([]counterIncr, []string) ([]int64, []int64, error) { + return nil, nil, errStoreDownStub +} + // TestHandleChat_HappyPath_NonStreaming_AccountsUsage drives a full // chat-completion request through Gateway.ServeHTTP against a fake // openai-shaped upstream, and asserts the response is forwarded verbatim @@ -1755,7 +1759,7 @@ func TestCacheCaptureWriter_ImplicitStatus_CapturesContentType(t *testing.T) { // builds a user AND a group limitScope unconditionally (v0.21 accounting // fix): whether either entity has Limits configured or not, both scopes // must always be present — a nil Limits field turns off ENFORCEMENT -// (evaluateScope's own nil check, limits.go) for that scope, never +// (checkAndCount's own nil check, limits.go) for that scope, never // accounting. Table-driven over every nil/non-nil combination of user and // group limits. func TestBuildLimitScopes_AlwaysBuildsBothScopes(t *testing.T) {