From bce9066a6adf8c20cb081432c4efcabb5705bd8a Mon Sep 17 00:00:00 2001 From: "Jakub A. W." Date: Fri, 30 Oct 2026 15:19:51 +0100 Subject: [PATCH 1/2] feat(ratelimit): enforce limits across replicas sharing a database --- internal/admin/handler_ratelimits_test.go | 2 +- internal/ratelimit/limiter.go | 107 +++++-- internal/ratelimit/peers.go | 202 ++++++++++++ internal/ratelimit/peers_test.go | 329 ++++++++++++++++++++ internal/ratelimit/persist.go | 58 +++- internal/ratelimit/persist_test.go | 34 +- internal/ratelimit/persist_writers_test.go | 100 ++++-- internal/ratelimit/service.go | 40 ++- internal/ratelimit/service_test.go | 40 ++- internal/ratelimit/snapshot.go | 250 ++++++++------- internal/ratelimit/store.go | 24 +- internal/ratelimit/store_counters_test.go | 76 ++++- internal/ratelimit/store_epoch_test.go | 2 +- internal/ratelimit/store_mongodb.go | 40 ++- internal/ratelimit/store_mongodb_migrate.go | 8 +- internal/ratelimit/store_mongodb_test.go | 18 +- internal/ratelimit/store_sql.go | 39 ++- internal/ratelimit/store_sql_test.go | 22 +- internal/ratelimit/types.go | 10 +- internal/server/ratelimit_support_test.go | 2 +- tests/perf/production_shape_test.go | 4 +- 21 files changed, 1105 insertions(+), 302 deletions(-) create mode 100644 internal/ratelimit/peers.go create mode 100644 internal/ratelimit/peers_test.go diff --git a/internal/admin/handler_ratelimits_test.go b/internal/admin/handler_ratelimits_test.go index 505be6d92..776938ca7 100644 --- a/internal/admin/handler_ratelimits_test.go +++ b/internal/admin/handler_ratelimits_test.go @@ -57,7 +57,7 @@ func (s *adminRateLimitStore) ReplaceConfigRules(ctx context.Context, rules []ra return s.UpsertRules(ctx, rules) } -func (s *adminRateLimitStore) LoadCounters(context.Context) ([]ratelimit.WindowSnapshot, error) { +func (s *adminRateLimitStore) LoadCounters(context.Context, int64) ([]ratelimit.WindowSnapshot, error) { return nil, nil } func (s *adminRateLimitStore) SaveCounters(context.Context, []ratelimit.WindowSnapshot) error { diff --git a/internal/ratelimit/limiter.go b/internal/ratelimit/limiter.go index 22a51fcf6..5bf837ba7 100644 --- a/internal/ratelimit/limiter.go +++ b/internal/ratelimit/limiter.go @@ -12,6 +12,11 @@ type windowCounter struct { windowStart int64 // unix seconds, aligned to the period current int64 previous int64 + // dirty marks a local counter that counted something since its last + // snapshot, so a flush writes only the rows that changed (see + // limiter.snapshot). Rolling the window forward is not a change: a + // reader rolls the saved row forward the same way. + dirty bool } // advance rolls the counter forward so windowStart covers now. @@ -38,6 +43,20 @@ func (w *windowCounter) estimate(now time.Time, periodSeconds int64) int64 { return int64(float64(w.previous)*weight) + w.current } +// rolledTo returns the counter as it stands in the window containing now. +// Unlike advance it never discards a counter saved in a later window than +// now's, as a replica whose clock runs slightly ahead can write: that row's +// counts are folded into now's window instead, so a boundary crossed a +// moment early elsewhere does not erase them here. +func (w windowCounter) rolledTo(now time.Time, periodSeconds int64) windowCounter { + start := now.Unix() - now.Unix()%periodSeconds + if w.windowStart > start { + return windowCounter{windowStart: start, current: w.current + w.previous} + } + w.advance(now, periodSeconds) + return w +} + // resetAfter returns the time until the current window rolls over. func (w *windowCounter) resetAfter(now time.Time, periodSeconds int64) time.Duration { remaining := max(w.windowStart+periodSeconds-now.Unix(), 1) @@ -81,12 +100,26 @@ func keyForRule(rule Rule) ruleKey { // limiter holds all live counters. Ordinary rules own one counter; per-child // rules own one counter per active direct child. Expired child partitions are // pruned below, and a single mutex keeps admission atomic across all matches. +// +// The local counters hold only what this process counted itself. The remote +// maps hold what every other writer sharing the store counted, summed per key +// as of the last load (see Service.loadPeers), and admission, status, and the +// routing probe check the sum of both. They stay empty while no other writer +// counted anything, so a lone instance pays only for lookups that miss. type limiter struct { mu sync.Mutex requests map[ruleKey]*windowCounter tokens map[ruleKey]*windowCounter inFlight map[ruleKey]int64 + remoteRequests map[ruleKey]windowCounter + remoteTokens map[ruleKey]windowCounter + remoteInFlight map[ruleKey]int64 + // publishedInFlight records the concurrency keys whose last snapshot + // published a nonzero gauge, so the next one publishes their drop to + // zero once instead of leaving peers to wait for the row to go stale. + publishedInFlight map[ruleKey]struct{} + expiries counterExpiryQueue expiryByCounter map[counterExpiryKey]*counterExpiry expiryWake chan struct{} @@ -99,12 +132,26 @@ type limiter struct { func newLimiter() *limiter { return &limiter{ - requests: make(map[ruleKey]*windowCounter), - tokens: make(map[ruleKey]*windowCounter), - inFlight: make(map[ruleKey]int64), + requests: make(map[ruleKey]*windowCounter), + tokens: make(map[ruleKey]*windowCounter), + inFlight: make(map[ruleKey]int64), + publishedInFlight: make(map[ruleKey]struct{}), } } +// withRemote returns the cluster-wide counter for key at now: the local +// counter, which the caller has advanced to now, plus every other writer's. +// The caller must hold l.mu. +func withRemote(remote map[ruleKey]windowCounter, key ruleKey, local *windowCounter, now time.Time, periodSeconds int64) windowCounter { + total := *local + if peers, ok := remote[key]; ok { + peers = peers.rolledTo(now, periodSeconds) + total.current += peers.current + total.previous += peers.previous + } + return total +} + func (l *limiter) counter(counters map[ruleKey]*windowCounter, key ruleKey) *windowCounter { counter, ok := counters[key] if !ok { @@ -147,7 +194,9 @@ func (l *limiter) admit(rules []Rule, now time.Time) (HeaderSnapshot, []ruleKey, if rule.MaxRequests != nil { counter := l.requests[key] counter.current++ - remaining := max(*rule.MaxRequests-counter.estimate(now, rule.PeriodSeconds), 0) + counter.dirty = true + total := withRemote(l.remoteRequests, key, counter, now, rule.PeriodSeconds) + remaining := max(*rule.MaxRequests-total.estimate(now, rule.PeriodSeconds), 0) if !headers.HasRequests || remaining < headers.RequestRemaining { headers.HasRequests = true headers.RequestLimit = *rule.MaxRequests @@ -157,7 +206,8 @@ func (l *limiter) admit(rules []Rule, now time.Time) (HeaderSnapshot, []ruleKey, } if rule.MaxTokens != nil { counter := l.tokens[key] - remaining := max(*rule.MaxTokens-counter.estimate(now, rule.PeriodSeconds), 0) + total := withRemote(l.remoteTokens, key, counter, now, rule.PeriodSeconds) + remaining := max(*rule.MaxTokens-total.estimate(now, rule.PeriodSeconds), 0) if !headers.HasTokens || remaining < headers.TokenRemaining { headers.HasTokens = true headers.TokenLimit = *rule.MaxTokens @@ -170,15 +220,16 @@ func (l *limiter) admit(rules []Rule, now time.Time) (HeaderSnapshot, []ruleKey, } // breach reports whether one rule currently rejects one more request, without -// consuming anything. The caller must hold l.mu. +// consuming anything. It counts every writer sharing the store, not only this +// one. The caller must hold l.mu. func (l *limiter) breach(rule Rule, now time.Time) *ExceededError { key := keyForRule(rule) if rule.PeriodSeconds == PeriodConcurrent { - if l.inFlight[key] >= *rule.MaxRequests { + if inFlight := l.inFlight[key] + l.remoteInFlight[key]; inFlight >= *rule.MaxRequests { return &ExceededError{ Rule: rule, Scope: ScopeConcurrency, - Observed: l.inFlight[key], + Observed: inFlight, Limit: *rule.MaxRequests, RetryAfter: time.Second, } @@ -189,13 +240,14 @@ func (l *limiter) breach(rule Rule, now time.Time) *ExceededError { counter := l.counter(l.requests, key) counter.advance(now, rule.PeriodSeconds) l.trackCounterExpiry(requestCounter, key, counter) - if used := counter.estimate(now, rule.PeriodSeconds); used >= *rule.MaxRequests { + total := withRemote(l.remoteRequests, key, counter, now, rule.PeriodSeconds) + if used := total.estimate(now, rule.PeriodSeconds); used >= *rule.MaxRequests { return &ExceededError{ Rule: rule, Scope: ScopeRequests, Observed: used, Limit: *rule.MaxRequests, - RetryAfter: counter.retryAfter(now, rule.PeriodSeconds, *rule.MaxRequests), + RetryAfter: total.retryAfter(now, rule.PeriodSeconds, *rule.MaxRequests), } } } @@ -203,13 +255,14 @@ func (l *limiter) breach(rule Rule, now time.Time) *ExceededError { counter := l.counter(l.tokens, key) counter.advance(now, rule.PeriodSeconds) l.trackCounterExpiry(tokenCounter, key, counter) - if used := counter.estimate(now, rule.PeriodSeconds); used >= *rule.MaxTokens { + total := withRemote(l.remoteTokens, key, counter, now, rule.PeriodSeconds) + if used := total.estimate(now, rule.PeriodSeconds); used >= *rule.MaxTokens { return &ExceededError{ Rule: rule, Scope: ScopeTokens, Observed: used, Limit: *rule.MaxTokens, - RetryAfter: counter.retryAfter(now, rule.PeriodSeconds, *rule.MaxTokens), + RetryAfter: total.retryAfter(now, rule.PeriodSeconds, *rule.MaxTokens), } } } @@ -258,11 +311,13 @@ func (l *limiter) recordTokens(rules []Rule, tokens int64, now time.Time) { counter := l.counter(l.tokens, keyForRule(rule)) counter.advance(now, rule.PeriodSeconds) counter.current += tokens + counter.dirty = true l.trackCounterExpiry(tokenCounter, keyForRule(rule), counter) } } -// status reports the live counter state for one rule. +// status reports the live counter state for one rule, summed over every +// writer sharing the store. func (l *limiter) status(rule Rule, now time.Time) Status { l.mu.Lock() defer l.mu.Unlock() @@ -270,7 +325,7 @@ func (l *limiter) status(rule Rule, now time.Time) Status { key := keyForRule(rule) status := Status{Rule: rule} if rule.PeriodSeconds == PeriodConcurrent { - status.InFlight = l.inFlight[key] + status.InFlight = l.inFlight[key] + l.remoteInFlight[key] if rule.MaxRequests != nil { remaining := max(*rule.MaxRequests-status.InFlight, 0) status.RequestsRemaining = &remaining @@ -285,7 +340,8 @@ func (l *limiter) status(rule Rule, now time.Time) Status { counter := l.counter(l.requests, key) counter.advance(now, rule.PeriodSeconds) l.trackCounterExpiry(requestCounter, key, counter) - status.RequestsUsed = counter.estimate(now, rule.PeriodSeconds) + total := withRemote(l.remoteRequests, key, counter, now, rule.PeriodSeconds) + status.RequestsUsed = total.estimate(now, rule.PeriodSeconds) remaining := max(*rule.MaxRequests-status.RequestsUsed, 0) status.RequestsRemaining = &remaining } @@ -293,18 +349,29 @@ func (l *limiter) status(rule Rule, now time.Time) Status { counter := l.counter(l.tokens, key) counter.advance(now, rule.PeriodSeconds) l.trackCounterExpiry(tokenCounter, key, counter) - status.TokensUsed = counter.estimate(now, rule.PeriodSeconds) + total := withRemote(l.remoteTokens, key, counter, now, rule.PeriodSeconds) + status.TokensUsed = total.estimate(now, rule.PeriodSeconds) remaining := max(*rule.MaxTokens-status.TokensUsed, 0) status.TokensRemaining = &remaining } return status } -// reset clears the window counters for one rule. In-flight gauges are left -// alone: they reflect requests that genuinely still hold slots. +// reset clears the window counters for one rule, this writer's and the +// other writers' alike: the reset applies to the whole cluster, and the next +// load sums only rows counted under the rule's new epoch. In-flight gauges +// are left alone, remote ones too: they reflect requests that genuinely still +// hold slots. func (l *limiter) reset(key ruleKey) { l.mu.Lock() defer l.mu.Unlock() + for _, remote := range []map[ruleKey]windowCounter{l.remoteRequests, l.remoteTokens} { + for candidate := range remote { + if candidate.sameDefinition(key) { + delete(remote, candidate) + } + } + } for candidate := range l.requests { if candidate.sameDefinition(key) { l.removeCounterExpiry(requestCounter, candidate) @@ -323,11 +390,13 @@ func (k ruleKey) sameDefinition(other ruleKey) bool { return k.scope == other.scope && k.subject == other.subject && k.periodSeconds == other.periodSeconds } -// resetAll clears every window counter. +// resetAll clears every window counter, local and remote. func (l *limiter) resetAll() { l.mu.Lock() defer l.mu.Unlock() l.requests = make(map[ruleKey]*windowCounter) l.tokens = make(map[ruleKey]*windowCounter) + l.remoteRequests = nil + l.remoteTokens = nil l.resetCounterExpiries() } diff --git a/internal/ratelimit/peers.go b/internal/ratelimit/peers.go new file mode 100644 index 000000000..972310f07 --- /dev/null +++ b/internal/ratelimit/peers.go @@ -0,0 +1,202 @@ +package ratelimit + +import ( + "context" + "time" +) + +const ( + // peerLoadOverlap is how far before the previous load an incremental + // load starts reading. A save stamps its rows before its transaction + // commits, up to persistTimeout later, and replicas' clocks differ a + // little; reading that far back again catches both. + peerLoadOverlap = 2 * persistTimeout + // peerFullLoadInterval is how often the incremental load is replaced by a + // full one. Nothing should need it, but it bounds what any row missed + // for whatever reason (a clock far off, a stalled commit) can cost, and + // drops rows other replicas deleted. + peerFullLoadInterval = time.Minute + // minInFlightLiveness is the floor of Service.inFlightLiveness. + minInFlightLiveness = 5 * time.Second +) + +// rowIdentity is the key every store gives a counter row. +type rowIdentity struct { + scope string + subject string + partition string + periodSeconds int64 + writer string +} + +func identityOf(snap WindowSnapshot) rowIdentity { + return rowIdentity{ + scope: snap.Scope, + subject: snap.Subject, + partition: snap.Partition, + periodSeconds: snap.PeriodSeconds, + writer: snap.Writer, + } +} + +// peerView caches the other writers' rows between loads, so a load only has +// to read the rows that changed since the previous one. Rows nobody rewrites +// any more, a stopped writer's, stay cached until their windows roll out. +// It is guarded by Service.persistMu. +type peerView struct { + rows map[rowIdentity]WindowSnapshot + // loadedAt is the clock reading taken just before the last successful + // load, and fullAt that of the last full one. + loadedAt time.Time + fullAt time.Time +} + +// merge folds a load into the cache. A full load replaces it. +func (v *peerView) merge(rows []WindowSnapshot, self string, full bool) { + if full || v.rows == nil { + v.rows = make(map[rowIdentity]WindowSnapshot, len(rows)) + } + for _, row := range rows { + if row.Writer == self { + continue + } + v.rows[identityOf(row)] = row + } +} + +// sum adds the cached rows up per key, as of now, for the rules in force: +// +// - A window row is rolled forward to the window containing now, as +// admission rolls a counter, and added to the key's request and token +// sums. Rolling forward is linear, so the sliding estimate of the sum is +// the sum of the rows' estimates. Rows counted under another epoch of the +// rule (see Rule.CounterEpoch) are skipped; rows whose windows rolled out +// are dropped from the cache. +// - An in-flight row counts only while its writer keeps rewriting it: one +// not saved within liveness is a stopped or crashed writer's, and its +// slots are free. It is dropped from the cache then, and once it reads +// zero. In-flight rows ignore the epoch, as a reset leaves in-flight +// gauges alone (see limiter.reset). +// +// A key nobody else counted gets no entry, and a cluster of one gets nil maps. +func (v *peerView) sum(rules []Rule, now time.Time, liveness time.Duration) (requests, tokens map[ruleKey]windowCounter, inFlight map[ruleKey]int64) { + index := indexRules(rules) + nowUnix := now.Unix() + for id, row := range v.rows { + key := ruleKey{ + scope: RuleScope(row.Scope), + subject: row.Subject, + partition: row.Partition, + periodSeconds: row.PeriodSeconds, + } + switch { + case row.PeriodSeconds == PeriodConcurrent: + if row.InFlight <= 0 || liveness <= 0 || nowUnix-row.UpdatedAt > int64(liveness/time.Second) { + delete(v.rows, id) + continue + } + if _, ok := index.lookup(key); !ok { + continue + } + if inFlight == nil { + inFlight = make(map[ruleKey]int64) + } + inFlight[key] += row.InFlight + case row.PeriodSeconds < 0: + delete(v.rows, id) + default: + latest := max(row.RequestsWindowStart, row.TokensWindowStart) + if latest+2*row.PeriodSeconds < nowUnix { + delete(v.rows, id) + continue + } + rule, ok := index.lookup(key) + if !ok || row.Epoch != rule.CounterEpoch { + continue + } + requests = addCounter(requests, key, windowCounter{ + windowStart: row.RequestsWindowStart, + current: row.RequestsCurrent, + previous: row.RequestsPrevious, + }, now) + tokens = addCounter(tokens, key, windowCounter{ + windowStart: row.TokensWindowStart, + current: row.TokensCurrent, + previous: row.TokensPrevious, + }, now) + } + } + return requests, tokens, inFlight +} + +// addCounter adds one row's counter, rolled to now, to the sum for key. A row +// that holds nothing at now adds no entry. +func addCounter(sums map[ruleKey]windowCounter, key ruleKey, row windowCounter, now time.Time) map[ruleKey]windowCounter { + if row.current == 0 && row.previous == 0 { + return sums + } + row = row.rolledTo(now, key.periodSeconds) + if row.current == 0 && row.previous == 0 { + return sums + } + if sums == nil { + sums = make(map[ruleKey]windowCounter) + } + sum := sums[key] + sum.windowStart = row.windowStart + sum.current += row.current + sum.previous += row.previous + sums[key] = sum + return sums +} + +// inFlightLiveness is how long a peer's in-flight row counts after its last +// save. Live writers rewrite it every flush, so three missed flushes, with a +// floor for very short intervals, mean the writer stopped. Without a flush +// loop nothing rewrites in-flight rows, or reloads them, so none count. +func (s *Service) inFlightLiveness() time.Duration { + if s.flushInterval <= 0 { + return 0 + } + return max(3*s.flushInterval, minInFlightLiveness) +} + +// loadPeers reads the rows other writers saved since the previous load (or +// all of them, see peerFullLoadInterval) and installs their sums as the +// limiter's remote counters. +func (s *Service) loadPeers(ctx context.Context) error { + startedAt := s.now() + s.persistMu.Lock() + full := s.peers.rows == nil || startedAt.Sub(s.peers.fullAt) >= peerFullLoadInterval + var since int64 + if !full { + since = s.peers.loadedAt.Add(-peerLoadOverlap).Unix() + } + s.persistMu.Unlock() + + // The read runs outside persistMu, so a slow store never holds up a + // Refresh. installPeers sums under the rules in force when it runs. + loadCtx, cancel := persistContext(ctx) + rows, err := s.store.LoadCounters(loadCtx, since) + cancel() + if err != nil { + return err + } + s.installPeers(rows, startedAt, full) + return nil +} + +// installPeers merges a load into the cache and installs the sums. It runs +// under persistMu, like Refresh's swap and the resets it implies, so sums +// taken under rules a Refresh just replaced never land after its resets. +func (s *Service) installPeers(rows []WindowSnapshot, loadedAt time.Time, full bool) { + s.persistMu.Lock() + defer s.persistMu.Unlock() + s.peers.merge(rows, s.writer, full) + s.peers.loadedAt = loadedAt + if full { + s.peers.fullAt = loadedAt + } + requests, tokens, inFlight := s.peers.sum(s.Rules(), s.now(), s.inFlightLiveness()) + s.limiter.setRemote(requests, tokens, inFlight) +} diff --git a/internal/ratelimit/peers_test.go b/internal/ratelimit/peers_test.go new file mode 100644 index 000000000..7b3be57a4 --- /dev/null +++ b/internal/ratelimit/peers_test.go @@ -0,0 +1,329 @@ +package ratelimit + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// fakeClock is a settable clock shared by services and their store, so +// tests age rows without sleeping. +type fakeClock struct { + mu sync.Mutex + now time.Time +} + +func (c *fakeClock) Now() time.Time { + c.mu.Lock() + defer c.mu.Unlock() + return c.now +} + +func (c *fakeClock) Advance(d time.Duration) { + c.mu.Lock() + defer c.mu.Unlock() + c.now = c.now.Add(d) +} + +// clusterStore is a memStore on a fake clock, recording the since of every +// load. +type clusterStore struct { + memStore + clock *fakeClock + loads []int64 + saves int +} + +func newClusterStore(t *testing.T, rules ...Rule) *clusterStore { + t.Helper() + clock := &fakeClock{now: windowBase.Add(10 * time.Second)} + store := &clusterStore{clock: clock} + store.now = clock.Now + require.NoError(t, store.UpsertRules(context.Background(), rules)) + return store +} + +func (s *clusterStore) LoadCounters(ctx context.Context, since int64) ([]WindowSnapshot, error) { + s.loads = append(s.loads, since) + return s.memStore.LoadCounters(ctx, since) +} + +func (s *clusterStore) SaveCounters(ctx context.Context, snapshots []WindowSnapshot) error { + s.saves++ + return s.memStore.SaveCounters(ctx, snapshots) +} + +// replica starts a service on store's clock. Its flush interval only sets +// the in-flight liveness bound: the loop never runs, tests call flush. +func replica(t *testing.T, store *clusterStore) *Service { + t.Helper() + service, err := NewService(context.Background(), store, WithFlushInterval(0)) + require.NoError(t, err) + t.Cleanup(service.Close) + service.now = store.clock.Now + service.Start(context.Background()) + service.flushInterval = time.Second + return service +} + +var teamOnOpenAI = Subjects{UserPath: "/team", Provider: "openai", Model: "openai/gpt-4o"} + +func TestReplicasEnforceOneLimitTogether(t *testing.T) { + ctx := context.Background() + store := newClusterStore(t, + Rule{Scope: ScopeUserPath, Subject: "/team", PeriodSeconds: PeriodHourSeconds, MaxRequests: new(int64(10)), MaxTokens: new(int64(1000)), Source: SourceManual}, + Rule{Scope: ScopeProvider, Subject: "openai", PeriodSeconds: PeriodHourSeconds, MaxRequests: new(int64(100)), Source: SourceManual}, + ) + a := replica(t, store) + b := replica(t, store) + now := store.clock.Now + + // Each stays under the limit on its own. + for range 6 { + _, err := a.Acquire(teamOnOpenAI, now()) + require.NoError(t, err) + } + for range 4 { + _, err := b.Acquire(teamOnOpenAI, now()) + require.NoError(t, err) + } + a.RecordTokens(onPath("/team"), 300, now()) + b.RecordTokens(onPath("/team"), 200, now()) + require.Equal(t, int64(6), requestsUsedFor(t, a, "/team", now())) + + // One flush cycle later both see the cluster's count and refuse. + a.flush(ctx) + b.flush(ctx) + a.flush(ctx) + for _, service := range []*Service{a, b} { + assert.Equal(t, int64(10), requestsUsedFor(t, service, "/team", now())) + statuses := service.StatusesForUserPath("/team", now()) + require.Len(t, statuses, 1) + assert.Equal(t, int64(500), statuses[0].TokensUsed) + _, err := service.Acquire(teamOnOpenAI, now()) + var exceeded *ExceededError + require.ErrorAs(t, err, &exceeded) + assert.Equal(t, int64(10), exceeded.Observed) + assert.Greater(t, exceeded.RetryAfter, time.Duration(0)) + } + + // The provider rule, with room left, sees the cluster too: its + // headers and the routing probe count both replicas. + store.rules = store.rules[1:] + require.NoError(t, a.Refresh(ctx)) + reservation, err := a.Acquire(teamOnOpenAI, now()) + require.NoError(t, err) + assert.Equal(t, int64(100-11), reservation.Headers().RequestRemaining) + assert.True(t, a.routeAvailableAt("openai", "openai/gpt-4o", now())) + + // Rows hold each replica's own counts only. + byWriter := make(map[string]int64) + for _, row := range store.counters { + if row.Subject == "/team" { + byWriter[row.Writer] = row.RequestsCurrent + } + } + assert.Equal(t, map[string]int64{a.writer: 6, b.writer: 4}, byWriter) +} + +func requestsUsedFor(t *testing.T, service *Service, path string, now time.Time) int64 { + t.Helper() + statuses := service.StatusesForUserPath(path, now) + require.Len(t, statuses, 1) + return statuses[0].RequestsUsed +} + +func TestResetClearsPeerCountsOnEveryReplica(t *testing.T) { + ctx := context.Background() + store := newClusterStore(t, teamRule(false, 100)) + a := replica(t, store) + b := replica(t, store) + now := store.clock.Now + a.SetChangeNotifier(func(ctx context.Context) { require.NoError(t, b.Refresh(ctx)) }) + + acquireN(t, a, 5, now()) + acquireN(t, b, 3, now()) + a.flush(ctx) + b.flush(ctx) + a.flush(ctx) + require.Equal(t, int64(8), requestsUsed(t, a, now())) + require.Equal(t, int64(8), requestsUsed(t, b, now())) + + require.NoError(t, a.ResetRule(ctx, ScopeUserPath, "/team", PeriodDaySeconds)) + assert.Equal(t, int64(0), requestsUsed(t, a, now()), "the replica that reset forgets the peer's count") + assert.Equal(t, int64(0), requestsUsed(t, b, now()), "the peer forgets the resetting replica's count") + + // Still cached rows of the old epoch never come back. + a.flush(ctx) + b.flush(ctx) + assert.Equal(t, int64(0), requestsUsed(t, a, now())) + assert.Equal(t, int64(0), requestsUsed(t, b, now())) + + // A mode round trip on one replica clears the other's view the same way. + acquireN(t, a, 2, now()) + a.flush(ctx) + b.flush(ctx) + require.Equal(t, int64(2), requestsUsed(t, b, now())) + require.NoError(t, a.UpsertRules(ctx, []Rule{teamRule(true, 100)})) + require.NoError(t, a.UpsertRules(ctx, []Rule{teamRule(false, 100)})) + assert.Equal(t, int64(0), requestsUsed(t, b, now())) + b.flush(ctx) + assert.Equal(t, int64(0), requestsUsed(t, b, now())) +} + +func TestPeerInFlightCountsWhileFreshAndExpires(t *testing.T) { + ctx := context.Background() + store := newClusterStore(t, Rule{ + Scope: ScopeUserPath, Subject: "/team", PeriodSeconds: PeriodConcurrent, + MaxRequests: new(int64(2)), Source: SourceManual, + }) + a := replica(t, store) + b := replica(t, store) + now := store.clock.Now + + first, err := a.Acquire(onPath("/team"), now()) + require.NoError(t, err) + second, err := a.Acquire(onPath("/team"), now()) + require.NoError(t, err) + a.flush(ctx) + b.flush(ctx) + statuses := b.StatusesForUserPath("/team", now()) + require.Len(t, statuses, 1) + assert.Equal(t, int64(2), statuses[0].InFlight) + _, err = b.Acquire(onPath("/team"), now()) + var exceeded *ExceededError + require.ErrorAs(t, err, &exceeded, "the peer's slots count") + assert.Equal(t, ScopeConcurrency, exceeded.Scope) + + // A release is published as it happens, not left to go stale. + first.Release() + a.flush(ctx) + b.flush(ctx) + held, err := b.Acquire(onPath("/team"), now()) + require.NoError(t, err) + held.Release() + + // A peer that keeps flushing keeps its slot... + store.clock.Advance(4 * time.Second) + a.flush(ctx) + store.clock.Advance(4 * time.Second) + b.flush(ctx) + held, err = b.Acquire(onPath("/team"), now()) + require.NoError(t, err) + _, err = b.Acquire(onPath("/team"), now()) + require.Error(t, err, "a live peer's slot still counts") + held.Release() + + // ...and one that stops, crashed, loses it once the row goes stale. + store.clock.Advance(minInFlightLiveness + time.Second) + b.flush(ctx) + for range 2 { + _, err = b.Acquire(onPath("/team"), now()) + require.NoError(t, err) + } + second.Release() +} + +func TestIncrementalLoadKeepsUnchangedRowsAndDropsRolledOutOnes(t *testing.T) { + ctx := context.Background() + store := newClusterStore(t, Rule{ + Scope: ScopeUserPath, Subject: "/team", PeriodSeconds: PeriodHourSeconds, + MaxRequests: new(int64(100)), Source: SourceManual, + }) + a := replica(t, store) + b := replica(t, store) + now := store.clock.Now + require.Equal(t, []int64{0, 0}, store.loads, "Start loads every row") + + acquireN(t, a, 3, now()) + a.flush(ctx) + b.flush(ctx) + require.Equal(t, int64(3), requestsUsed(t, b, now())) + since := store.loads[len(store.loads)-1] + assert.Equal(t, now().Add(-peerLoadOverlap).Unix(), since, "a flush loads only rows saved since the last load") + + // a stops counting: its row is not rewritten, so loads stop returning + // it, and b keeps counting it from the cache. + saves := store.saves + a.flush(ctx) + assert.Equal(t, saves, store.saves, "nothing changed, nothing is written") + store.clock.Advance(30 * time.Second) + b.flush(ctx) + b.flush(ctx) + rows, err := store.memStore.LoadCounters(ctx, store.loads[len(store.loads)-1]) + require.NoError(t, err) + require.Empty(t, rows) + assert.Equal(t, int64(3), requestsUsed(t, b, now())) + + // A change is picked up by the next load. + acquireN(t, a, 2, now()) + a.flush(ctx) + b.flush(ctx) + assert.Equal(t, int64(5), requestsUsed(t, b, now())) + + // Once a minute the load is a full one. + store.clock.Advance(peerFullLoadInterval) + b.flush(ctx) + assert.Equal(t, int64(0), store.loads[len(store.loads)-1]) + assert.Equal(t, int64(5), requestsUsed(t, b, now())) + + // Two periods on, the row holds nothing and leaves the cache. + store.clock.Advance(2 * time.Hour) + b.flush(ctx) + assert.Equal(t, int64(0), requestsUsed(t, b, now())) + b.persistMu.Lock() + assert.Empty(t, b.peers.rows) + b.persistMu.Unlock() +} + +// failingSaveStore fails SaveCounters while err is set. +type failingSaveStore struct { + *clusterStore + err error +} + +func (s *failingSaveStore) SaveCounters(ctx context.Context, snapshots []WindowSnapshot) error { + if s.err != nil { + return s.err + } + return s.clusterStore.SaveCounters(ctx, snapshots) +} + +// Rows are written only when they change, so a failed save must leave them +// to the next one rather than lose them. +func TestFailedSaveIsRetried(t *testing.T) { + ctx := context.Background() + store := &failingSaveStore{clusterStore: newClusterStore(t, + teamRule(false, 100), + Rule{Scope: ScopeUserPath, Subject: "/team", PeriodSeconds: PeriodConcurrent, MaxRequests: new(int64(5)), Source: SourceManual}, + )} + service, err := NewService(ctx, store, WithFlushInterval(0)) + require.NoError(t, err) + t.Cleanup(service.Close) + service.now = store.clock.Now + service.Start(ctx) + + held, err := service.Acquire(onPath("/team"), store.clock.Now()) + require.NoError(t, err) + service.flush(ctx) + held.Release() + store.err = errors.New("store down") + held, err = service.Acquire(onPath("/team"), store.clock.Now()) + require.NoError(t, err) + held.Release() + service.flush(ctx) + store.err = nil + service.flush(ctx) + + byPeriod := make(map[int64]WindowSnapshot) + for _, row := range store.counters { + byPeriod[row.PeriodSeconds] = row + } + assert.Equal(t, int64(2), byPeriod[PeriodDaySeconds].RequestsCurrent) + assert.Equal(t, int64(0), byPeriod[PeriodConcurrent].InFlight, "the released slot is published after all") +} diff --git a/internal/ratelimit/persist.go b/internal/ratelimit/persist.go index dc6fa9817..7c1773a9a 100644 --- a/internal/ratelimit/persist.go +++ b/internal/ratelimit/persist.go @@ -18,9 +18,10 @@ const ( persistClosed ) -// WithFlushInterval sets how often an active generation writes window -// snapshots. Zero disables the periodic loop; Start still loads and Close -// of an active generation still writes once. +// WithFlushInterval sets how often an active generation writes its changed +// counters and reads the other writers'. Zero disables the periodic loop; +// Start still loads and Close of an active generation still writes once, but +// nothing counted elsewhere after Start is seen. func WithFlushInterval(interval time.Duration) ServiceOption { return func(service *Service) { if interval < 0 { @@ -49,13 +50,19 @@ func persistContext(parent context.Context) (context.Context, context.CancelFunc return context.WithTimeout(context.Background(), persistTimeout) } -// Start restores persisted windows and, from then on, keeps writing them. -// It is separate from NewService because a reload builds the next generation -// while the current one still serves: only the generation that gets to serve -// may touch the snapshot, or a replacement that is built and then discarded -// would flush its empty windows over the live ones. Start is idempotent, and -// a failed load leaves the generation idle rather than persisting from a -// blank slate. +// Start loads the counters every writer saved, earlier runs of this one +// included, and from then on keeps writing this generation's and reading the +// others'. It is separate from NewService because a reload builds the next +// generation while the current one still serves: only the generation that +// gets to serve may touch the snapshot, or a replacement that is built and +// then discarded would publish its counters too. Start is idempotent, and a +// failed load leaves the generation idle rather than persisting from a blank +// slate. +// +// The loaded rows become the limiter's remote counters, not its own: this +// generation's rows only ever hold what it counted itself (see +// WindowSnapshot), so a restart counts its predecessor's requests once, +// through their rows, rather than again through its own. func (s *Service) Start(ctx context.Context) { if s == nil || s.store == nil { return @@ -69,8 +76,9 @@ func (s *Service) Start(ctx context.Context) { s.persistState = persistStarting s.lifeMu.Unlock() + loadedAt := s.now() loadCtx, cancel := persistContext(ctx) - snapshots, err := s.store.LoadCounters(loadCtx) + snapshots, err := s.store.LoadCounters(loadCtx, 0) cancel() s.lifeMu.Lock() @@ -83,7 +91,7 @@ func (s *Service) Start(ctx context.Context) { slog.Warn("rate limit counters: load failed; not persisting this generation", "error", err) return } - s.limiter.restore(snapshots, s.Rules(), time.Now().UTC()) + s.installPeers(snapshots, loadedAt, true) s.startFlushLoop() s.persistState = persistActive if len(snapshots) > 0 { @@ -107,9 +115,7 @@ func (s *Service) startFlushLoop() { for { select { case <-ticker.C: - flushCtx, cancel := persistContext(context.Background()) - s.flush(flushCtx) - cancel() + s.flush(context.Background()) case <-stop: return } @@ -117,17 +123,37 @@ func (s *Service) startFlushLoop() { }() } +// flush is one tick of the flush loop: it saves this writer's changes, then +// reads the other writers'. The read comes second so that, within one tick, +// peers see this writer's counts no later than it sees theirs. func (s *Service) flush(ctx context.Context) { + if s == nil || s.store == nil { + return + } + s.save(ctx) + if err := s.loadPeers(ctx); err != nil { + slog.Warn("rate limit counters: load failed", "error", err) + } +} + +// save writes the rows of the counters that changed since the last save. +func (s *Service) save(ctx context.Context) { if s == nil || s.store == nil { return } s.persistMu.Lock() defer s.persistMu.Unlock() snapshots := s.limiter.snapshot(s.Rules()) + if len(snapshots) == 0 { + return + } for i := range snapshots { snapshots[i].Writer = s.writer } - if err := s.store.SaveCounters(ctx, snapshots); err != nil { + saveCtx, cancel := persistContext(ctx) + defer cancel() + if err := s.store.SaveCounters(saveCtx, snapshots); err != nil { + s.limiter.unsaved(snapshots) slog.Warn("rate limit counters: flush failed", "error", err) } } diff --git a/internal/ratelimit/persist_test.go b/internal/ratelimit/persist_test.go index 58bf7a050..4a37681cb 100644 --- a/internal/ratelimit/persist_test.go +++ b/internal/ratelimit/persist_test.go @@ -32,14 +32,14 @@ func TestSnapshotRoundTripPreservesEstimate(t *testing.T) { require.Empty(t, snaps[0].Partition) dst := newLimiter() - dst.restore(snaps, []Rule{rule}, now) + remoteOf(dst, snaps, []Rule{rule}, now) got := dst.status(rule, now).RequestsUsed require.Equal(t, wantReq, got) got = dst.status(rule, now).TokensUsed require.Equal(t, wantTok, got) } -func TestSnapshotSkipsConcurrentAndExpiredChild(t *testing.T) { +func TestSnapshotPublishesInFlightAndDropsExpiredChild(t *testing.T) { now := time.Unix(1_700_000_000, 0).UTC() shared := Rule{ Scope: ScopeUserPath, Subject: "/team", PeriodSeconds: PeriodConcurrent, @@ -53,10 +53,10 @@ func TestSnapshotSkipsConcurrentAndExpiredChild(t *testing.T) { require.True(t, ok) src := newLimiter() - _, _, err := src.admit([]Rule{shared, child}, now) + _, held, err := src.admit([]Rule{shared, child}, now) require.Nil(t, err) - // Force the child window into the distant past so restore drops it. + // Force the child window into the distant past so the peer view drops it. src.mu.Lock() for _, counter := range src.requests { if counter != nil { @@ -66,14 +66,24 @@ func TestSnapshotSkipsConcurrentAndExpiredChild(t *testing.T) { src.mu.Unlock() snaps := src.snapshot([]Rule{shared, template}) + var gauges []WindowSnapshot for _, snap := range snaps { - require.NotEqual(t, PeriodConcurrent, snap.PeriodSeconds) + if snap.PeriodSeconds == PeriodConcurrent { + gauges = append(gauges, snap) + } } + require.Equal(t, []WindowSnapshot{{Scope: string(ScopeUserPath), Subject: "/team", InFlight: 1}}, gauges) dst := newLimiter() - dst.restore(snaps, []Rule{template}, now) + remoteOf(dst, snaps, []Rule{template}, now) got := dst.status(child, now).RequestsUsed require.Equal(t, int64(0), got) + + // Released, the gauge is published once more, as zero, then no longer. + src.release(held) + snaps = src.snapshot([]Rule{shared, template}) + require.Equal(t, []WindowSnapshot{{Scope: string(ScopeUserPath), Subject: "/team"}}, snaps) + require.Empty(t, src.snapshot([]Rule{shared, template})) } func TestSnapshotIsolatesPerChildPartitions(t *testing.T) { @@ -95,7 +105,7 @@ func TestSnapshotIsolatesPerChildPartitions(t *testing.T) { require.Len(t, snaps, 2) dst := newLimiter() - dst.restore(snaps, []Rule{template}, now) + remoteOf(dst, snaps, []Rule{template}, now) _, _, err = dst.admit([]Rule{alice}, now) require.NotNil(t, err) _, _, err = dst.admit([]Rule{bob}, now) @@ -182,14 +192,14 @@ func TestResetClearsPersistedWindow(t *testing.T) { require.NoError(t, err) } -func TestRestoreIgnoresSharedRowOnPerChildRule(t *testing.T) { +func TestPeerSumIgnoresSharedRowOnPerChildRule(t *testing.T) { now := time.Unix(1_700_000_000, 0).UTC() template := Rule{ Scope: ScopeUserPath, Subject: "/customers", PerChild: true, PeriodSeconds: PeriodHourSeconds, MaxRequests: new(int64(1)), } dst := newLimiter() - dst.restore([]WindowSnapshot{{ + remoteOf(dst, []WindowSnapshot{{ Scope: string(ScopeUserPath), Subject: "/customers", Partition: "", PeriodSeconds: PeriodHourSeconds, RequestsWindowStart: now.Unix(), RequestsCurrent: 1, }}, []Rule{template}, now) @@ -359,7 +369,7 @@ type failLoadStore struct { err error } -func (s *failLoadStore) LoadCounters(context.Context) ([]WindowSnapshot, error) { +func (s *failLoadStore) LoadCounters(context.Context, int64) ([]WindowSnapshot, error) { return nil, s.err } @@ -378,10 +388,10 @@ type blockingLoadStore struct { release chan struct{} } -func (s *blockingLoadStore) LoadCounters(context.Context) ([]WindowSnapshot, error) { +func (s *blockingLoadStore) LoadCounters(_ context.Context, since int64) ([]WindowSnapshot, error) { close(s.started) <-s.release - return s.memStore.LoadCounters(context.Background()) + return s.memStore.LoadCounters(context.Background(), since) } // TestDeleteRuleStopsEnforcingWhenSnapshotDeleteFails: the rule row is already diff --git a/internal/ratelimit/persist_writers_test.go b/internal/ratelimit/persist_writers_test.go index 2478b0b57..d2f40bce6 100644 --- a/internal/ratelimit/persist_writers_test.go +++ b/internal/ratelimit/persist_writers_test.go @@ -10,21 +10,27 @@ import ( "github.com/stretchr/testify/require" ) -func TestRestoreMergesWritersByMax(t *testing.T) { +// remoteOf installs rows as another writer's, summed as a load does. +func remoteOf(l *limiter, rows []WindowSnapshot, rules []Rule, now time.Time) { + var view peerView + view.merge(rows, "self", true) + l.setRemote(view.sum(rules, now, time.Minute)) +} + +func TestPeerSumAddsWritersRolledForward(t *testing.T) { period := PeriodHourSeconds start := int64(1_700_000_000) - int64(1_700_000_000)%period now := time.Unix(start+60, 0).UTC() rule := Rule{ Scope: ScopeUserPath, Subject: "/team", PeriodSeconds: period, - MaxRequests: new(int64(100)), MaxTokens: new(int64(1000)), + MaxRequests: new(int64(100)), MaxTokens: new(int64(1000)), CounterEpoch: "e1", } row := func(writer string) WindowSnapshot { return WindowSnapshot{ - Scope: string(ScopeUserPath), Subject: "/team", PeriodSeconds: period, Writer: writer, + Scope: string(ScopeUserPath), Subject: "/team", PeriodSeconds: period, Writer: writer, Epoch: "e1", } } - // Same window: the larger current and the larger previous win, - // independently and per counter kind. + // Same window: added field by field, per counter kind. a := row("a") a.RequestsWindowStart, a.RequestsCurrent, a.RequestsPrevious = start, 5, 2 a.TokensWindowStart, a.TokensCurrent, a.TokensPrevious = start, 10, 90 @@ -32,37 +38,73 @@ func TestRestoreMergesWritersByMax(t *testing.T) { b.RequestsWindowStart, b.RequestsCurrent, b.RequestsPrevious = start, 3, 7 b.TokensWindowStart, b.TokensCurrent, b.TokensPrevious = start, 40, 20 // One period behind: rolled forward first, so its current becomes - // previous (9 beats 7) and its previous falls out. + // previous and its previous falls out. c := row("c") c.RequestsWindowStart, c.RequestsCurrent, c.RequestsPrevious = start-period, 9, 50 - // Two periods behind: holds nothing at now, whatever it counted. + // Two periods behind: holds nothing at now, whatever it counted, and + // leaves the cache. d := row("d") d.RequestsWindowStart, d.RequestsCurrent = start-2*period, 80 d.TokensWindowStart, d.TokensCurrent = start-2*period, 800 - - l := newLimiter() - l.restore([]WindowSnapshot{a, b, c, d}, []Rule{rule}, now) + // Another epoch: skipped, but kept, as this replica may not have + // reloaded the rule yet. + e := row("e") + e.Epoch = "e0" + e.RequestsWindowStart, e.RequestsCurrent = start, 1000 + // This writer's own row: already counted locally. + self := row("self") + self.RequestsWindowStart, self.RequestsCurrent = start, 1000 + + var view peerView + view.merge([]WindowSnapshot{a, b, c, d, e, self}, "self", true) + requests, tokens, inFlight := view.sum([]Rule{rule}, now, time.Minute) key := keyForRule(rule) - require.Contains(t, l.requests, key) - assert.Equal(t, windowCounter{windowStart: start, current: 5, previous: 9}, *l.requests[key]) - require.Contains(t, l.tokens, key) - assert.Equal(t, windowCounter{windowStart: start, current: 40, previous: 90}, *l.tokens[key]) + assert.Equal(t, map[ruleKey]windowCounter{key: {windowStart: start, current: 8, previous: 18}}, requests) + assert.Equal(t, map[ruleKey]windowCounter{key: {windowStart: start, current: 50, previous: 110}}, tokens) + assert.Nil(t, inFlight) + assert.Len(t, view.rows, 4, "the rolled-out row leaves the cache, the other epoch's stays") + assert.NotContains(t, view.rows, identityOf(d)) + + // The limiter checks the local count plus the sum. + l := newLimiter() + l.setRemote(requests, tokens, inFlight) + _, _, exceeded := l.admit([]Rule{rule}, now) + require.Nil(t, exceeded) + assert.Equal(t, int64(9+18*(period-60)/period), l.status(rule, now).RequestsUsed) } -func TestRestoreSkipsKeysThatHoldNothing(t *testing.T) { +func TestPeerSumSkipsKeysThatHoldNothing(t *testing.T) { period := PeriodHourSeconds start := int64(1_700_000_000) - int64(1_700_000_000)%period now := time.Unix(start, 0).UTC() rule := Rule{Scope: ScopeUserPath, Subject: "/team", PeriodSeconds: period, MaxRequests: new(int64(1))} - l := newLimiter() - l.restore([]WindowSnapshot{{ + var view peerView + view.merge([]WindowSnapshot{{ Scope: string(ScopeUserPath), Subject: "/team", PeriodSeconds: period, Writer: "a", RequestsWindowStart: start - 2*period, RequestsCurrent: 5, - }}, []Rule{rule}, now) - assert.Empty(t, l.requests) - assert.Empty(t, l.tokens) + }, { + Scope: string(ScopeUserPath), Subject: "/team", PeriodSeconds: period, Writer: "b", + RequestsWindowStart: start, + }}, "self", true) + requests, tokens, inFlight := view.sum([]Rule{rule}, now, time.Minute) + assert.Nil(t, requests) + assert.Nil(t, tokens) + assert.Nil(t, inFlight) +} + +// A replica whose clock runs a little ahead can save a row in the next +// window before this one reaches it. Its counts are folded into this +// window rather than dropped. +func TestRolledToFoldsCounterFromALaterWindow(t *testing.T) { + period := PeriodMinuteSeconds + start := int64(1_700_000_000) - int64(1_700_000_000)%period + now := time.Unix(start+period-1, 0) + ahead := windowCounter{windowStart: start + period, current: 2, previous: 5} + assert.Equal(t, windowCounter{windowStart: start, current: 7}, ahead.rolledTo(now, period)) + behind := windowCounter{windowStart: start - period, current: 3, previous: 4} + assert.Equal(t, windowCounter{windowStart: start, previous: 3}, behind.rolledTo(now, period)) } // startedService builds and starts a service on store, closing it at cleanup @@ -132,9 +174,21 @@ func TestRestartUnderNewWriterRestoresOwnCounts(t *testing.T) { acquireN(t, second, 1, now) second.Close() require.Len(t, store.counters, 2, "each run keeps its own row") + byWriter := make(map[string]int64, len(store.counters)) + for _, snap := range store.counters { + byWriter[snap.Writer] = snap.RequestsCurrent + } + require.Equal(t, map[string]int64{first.writer: 2, second.writer: 1}, byWriter, + "a row holds only what its run counted, never what it loaded") third := startedService(t, store) assert.Equal(t, int64(3), requestsUsed(t, third, now)) + acquireN(t, third, 4, now) + third.Close() + + // Each generation is counted once, however many came before it. + fourth := startedService(t, store) + assert.Equal(t, int64(7), requestsUsed(t, fourth, now)) } func TestReplicasDoNotOverwriteEachOther(t *testing.T) { @@ -154,10 +208,10 @@ func TestReplicasDoNotOverwriteEachOther(t *testing.T) { } require.Equal(t, map[string]int64{a.writer: 4, b.writer: 1}, byWriter) - // A replica restarting takes the busiest peer's count as its own share - // rather than the sum of all of them. + // A replica restarting sees what the cluster counted: every replica's + // row, each once. restarted := startedService(t, store) - assert.Equal(t, int64(4), requestsUsed(t, restarted, now)) + assert.Equal(t, int64(5), requestsUsed(t, restarted, now)) } func TestResetDeletesEveryWritersRows(t *testing.T) { diff --git a/internal/ratelimit/service.go b/internal/ratelimit/service.go index 61d121eb2..a3c0ada1f 100644 --- a/internal/ratelimit/service.go +++ b/internal/ratelimit/service.go @@ -43,15 +43,20 @@ type Service struct { // writer keys this generation's window snapshot rows, so replicas // sharing a store never overwrite each other. Every Service draws a new - // one, so a restart or reload writes beside its predecessor's rows and - // restore merges them (see limiter.restore). + // one, so a restart or reload writes beside its predecessor's rows, and + // counts them as another writer's (see peerView.sum). writer string flushInterval time.Duration - persistMu sync.Mutex - lifeMu sync.Mutex - persistState persistState - flushStop chan struct{} - flushDone chan struct{} + // persistMu serializes saves, resets of saved rows, and installing the + // peer view with Refresh's swap; it also guards peers. + persistMu sync.Mutex + lifeMu sync.Mutex + persistState persistState + flushStop chan struct{} + flushDone chan struct{} + peers peerView + // now is the clock the peer view is kept by. Tests replace it. + now func() time.Time // changed is called after every successful rule write; see // SetChangeNotifier. @@ -89,6 +94,7 @@ func NewService(ctx context.Context, store Store, options ...ServiceOption) (*Se limiter: newLimiter(), quotaTemplates: true, writer: uuid.NewString(), + now: time.Now, } for _, option := range options { if option != nil { @@ -101,8 +107,8 @@ func NewService(ctx context.Context, store Store, options ...ServiceOption) (*Se return service, nil } -// Close stops the flush loop, writes a final snapshot if this generation -// was started, and stops the in-memory expiry cleanup worker. +// Close stops the flush loop, saves what changed since the last flush if this +// generation was started, and stops the in-memory expiry cleanup worker. func (s *Service) Close() { if s == nil { return @@ -125,9 +131,7 @@ func (s *Service) Close() { <-done } if wasActive { - flushCtx, cancel := persistContext(context.Background()) - s.flush(flushCtx) - cancel() + s.save(context.Background()) } if s.limiter != nil { s.limiter.close() @@ -158,7 +162,8 @@ func (s *Service) Refresh(ctx context.Context) error { }) // The swap and the counter resets it implies happen under persistMu, so // a flush never pairs the new rules, and with them a new counter epoch, - // with counts the change clears. + // with counts the change clears, and a peer load summed under the old + // rules never lands after the resets. s.persistMu.Lock() defer s.persistMu.Unlock() s.mu.Lock() @@ -171,10 +176,11 @@ func (s *Service) Refresh(ctx context.Context) error { // Removed definitions, shared/per-child mode changes, and definitions // reset, or deleted and re-added, elsewhere (a new counter epoch) drop - // their in-memory counters. Their saved rows need no delete: they carry the old - // epoch, which restore skips (see Rule.CounterEpoch), and the store - // collects them once no writer refreshes them. Deleting them here could - // also take rows another replica already saved under the new rule. + // their in-memory counters, this writer's and the other writers' sums. + // Their saved rows need no delete: they carry the old epoch, which the + // peer view skips (see Rule.CounterEpoch), and the store collects them + // once no writer refreshes them. Deleting them here could also take rows + // another replica already saved under the new rule. next := make(map[ruleKey]Rule, len(rules)) for _, rule := range rules { next[keyForRule(rule)] = rule diff --git a/internal/ratelimit/service_test.go b/internal/ratelimit/service_test.go index f14f8c061..663e1c9fb 100644 --- a/internal/ratelimit/service_test.go +++ b/internal/ratelimit/service_test.go @@ -15,8 +15,16 @@ import ( type memStore struct { rules []Rule counters []WindowSnapshot - // written stamps each row the way a real store's updated_at column does. - written map[string]int64 + // now stamps saved rows, as the saving replica's clock does for a real + // store; nil means time.Now. + now func() time.Time +} + +func (m *memStore) clock() time.Time { + if m.now != nil { + return m.now() + } + return time.Now() } func (m *memStore) ListRules(context.Context) ([]Rule, error) { @@ -86,8 +94,14 @@ func (m *memStore) ReplaceConfigRules(ctx context.Context, rules []Rule) error { func (m *memStore) Close() error { return nil } -func (m *memStore) LoadCounters(context.Context) ([]WindowSnapshot, error) { - return append([]WindowSnapshot(nil), m.counters...), nil +func (m *memStore) LoadCounters(_ context.Context, since int64) ([]WindowSnapshot, error) { + var rows []WindowSnapshot + for _, snap := range m.counters { + if snap.UpdatedAt >= since { + rows = append(rows, snap) + } + } + return rows, nil } // snapshotIdentity is the primary key every store keys a window row by. @@ -95,15 +109,13 @@ func snapshotIdentity(s WindowSnapshot) string { return s.Scope + "\x00" + s.Subject + "\x00" + s.Partition + "\x00" + strconv.FormatInt(s.PeriodSeconds, 10) + "\x00" + s.Writer } -// SaveCounters mirrors the real stores: upsert by identity, then collect rows -// that went two periods without a write. +// SaveCounters mirrors the real stores: stamp and upsert by identity, then +// collect window rows that went two periods without a write and in-flight +// rows that went inFlightRowRetention. func (m *memStore) SaveCounters(_ context.Context, snapshots []WindowSnapshot) error { - now := time.Now().Unix() - if m.written == nil { - m.written = make(map[string]int64) - } + now := m.clock().Unix() for _, snap := range snapshots { - m.written[snapshotIdentity(snap)] = now + snap.UpdatedAt = now replaced := false for i, existing := range m.counters { if snapshotIdentity(existing) == snapshotIdentity(snap) { @@ -118,7 +130,11 @@ func (m *memStore) SaveCounters(_ context.Context, snapshots []WindowSnapshot) e } kept := m.counters[:0] for _, snap := range m.counters { - if m.written[snapshotIdentity(snap)]+2*snap.PeriodSeconds >= now { + retention := 2 * snap.PeriodSeconds + if snap.PeriodSeconds == PeriodConcurrent { + retention = inFlightRowRetention + } + if snap.UpdatedAt+retention >= now { kept = append(kept, snap) } } diff --git a/internal/ratelimit/snapshot.go b/internal/ratelimit/snapshot.go index 2b679422d..cb6cbdd9e 100644 --- a/internal/ratelimit/snapshot.go +++ b/internal/ratelimit/snapshot.go @@ -1,13 +1,17 @@ package ratelimit -import "time" - -// WindowSnapshot is one persisted request/token sliding window. -// Partition is empty for a shared rule and the child path for a per-child -// template (Rule.EffectiveSubject). Writer is the random ID of the process -// that wrote the row, so replicas sharing a store keep separate rows instead -// of overwriting each other's; rows imported from before writer IDs existed -// carry an empty writer. +// WindowSnapshot is one writer's persisted counters for one key: its +// request/token sliding windows, or for a concurrency rule its in-flight +// gauge. Partition is empty for a shared rule and the child path for a +// per-child template (Rule.EffectiveSubject). Writer is the random ID of the +// process that wrote the row, so replicas sharing a store keep separate rows +// instead of overwriting each other's; rows imported from before writer IDs +// existed carry an empty writer. +// +// A row holds only what its writer counted itself, never counts it loaded +// from other rows, so the cluster's count for a key is the sum of its rows +// (see peerView.sum). A restarted process writes under a new ID, and its +// earlier run's rows are summed like any other writer's. type WindowSnapshot struct { Scope string `bson:"scope"` Subject string `bson:"subject"` @@ -22,162 +26,156 @@ type WindowSnapshot struct { TokensWindowStart int64 `bson:"tokens_window_start"` TokensCurrent int64 `bson:"tokens_current"` TokensPrevious int64 `bson:"tokens_previous"` + // InFlight is the writer's in-flight gauge on a concurrency rule + // (PeriodConcurrent), whose window fields stay zero. + InFlight int64 `bson:"in_flight"` + // UpdatedAt is when the writer last saved the row, in unix seconds by + // its own clock. Stores stamp it on save, whatever the snapshot holds, + // and report it on load. + UpdatedAt int64 `bson:"updated_at"` } func definitionKey(scope RuleScope, subject string, periodSeconds int64) ruleKey { return ruleKey{scope: scope, subject: subject, periodSeconds: periodSeconds} } +// ruleIndex finds the rule definition a counter key or stored row belongs to. +type ruleIndex map[ruleKey]Rule + +func indexRules(rules []Rule) ruleIndex { + index := make(ruleIndex, len(rules)) + for _, rule := range rules { + index[definitionKey(rule.Scope, rule.Subject, rule.PeriodSeconds)] = rule + } + return index +} + +// lookup returns the rule for key, provided the key's partition fits the +// rule's mode: a shared rule owns only the empty partition, a per-child rule +// only child partitions. +func (index ruleIndex) lookup(key ruleKey) (Rule, bool) { + rule, ok := index[definitionKey(key.scope, key.subject, key.periodSeconds)] + if !ok || rule.PerChild != (key.partition != "") { + return Rule{}, false + } + return rule, true +} + +func newSnapshot(key ruleKey, rule Rule) WindowSnapshot { + return WindowSnapshot{ + Scope: string(key.scope), + Subject: key.subject, + Partition: key.partition, + PeriodSeconds: key.periodSeconds, + Epoch: rule.CounterEpoch, + } +} + +// snapshot returns the rows this writer has to save: one per key whose +// request or token counter changed since the previous snapshot, and one per +// concurrency key whose gauge is nonzero, or dropped to zero since the +// previous snapshot published it. A row of an unchanged key is left as it is: +// it already holds the same counts, and a reader rolls it forward as admission +// rolls the counter. In-flight rows are republished every time, so peers can +// tell a live writer's slots from a crashed one's (see Service.inFlightLiveness). +// +// It clears the change marks; a save that fails hands its rows back through +// unsaved. func (l *limiter) snapshot(rules []Rule) []WindowSnapshot { l.mu.Lock() defer l.mu.Unlock() - defs := make(map[ruleKey]Rule, len(rules)) - for _, rule := range rules { - if rule.PeriodSeconds <= 0 { - continue + index := indexRules(rules) + var changed []ruleKey + for _, counters := range []map[ruleKey]*windowCounter{l.requests, l.tokens} { + for key, counter := range counters { + if !counter.dirty { + continue + } + counter.dirty = false + if key.periodSeconds > 0 { + changed = append(changed, key) + } } - defs[definitionKey(rule.Scope, rule.Subject, rule.PeriodSeconds)] = rule } - byKey := make(map[ruleKey]*WindowSnapshot) - add := func(key ruleKey, kind counterKind, counter windowCounter) { - rule, ok := defs[definitionKey(key.scope, key.subject, key.periodSeconds)] - if !ok { - return - } - if rule.PerChild != (key.partition != "") { - return - } - snap := byKey[key] - if snap == nil { - snap = &WindowSnapshot{ - Scope: string(key.scope), - Subject: key.subject, - Partition: key.partition, - PeriodSeconds: key.periodSeconds, - Epoch: rule.CounterEpoch, - } - byKey[key] = snap + var out []WindowSnapshot + seen := make(map[ruleKey]bool, len(changed)) + for _, key := range changed { + rule, ok := index.lookup(key) + if !ok || seen[key] { + continue } - switch kind { - case requestCounter: + seen[key] = true + snap := newSnapshot(key, rule) + if counter := l.requests[key]; counter != nil { snap.RequestsWindowStart = counter.windowStart snap.RequestsCurrent = counter.current snap.RequestsPrevious = counter.previous - case tokenCounter: + } + if counter := l.tokens[key]; counter != nil { snap.TokensWindowStart = counter.windowStart snap.TokensCurrent = counter.current snap.TokensPrevious = counter.previous } + out = append(out, snap) } - for key, counter := range l.requests { - add(key, requestCounter, *counter) - } - for key, counter := range l.tokens { - add(key, tokenCounter, *counter) + for key, inFlight := range l.inFlight { + rule, ok := index.lookup(key) + if !ok || key.periodSeconds != PeriodConcurrent { + continue + } + snap := newSnapshot(key, rule) + snap.InFlight = inFlight + out = append(out, snap) + l.publishedInFlight[key] = struct{}{} } - - out := make([]WindowSnapshot, 0, len(byKey)) - for _, snap := range byKey { - out = append(out, *snap) + for key := range l.publishedInFlight { + if l.inFlight[key] > 0 { + continue + } + delete(l.publishedInFlight, key) + if rule, ok := index.lookup(key); ok { + out = append(out, newSnapshot(key, rule)) + } } return out } -// restore loads persisted windows into the limiter. A key can have one row -// per writer: every replica sharing the store writes its own, and a restarted -// process comes back under a new writer ID next to the rows of its previous -// run. Rows for the same key are merged per counter (requests and tokens -// separately): -// -// 1. Each row is rolled forward to the window containing now, exactly as -// admission would roll it (windowCounter.advance): a row one period behind -// shifts current into previous, an older one drops to zero. -// 2. The merged counter takes the maximum current and the maximum previous -// across those normalized rows. -// -// The maximum, not the sum, is deliberate. A single instance restarting gets -// its own counts back: its last run started from the run before it, so its -// rows dominate every older one field by field. With N replicas each serving -// about 1/N of the traffic, the largest peer row approximates this replica's -// own share, where a sum would charge it for all N. -func (l *limiter) restore(snapshots []WindowSnapshot, rules []Rule, now time.Time) { +// unsaved marks the rows of a failed save for the next snapshot again. A +// counter reset since the snapshot has nothing left to mark. +func (l *limiter) unsaved(snapshots []WindowSnapshot) { l.mu.Lock() defer l.mu.Unlock() - - byDef := make(map[ruleKey]Rule, len(rules)) - for _, rule := range rules { - byDef[definitionKey(rule.Scope, rule.Subject, rule.PeriodSeconds)] = rule - } - - nowUnix := now.Unix() - requests := make(map[ruleKey]*windowCounter) - tokens := make(map[ruleKey]*windowCounter) for _, snap := range snapshots { - if snap.PeriodSeconds <= 0 { - continue - } - rule, ok := byDef[definitionKey(RuleScope(snap.Scope), snap.Subject, snap.PeriodSeconds)] - if !ok { - continue - } - if rule.PerChild != (snap.Partition != "") { - continue - } - // Counted under an earlier definition of the rule: a mode change or - // a delete and re-add cleared these counts. - if snap.Epoch != rule.CounterEpoch { - continue - } - latest := max(snap.RequestsWindowStart, snap.TokensWindowStart) - if latest > 0 && latest+2*snap.PeriodSeconds < nowUnix { - continue - } - key := ruleKey{ scope: RuleScope(snap.Scope), subject: snap.Subject, partition: snap.Partition, periodSeconds: snap.PeriodSeconds, } - mergeCounter(requests, key, windowCounter{ - windowStart: snap.RequestsWindowStart, - current: snap.RequestsCurrent, - previous: snap.RequestsPrevious, - }, now) - mergeCounter(tokens, key, windowCounter{ - windowStart: snap.TokensWindowStart, - current: snap.TokensCurrent, - previous: snap.TokensPrevious, - }, now) - } - for key, counter := range requests { - l.requests[key] = counter - l.trackCounterExpiry(requestCounter, key, counter) - } - for key, counter := range tokens { - l.tokens[key] = counter - l.trackCounterExpiry(tokenCounter, key, counter) + if snap.PeriodSeconds == PeriodConcurrent { + // A nonzero gauge goes out with every snapshot anyway; a drop + // to zero has to be retried. + l.publishedInFlight[key] = struct{}{} + continue + } + if counter := l.requests[key]; counter != nil { + counter.dirty = true + } + if counter := l.tokens[key]; counter != nil { + counter.dirty = true + } } } -// mergeCounter folds one persisted row into the merged counter for key, as -// described on restore. A row that holds nothing at now leaves no counter. -func mergeCounter(merged map[ruleKey]*windowCounter, key ruleKey, row windowCounter, now time.Time) { - if row.current == 0 && row.previous == 0 { - return - } - row.advance(now, key.periodSeconds) - if row.current == 0 && row.previous == 0 { - return - } - existing := merged[key] - if existing == nil { - merged[key] = &row - return - } - existing.current = max(existing.current, row.current) - existing.previous = max(existing.previous, row.previous) +// setRemote installs the other writers' summed counters, replacing the +// previous load's. +func (l *limiter) setRemote(requests, tokens map[ruleKey]windowCounter, inFlight map[ruleKey]int64) { + l.mu.Lock() + defer l.mu.Unlock() + l.remoteRequests = requests + l.remoteTokens = tokens + l.remoteInFlight = inFlight } diff --git a/internal/ratelimit/store.go b/internal/ratelimit/store.go index dfb679a43..e534430aa 100644 --- a/internal/ratelimit/store.go +++ b/internal/ratelimit/store.go @@ -9,25 +9,33 @@ import ( var ErrNotFound = errors.New("rate limit rule not found") +// inFlightRowRetention is how long a store keeps an in-flight row nobody +// saves any more. Peers stop counting it after a few flush intervals (see +// Service.inFlightLiveness); this only bounds how long it takes space. +const inFlightRowRetention = PeriodDaySeconds + // Store persists rate limit rule definitions and optional window snapshots. type Store interface { ListRules(ctx context.Context) ([]Rule, error) UpsertRules(ctx context.Context, rules []Rule) error DeleteRule(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error ReplaceConfigRules(ctx context.Context, rules []Rule) error - // LoadCounters returns every writer's rows; restore merges them per key. - LoadCounters(ctx context.Context) ([]WindowSnapshot, error) + // LoadCounters returns every writer's rows saved at or after since, in + // unix seconds by the writer's clock (WindowSnapshot.UpdatedAt); since 0 + // returns them all. + LoadCounters(ctx context.Context, since int64) ([]WindowSnapshot, error) // SaveCounters upserts the given snapshots, keyed by their Writer as well - // as their window, so writers never overwrite each other's rows. It also - // collects rows no writer has refreshed for two of their own periods — the - // same staleness bound a load applies, so it only ever drops rows that - // could no longer be restored. + // as their window, so writers never overwrite each other's rows, and + // stamps them with the current time. It also collects window rows no + // writer has saved for two of their own periods — by then every count + // they hold has rolled out — and in-flight rows no writer has saved for + // a day, long after peers stopped counting them. // Nothing is deleted to make room for a write, so a crash mid-save costs at // most the flush interval rather than a whole window. SaveCounters(ctx context.Context, snapshots []WindowSnapshot) error // ResetCounter gives the rule a new counter epoch, so every replica's - // next Refresh clears its counts and restore skips every row saved - // before, and removes the rows of every writer. ResetAllCounters does + // next Refresh clears its counts and loads skip every row saved before, + // and removes the rows of every writer. ResetAllCounters does // the same for every rule. ResetCounter(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error ResetAllCounters(ctx context.Context) error diff --git a/internal/ratelimit/store_counters_test.go b/internal/ratelimit/store_counters_test.go index e3723eafa..f5ba1da41 100644 --- a/internal/ratelimit/store_counters_test.go +++ b/internal/ratelimit/store_counters_test.go @@ -8,6 +8,17 @@ import ( "github.com/stretchr/testify/require" ) +// unstamped returns rows without the save stamp a store adds, for comparing +// them with the snapshots that were saved. +func unstamped(rows []WindowSnapshot) []WindowSnapshot { + out := make([]WindowSnapshot, len(rows)) + for i, row := range rows { + row.UpdatedAt = 0 + out[i] = row + } + return out +} + // runCounterStoreSuite asserts the counter half of the Store contract. Every // backend runs the same body, so a window that round-trips on SQLite has to // round-trip on PostgreSQL and MongoDB too. seedStale writes a row directly, @@ -15,6 +26,7 @@ import ( func runCounterStoreSuite(t *testing.T, store Store, seedStale func(t *testing.T, snap WindowSnapshot, updatedAt int64)) { t.Helper() ctx := context.Background() + saveStart := time.Now().Unix() live := []WindowSnapshot{ { Scope: string(ScopeUserPath), Subject: "/customers", Partition: "/customers/alice", @@ -31,7 +43,7 @@ func runCounterStoreSuite(t *testing.T, store Store, seedStale func(t *testing.T } loaded := func(t *testing.T, step string) map[string]WindowSnapshot { t.Helper() - got, err := store.LoadCounters(ctx) + got, err := store.LoadCounters(ctx, 0) require.NoError(t, err, "LoadCounters %s: %v", step, err) byPartition := make(map[string]WindowSnapshot, len(got)) @@ -40,20 +52,70 @@ func runCounterStoreSuite(t *testing.T, store Store, seedStale func(t *testing.T } return byPartition } - // Every field survives the round trip, per partition. - err := store.SaveCounters(ctx, live) + // Every field survives the round trip, per partition, and the store + // stamps the save whatever the snapshot carried. + stamped := append([]WindowSnapshot(nil), live...) + stamped[0].UpdatedAt = 1 + err := store.SaveCounters(ctx, stamped) require.NoError(t, err) got := loaded(t, "after save") for _, want := range live { - require.Equal(t, want, got[want.Partition]) + row := got[want.Partition] + require.GreaterOrEqual(t, row.UpdatedAt, saveStart, "the save stamps the row") + require.LessOrEqual(t, row.UpdatedAt, time.Now().Unix()) + require.Equal(t, want, unstamped([]WindowSnapshot{row})[0]) } + + // A load since a time returns only the rows saved since: an incremental + // load must see every save, and skip rows nobody saved since. + seedStale(t, WindowSnapshot{ + Scope: string(ScopeUserPath), Subject: "/customers", Partition: "/customers/carol", + PeriodSeconds: PeriodHourSeconds, Writer: "writer-a", + RequestsWindowStart: 1700000000, RequestsCurrent: 5, + }, saveStart-60) + recent, err := store.LoadCounters(ctx, saveStart) + require.NoError(t, err) + require.ElementsMatch(t, live, unstamped(recent)) + older, err := store.LoadCounters(ctx, saveStart-60) + require.NoError(t, err) + require.Len(t, older, 3) + none, err := store.LoadCounters(ctx, time.Now().Unix()+60) + require.NoError(t, err) + require.Empty(t, none) // Omitting a row does not delete it: it stays restorable until it goes // two periods without a write. err = store.SaveCounters(ctx, live[:1]) require.NoError(t, err) got = loaded(t, "after partial save") - require.Len(t, got, 2) + require.Len(t, got, 3) + require.NoError(t, store.ResetCounter(ctx, ScopeUserPath, "/customers", PeriodHourSeconds)) + require.NoError(t, store.SaveCounters(ctx, live)) + + // An in-flight row round-trips its gauge. It has no window, so the store + // keeps it a day after its last save, and drops it after that. + inFlight := WindowSnapshot{ + Scope: string(ScopeProvider), Subject: "openai", PeriodSeconds: PeriodConcurrent, + Writer: "writer-a", Epoch: "e1", InFlight: 3, + } + require.NoError(t, store.SaveCounters(ctx, []WindowSnapshot{inFlight})) + idle := inFlight + idle.Writer, idle.InFlight = "writer-idle", 1 + seedStale(t, idle, time.Now().Unix()-PeriodHourSeconds) + gone := inFlight + gone.Writer = "writer-gone" + seedStale(t, gone, time.Now().Unix()-inFlightRowRetention-60) + require.NoError(t, store.SaveCounters(ctx, live)) + rows, err := store.LoadCounters(ctx, 0) + require.NoError(t, err) + var gauges []WindowSnapshot + for _, row := range rows { + if row.PeriodSeconds == PeriodConcurrent { + gauges = append(gauges, row) + } + } + require.ElementsMatch(t, []WindowSnapshot{inFlight, idle}, unstamped(gauges)) + require.NoError(t, store.ResetCounter(ctx, ScopeProvider, "openai", PeriodConcurrent)) // Two periods without a write makes it collectable by the next save. seedStale(t, WindowSnapshot{ @@ -74,9 +136,9 @@ func runCounterStoreSuite(t *testing.T, store Store, seedStale func(t *testing.T } err = store.SaveCounters(ctx, peer) require.NoError(t, err) - all, err := store.LoadCounters(ctx) + all, err := store.LoadCounters(ctx, 0) require.NoError(t, err) - require.ElementsMatch(t, append(append([]WindowSnapshot(nil), live...), peer...), all) + require.ElementsMatch(t, append(append([]WindowSnapshot(nil), live...), peer...), unstamped(all)) // A reset clears every partition and every writer of the definition, and // gives the rule a new counter epoch, leaving other rules alone. err = store.UpsertRules(ctx, []Rule{ diff --git a/internal/ratelimit/store_epoch_test.go b/internal/ratelimit/store_epoch_test.go index 984d14f81..9be54b772 100644 --- a/internal/ratelimit/store_epoch_test.go +++ b/internal/ratelimit/store_epoch_test.go @@ -97,7 +97,7 @@ func runCounterEpochSuite(t *testing.T, store Store) { require.NoError(t, store.SaveCounters(ctx, []WindowSnapshot{snap})) snap.Epoch = "e2" require.NoError(t, store.SaveCounters(ctx, []WindowSnapshot{snap})) - rows, err := store.LoadCounters(ctx) + rows, err := store.LoadCounters(ctx, 0) require.NoError(t, err) require.Len(t, rows, 1) assert.Equal(t, "e2", rows[0].Epoch, "a save rewrites the row's epoch") diff --git a/internal/ratelimit/store_mongodb.go b/internal/ratelimit/store_mongodb.go index 1d3bdc8fa..1b99e8fb7 100644 --- a/internal/ratelimit/store_mongodb.go +++ b/internal/ratelimit/store_mongodb.go @@ -55,6 +55,13 @@ func NewMongoDBStore(ctx context.Context, database *mongo.Database) (*MongoDBSto if err != nil { return nil, fmt.Errorf("create rate limit counter indexes: %w", err) } + // Every replica's flush loads the rows saved since its previous load. + _, err = store.counters.Indexes().CreateOne(ctx, mongo.IndexModel{ + Keys: bson.D{{Key: "updated_at", Value: 1}}, + }) + if err != nil { + return nil, fmt.Errorf("create rate limit counter indexes: %w", err) + } if err := store.importLegacyCounters(ctx); err != nil { return nil, err } @@ -394,8 +401,12 @@ func (s *MongoDBStore) configRulesWithoutManualCollisions(ctx context.Context, r return filtered, nil } -func (s *MongoDBStore) LoadCounters(ctx context.Context) ([]WindowSnapshot, error) { - cursor, err := s.counters.Find(ctx, bson.D{}) +func (s *MongoDBStore) LoadCounters(ctx context.Context, since int64) ([]WindowSnapshot, error) { + filter := bson.D{} + if since > 0 { + filter = bson.D{{Key: "updated_at", Value: bson.D{{Key: "$gte", Value: since}}}} + } + cursor, err := s.counters.Find(ctx, filter) if err != nil { return nil, fmt.Errorf("list rate limit counters: %w", err) } @@ -416,35 +427,30 @@ func (s *MongoDBStore) LoadCounters(ctx context.Context) ([]WindowSnapshot, erro return snapshots, nil } -// counterDocument is a window row as stored: the snapshot plus the write -// stamp that drives staleness collection. The stamp is the store's own -// bookkeeping, so it stays out of WindowSnapshot. -type counterDocument struct { - WindowSnapshot `bson:",inline"` - UpdatedAt int64 `bson:"updated_at"` -} - func (s *MongoDBStore) SaveCounters(ctx context.Context, snapshots []WindowSnapshot) error { now := time.Now().Unix() if len(snapshots) > 0 { writes := make([]mongo.WriteModel, 0, len(snapshots)) for _, snap := range snapshots { + snap.UpdatedAt = now writes = append(writes, mongo.NewUpdateOneModel(). SetFilter(counterIdentityFilter(snap)). - SetUpdate(bson.D{{Key: "$set", Value: counterDocument{WindowSnapshot: snap, UpdatedAt: now}}}). + SetUpdate(bson.D{{Key: "$set", Value: snap}}). SetUpsert(true)) } if _, err := s.counters.BulkWrite(ctx, writes, options.BulkWrite().SetOrdered(false)); err != nil { return fmt.Errorf("upsert rate limit counters: %w", err) } } - // Collect the rows nobody writes any more, on the same staleness bound - // restore applies, so this only ever deletes rows a load would discard. + // Collect the rows nobody writes any more, on the bounds the SQL stores + // use (see pruneCountersSQL). + retention := bson.D{{Key: "$cond", Value: bson.A{ + bson.D{{Key: "$eq", Value: bson.A{"$period_seconds", PeriodConcurrent}}}, + inFlightRowRetention, + bson.D{{Key: "$multiply", Value: bson.A{2, "$period_seconds"}}}, + }}} stale := bson.D{{Key: "$expr", Value: bson.D{{Key: "$lt", Value: bson.A{ - bson.D{{Key: "$add", Value: bson.A{ - "$updated_at", - bson.D{{Key: "$multiply", Value: bson.A{2, "$period_seconds"}}}, - }}}, + bson.D{{Key: "$add", Value: bson.A{"$updated_at", retention}}}, now, }}}}} if _, err := s.counters.DeleteMany(ctx, stale); err != nil { diff --git a/internal/ratelimit/store_mongodb_migrate.go b/internal/ratelimit/store_mongodb_migrate.go index d2d675dde..413f730d5 100644 --- a/internal/ratelimit/store_mongodb_migrate.go +++ b/internal/ratelimit/store_mongodb_migrate.go @@ -24,7 +24,7 @@ import ( // the reset's delete of a document read here makes the transaction conflict // and retry from a fresh read. A standalone server has no transactions, so // there the counts can be written back, but harmlessly: imported documents -// carry the empty epoch, and the reset gave its rule a new one, which restore +// carry the empty epoch, and the reset gave its rule a new one, which a load // requires (see Rule.CounterEpoch). func (s *MongoDBStore) importLegacyCounters(ctx context.Context) error { return s.inTransaction(ctx, "import legacy rate limit counters", s.moveLegacyCounters) @@ -46,14 +46,14 @@ func (s *MongoDBStore) moveLegacyCounters(ctx context.Context) error { // A malformed document is cleared with the rest: no release can // restore it, and keeping it would warn on every start. imported = append(imported, cursor.Current.Lookup("_id")) - var doc counterDocument + var doc WindowSnapshot if err := cursor.Decode(&doc); err != nil { slog.Warn("rate limit counters: skipping malformed legacy snapshot", "error", err) continue } - doc.Writer, doc.Epoch = "", "" + doc.Writer, doc.Epoch, doc.InFlight = "", "", 0 writes = append(writes, mongo.NewUpdateOneModel(). - SetFilter(counterIdentityFilter(doc.WindowSnapshot)). + SetFilter(counterIdentityFilter(doc)). SetUpdate(bson.D{{Key: "$set", Value: doc}}). SetUpsert(true)) } diff --git a/internal/ratelimit/store_mongodb_test.go b/internal/ratelimit/store_mongodb_test.go index 989232605..767ffc0f6 100644 --- a/internal/ratelimit/store_mongodb_test.go +++ b/internal/ratelimit/store_mongodb_test.go @@ -270,8 +270,8 @@ func TestMongoDBStoreCounterRoundTrip(t *testing.T) { runCounterStoreSuite(t, store, func(t *testing.T, snap WindowSnapshot, updatedAt int64) { t.Helper() - doc := counterDocument{WindowSnapshot: snap, UpdatedAt: updatedAt} - _, err := db.Collection("rate_limit_window_snapshots").InsertOne(context.Background(), doc) + snap.UpdatedAt = updatedAt + _, err := db.Collection("rate_limit_window_snapshots").InsertOne(context.Background(), snap) require.NoError(t, err) }) }) @@ -301,10 +301,10 @@ func TestMongoDBStoreLoadCountersSkipsMalformedDocument(t *testing.T) { }) require.NoError(t, err) - got, err := store.LoadCounters(ctx) + got, err := store.LoadCounters(ctx, 0) require.NoError(t, err) require.Len(t, got, 1) - require.Equal(t, good, got[0]) + require.Equal(t, good, unstamped(got)[0]) }) } @@ -354,9 +354,9 @@ func TestMongoDBStoreImportsLegacyCounters(t *testing.T) { store, err := NewMongoDBStore(ctx, db) require.NoError(t, err) - got, err := store.LoadCounters(ctx) + got, err := store.LoadCounters(ctx, 0) require.NoError(t, err) - require.Equal(t, []WindowSnapshot{legacy}, got, "the legacy document loads under the empty writer") + require.Equal(t, []WindowSnapshot{legacy}, unstamped(got), "the legacy document loads under the empty writer") remaining, err := legacyColl.CountDocuments(ctx, bson.D{}) require.NoError(t, err) require.Zero(t, remaining, "imported documents leave the legacy collection") @@ -364,7 +364,7 @@ func TestMongoDBStoreImportsLegacyCounters(t *testing.T) { legacyWrite(t, 7) store, err = NewMongoDBStore(ctx, db) require.NoError(t, err) - got, err = store.LoadCounters(ctx) + got, err = store.LoadCounters(ctx, 0) require.NoError(t, err) require.Len(t, got, 1) require.Equal(t, int64(7), got[0].RequestsCurrent) @@ -374,7 +374,7 @@ func TestMongoDBStoreImportsLegacyCounters(t *testing.T) { require.NoError(t, err) store, err = NewMongoDBStore(ctx, db) require.NoError(t, err) - got, err = store.LoadCounters(ctx) + got, err = store.LoadCounters(ctx, 0) require.NoError(t, err) require.Empty(t, got) }) @@ -400,7 +400,7 @@ func TestMongoDBStoreLegacyCountsWrittenBackAfterResetAreNotRestored(t *testing. require.NoError(t, err) importLegacy := func(t *testing.T) { t.Helper() - _, err := store.legacyCounters.InsertOne(ctx, counterDocument{ + _, err := store.legacyCounters.InsertOne(ctx, WindowSnapshot{ Scope: string(ScopeUserPath), Subject: "/team", PeriodSeconds: PeriodHourSeconds, RequestsWindowStart: now.Unix() - now.Unix()%PeriodHourSeconds, RequestsCurrent: 90, UpdatedAt: now.Unix(), diff --git a/internal/ratelimit/store_sql.go b/internal/ratelimit/store_sql.go index 6c19845c9..f3d9ddecf 100644 --- a/internal/ratelimit/store_sql.go +++ b/internal/ratelimit/store_sql.go @@ -49,6 +49,7 @@ const sqlRateLimitWindowSnapshotsSchema = ` tokens_window_start ` + sqlx.TypeInt64 + ` NOT NULL DEFAULT 0, tokens_current ` + sqlx.TypeInt64 + ` NOT NULL DEFAULT 0, tokens_previous ` + sqlx.TypeInt64 + ` NOT NULL DEFAULT 0, + in_flight ` + sqlx.TypeInt64 + ` NOT NULL DEFAULT 0, updated_at ` + sqlx.TypeInt64 + ` NOT NULL, PRIMARY KEY (scope, subject, partition, period_seconds, writer) )` @@ -105,6 +106,15 @@ func NewSQLStore(ctx context.Context, db sqlx.DB) (*SQLStore, error) { if err := db.Schema(ctx, sqlRateLimitWindowSnapshotsSchema); err != nil { return nil, fmt.Errorf("failed to create rate_limit_window_snapshots table: %w", err) } + if err := sqlx.AddColumns(ctx, db, + `ALTER TABLE rate_limit_window_snapshots ADD COLUMN in_flight `+sqlx.TypeInt64+` NOT NULL DEFAULT 0`, + ); err != nil { + return nil, fmt.Errorf("failed to migrate rate_limit_window_snapshots table: %w", err) + } + // Every replica's flush loads the rows saved since its previous load. + if err := db.Schema(ctx, `CREATE INDEX IF NOT EXISTS idx_rate_limit_window_snapshots_updated_at ON rate_limit_window_snapshots(updated_at)`); err != nil { + return nil, fmt.Errorf("failed to create rate limit counter index: %w", err) + } legacy, err := importLegacyCounters(ctx, db) if err != nil { return nil, err @@ -197,13 +207,15 @@ func (s *SQLStore) ReplaceConfigRules(ctx context.Context, rules []Rule) error { }) } -func (s *SQLStore) LoadCounters(ctx context.Context) ([]WindowSnapshot, error) { +func (s *SQLStore) LoadCounters(ctx context.Context, since int64) ([]WindowSnapshot, error) { rows, err := s.db.Query(ctx, ` SELECT scope, subject, partition, period_seconds, writer, epoch, requests_window_start, requests_current, requests_previous, - tokens_window_start, tokens_current, tokens_previous + tokens_window_start, tokens_current, tokens_previous, + in_flight, updated_at FROM rate_limit_window_snapshots - `) + WHERE updated_at >= ? + `, since) if err != nil { return nil, fmt.Errorf("list rate limit counters: %w", err) } @@ -216,6 +228,7 @@ func (s *SQLStore) LoadCounters(ctx context.Context) ([]WindowSnapshot, error) { &snap.Scope, &snap.Subject, &snap.Partition, &snap.PeriodSeconds, &snap.Writer, &snap.Epoch, &snap.RequestsWindowStart, &snap.RequestsCurrent, &snap.RequestsPrevious, &snap.TokensWindowStart, &snap.TokensCurrent, &snap.TokensPrevious, + &snap.InFlight, &snap.UpdatedAt, ); err != nil { // One unreadable row must not cost every other window its // restore: Start treats a load error as "do not persist this @@ -235,8 +248,8 @@ const upsertCounterSQL = ` INSERT INTO rate_limit_window_snapshots ( scope, subject, partition, period_seconds, writer, epoch, requests_window_start, requests_current, requests_previous, - tokens_window_start, tokens_current, tokens_previous, updated_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + tokens_window_start, tokens_current, tokens_previous, in_flight, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(scope, subject, partition, period_seconds, writer) DO UPDATE SET epoch = excluded.epoch, requests_window_start = excluded.requests_window_start, @@ -245,6 +258,7 @@ const upsertCounterSQL = ` tokens_window_start = excluded.tokens_window_start, tokens_current = excluded.tokens_current, tokens_previous = excluded.tokens_previous, + in_flight = excluded.in_flight, updated_at = excluded.updated_at ` @@ -255,12 +269,12 @@ func (s *SQLStore) SaveCounters(ctx context.Context, snapshots []WindowSnapshot) if _, err := q.Exec(ctx, upsertCounterSQL, snap.Scope, snap.Subject, snap.Partition, snap.PeriodSeconds, snap.Writer, snap.Epoch, snap.RequestsWindowStart, snap.RequestsCurrent, snap.RequestsPrevious, - snap.TokensWindowStart, snap.TokensCurrent, snap.TokensPrevious, now, + snap.TokensWindowStart, snap.TokensCurrent, snap.TokensPrevious, snap.InFlight, now, ); err != nil { return fmt.Errorf("upsert rate limit counter: %w", err) } } - if _, err := q.Exec(ctx, pruneCountersSQL, now); err != nil { + if _, err := q.Exec(ctx, pruneCountersSQL, PeriodConcurrent, inFlightRowRetention, now); err != nil { return fmt.Errorf("prune expired rate limit counters: %w", err) } return nil @@ -268,13 +282,14 @@ func (s *SQLStore) SaveCounters(ctx context.Context, snapshots []WindowSnapshot) } // pruneCountersSQL collects the rows nobody writes any more: a rule that was -// deleted, a window that fell out of memory, or a writer that stopped — every -// restart leaves its predecessor's rows behind. Two periods without a write is -// the same staleness bound restore applies, so this only ever deletes rows a -// load would have discarded. +// deleted, a window that went unchanged, or a writer that stopped — every +// restart leaves its predecessor's rows behind. A window row is saved +// whenever it counts something, so two periods without a write means every +// count it holds has rolled out, the same bound a load drops it on. An +// in-flight row has no window: it is kept for inFlightRowRetention. const pruneCountersSQL = ` DELETE FROM rate_limit_window_snapshots - WHERE updated_at + 2 * period_seconds < ? + WHERE updated_at + CASE WHEN period_seconds = ? THEN ? ELSE 2 * period_seconds END < ? ` // ResetCounter draws a new counter epoch for the rule and removes every diff --git a/internal/ratelimit/store_sql_test.go b/internal/ratelimit/store_sql_test.go index 5d278ec50..47778f479 100644 --- a/internal/ratelimit/store_sql_test.go +++ b/internal/ratelimit/store_sql_test.go @@ -268,7 +268,7 @@ func TestSQLStoreCounterRoundTrip(t *testing.T) { _, err := store.db.Exec(context.Background(), upsertCounterSQL, snap.Scope, snap.Subject, snap.Partition, snap.PeriodSeconds, snap.Writer, snap.Epoch, snap.RequestsWindowStart, snap.RequestsCurrent, snap.RequestsPrevious, - snap.TokensWindowStart, snap.TokensCurrent, snap.TokensPrevious, updatedAt, + snap.TokensWindowStart, snap.TokensCurrent, snap.TokensPrevious, snap.InFlight, updatedAt, ) require.NoError(t, err) }) @@ -293,14 +293,14 @@ func TestSQLStoreLoadCountersSkipsMalformedRow(t *testing.T) { require.NoError(t, err) _, err = db.Exec(ctx, upsertCounterSQL, string(ScopeUserPath), "/broken", "", PeriodHourSeconds, "w1", "", - 0, "not-a-number", 0, 0, 0, 0, time.Now().Unix(), + 0, "not-a-number", 0, 0, 0, 0, 0, time.Now().Unix(), ) require.NoError(t, err) - got, err := store.LoadCounters(ctx) + got, err := store.LoadCounters(ctx, 0) require.NoError(t, err) require.Len(t, got, 1) - require.Equal(t, good, got[0]) + require.Equal(t, good, unstamped(got)[0]) } // legacyCountersSchema is rate_limit_counters as releases before writer IDs @@ -353,9 +353,9 @@ func TestSQLStoreImportsLegacyCounters(t *testing.T) { store, err := NewSQLStore(ctx, db) require.NoError(t, err) - got, err := store.LoadCounters(ctx) + got, err := store.LoadCounters(ctx, 0) require.NoError(t, err) - require.Equal(t, []WindowSnapshot{legacy}, got, "the legacy row loads under the empty writer") + require.Equal(t, []WindowSnapshot{legacy}, unstamped(got), "the legacy row loads under the empty writer") var remaining int64 err = db.QueryRow(ctx, `SELECT COUNT(*) FROM rate_limit_counters`).Scan(&remaining) require.NoError(t, err) @@ -364,9 +364,9 @@ func TestSQLStoreImportsLegacyCounters(t *testing.T) { // Re-opening imports nothing new and keeps the imported row. _, err = NewSQLStore(ctx, db) require.NoError(t, err) - got, err = store.LoadCounters(ctx) + got, err = store.LoadCounters(ctx, 0) require.NoError(t, err) - require.Equal(t, []WindowSnapshot{legacy}, got) + require.Equal(t, []WindowSnapshot{legacy}, unstamped(got)) // Rolling upgrade: an older replica still writes the legacy shape. The // next start imports its newer count over the earlier import. @@ -378,7 +378,7 @@ func TestSQLStoreImportsLegacyCounters(t *testing.T) { require.NoError(t, err) store, err = NewSQLStore(ctx, db) require.NoError(t, err) - got, err = store.LoadCounters(ctx) + got, err = store.LoadCounters(ctx, 0) require.NoError(t, err) require.Len(t, got, 1) require.Equal(t, int64(7), got[0].RequestsCurrent) @@ -394,7 +394,7 @@ func TestSQLStoreImportsLegacyCounters(t *testing.T) { require.NoError(t, err) store, err = NewSQLStore(ctx, db) require.NoError(t, err) - got, err = store.LoadCounters(ctx) + got, err = store.LoadCounters(ctx, 0) require.NoError(t, err) require.Empty(t, got) }) @@ -484,7 +484,7 @@ func TestSQLStoreLegacyImportCannotUndoConcurrentReset(t *testing.T) { require.NoError(t, <-resetDone) require.NoError(t, <-importDone) - got, err := store.LoadCounters(ctx) + got, err := store.LoadCounters(ctx, 0) require.NoError(t, err) require.Empty(t, got, "the import must not bring back the count the reset deleted") }) diff --git a/internal/ratelimit/types.go b/internal/ratelimit/types.go index e001b63b8..9f75a85aa 100644 --- a/internal/ratelimit/types.go +++ b/internal/ratelimit/types.go @@ -1,8 +1,10 @@ // Package ratelimit enforces request, token, and concurrency limits for the // AI gateway. Rules are scoped to a consumer user-path subtree, a provider, or -// a model. Rule definitions are persisted. Request and token windows are -// snapshotted to the same store so they survive restart and reload; live -// admission stays in memory and per instance. +// a model. Rule definitions are persisted. Every instance saves what it +// counted to the same store, so counts survive restart and reload, and reads +// what the others counted, so instances sharing a store enforce one limit +// together. Admission itself stays in memory: it checks this instance's live +// counters plus the others' as of the last load, never the store. package ratelimit import ( @@ -68,7 +70,7 @@ type Rule struct { // CounterEpoch names the counters this rule definition owns. The store // draws a new one when the rule is created, switches between shared and // per-child, or has its counters reset, and keeps it on any other edit. Saved window rows carry - // the epoch they were counted under, and restore skips rows of any other + // the epoch they were counted under, and loads skip rows of any other // epoch, so counts a change cleared never come back, whether or not // deleting their rows succeeded. Rules stored before it existed have the // empty epoch, as do the rows saved for them. diff --git a/internal/server/ratelimit_support_test.go b/internal/server/ratelimit_support_test.go index 087745f44..b7399ad5f 100644 --- a/internal/server/ratelimit_support_test.go +++ b/internal/server/ratelimit_support_test.go @@ -34,7 +34,7 @@ func (s *staticRuleStore) DeleteRule(context.Context, ratelimit.RuleScope, strin func (s *staticRuleStore) ReplaceConfigRules(context.Context, []ratelimit.Rule) error { return nil } -func (s *staticRuleStore) LoadCounters(context.Context) ([]ratelimit.WindowSnapshot, error) { +func (s *staticRuleStore) LoadCounters(context.Context, int64) ([]ratelimit.WindowSnapshot, error) { return nil, nil } func (s *staticRuleStore) SaveCounters(context.Context, []ratelimit.WindowSnapshot) error { diff --git a/tests/perf/production_shape_test.go b/tests/perf/production_shape_test.go index a4d0a4f32..97590939c 100644 --- a/tests/perf/production_shape_test.go +++ b/tests/perf/production_shape_test.go @@ -39,7 +39,7 @@ func (s *benchRateLimitStore) DeleteRule(context.Context, ratelimit.RuleScope, s func (s *benchRateLimitStore) ReplaceConfigRules(context.Context, []ratelimit.Rule) error { return nil } -func (s *benchRateLimitStore) LoadCounters(context.Context) ([]ratelimit.WindowSnapshot, error) { +func (s *benchRateLimitStore) LoadCounters(context.Context, int64) ([]ratelimit.WindowSnapshot, error) { return nil, nil } func (s *benchRateLimitStore) SaveCounters(context.Context, []ratelimit.WindowSnapshot) error { @@ -49,7 +49,7 @@ func (s *benchRateLimitStore) ResetCounter(context.Context, ratelimit.RuleScope, return nil } func (s *benchRateLimitStore) ResetAllCounters(context.Context) error { return nil } -func (s *benchRateLimitStore) Close() error { return nil } +func (s *benchRateLimitStore) Close() error { return nil } // newBenchRateLimiter builds a real ratelimit.Service with one high user-path // request limit that matches every request but never rejects. From 802a0512e164111f477f1a64c002f3aea8f98320 Mon Sep 17 00:00:00 2001 From: "Jakub A. W." Date: Fri, 30 Oct 2026 15:20:34 +0100 Subject: [PATCH 2/2] docs(ratelimit): describe cluster-wide limits and their overshoot bound --- config/config.example.yaml | 2 +- config/ratelimit.go | 8 ++-- docs/features/rate-limits.mdx | 76 ++++++++++++++++++++++++----------- docs/guides/production.mdx | 21 ++++++---- 4 files changed, 71 insertions(+), 36 deletions(-) diff --git a/config/config.example.yaml b/config/config.example.yaml index 0d00646ae..df6223505 100644 --- a/config/config.example.yaml +++ b/config/config.example.yaml @@ -338,7 +338,7 @@ budgets: rate_limits: enabled: true # env: RATE_LIMITS_ENABLED; with no configured rules this has no effect - flush_interval: 1 # seconds; env: RATE_LIMITS_FLUSH_INTERVAL; 0 = no periodic snapshot (still load + shutdown flush) + flush_interval: 1 # seconds; env: RATE_LIMITS_FLUSH_INTERVAL; how often replicas share counts; 0 = no periodic sharing (still load + shutdown flush) user_paths: # Env equivalent: # SET_RATE_LIMIT_USER__PATH__EXAMPLE="rpm=100,tpm=50000,rpd=10000,concurrent=10" diff --git a/config/ratelimit.go b/config/ratelimit.go index 363b49a1a..95328a546 100644 --- a/config/ratelimit.go +++ b/config/ratelimit.go @@ -15,9 +15,11 @@ type RateLimitsConfig struct { // Default: true. With no rules configured the check is a no-op. Enabled bool `yaml:"enabled" env:"RATE_LIMITS_ENABLED"` - // FlushInterval is how often live request/token windows are written to - // storage, in seconds. Default 1. 0 disables the periodic loop; Start - // still loads and Close of an active generation still writes once. + // FlushInterval is how often each instance writes its changed counters to + // storage and reads the other instances', in seconds. Default 1. 0 + // disables the periodic loop; Start still loads and Close of an active + // generation still writes once, but instances stop seeing each other's + // traffic after start. FlushInterval int `yaml:"flush_interval" env:"RATE_LIMITS_FLUSH_INTERVAL"` // UserPaths declares rate limit rules by tracked user path. diff --git a/docs/features/rate-limits.mdx b/docs/features/rate-limits.mdx index 76fdeebd4..35fdf3139 100644 --- a/docs/features/rate-limits.mdx +++ b/docs/features/rate-limits.mdx @@ -32,7 +32,8 @@ embeddings, audio, passthrough, realtime sessions, and batch submission. Use rate limits for traffic control and budgets for spend control: budgets cap dollars over hours-to-months from durable usage records, rate limits cap -request and token velocity over seconds-to-a-day with fast in-memory counters. +request and token velocity over seconds-to-a-day with fast in-memory counters +that replicas sharing a database keep in step. ## Enable rate limits @@ -229,30 +230,57 @@ Details worth knowing: - A realtime session counts as one request and holds one concurrency slot for the whole session. A batch submission counts as one request and holds no concurrency slot. -- Request and token windows live in memory for admission and are snapshotted - to the configured store about once a second (and on shutdown), so they - survive a restart and `gomodel --reload`. Concurrency gauges are not - persisted: after a bounce nothing is in flight. Counters are per replica: - each instance writes its own snapshot rows, so replicas sharing a database - never overwrite each other. On start an instance rolls every stored row - forward to the current window and takes, for requests and tokens - separately, the largest current and the largest previous bucket. Those can - come from different replicas, so the restored count can be a little above - any single replica's, but it is never the sum. A reset applies to every - replica: the replica that serves it clears its counters at once, and the - others within a few seconds, when they +- Request and token windows live in memory for admission and are saved to + the configured store about once a second (and on shutdown), so they + survive a restart and `gomodel --reload`. A restart counts what the + previous run counted exactly once. Tune the interval with + `RATE_LIMITS_FLUSH_INTERVAL` (seconds; default `1`). `0` skips the + periodic write: counts still load on start and save on a clean shutdown, + but replicas no longer see each other's traffic after start. +- A reset applies to every replica: the replica that serves it clears its + counters at once, and the others within a few seconds, when they [reload the rules](/guides/production#config-changes-reach-every-replica-within-seconds). - Switching a rule between shared and per-child, or deleting and re-creating - it, starts its counts from zero the same way. Stored rows from before a - reset or such a change are never restored, even after a restart or after - switching back. Editing a limit keeps the counts. Finish a rolling upgrade - from a release before per-replica snapshots before resetting or deleting - rules: a reset or delete served by an older replica can come back after a - restart. With N replicas the - effective limit is still about N times the configured value — use budgets - for durable, cross-instance control. Tune the snapshot interval with - `RATE_LIMITS_FLUSH_INTERVAL` (seconds; default `1`; `0` skips the periodic - write but still loads on start and flushes on a clean shutdown). + Switching a rule between shared and per-child, or deleting and + re-creating it, starts its counts from zero the same way. Stored counts + from before a reset or such a change never come back, even after a + restart or after switching back. Editing a limit keeps the counts. Finish + a rolling upgrade from a release before per-replica snapshots before + resetting or deleting rules: a reset or delete served by an older replica can come back after a + restart. + +## Several replicas share one limit + +Replicas that share a database enforce each rule together: a `100 rpm` rule +admits about 100 requests per minute across all of them, not 100 per +replica. Each replica saves what it counted itself, and on every flush reads +what the others saved since its previous read. Admission checks the +replica's own live count plus the others' last saved counts, entirely in +memory: no request waits on the database. Remaining-request headers, the +admin and dashboard usage, `GET /v1/usage`, and routing around a saturated +provider or model all report the cluster-wide count. + +The limit is approximate, not exact: + +- Another replica's requests reach this one within about two flush + intervals: one for that replica to save them, one for this replica to + read them. Until then each replica admits as if those requests had not + happened, so the cluster can overshoot by up to about two seconds of the + other replicas' traffic at the default interval. Requests that several + replicas admit at the same instant can all take the last slot. +- Concurrency caps are shared the same way: each replica publishes its + in-flight count every flush. A replica that crashes holds its slots until + its count goes stale, after three flush intervals and at least 5 seconds, + and then they free on every other replica. Freshness compares the time a + replica saved its count with the reader's clock, so keep replica clocks + synchronized (NTP). +- A cap that must never be exceeded, such as `concurrent=1` across + replicas, needs a per-request check against shared state, which GoModel + does not offer yet. Use sticky routing or a single replica for it. + +Each flush writes only the counters that changed since the last one, and +reads only the rows written since its previous read, with a full read once +a minute. [Budgets](/features/budgets) remain the exact, durable control for +spend. ## Client behavior diff --git a/docs/guides/production.mdx b/docs/guides/production.mdx index cac994c9d..48ac76f24 100644 --- a/docs/guides/production.mdx +++ b/docs/guides/production.mdx @@ -96,8 +96,8 @@ Redis when several replicas should share discovery results and cache hits. Take that into account when ordering service startup. -Redis does not coordinate rate limits or budgets. It is a cache, not a -distributed lock. +Redis does not coordinate rate limits or budgets: replicas share both through +the database. Redis is a cache, not a distributed lock. ## Provider credentials @@ -405,17 +405,22 @@ coordinate across instances. | [Usage and cost records](/features/cost-tracking) | Durable (database) | Aggregated correctly. | | Audit logs | Durable (database) | Complete across replicas. | | [SSO](/pro/sso) sign-out (Pro) | Durable (database) | A signed-out session is rejected on every replica within about 2 seconds. | -| [Rate limits](/features/rate-limits) | In-memory, per instance | Effective limit is about `N x` the configured limit. Counters survive a restart, and a reset clears every replica. | +| [Rate limits](/features/rate-limits) | In memory, shared through the database | Replicas enforce one limit together, approximately: each sees the others' counts within about two flush intervals. Counters survive a restart, and a reset clears every replica. | | Circuit breakers | In-memory, per process | Each replica learns an outage separately, so upstream sees up to `N x failure_threshold` failures before any breaker opens. | | [Session affinity](/features/session-keeping) pins | In-memory, per instance | Sessions stick per replica, not globally, unless the load balancer is sticky. | | Live log buffer | In-memory, per instance | A connected dashboard sees only that replica's live traffic, roughly `1/N` of it. History still comes from the database. | | MCP downstream sessions | In-memory, per instance | An `Mcp-Session-Id` from one replica returns 404 on another. [MCP](/features/mcp-gateway) clients need sticky sessions. | | Provider request health | In-memory, per instance | Dashboard provider status reflects the replica that answered. | -The one that surprises people is rate limits: they are an in-process control, so -three replicas with a 100 rpm rule admit up to 300 rpm in aggregate. Size rules -per replica, and use [budgets](/features/budgets) as the durable cross-instance -control since they read from the shared database. +The one that surprises people is rate limits: admission is in memory, so a +replica learns the others' traffic only when it reads their saved counts, about +once a second. Three replicas with a 100 rpm rule admit about 100 rpm in +aggregate, but can overshoot by about two seconds of the other replicas' +traffic, and a crashed replica's concurrency slots stay taken for a few +seconds. Do not rely on rate limits for a cap that must never be exceeded, such +as one concurrent request across replicas, and use +[budgets](/features/budgets) as the exact cross-instance control for spend. See +[several replicas share one limit](/features/rate-limits#several-replicas-share-one-limit). ### Config changes reach every replica within seconds @@ -568,7 +573,7 @@ if the file size matters. with SQLite, mount a volume and set `SQLITE_PATH`. - Declare providers in environment variables or `config.yaml`, with values from a secret manager. -- Size rate limit rules per replica when running several replicas. +- Keep replica clocks synchronized (NTP); shared concurrency limits rely on it. - Decide on `LOGGING_LOG_BODIES` and set retention windows deliberately. - Point readiness at `/health/ready` and set `terminationGracePeriodSeconds` above 30.