From ed1b1d7b2f59eb22f92f73aadf35db4293ced9d9 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Thu, 20 Aug 2026 01:39:41 +0100 Subject: [PATCH] fix(limits): cap unknown-model warnings, time-gated sweep, rounded micro-USD conversion --- limits.go | 50 +++++++++++++++++------- limits_test.go | 91 ++++++++++++++++++++++++++++-------------- pricing.go | 79 +++++++++++++++++++++++++++++-------- pricing_test.go | 102 +++++++++++++++++++++++++++++++++++++++++------- 4 files changed, 247 insertions(+), 75 deletions(-) diff --git a/limits.go b/limits.go index 2535b00..d7ae3ed 100644 --- a/limits.go +++ b/limits.go @@ -38,6 +38,15 @@ const ( // counterStore. const usdToMicroFactor = 1_000_000 +// usdToMicros converts a USD amount to micro-USD, rounded to the nearest +// micro-USD. A bare int64(usd*usdToMicroFactor) truncates rather than +// rounds — float64 can represent a value like 8.2 as very slightly under +// its true value, so a truncating cast turns an intended $8.20 limit into +// 8_199_999 micro-USD instead of 8_200_000, one micro-USD too strict. +func usdToMicros(usd float64) int64 { + return int64(math.Round(usd * usdToMicroFactor)) +} + // 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. @@ -57,37 +66,48 @@ type memoryEntry struct { value int64 } -// memoryStoreSweepThreshold is the key count above which incrBy -// opportunistically deletes expired entries before writing. Kept well -// above typical live-scope counts so the sweep is rare on a healthy -// gateway and only fires when stale keys are actually accumulating. -const memoryStoreSweepThreshold = 1024 +// sweepEvery is the minimum real time between incrBy's opportunistic +// expired-entry sweeps, regardless of how many keys memoryStore holds. A +// size-gated sweep (only above N keys) was measured to scan the whole map +// on every incrBy once live keys exceeded that threshold, since a sweep +// that frees nothing (all keys still live) never brings the count back +// down — a 167x incrBy cliff. Gating on elapsed time instead bounds sweep +// frequency independent of key count: worst case is one full-map scan +// every 30s. +const sweepEvery = 30 * time.Second // 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 // started here would leak on rebuild instead of being collected with the // rest of the old instance. Expired entries are instead swept -// opportunistically from incrBy once the map grows past -// memoryStoreSweepThreshold keys. +// opportunistically from incrBy, at most once every sweepEvery. type memoryStore struct { - data map[string]*memoryEntry - mu sync.Mutex + data map[string]*memoryEntry + nowFn func() time.Time // injected for tests; defaults to time.Now + lastSweep time.Time // guarded by mu; zero value sweeps on the first incrBy + mu sync.Mutex } // newMemoryStore returns an empty memoryStore. func newMemoryStore() *memoryStore { - return &memoryStore{data: make(map[string]*memoryEntry)} + return &memoryStore{data: make(map[string]*memoryEntry), nowFn: time.Now} +} + +// now returns m's current time, via nowFn. +func (m *memoryStore) now() time.Time { + return m.nowFn() } // incrBy implements counterStore. func (m *memoryStore) incrBy(key string, n int64, ttl time.Duration) (int64, error) { - now := time.Now() + now := m.now() m.mu.Lock() defer m.mu.Unlock() - if len(m.data) > memoryStoreSweepThreshold { + if now.Sub(m.lastSweep) >= sweepEvery { m.sweepLocked(now) + m.lastSweep = now } e, ok := m.data[key] @@ -101,7 +121,7 @@ func (m *memoryStore) incrBy(key string, n int64, ttl time.Duration) (int64, err // get implements counterStore. func (m *memoryStore) get(key string) (int64, error) { - now := time.Now() + now := m.now() m.mu.Lock() defer m.mu.Unlock() @@ -314,10 +334,10 @@ func (l *limiter) evaluateScope(sc limitScope, minCount, dayCount int64, now tim if v := l.budgetViolation(sc, metricTok, "tokens-per-month", lim.TokensPerMonth, windowMonth, now); v != nil { return v } - if v := l.budgetViolation(sc, metricCost, "cost-per-day", int64(lim.CostPerDayUSD*usdToMicroFactor), windowDay, now); v != nil { + 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", int64(lim.CostPerMonthUSD*usdToMicroFactor), windowMonth, now); v != nil { + if v := l.budgetViolation(sc, metricCost, "cost-per-month", usdToMicros(lim.CostPerMonthUSD), windowMonth, now); v != nil { return v } return nil diff --git a/limits_test.go b/limits_test.go index 2335aac..69ccd5d 100644 --- a/limits_test.go +++ b/limits_test.go @@ -1,7 +1,6 @@ package traefikllmgateway import ( - "fmt" "testing" "time" ) @@ -227,40 +226,74 @@ func TestMemoryStoreExpiry(t *testing.T) { } } -// TestMemoryStoreSweepOnAccess grows the store past -// memoryStoreSweepThreshold, backdates every entry's expiry directly -// (rather than racing a short TTL against the insertion loop's own real -// wall-clock time, which is inherently flaky), and asserts the next -// incrBy's opportunistic sweep reclaims every expired entry. -func TestMemoryStoreSweepOnAccess(t *testing.T) { +// 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 +// key physically in place (no scan happened), and the first incrBy at or +// past the gate reclaims it. Time is injected via m.nowFn so the test +// needs no real sleep and no size threshold to trigger the behavior. +func TestMemoryStoreSweepTimeGated(t *testing.T) { m := newMemoryStore() - for i := 0; i < memoryStoreSweepThreshold+10; i++ { - key := fmt.Sprintf("k%d", i) - if _, err := m.incrBy(key, 1, time.Hour); err != nil { - t.Fatal(err) - } - } + fake := time.Date(2026, 8, 20, 0, 0, 0, 0, time.UTC) + m.nowFn = func() time.Time { return fake } - m.mu.Lock() - past := time.Now().Add(-time.Hour) - for _, e := range m.data { - e.expiry = past - } - setupCount := len(m.data) - m.mu.Unlock() - if setupCount != memoryStoreSweepThreshold+10 { - t.Fatalf("setup: got %d keys, want %d", setupCount, memoryStoreSweepThreshold+10) - } - - if _, err := m.incrBy("trigger-sweep", 1, time.Minute); err != nil { + // lastSweep starts at the zero time, so this first incrBy always + // sweeps (a no-op on an empty store) and sets lastSweep to fake. + if _, err := m.incrBy("k1", 1, time.Nanosecond); err != nil { t.Fatal(err) } + // k1 is already expired (ttl was 1ns), but stay well inside the + // sweepEvery gate — the next incrBy must not scan, so k1 survives as + // a physical (if stale) map entry. + fake = fake.Add(time.Millisecond) + if _, err := m.incrBy("k2", 1, time.Minute); err != nil { + t.Fatal(err) + } + if !memoryStoreHasKey(m, "k1") { + t.Fatal("k1 must still be present: the sweep gate has not elapsed yet") + } + + // Cross the sweepEvery boundary from the last sweep: this incrBy must + // scan and reclaim k1. + fake = fake.Add(sweepEvery) + if _, err := m.incrBy("k3", 1, time.Minute); err != nil { + t.Fatal(err) + } + if memoryStoreHasKey(m, "k1") { + t.Error("k1 should have been swept once sweepEvery elapsed") + } +} + +// memoryStoreHasKey reports whether key is physically present in m.data, +// bypassing get's expiry check — used to observe sweep timing directly. +func memoryStoreHasKey(m *memoryStore, key string) bool { m.mu.Lock() - n := len(m.data) - m.mu.Unlock() - if n != 1 { - t.Errorf("expected sweep to reclaim every expired entry, %d keys remain (want 1: trigger-sweep)", n) + defer m.mu.Unlock() + _, ok := m.data[key] + return ok +} + +// TestLimiterCostBudgetRoundsUSDConversion pins usdToMicros' rounding +// against a value (8.2) where float64 representation sits just under the +// true value: a truncating int64(8.2*1e6) would undershoot to 8_199_999, +// making a spend of exactly $8.199999 look already over an intended +// $8.20 budget. With rounding, 8_199_999 must stay under budget and only +// the exact $8.20 spend violates. +func TestLimiterCostBudgetRoundsUSDConversion(t *testing.T) { + now := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC) + l := newLimiter(nil, true) + l.nowFn = func() time.Time { return now } + scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{CostPerDayUSD: 8.2}}} + + l.account(scopes, usage{}, 8_199_999) + if v := l.checkAndCount(scopes); v != nil { + t.Fatalf("spending $8.199999 against an $8.20 budget must not violate, got %+v", v) + } + + l.account(scopes, usage{}, 1) // now exactly at $8.20 + if v := l.checkAndCount(scopes); v == nil { + t.Fatal("spending exactly $8.20 against an $8.20 budget must violate") } } diff --git a/pricing.go b/pricing.go index 1d001b8..c71e762 100644 --- a/pricing.go +++ b/pricing.go @@ -1,6 +1,9 @@ package traefikllmgateway -import "sync" +import ( + "math" + "sync" +) // tokensPerMillion is the denominator for the providers' per-1M-token // pricing unit. @@ -22,33 +25,55 @@ var builtinPricing = map[string]ModelPricing{ "grok-4": {InputPerM: 3, OutputPerM: 15}, } +// unknownModelWarnCap bounds warnedModels. The model id in costMicros' +// calls comes from client-supplied request bodies (a group with an empty +// Models allowlist accepts any model string), so without a cap an attacker +// or a buggy client cycling through model names could grow warnedModels +// without bound. Once the cap is reached, costMicros stops recording new +// names and pricingWarnFn fires once more with a summary message instead +// of a per-model one. +const unknownModelWarnCap = 128 + +// warnCapMessage is the single summary warning fired once the dedup set +// in warnedModels reaches unknownModelWarnCap. +const warnCapMessage = "pricing: unknown-model warning cap reached; further unknown models unwarned" + // pricingWarnFn is called once per unknown model id encountered by -// costMicros, with the model id. The default is a no-op so pricing.go has -// no logging dependency of its own; a later task wires this to the -// gateway's logf, and tests can substitute their own func to capture -// warnings. +// costMicros, with the model id — or, once unknownModelWarnCap is +// reached, once more with warnCapMessage. The default is a no-op so +// pricing.go has no logging dependency of its own; a later task wires +// this to the gateway's logf via setPricingWarnFn, and tests can +// substitute their own func to capture warnings the same way. Read inside +// warnUnknownModel under warnedModelsMu — always assign through +// setPricingWarnFn, never by a bare `pricingWarnFn = …`, or a concurrent +// costMicros call can race the assignment. var pricingWarnFn = func(string) {} // warnedModels tracks which unknown model ids have already triggered // pricingWarnFn, so a gateway processing many requests for the same // unpriced model logs the warning once per process lifetime rather than -// once per request. Package-level and mutex-guarded by design: costMicros -// is called from the request path with no natural place to thread a -// longer-lived cache through, and the dedup set must survive across -// requests and across limiter instances, not reset per call. +// once per request. warnCapNotified tracks whether the one-time summary +// warning has already fired once the cap is reached. Package-level and +// mutex-guarded by design: costMicros is called from the request path +// with no natural place to thread a longer-lived cache through, and the +// dedup set must survive across requests and across limiter instances, +// not reset per call. var ( - warnedModelsMu sync.Mutex - warnedModels = map[string]bool{} + warnedModelsMu sync.Mutex + warnedModels = map[string]bool{} + warnCapNotified bool ) // costMicros computes the micro-USD cost of u's token usage for model, // looking up its price in overrides first (exact model id match), then // builtinPricing. An unknown model costs 0 and triggers one pricingWarnFn -// call per model id. +// call per model id, up to unknownModelWarnCap distinct names. // // The math stays in integer micro-USD throughout to avoid float drift: -// priceMicrosPerM is the per-1M-token price scaled to micro-USD, and each -// of input/output cost is tokens*priceMicrosPerM/1e6, summed. +// priceMicrosPerM is the per-1M-token price scaled to micro-USD and +// rounded to the nearest micro-USD (a bare truncating cast can undershoot +// a price like 8.2 by one micro-USD due to float64 representation), and +// each of input/output cost is tokens*priceMicrosPerM/1e6, summed. func costMicros(model string, u usage, overrides map[string]*ModelPricing) int64 { price, ok := lookupPricing(model, overrides) if !ok { @@ -56,8 +81,8 @@ func costMicros(model string, u usage, overrides map[string]*ModelPricing) int64 return 0 } - inputMicrosPerM := int64(price.InputPerM * tokensPerMillion) - outputMicrosPerM := int64(price.OutputPerM * tokensPerMillion) + inputMicrosPerM := int64(math.Round(price.InputPerM * tokensPerMillion)) + outputMicrosPerM := int64(math.Round(price.OutputPerM * tokensPerMillion)) inputCost := u.prompt * inputMicrosPerM / tokensPerMillion outputCost := u.completion * outputMicrosPerM / tokensPerMillion @@ -77,14 +102,34 @@ func lookupPricing(model string, overrides map[string]*ModelPricing) (ModelPrici return ModelPricing{}, false } +// setPricingWarnFn replaces pricingWarnFn, guarded by the same mutex +// warnUnknownModel reads it under. A later task calls this at New() to +// wire warnings to the gateway's logf; tests call it to capture warnings. +// Always use this instead of assigning pricingWarnFn directly. +func setPricingWarnFn(fn func(string)) { + warnedModelsMu.Lock() + defer warnedModelsMu.Unlock() + pricingWarnFn = fn +} + // warnUnknownModel calls pricingWarnFn for model, once per process -// lifetime. +// lifetime. Once warnedModels reaches unknownModelWarnCap, it stops +// recording new model ids and instead fires a single one-time summary +// warning via warnCapMessage — after that, further unknown models are +// silently unwarned rather than growing the set forever. func warnUnknownModel(model string) { warnedModelsMu.Lock() defer warnedModelsMu.Unlock() if warnedModels[model] { return } + if len(warnedModels) >= unknownModelWarnCap { + if !warnCapNotified { + warnCapNotified = true + pricingWarnFn(warnCapMessage) + } + return + } warnedModels[model] = true pricingWarnFn(model) } diff --git a/pricing_test.go b/pricing_test.go index a4c7c0a..955727a 100644 --- a/pricing_test.go +++ b/pricing_test.go @@ -1,6 +1,10 @@ package traefikllmgateway -import "testing" +import ( + "strconv" + "sync" + "testing" +) // TestCostMicrosBuiltin exercises the integer micro-USD math against // builtinPricing entries directly, so a regression to float-based pricing @@ -59,19 +63,42 @@ func TestCostMicrosOverrideNilEntryFallsBack(t *testing.T) { } } +// resetPricingWarnState snapshots and clears the package-level unknown- +// model dedup state (warnedModels, warnCapNotified) and pricingWarnFn, and +// registers a t.Cleanup to restore them — so tests exercising the +// warn-once/warn-cap behavior are independent of each other and of test +// order. All access goes through setPricingWarnFn / warnedModelsMu so +// these tests stay race-clean. +func resetPricingWarnState(t *testing.T) *[]string { + t.Helper() + warnedModelsMu.Lock() + prevWarn, prevWarned, prevCapNotified := pricingWarnFn, warnedModels, warnCapNotified + warnedModels = map[string]bool{} + warnCapNotified = false + warnedModelsMu.Unlock() + + t.Cleanup(func() { + warnedModelsMu.Lock() + warnedModels, warnCapNotified = prevWarned, prevCapNotified + warnedModelsMu.Unlock() + setPricingWarnFn(prevWarn) + }) + + var warned []string + var mu sync.Mutex + setPricingWarnFn(func(model string) { + mu.Lock() + defer mu.Unlock() + warned = append(warned, model) + }) + return &warned +} + // TestCostMicrosUnknownModelWarnsOnce asserts an unpriced model costs 0 // and fires pricingWarnFn exactly once even across repeated calls, then // resumes warning for a second distinct unknown model. func TestCostMicrosUnknownModelWarnsOnce(t *testing.T) { - prevWarn, prevWarned := pricingWarnFn, warnedModels - t.Cleanup(func() { - pricingWarnFn = prevWarn - warnedModels = prevWarned - }) - warnedModels = map[string]bool{} - - var warned []string - pricingWarnFn = func(model string) { warned = append(warned, model) } + warned := resetPricingWarnState(t) for i := 0; i < 3; i++ { if got := costMicros("totally-unpriced-model", usage{prompt: 100, completion: 100}, nil); got != 0 { @@ -83,16 +110,63 @@ func TestCostMicrosUnknownModelWarnsOnce(t *testing.T) { } want := []string{"totally-unpriced-model", "another-unpriced-model"} - if len(warned) != len(want) { - t.Fatalf("warned = %v, want %v", warned, want) + if len(*warned) != len(want) { + t.Fatalf("warned = %v, want %v", *warned, want) } for i, m := range want { - if warned[i] != m { - t.Errorf("warned[%d] = %q, want %q", i, warned[i], m) + if (*warned)[i] != m { + t.Errorf("warned[%d] = %q, want %q", i, (*warned)[i], m) } } } +// TestCostMicrosUnknownModelWarnCap asserts warnedModels never grows past +// unknownModelWarnCap, however many distinct unknown model ids costMicros +// sees — the model id comes from client-controlled request bodies (a +// group with an empty Models allowlist accepts any string), so an +// unbounded dedup set would be a memory-growth vector. Once the cap is +// reached, one summary warning fires and no more do. +func TestCostMicrosUnknownModelWarnCap(t *testing.T) { + warned := resetPricingWarnState(t) + + const distinctModels = 200 + for i := 0; i < distinctModels; i++ { + model := "unpriced-model-" + strconv.Itoa(i) + if got := costMicros(model, usage{prompt: 1}, nil); got != 0 { + t.Fatalf("costMicros(%q, ...) = %d, want 0", model, got) + } + } + + warnedModelsMu.Lock() + setSize := len(warnedModels) + warnedModelsMu.Unlock() + if setSize > unknownModelWarnCap { + t.Errorf("warnedModels size = %d, want <= %d", setSize, unknownModelWarnCap) + } + + if got, max := len(*warned), unknownModelWarnCap+1; got > max { + t.Errorf("warn count = %d, want <= %d (cap distinct warnings + 1 summary)", got, max) + } + + last := (*warned)[len(*warned)-1] + if last != warnCapMessage { + t.Errorf("last warning = %q, want the cap-reached summary %q", last, warnCapMessage) + } +} + +// TestCostMicrosRoundsPriceConversion pins costMicros' per-1M-token price +// conversion to round-to-nearest rather than truncate: at $8.20/M, a +// truncating int64(8.2*1e6) undershoots to 8_199_999 (float64 can +// represent 8.2 as very slightly under its true value); rounding must +// yield the intended 8_200_000. +func TestCostMicrosRoundsPriceConversion(t *testing.T) { + overrides := map[string]*ModelPricing{"round-test": {InputPerM: 8.2, OutputPerM: 8.2}} + u := usage{prompt: 1_000_000, completion: 1_000_000} + if got, want := costMicros("round-test", u, overrides), int64(16_400_000); got != want { + t.Errorf("costMicros with an 8.2 price = %d, want %d", got, want) + } +} + // TestBuiltinPricingSeedComplete asserts every model the task brief seeds // is present in builtinPricing with a positive price on both sides — a // missing or zeroed entry would silently price real traffic at $0.