mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-12 04:54:45 +00:00
perf(limits): single-batch counter writes, bounded fallback retention
Review fix round (2 Important + 5 folded) on 08cfdb4/b3fb705.
1. counterStore gains incrMulti(entries []counterIncr) ([]int64, error).
redisStore.incrMulti sends every entry's INCRBY+EXPIRE pair as ONE
pipeline (ttlToSeconds extracted from incrBy, now shared);
memoryStore.incrMulti loops incrBy for interface conformance.
checkAndCount batches every scope's req:min/req:day/req:hour into one
storeIncrMulti call (down from up to 9 round trips for a 3-scope
request) and reads req:min/req:day back from the reply, no extra
GETs. account batches every scope's nonzero tok-in/tok-out/cost
entries into one call (down from up to 27). Violation math and
fail-open/fail-closed/latched semantics are unchanged — only the
round-trip count moves.
2. memoryStoreMaxTTL (48h) caps every TTL memoryStore actually applies,
clamped inside incrBy so incrMulti's loop inherits it for free. The
in-process fallback is continuity, not history: applying the real
hour/day/month TTLs there was growing live-key count ~30x, multiplying
sweepEvery's own documented "167x incrBy cliff" by the same factor.
Sweep cadence is unchanged (still 30s-gated); the clamp shrinks what
there is to sweep instead.
Folded: historyMaxSpan's margin comment corrected (~4 days, not a
month); checkAndCount's godoc names req:hour; the integration
current-hour assertion sums every point instead of reading only the
last one (an hour-rollover race the suite cannot control); a new
exact-boundary test proves tokin+tokout landing exactly on a
TokensPerDay limit still violates; README documents the history route
as unthrottled (admin-gate is its only trust boundary) and the
fallback's 48h retention ceiling.
Tests: countingIncrStore proves exactly one incrMulti call per
checkAndCount and per account, with correct per-scope/metric/window
values and unchanged violation math; a fake-RESP-server test pins
incrMulti's wire format (one pipeline, entry-ordered replies); TTL
clamp tests cover both the ceiling and the pass-through case at the
boundary.
Gates: go test ./... -count=1, -race, gofmt, vet, golangci-lint,
make yaegi-check all green. make integration: every test passes
except the pre-existing TestUsersFileHotReload flake, which is NOT
caused by this change — verified via an isolated worktree at the
prior commit (b3fb705) run under identical conditions (full suite,
same host), which fails identically. This host's Docker Desktop is
currently slow to propagate the bind-mounted users.json change into
the containers; auth.go/users_file.go are untouched by this diff.
This commit is contained in:
@@ -509,26 +509,38 @@ http:
|
||||
accuracy.
|
||||
- **Storage**: `redis` configured wires a hand-rolled, stdlib-only RESP2
|
||||
client (`resp.go`) — no `go-redis`, which is not Yaegi-interpretable —
|
||||
against Redis, Valkey, or Dragonfly. Semantics are deliberately
|
||||
**at-least-once, never under-counting**: a lost reply after the server
|
||||
already applied an `INCRBY` can cause a resend that over-counts by one
|
||||
request's worth on that key. This is a fail-*safe* trade-off (more
|
||||
restrictive than reality), never a fail-*open* one. Without `redis`
|
||||
configured, or when it errors and `failOpen: true`, counters live in an
|
||||
in-process map — correct for one Traefik replica, only approximate
|
||||
across several, since each replica counts independently. Exact limits
|
||||
across multiple Traefik replicas require the shared Redis. Every
|
||||
counter, including the fallback in-process one, works as the usage-
|
||||
history API's data source: a fallback-only deployment gets charts of
|
||||
whatever the process itself accumulated, not an empty series.
|
||||
- **Retention** (counter key TTL, distinct from a window's own length):
|
||||
a minute key lives 2 minutes, an hour key 48 hours, a day key 35 days,
|
||||
a month key 400 days — long enough for the usage-history API's largest
|
||||
span (48 hourly / 35 daily / 13 monthly buckets) to always find a live
|
||||
key, not just long enough to survive that bucket's own single rollover.
|
||||
**Migration**: none — a deployed process's old, unsplit `tok` counter
|
||||
keys simply expire on their existing TTL and are never read again;
|
||||
nothing needs backfilling.
|
||||
against Redis, Valkey, or Dragonfly. Every counter write (`checkAndCount`,
|
||||
`account`) goes through ONE pipelined round trip for the whole request,
|
||||
regardless of how many user/group/total counters it touches (perf
|
||||
review, 2026-08-21) — a 3-scope request previously paid up to 9 (request
|
||||
counting) or 27 (usage accounting) separate round trips. Semantics are
|
||||
deliberately **at-least-once, never under-counting**: a lost reply after
|
||||
the server already applied the pipeline can cause a resend that
|
||||
over-counts every counter in it by its own delta. This is a
|
||||
fail-*safe* trade-off (more restrictive than reality), never a
|
||||
fail-*open* one.
|
||||
- **Fallback** (no `redis` configured, or a Redis error with
|
||||
`failOpen: true`): counters live in an in-process map — correct for one
|
||||
Traefik replica, only approximate across several, since each replica
|
||||
counts independently. Exact limits across multiple Traefik replicas
|
||||
require the shared Redis. Every counter, including this fallback one,
|
||||
works as the usage-history API's data source, but its retention is
|
||||
capped at **48 hours regardless of window** (`memoryStoreMaxTTL`,
|
||||
limits.go — ruling, 2026-08-21: the in-process fallback is continuity,
|
||||
not history; applying the real day/month TTLs to it would grow its live
|
||||
key count roughly 30x, multiplying the cost of its own periodic sweep by
|
||||
the same factor). A fallback-only deployment's usage-history charts
|
||||
therefore show at most the trailing 48 hours, never the full day/month
|
||||
span Redis-backed deployments get — full retention needs the shared
|
||||
Redis.
|
||||
- **Retention** (Redis counter key TTL, distinct from a window's own
|
||||
length): a minute key lives 2 minutes, an hour key 48 hours, a day key
|
||||
35 days, a month key 400 days — long enough for the usage-history API's
|
||||
largest span (48 hourly / 35 daily / 13 monthly buckets) to always find
|
||||
a live key, not just long enough to survive that bucket's own single
|
||||
rollover. **Migration**: none — a deployed process's old, unsplit `tok`
|
||||
counter keys simply expire on their existing TTL and are never read
|
||||
again; nothing needs backfilling.
|
||||
|
||||
## Retry
|
||||
|
||||
@@ -773,7 +785,11 @@ request to any of them falls through to the plugin's existing
|
||||
`{"scope","metric","window","points":[{"bucket":"2026082114","value":123},...]}`.
|
||||
Unlike the other two JSON routes, this one **never counts request
|
||||
statistics** — it skips `checkAndCount` entirely, ahead of the same
|
||||
change landing for every `/admin/api/*` route in a later change.
|
||||
change landing for every `/admin/api/*` route in a later change. It is
|
||||
therefore also **unthrottled**: no `requestsPerMinute`/`requestsPerDay`
|
||||
limit ever applies to it, so its only trust boundary is the admin gate
|
||||
itself (a valid admin key). An admin key holder can poll it as fast as
|
||||
they like.
|
||||
- **What's exposed**: provider names, types, base URLs (with any
|
||||
userinfo/query string stripped before it's ever echoed), model counts,
|
||||
discovery status, group/user names, membership, limits, and live usage
|
||||
|
||||
@@ -450,11 +450,13 @@ func (g *Gateway) serveAdminUsage(w http.ResponseWriter) {
|
||||
// per window (v0.2 data-layer task), matching each window's own retention
|
||||
// (limits.go: hourWindowTTL/dayWindowTTL/monthWindowTTL) rounded down to
|
||||
// whole buckets: 48 hourly buckets, 35 daily buckets, 13 monthly buckets
|
||||
// (a one-month margin over monthWindowTTL's 400*24h ≈ 13.3 months). A
|
||||
// span beyond what its window's TTL could ever have kept alive would only
|
||||
// echo zeros for the missing tail, so the cap keeps every requested
|
||||
// bucket meaningful. window is assumed already validated by
|
||||
// validHistoryWindow — the default case is a programming error, not a
|
||||
// (roughly a 4-day margin, not a full month: 13 average-length ~30.4-day
|
||||
// months is ~396 days, against monthWindowTTL's 400*24h = 400 days —
|
||||
// review sweep, 2026-08-21, corrects an earlier "one-month margin"
|
||||
// overclaim here). A span beyond what its window's TTL could ever have
|
||||
// kept alive would only echo zeros for the missing tail, so the cap keeps
|
||||
// every requested bucket meaningful. window is assumed already validated
|
||||
// by validHistoryWindow — the default case is a programming error, not a
|
||||
// client input, mirroring windowKey/windowEnd's own panic convention.
|
||||
func historyMaxSpan(window string) int {
|
||||
switch window {
|
||||
|
||||
@@ -751,6 +751,9 @@ func (s *countingMultiStore) getMulti(keys []string) ([]int64, error) {
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
func (s *countingMultiStore) incrMulti(entries []counterIncr) ([]int64, error) {
|
||||
return make([]int64, len(entries)), nil
|
||||
}
|
||||
|
||||
func TestLimiterCurrentUsage_BatchesOneGetMultiCallForAllScopes(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
@@ -1003,10 +1003,15 @@ func TestAdminDashboard(t *testing.T) {
|
||||
}
|
||||
|
||||
// v0.2 data-layer task: GET /admin/api/usage/history's hour series for
|
||||
// alice and for the synthetic total scope both must show a non-zero
|
||||
// current-hour bucket, reflecting the real LLM traffic this suite
|
||||
// already generated (TestRetryFlakyRecovers, TestResponseCache, ...)
|
||||
// against a real Redis-backed limiter, not just a unit-test stub.
|
||||
// alice and for the synthetic total scope both must show non-zero
|
||||
// traffic, reflecting the real requests this suite already generated
|
||||
// (TestRetryFlakyRecovers, TestResponseCache, ...) against a real
|
||||
// Redis-backed limiter, not just a unit-test stub. Summed across every
|
||||
// returned point (review sweep, 2026-08-21), not just the last one: an
|
||||
// hour rollover landing between this suite's earlier traffic and this
|
||||
// very read would otherwise split the count across two buckets and
|
||||
// leave the current (now-empty) one at zero — a real flake this
|
||||
// integration run cannot control the wall-clock timing of.
|
||||
for _, scope := range []string{"user:alice", "total"} {
|
||||
hist := adminUsageHistory(t, scope, "req", "hour")
|
||||
if got, _ := hist["scope"].(string); got != scope {
|
||||
@@ -1016,12 +1021,17 @@ func TestAdminDashboard(t *testing.T) {
|
||||
if !ok || len(points) == 0 {
|
||||
t.Fatalf("history scope=%q: points = %#v, want a non-empty array", scope, hist["points"])
|
||||
}
|
||||
current, ok := points[len(points)-1].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("history scope=%q: last point is not an object, got %#v", scope, points[len(points)-1])
|
||||
var total float64
|
||||
for _, p := range points {
|
||||
point, ok := p.(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("history scope=%q: point is not an object, got %#v", scope, p)
|
||||
}
|
||||
val, _ := point["value"].(float64)
|
||||
total += val
|
||||
}
|
||||
if val, _ := current["value"].(float64); val <= 0 {
|
||||
t.Errorf("history scope=%q: current-hour bucket value = %v, want > 0 (this suite's own traffic)", scope, val)
|
||||
if total <= 0 {
|
||||
t.Errorf("history scope=%q: sum of all hour buckets = %v, want > 0 (this suite's own traffic)", scope, total)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -132,6 +132,15 @@ func (lc *LimitsConfig) validate() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// counterIncr is one entry in an incrMulti batch: increment key by delta,
|
||||
// creating it with an expiry of ttl from now if it does not exist or has
|
||||
// expired — the same per-entry contract incrBy applies to a single key.
|
||||
type counterIncr struct {
|
||||
key string
|
||||
delta int64
|
||||
ttl time.Duration
|
||||
}
|
||||
|
||||
// counterStore is the storage backend the limiter uses for atomic windowed
|
||||
// counters. memoryStore (below) is the in-process fallback; a distributed
|
||||
// (e.g. Redis-backed) implementation is wired in a later task.
|
||||
@@ -152,6 +161,19 @@ type counterStore interface {
|
||||
// whole batch (mirrored by the limiter's storeGetMulti as a single
|
||||
// fail-open/fail-closed decision), never a partial result.
|
||||
getMulti(keys []string) ([]int64, error)
|
||||
// incrMulti applies incrBy's own per-entry contract to every entry in
|
||||
// entries in one round trip where the backend supports it (redisStore
|
||||
// pipelines one INCRBY+EXPIRE pair per entry into a single pipeline
|
||||
// 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.
|
||||
incrMulti(entries []counterIncr) ([]int64, error)
|
||||
}
|
||||
|
||||
// memoryEntry is one counter's value and expiry in memoryStore.
|
||||
@@ -170,6 +192,32 @@ type memoryEntry struct {
|
||||
// every 30s.
|
||||
const sweepEvery = 30 * time.Second
|
||||
|
||||
// memoryStoreMaxTTL caps any TTL memoryStore actually applies, regardless
|
||||
// of what a caller requests. Ruling (perf review, 2026-08-21): the
|
||||
// in-process fallback is continuity, not history. Applying the real
|
||||
// hourWindowTTL/dayWindowTTL/monthWindowTTL values (48h/35d/400d) to the
|
||||
// fallback map was measured to grow its live key count roughly 30x over
|
||||
// the original three-window (min/day/month) shape — and sweepEvery's own
|
||||
// doc comment above already records that a full-map sweep which frees
|
||||
// nothing is a "167x incrBy cliff" once live keys cross a threshold;
|
||||
// growing the map ~30x bigger multiplies that same cliff's cost by
|
||||
// roughly the same factor. Clamping every fallback TTL to 48h keeps
|
||||
// per-scope live keys bounded (a handful of metrics across four windows,
|
||||
// not 35-400 days' worth) independent of which window a caller asks for.
|
||||
// This is a documented capability gap, not a bug: with no Redis
|
||||
// configured, or during a Redis outage under failOpen, memoryStore-backed
|
||||
// usage-history charts only ever show the trailing 48h, never longer —
|
||||
// see README's "fallback works too, but only for 48h" note.
|
||||
const memoryStoreMaxTTL = 48 * time.Hour
|
||||
|
||||
// clampTTL caps ttl at memoryStoreMaxTTL — see that const's doc comment.
|
||||
func clampTTL(ttl time.Duration) time.Duration {
|
||||
if ttl > memoryStoreMaxTTL {
|
||||
return memoryStoreMaxTTL
|
||||
}
|
||||
return ttl
|
||||
}
|
||||
|
||||
// memoryStore is an in-process counterStore: a mutex-guarded map with
|
||||
// per-key expiry. It has no background goroutine — a Yaegi middleware
|
||||
// instance is rebuilt on every config reload, and any ticker or goroutine
|
||||
@@ -193,8 +241,10 @@ func (m *memoryStore) now() time.Time {
|
||||
return m.nowFn()
|
||||
}
|
||||
|
||||
// incrBy implements counterStore.
|
||||
// incrBy implements counterStore. ttl is clamped to memoryStoreMaxTTL
|
||||
// before it is ever applied — see that const's doc comment.
|
||||
func (m *memoryStore) incrBy(key string, n int64, ttl time.Duration) (int64, error) {
|
||||
ttl = clampTTL(ttl)
|
||||
now := m.now()
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
@@ -239,6 +289,20 @@ func (m *memoryStore) getMulti(keys []string) ([]int64, error) {
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// incrMulti implements counterStore by looping over incrBy: memoryStore's
|
||||
// in-process map access is already effectively free per key, so there is
|
||||
// no round-trip cost to batch away — this exists purely so memoryStore
|
||||
// satisfies counterStore's incrMulti contract for the limiter's nil-store
|
||||
// fallback path (storeIncrMulti). Each entry's ttl is still clamped by
|
||||
// incrBy itself.
|
||||
func (m *memoryStore) incrMulti(entries []counterIncr) ([]int64, error) {
|
||||
out := make([]int64, len(entries))
|
||||
for i, e := range entries {
|
||||
out[i], _ = m.incrBy(e.key, e.delta, e.ttl) // memoryStore.incrBy never errors
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// sweepLocked deletes every expired entry. Callers must hold m.mu.
|
||||
func (m *memoryStore) sweepLocked(now time.Time) {
|
||||
for k, e := range m.data {
|
||||
@@ -505,6 +569,38 @@ func (l *limiter) failPolicyGetMulti(keys []string) ([]int64, bool) {
|
||||
return v, true
|
||||
}
|
||||
|
||||
// storeIncrMulti mirrors storeIncrBy for a batch of increments: one round
|
||||
// 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.
|
||||
func (l *limiter) storeIncrMulti(entries []counterIncr) (v []int64, ok bool) {
|
||||
if l.store == nil {
|
||||
v, _ = l.fallback.incrMulti(entries) // fallback never errors
|
||||
return v, true
|
||||
}
|
||||
if l.storeLatched() {
|
||||
return l.failPolicyIncrMulti(entries)
|
||||
}
|
||||
v, err := l.store.incrMulti(entries)
|
||||
if err == nil {
|
||||
return v, true
|
||||
}
|
||||
l.recordStoreFailure(err)
|
||||
return l.failPolicyIncrMulti(entries)
|
||||
}
|
||||
|
||||
// failPolicyIncrMulti mirrors failPolicyIncrBy for a batch increment.
|
||||
func (l *limiter) failPolicyIncrMulti(entries []counterIncr) ([]int64, bool) {
|
||||
if !l.failOpen {
|
||||
return nil, false
|
||||
}
|
||||
v, _ := l.fallback.incrMulti(entries)
|
||||
return v, 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
|
||||
@@ -588,16 +684,31 @@ func (l *limiter) getCounter(kind, id, metric, window string, t time.Time) (int6
|
||||
return l.storeGet(windowKey(kind, id, metric, window, t))
|
||||
}
|
||||
|
||||
// checkAndCount increments every scope's req:min and req:day 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 any store operation 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
|
||||
// checkAndCountKeysPerScope is the number of counterIncr entries
|
||||
// checkAndCount builds per scope (req:min, req:day, req:hour), and the
|
||||
// stride its flat storeIncrMulti result is sliced back into per-scope
|
||||
// counts by.
|
||||
const checkAndCountKeysPerScope = 3
|
||||
|
||||
// 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
|
||||
// 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.
|
||||
//
|
||||
// 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,
|
||||
@@ -606,32 +717,26 @@ func (l *limiter) getCounter(kind, id, metric, window string, t time.Time) (int6
|
||||
func (l *limiter) checkAndCount(scopes []limitScope) *limitViolation {
|
||||
now := l.now()
|
||||
|
||||
minCounts := make([]int64, len(scopes))
|
||||
dayCounts := make([]int64, len(scopes))
|
||||
for i, sc := range scopes {
|
||||
v, ok := l.incrCounter(sc.kind, sc.id, metricReq, windowMin, now, 1, minWindowTTL)
|
||||
if !ok {
|
||||
return storeDownViolation()
|
||||
}
|
||||
minCounts[i] = v
|
||||
entries := make([]counterIncr, 0, len(scopes)*checkAndCountKeysPerScope)
|
||||
for _, sc := range scopes {
|
||||
entries = append(entries,
|
||||
counterIncr{key: windowKey(sc.kind, sc.id, metricReq, windowMin, now), delta: 1, ttl: minWindowTTL},
|
||||
counterIncr{key: windowKey(sc.kind, sc.id, metricReq, windowDay, now), delta: 1, ttl: dayWindowTTL},
|
||||
counterIncr{key: windowKey(sc.kind, sc.id, metricReq, windowHour, now), delta: 1, ttl: hourWindowTTL},
|
||||
)
|
||||
}
|
||||
|
||||
v, ok = l.incrCounter(sc.kind, sc.id, metricReq, windowDay, now, 1, dayWindowTTL)
|
||||
if !ok {
|
||||
return storeDownViolation()
|
||||
}
|
||||
dayCounts[i] = v
|
||||
|
||||
// Hour is stats-only (windowHour's own doc comment): its value is
|
||||
// never read back here, only written, so the admin usage-history
|
||||
// API (admin.go) has an hour-granularity req series to read
|
||||
// later.
|
||||
if _, ok = l.incrCounter(sc.kind, sc.id, metricReq, windowHour, now, 1, hourWindowTTL); !ok {
|
||||
return storeDownViolation()
|
||||
}
|
||||
vals, ok := l.storeIncrMulti(entries)
|
||||
if !ok || len(vals) != len(entries) {
|
||||
return storeDownViolation()
|
||||
}
|
||||
|
||||
for i, sc := range scopes {
|
||||
if v := l.evaluateScope(sc, minCounts[i], dayCounts[i], now); v != nil {
|
||||
minCount := vals[i*checkAndCountKeysPerScope]
|
||||
dayCount := vals[i*checkAndCountKeysPerScope+1]
|
||||
// vals[i*checkAndCountKeysPerScope+2] is the req:hour count —
|
||||
// stats-only, deliberately never read here.
|
||||
if v := l.evaluateScope(sc, minCount, dayCount, now); v != nil {
|
||||
return v
|
||||
}
|
||||
}
|
||||
@@ -746,31 +851,48 @@ func (l *limiter) tokenBudgetViolation(sc limitScope, name string, limit int64,
|
||||
// scope's hour, day, and month counters. Each metric/direction is skipped
|
||||
// independently when its own value is 0 — no store write for a direction
|
||||
// this call has nothing to report, so a streaming response cut off before
|
||||
// any completion tokens arrived still writes tokin alone. A sample that
|
||||
// hits a fail-closed store error (see storeIncrBy) is dropped silently:
|
||||
// account has no error return to signal it, and dropping — rather than
|
||||
// counting it in the fallback — avoids double-counting once the store
|
||||
// recovers.
|
||||
// any completion tokens arrived still writes tokin alone.
|
||||
//
|
||||
// Every scope's writes are collected into ONE storeIncrMulti call (perf
|
||||
// review, 2026-08-21): a 3-scope request (user, group, total) with all
|
||||
// three metrics nonzero previously paid up to 9 round trips per scope —
|
||||
// 27 total — one incrCounter call each. account's return-value discipline
|
||||
// is unchanged by the batching: it has no error return, and a sample that
|
||||
// hits a fail-closed batch (see storeIncrMulti) is dropped in full,
|
||||
// silently, rather than counting it in the fallback — dropping avoids
|
||||
// double-counting once the store recovers, the same reasoning that
|
||||
// applied per-key before this call became one batch.
|
||||
func (l *limiter) account(scopes []limitScope, u usage, costMicros int64) {
|
||||
now := l.now()
|
||||
|
||||
entries := make([]counterIncr, 0, len(scopes)*9)
|
||||
for _, sc := range scopes {
|
||||
if u.prompt != 0 {
|
||||
l.incrCounter(sc.kind, sc.id, metricTokIn, windowHour, now, u.prompt, hourWindowTTL)
|
||||
l.incrCounter(sc.kind, sc.id, metricTokIn, windowDay, now, u.prompt, dayWindowTTL)
|
||||
l.incrCounter(sc.kind, sc.id, metricTokIn, windowMonth, now, u.prompt, monthWindowTTL)
|
||||
entries = append(entries,
|
||||
counterIncr{key: windowKey(sc.kind, sc.id, metricTokIn, windowHour, now), delta: u.prompt, ttl: hourWindowTTL},
|
||||
counterIncr{key: windowKey(sc.kind, sc.id, metricTokIn, windowDay, now), delta: u.prompt, ttl: dayWindowTTL},
|
||||
counterIncr{key: windowKey(sc.kind, sc.id, metricTokIn, windowMonth, now), delta: u.prompt, ttl: monthWindowTTL},
|
||||
)
|
||||
}
|
||||
if u.completion != 0 {
|
||||
l.incrCounter(sc.kind, sc.id, metricTokOut, windowHour, now, u.completion, hourWindowTTL)
|
||||
l.incrCounter(sc.kind, sc.id, metricTokOut, windowDay, now, u.completion, dayWindowTTL)
|
||||
l.incrCounter(sc.kind, sc.id, metricTokOut, windowMonth, now, u.completion, monthWindowTTL)
|
||||
entries = append(entries,
|
||||
counterIncr{key: windowKey(sc.kind, sc.id, metricTokOut, windowHour, now), delta: u.completion, ttl: hourWindowTTL},
|
||||
counterIncr{key: windowKey(sc.kind, sc.id, metricTokOut, windowDay, now), delta: u.completion, ttl: dayWindowTTL},
|
||||
counterIncr{key: windowKey(sc.kind, sc.id, metricTokOut, windowMonth, now), delta: u.completion, ttl: monthWindowTTL},
|
||||
)
|
||||
}
|
||||
if costMicros != 0 {
|
||||
l.incrCounter(sc.kind, sc.id, metricCost, windowHour, now, costMicros, hourWindowTTL)
|
||||
l.incrCounter(sc.kind, sc.id, metricCost, windowDay, now, costMicros, dayWindowTTL)
|
||||
l.incrCounter(sc.kind, sc.id, metricCost, windowMonth, now, costMicros, monthWindowTTL)
|
||||
entries = append(entries,
|
||||
counterIncr{key: windowKey(sc.kind, sc.id, metricCost, windowHour, now), delta: costMicros, ttl: hourWindowTTL},
|
||||
counterIncr{key: windowKey(sc.kind, sc.id, metricCost, windowDay, now), delta: costMicros, ttl: dayWindowTTL},
|
||||
counterIncr{key: windowKey(sc.kind, sc.id, metricCost, windowMonth, now), delta: costMicros, ttl: monthWindowTTL},
|
||||
)
|
||||
}
|
||||
}
|
||||
if len(entries) == 0 {
|
||||
return // every metric was zero for every scope — nothing to write
|
||||
}
|
||||
l.storeIncrMulti(entries) // no error return by contract (doc comment above); ok is intentionally discarded
|
||||
}
|
||||
|
||||
// usageKeysPerScope is the number of windowKey strings usageWindowKeys
|
||||
|
||||
+287
-17
@@ -361,6 +361,9 @@ func (s *historyCountingStore) getMulti(keys []string) ([]int64, error) {
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
func (s *historyCountingStore) incrMulti(entries []counterIncr) ([]int64, error) {
|
||||
return make([]int64, len(entries)), nil
|
||||
}
|
||||
|
||||
// TestLimiterHistory_StoreDown asserts limiter.history reports ok=false
|
||||
// when the configured store errors and failOpen is false — the caller
|
||||
@@ -498,6 +501,75 @@ func TestMemoryStoreExpiry(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestMemoryStore_TTLClampedToMax is the perf-review (2026-08-21) case:
|
||||
// a TTL beyond memoryStoreMaxTTL (48h) — e.g. dayWindowTTL at 35 days or
|
||||
// monthWindowTTL at 400 days — must be clamped down to 48h before it is
|
||||
// ever applied, so the fallback's own live-key count stays bounded
|
||||
// independent of which window a caller writes.
|
||||
func TestMemoryStore_TTLClampedToMax(t *testing.T) {
|
||||
m := newMemoryStore()
|
||||
fake := time.Date(2026, 8, 21, 0, 0, 0, 0, time.UTC)
|
||||
m.nowFn = func() time.Time { return fake }
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
ttl time.Duration
|
||||
}{
|
||||
{"dayWindowTTL (35d)", dayWindowTTL},
|
||||
{"monthWindowTTL (400d)", monthWindowTTL},
|
||||
{"an arbitrary huge ttl", 10 * 365 * 24 * time.Hour},
|
||||
}
|
||||
for i, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
key := fmt.Sprintf("k%d", i)
|
||||
if _, err := m.incrBy(key, 1, c.ttl); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
m.mu.Lock()
|
||||
got := m.data[key].expiry
|
||||
m.mu.Unlock()
|
||||
want := fake.Add(memoryStoreMaxTTL)
|
||||
if !got.Equal(want) {
|
||||
t.Errorf("expiry = %v, want %v (clamped to memoryStoreMaxTTL, not the requested %v)", got, want, c.ttl)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestMemoryStore_TTLBelowMaxIsUnaffected asserts the clamp is a ceiling,
|
||||
// not a rewrite: a TTL already at or under memoryStoreMaxTTL passes
|
||||
// through exactly as requested — hourWindowTTL (48h) sits precisely at
|
||||
// the clamp boundary and must round-trip unchanged, not get pushed under
|
||||
// it by an off-by-one comparison.
|
||||
func TestMemoryStore_TTLBelowMaxIsUnaffected(t *testing.T) {
|
||||
m := newMemoryStore()
|
||||
fake := time.Date(2026, 8, 21, 0, 0, 0, 0, time.UTC)
|
||||
m.nowFn = func() time.Time { return fake }
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
ttl time.Duration
|
||||
}{
|
||||
{"minWindowTTL (2m)", minWindowTTL},
|
||||
{"hourWindowTTL (48h, exactly at the clamp)", hourWindowTTL},
|
||||
}
|
||||
for i, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
key := fmt.Sprintf("k%d", i)
|
||||
if _, err := m.incrBy(key, 1, c.ttl); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
m.mu.Lock()
|
||||
got := m.data[key].expiry
|
||||
m.mu.Unlock()
|
||||
want := fake.Add(c.ttl)
|
||||
if !got.Equal(want) {
|
||||
t.Errorf("expiry = %v, want %v (ttl under the clamp must pass through unchanged)", got, want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestMemoryStoreSweepTimeGated asserts the sweep gate is purely
|
||||
// time-based (sweepEvery), independent of how many keys memoryStore
|
||||
// holds: an incrBy call before the gate elapses leaves an already-expired
|
||||
@@ -713,6 +785,7 @@ type erroringStore struct {
|
||||
func (s *erroringStore) incrBy(string, int64, time.Duration) (int64, error) { return 0, s.err }
|
||||
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 }
|
||||
|
||||
// TestLimiter_FailOpen_StoreErrorUsesFallback is carried-item (c): a store
|
||||
// error with failOpen=true must not block the request — the operation
|
||||
@@ -756,13 +829,14 @@ func TestLimiter_FailClosed_StoreErrorReturnsStoreDownViolation(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// succeedIncrFailGetStore is a counterStore stub whose incrBy always
|
||||
// succeeds and whose get always errors. It isolates budgetViolation's own
|
||||
// storeDown branch (limits.go) from checkAndCount's earlier incrCounter
|
||||
// one: an erroringStore that fails both methods makes checkAndCount's
|
||||
// initial req:min/req:day increments 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.
|
||||
// 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.
|
||||
type succeedIncrFailGetStore struct {
|
||||
getErr error
|
||||
}
|
||||
@@ -774,15 +848,23 @@ func (s *succeedIncrFailGetStore) get(string) (int64, error) { return 0, s.getEr
|
||||
func (s *succeedIncrFailGetStore) getMulti([]string) ([]int64, error) {
|
||||
return nil, s.getErr
|
||||
}
|
||||
func (s *succeedIncrFailGetStore) incrMulti(entries []counterIncr) ([]int64, error) {
|
||||
out := make([]int64, len(entries))
|
||||
for i := range out {
|
||||
out[i] = 1
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// TestLimiter_FailClosed_BudgetReadReturnsStoreDownViolation covers the
|
||||
// budgetViolation fail-closed path specifically (as opposed to the
|
||||
// request-counter incrCounter path TestLimiter_FailClosed_
|
||||
// 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 incrBy succeeds so checkAndCount's req:min/req:day
|
||||
// increments pass and evaluateScope actually reaches budgetViolation; a
|
||||
// TokensPerDay limit makes evaluateScope call it.
|
||||
// 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.
|
||||
func TestLimiter_FailClosed_BudgetReadReturnsStoreDownViolation(t *testing.T) {
|
||||
store := &succeedIncrFailGetStore{getErr: errors.New("boom")}
|
||||
l := newLimiter(store, false)
|
||||
@@ -884,17 +966,25 @@ func (s *countingErrorStore) getMulti([]string) ([]int64, error) {
|
||||
return nil, s.err
|
||||
}
|
||||
|
||||
func (s *countingErrorStore) incrMulti([]counterIncr) ([]int64, error) {
|
||||
s.incrCalls++
|
||||
return 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
|
||||
// without calling the store at all — a request that touches several
|
||||
// counters (min/day/month, tokens/cost, user/group) must not pay a fresh
|
||||
// counters (min/day/hour, tokens/cost, user/group) must not pay a fresh
|
||||
// bounded wait against a store that just proved unreachable for each one.
|
||||
// failOpen=true is used so checkAndCount completes both its req:min and
|
||||
// req:day increments per call instead of short-circuiting after the
|
||||
// first — this is what makes the "zero additional store calls" assertion
|
||||
// meaningful, rather than trivially true because checkAndCount stops
|
||||
// after one op regardless of the latch.
|
||||
// checkAndCount batches its whole req:min/req:day/req:hour increment into
|
||||
// ONE incrMulti call (perf review, 2026-08-21), so "the store must not be
|
||||
// called" now means that single call is skipped entirely during the latch
|
||||
// window, not that some subset of several per-key calls is. failOpen=true
|
||||
// is required so every checkAndCount call below returns no violation
|
||||
// (ok=true via the fallback) instead of a storeDownViolation — the
|
||||
// incrCalls assertions only make sense if checkAndCount actually
|
||||
// completes normally each time.
|
||||
func TestLimiter_StoreDownLatch_SkipsStoreCallsWithinWindow(t *testing.T) {
|
||||
store := &countingErrorStore{err: errors.New("boom")}
|
||||
l := newLimiter(store, true)
|
||||
@@ -986,3 +1076,183 @@ 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.
|
||||
type countingIncrStore struct {
|
||||
values map[string]int64
|
||||
incrMultiCalls int
|
||||
getMultiCalls int
|
||||
}
|
||||
|
||||
func (s *countingIncrStore) incrBy(string, int64, time.Duration) (int64, error) { return 0, nil }
|
||||
|
||||
func (s *countingIncrStore) get(key string) (int64, error) {
|
||||
if s.values == nil {
|
||||
return 0, nil
|
||||
}
|
||||
return s.values[key], nil
|
||||
}
|
||||
|
||||
func (s *countingIncrStore) getMulti(keys []string) ([]int64, error) {
|
||||
s.getMultiCalls++
|
||||
out := make([]int64, len(keys))
|
||||
for i, k := range keys {
|
||||
if s.values != nil {
|
||||
out[i] = s.values[k]
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *countingIncrStore) incrMulti(entries []counterIncr) ([]int64, error) {
|
||||
s.incrMultiCalls++
|
||||
if s.values == nil {
|
||||
s.values = make(map[string]int64)
|
||||
}
|
||||
out := make([]int64, len(entries))
|
||||
for i, e := range entries {
|
||||
s.values[e.key] += e.delta
|
||||
out[i] = s.values[e.key]
|
||||
}
|
||||
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) {
|
||||
store := &countingIncrStore{}
|
||||
l := newLimiter(store, true)
|
||||
now := time.Date(2026, 8, 21, 9, 0, 0, 0, time.UTC)
|
||||
l.nowFn = func() time.Time { return now }
|
||||
|
||||
scopes := []limitScope{
|
||||
{kind: "user", id: "u", limits: &LimitsConfig{RequestsPerMinute: 2}},
|
||||
{kind: "group", id: "g", limits: nil},
|
||||
{kind: totalScopeKind, id: totalScopeID, limits: nil},
|
||||
}
|
||||
|
||||
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 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)
|
||||
}
|
||||
|
||||
v := l.checkAndCount(scopes)
|
||||
if v == nil {
|
||||
t.Fatal("3rd call should violate u's requestsPerMinute:2")
|
||||
}
|
||||
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)
|
||||
}
|
||||
|
||||
// Values correct: batching three scopes into one call must not
|
||||
// corrupt any one scope's own counter.
|
||||
uMin, ok := l.getCounter("user", "u", metricReq, windowMin, now)
|
||||
if !ok || uMin != 3 {
|
||||
t.Errorf("user req:min = %d, ok=%v, want 3", uMin, ok)
|
||||
}
|
||||
totalMin, ok := l.getCounter(totalScopeKind, totalScopeID, metricReq, windowMin, now)
|
||||
if !ok || totalMin != 3 {
|
||||
t.Errorf("total req:min = %d, ok=%v, want 3", totalMin, ok)
|
||||
}
|
||||
gDay, ok := l.getCounter("group", "g", metricReq, windowDay, now)
|
||||
if !ok || gDay != 3 {
|
||||
t.Errorf("group req:day = %d, ok=%v, want 3", gDay, ok)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAccount_OneIncrMultiCallRegardlessOfScopeOrMetricCount is the
|
||||
// perf-review (2026-08-21) case for account: a 3-scope request with all
|
||||
// three metrics nonzero (27 counters: 3 scopes x 3 metrics x 3 windows)
|
||||
// must cost exactly ONE incrMulti round trip, and the batched reply must
|
||||
// land every value in the right scope/metric/window slot.
|
||||
func TestAccount_OneIncrMultiCallRegardlessOfScopeOrMetricCount(t *testing.T) {
|
||||
store := &countingIncrStore{}
|
||||
l := newLimiter(store, true)
|
||||
now := time.Date(2026, 8, 21, 9, 0, 0, 0, time.UTC)
|
||||
l.nowFn = func() time.Time { return now }
|
||||
|
||||
scopes := []limitScope{
|
||||
{kind: "user", id: "u"},
|
||||
{kind: "group", id: "g"},
|
||||
{kind: totalScopeKind, id: totalScopeID},
|
||||
}
|
||||
l.account(scopes, usage{prompt: 10, completion: 5}, 700)
|
||||
|
||||
if store.incrMultiCalls != 1 {
|
||||
t.Errorf("incrMultiCalls = %d, want 1 (one batch for 27 counters)", store.incrMultiCalls)
|
||||
}
|
||||
|
||||
uTokInDay, ok := l.getCounter("user", "u", metricTokIn, windowDay, now)
|
||||
if !ok || uTokInDay != 10 {
|
||||
t.Errorf("user tokin:day = %d, ok=%v, want 10", uTokInDay, ok)
|
||||
}
|
||||
totalTokOutHour, ok := l.getCounter(totalScopeKind, totalScopeID, metricTokOut, windowHour, now)
|
||||
if !ok || totalTokOutHour != 5 {
|
||||
t.Errorf("total tokout:hour = %d, ok=%v, want 5", totalTokOutHour, ok)
|
||||
}
|
||||
gCostMonth, ok := l.getCounter("group", "g", metricCost, windowMonth, now)
|
||||
if !ok || gCostMonth != 700 {
|
||||
t.Errorf("group cost:month = %d, ok=%v, want 700", gCostMonth, ok)
|
||||
}
|
||||
|
||||
// A second call accumulates rather than overwrites, and still costs
|
||||
// exactly one more round trip.
|
||||
l.account(scopes, usage{prompt: 2}, 0)
|
||||
if store.incrMultiCalls != 2 {
|
||||
t.Errorf("incrMultiCalls after 2nd account = %d, want 2", store.incrMultiCalls)
|
||||
}
|
||||
uTokInDay, ok = l.getCounter("user", "u", metricTokIn, windowDay, now)
|
||||
if !ok || uTokInDay != 12 {
|
||||
t.Errorf("user tokin:day after 2nd account = %d, ok=%v, want 12", uTokInDay, ok)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenBudget_ExactBoundaryViolates asserts budgetViolation/
|
||||
// tokenBudgetViolation's ">= limit" semantics at the exact boundary: a
|
||||
// scope whose accumulated tokin+tokout lands EXACTLY on the limit (not
|
||||
// merely over it) must still violate. budgetViolation's own "used <
|
||||
// limit passes" check makes "used == limit" the smallest value that
|
||||
// violates; this pins that boundary directly rather than only exercising
|
||||
// values already over it (TestLimiterTokenBudget, TestTokenBudgetSumsInAndOut).
|
||||
func TestTokenBudget_ExactBoundaryViolates(t *testing.T) {
|
||||
l := newLimiter(nil, true)
|
||||
scopes := []limitScope{{kind: "group", id: "g", limits: &LimitsConfig{TokensPerDay: 100}}}
|
||||
|
||||
l.account(scopes, usage{prompt: 63, completion: 36}, 0) // 99 < 100: must not violate yet
|
||||
if v := l.checkAndCount(scopes); v != nil {
|
||||
t.Fatalf("99 < 100 must not violate, got %+v", v)
|
||||
}
|
||||
|
||||
l.account(scopes, usage{completion: 1}, 0) // pushes tokout to 37: total now exactly 100
|
||||
v := l.checkAndCount(scopes)
|
||||
if v == nil {
|
||||
t.Fatal("tokin+tokout landing exactly on the 100 budget must violate (>= limit, not only > limit)")
|
||||
}
|
||||
if v.storeDown {
|
||||
t.Error("storeDown must be false for a real limit breach")
|
||||
}
|
||||
}
|
||||
|
||||
+69
-16
@@ -18,14 +18,26 @@ func newRedisStore(client *respClient) *redisStore {
|
||||
return &redisStore{client: client}
|
||||
}
|
||||
|
||||
// incrBy implements counterStore: INCRBY key n, then EXPIRE key
|
||||
// ttlSeconds, sent as one pipeline so the two commands share a single round
|
||||
// trip. It returns INCRBY's resulting counter value; EXPIRE's reply is not
|
||||
// otherwise inspected — a failed EXPIRE right after a successful INCRBY on
|
||||
// the same key would only leave that key without a fresh TTL, not corrupt
|
||||
// the count. ttl is rounded up to whole seconds (Redis EXPIRE's unit),
|
||||
// ttlToSeconds converts ttl to whole Redis EXPIRE seconds, rounded up,
|
||||
// with a floor of 1s so a sub-second ttl never turns into EXPIRE 0 (an
|
||||
// immediate delete).
|
||||
// immediate delete). Shared by incrBy and incrMulti.
|
||||
func ttlToSeconds(ttl time.Duration) int64 {
|
||||
s := int64(ttl / time.Second)
|
||||
if ttl%time.Second != 0 {
|
||||
s++
|
||||
}
|
||||
if s < 1 {
|
||||
s = 1
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// incrBy implements counterStore: INCRBY key n, then EXPIRE key
|
||||
// ttlToSeconds(ttl), sent as one pipeline so the two commands share a
|
||||
// single round trip. It returns INCRBY's resulting counter value;
|
||||
// EXPIRE's reply is not otherwise inspected — a failed EXPIRE right after
|
||||
// a successful INCRBY on the same key would only leave that key without a
|
||||
// fresh TTL, not corrupt the count.
|
||||
//
|
||||
// Semantics are deliberately at-least-once, not exactly-once, both in the
|
||||
// conservative direction (never under-counts): respClient.pipeline's
|
||||
@@ -39,17 +51,9 @@ func newRedisStore(client *respClient) *redisStore {
|
||||
// that occasionally over-counts by one request's worth is fail-safe (more
|
||||
// restrictive than reality), never fail-open in the unsafe direction.
|
||||
func (s *redisStore) incrBy(key string, n int64, ttl time.Duration) (int64, error) {
|
||||
ttlSeconds := int64(ttl / time.Second)
|
||||
if ttl%time.Second != 0 {
|
||||
ttlSeconds++
|
||||
}
|
||||
if ttlSeconds < 1 {
|
||||
ttlSeconds = 1
|
||||
}
|
||||
|
||||
replies, err := s.client.pipeline([][]string{
|
||||
{"INCRBY", key, strconv.FormatInt(n, 10)},
|
||||
{"EXPIRE", key, strconv.FormatInt(ttlSeconds, 10)},
|
||||
{"EXPIRE", key, strconv.FormatInt(ttlToSeconds(ttl), 10)},
|
||||
})
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("redisStore: incrBy %q: %w", key, err)
|
||||
@@ -67,6 +71,55 @@ func (s *redisStore) incrBy(key string, n int64, ttl time.Duration) (int64, erro
|
||||
return v, nil
|
||||
}
|
||||
|
||||
// incrMulti implements counterStore: every entry's INCRBY+EXPIRE pair sent
|
||||
// as ONE pipeline (perf review, 2026-08-21) — 2*len(entries) commands, one
|
||||
// round trip regardless of len(entries), extending incrBy's own
|
||||
// single-key pipelining to a whole batch of counters at once. Reply i*2
|
||||
// is entry i's INCRBY result (returned in out[i]); reply i*2+1 is its
|
||||
// EXPIRE result, not otherwise inspected — same reasoning as incrBy's own
|
||||
// doc comment: a failed EXPIRE right after a successful INCRBY only
|
||||
// leaves that one key without a fresh TTL, not a corrupted count. The
|
||||
// same at-least-once semantics as incrBy (see its own doc comment) apply
|
||||
// here too, extended to the whole batch: a lost reply after the server
|
||||
// already applied the pipeline can cause respClient.pipeline's
|
||||
// reconnect-once retry to re-send it, over-counting every entry in it by
|
||||
// its own delta.
|
||||
func (s *redisStore) incrMulti(entries []counterIncr) ([]int64, error) {
|
||||
if len(entries) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
cmds := make([][]string, 0, len(entries)*2)
|
||||
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)},
|
||||
)
|
||||
}
|
||||
|
||||
replies, err := s.client.pipeline(cmds)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("redisStore: incrMulti: %w", err)
|
||||
}
|
||||
if len(replies) != len(cmds) {
|
||||
return nil, fmt.Errorf("redisStore: incrMulti: expected %d replies, got %d", len(cmds), len(replies))
|
||||
}
|
||||
|
||||
out := make([]int64, len(entries))
|
||||
for i, e := range entries {
|
||||
reply := replies[i*2]
|
||||
if re, ok := reply.(error); ok {
|
||||
return nil, fmt.Errorf("redisStore: incrMulti %q: INCRBY failed: %w", e.key, re)
|
||||
}
|
||||
v, ok := reply.(int64)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("redisStore: incrMulti %q: unexpected INCRBY reply type %T", e.key, reply)
|
||||
}
|
||||
out[i] = v
|
||||
}
|
||||
return out, 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
|
||||
|
||||
@@ -29,6 +29,47 @@ func TestRedisStore_IncrBy_PipelinesIncrbyAndExpire(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestRedisStore_IncrMulti_PipelinesAllPairsInOneCall is the perf-review
|
||||
// (2026-08-21) case: incrMulti sends every entry's INCRBY+EXPIRE pair as
|
||||
// ONE pipeline — no SELECT (or any other command) interleaved between
|
||||
// entries, which would only happen if incrMulti dialled a fresh
|
||||
// connection or made a separate call per entry instead of the single
|
||||
// s.client.pipeline call it actually makes — and returns each entry's
|
||||
// INCRBY reply in the same order as entries, not the reply's position in
|
||||
// the full 2*N command/reply stream.
|
||||
func TestRedisStore_IncrMulti_PipelinesAllPairsInOneCall(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:hour:2026082014", "1"}, reply: []byte(":11\r\n")},
|
||||
{wantArgs: []string{"EXPIRE", "llmgw:user:a:req:hour:2026082014", "172800"}, reply: []byte(":1\r\n")},
|
||||
{wantArgs: []string{"INCRBY", "llmgw:user:a:tokin:day:20260820", "40"}, reply: []byte(":140\r\n")},
|
||||
{wantArgs: []string{"EXPIRE", "llmgw:user:a:tokin:day:20260820", "3024000"}, reply: []byte(":1\r\n")},
|
||||
{wantArgs: []string{"INCRBY", "llmgw:total:all:cost:month:202608", "500"}, reply: []byte(":9500\r\n")},
|
||||
{wantArgs: []string{"EXPIRE", "llmgw:total:all:cost:month:202608", "34560000"}, reply: []byte(":1\r\n")},
|
||||
})
|
||||
|
||||
store := newRedisStore(newRESPClient(ln.Addr().String(), "", 0))
|
||||
entries := []counterIncr{
|
||||
{key: "llmgw:user:a:req:hour:2026082014", delta: 1, ttl: hourWindowTTL},
|
||||
{key: "llmgw:user:a:tokin:day:20260820", delta: 40, ttl: dayWindowTTL},
|
||||
{key: "llmgw:total:all:cost:month:202608", delta: 500, ttl: monthWindowTTL},
|
||||
}
|
||||
got, err := store.incrMulti(entries)
|
||||
if err != nil {
|
||||
t.Fatalf("incrMulti: %v", err)
|
||||
}
|
||||
want := []int64{11, 140, 9500}
|
||||
if len(got) != len(want) {
|
||||
t.Fatalf("len(got) = %d, want %d", len(got), len(want))
|
||||
}
|
||||
for i, w := range want {
|
||||
if got[i] != w {
|
||||
t.Errorf("got[%d] = %d, want %d (entry order, not reply-stream position)", i, got[i], w)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 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.
|
||||
|
||||
@@ -83,6 +83,10 @@ func (alwaysErrStore) getMulti([]string) ([]int64, error) {
|
||||
return nil, errStoreDownStub
|
||||
}
|
||||
|
||||
func (alwaysErrStore) incrMulti([]counterIncr) ([]int64, error) {
|
||||
return 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
|
||||
|
||||
Reference in New Issue
Block a user