ratelimit: key limiter state by the rule's prefix, not its index (#358)
Buckets and held slots were keyed by the rule's index in the sorted
Effective.Paths. Inserting /checkout/ ahead of an existing /api/ rule
handed /checkout/ the tokens and held slots /api/ had, and gave /api/ a
fresh entry: checkout could be refused for API's use and API's allowance
was briefly replenished.
The key now names the rule by its prefix ("" for the base rule, which
Validate guarantees no path rule uses), and RuleFor returns the prefix
instead of the index.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
4 files changed,
+78
−25
internal/ratelimit/limiter.go+6 −4internal/ratelimit/limiter_test.go+54 −1internal/ratelimit/settings.go+7 −6internal/ratelimit/settings_test.go+11 −14
diff --git a/internal/ratelimit/limiter.go b/internal/ratelimit/limiter.go| index dacca5d42..d3e5ab9fc 100644 |
| --- a/internal/ratelimit/limiter.go |
| +++ b/internal/ratelimit/limiter.go |
| @@ -37,10 +37,12 @@ const nano = int64(time.Second) |
| // key is one client's allowance under one rule of one host. A path rule's |
| // requests count against its own key, so heavy API use leaves the client's |
| -// base allowance alone. |
| +// base allowance alone. The rule is named by its prefix, not its place in |
| +// Effective.Paths, so adding or removing another rule never moves live |
| +// tokens or held slots between prefixes. |
| type key struct { |
| host string |
| - rule int // index into Effective.Paths, -1 for the base rule |
| + rule string // the path rule's prefix, "" for the base rule |
| client netip.Prefix |
| } |
| @@ -134,8 +136,8 @@ func (t *Ticket) Release() { |
| // second otherwise. A rule field of zero or less turns that check off, so a |
| // zero Effective refuses nothing. |
| func (l *Limiter) Admit(host string, eff Effective, client netip.Prefix, path string, holdSlot, observe bool) (t *Ticket, d Decision, retryAfter time.Duration) { |
| - idx, rule := eff.RuleFor(path) |
| - k := key{host: host, rule: idx, client: client} |
| + prefix, rule := eff.RuleFor(path) |
| + k := key{host: host, rule: prefix, client: client} |
| now := int64(l.now().Sub(l.epoch)) |
| l.mu.Lock() |
diff --git a/internal/ratelimit/limiter_test.go b/internal/ratelimit/limiter_test.go| index 8767d9168..e9ac924d2 100644 |
| --- a/internal/ratelimit/limiter_test.go |
| +++ b/internal/ratelimit/limiter_test.go |
| @@ -53,7 +53,7 @@ func (l *Limiter) tracked() int { |
| func (l *Limiter) heldBy(client netip.Prefix) int { |
| l.mu.Lock() |
| defer l.mu.Unlock() |
| - e := l.entries[key{host: host, rule: -1, client: client}] |
| + e := l.entries[key{host: host, rule: "", client: client}] |
| if e == nil { |
| return -1 |
| } |
| @@ -192,6 +192,59 @@ func TestPathRuleSeparateAllowance(t *testing.T) { |
| } |
| } |
| +// A path rule's state belongs to its prefix, not to where the rule sorts: |
| +// inserting /checkout/ ahead of /api/ while /api/ requests are outstanding |
| +// leaves /api/'s held slots and spent tokens with /api/, gives /checkout/ a |
| +// fresh allowance, and every release still balances. |
| +func TestPathRuleStateFollowsPrefix(t *testing.T) { |
| + clk := newClock() |
| + l := NewLimiter(clk.now, 100) |
| + a := v4("203.0.113.4") |
| + apiOnly := Resolve(Settings{BusyAt: 1, Paths: []PathRule{ |
| + {Prefix: "/api/", Rule: Rule{ClientShare: 1, Rate: 1, Burst: 1}}, |
| + }}, Ceiling{}) |
| + withCheckout := Resolve(Settings{BusyAt: 1, Paths: []PathRule{ |
| + {Prefix: "/api/", Rule: Rule{ClientShare: 1, Rate: 1, Burst: 1}}, |
| + {Prefix: "/checkout/", Rule: Rule{ClientShare: 1, Rate: 1, Burst: 1}}, |
| + }}, Ceiling{}) |
| + if withCheckout.Paths[0].Prefix != "/checkout/" { |
| + t.Fatalf("paths = %+v, want /checkout/ sorted ahead of /api/", withCheckout.Paths) |
| + } |
| + |
| + // /api/ holds a slot and has spent its one token. |
| + api, d, _ := admit(t, l, apiOnly, a, "/api/x") |
| + if d != Allow || api == nil { |
| + t.Fatalf("/api/ first: %v ticket %v, want Allow", d, api) |
| + } |
| + |
| + // /checkout/ starts fresh: no slot held, a full bucket. |
| + checkout, d, _ := admit(t, l, withCheckout, a, "/checkout/pay") |
| + if d != Allow || checkout == nil { |
| + t.Fatalf("/checkout/ after the rule was inserted: %v, want Allow: it must not inherit /api/'s slot or tokens", d) |
| + } |
| + // /api/ keeps its held slot (shed while busy) and its empty bucket. |
| + if _, d, _ := admit(t, l, withCheckout, a, "/api/y"); d != OverCeiling { |
| + t.Fatalf("/api/ after /checkout/ was inserted: %v, want OverCeiling: its spent token must stay spent", d) |
| + } |
| + clk.add(time.Second) |
| + if _, d, _ := admit(t, l, withCheckout, a, "/api/y"); d != Shed { |
| + t.Fatalf("/api/ with a refilled token while it holds its share: %v, want Shed: its held slot must stay with it", d) |
| + } |
| + |
| + api.Release() |
| + checkout.Release() |
| + if got := l.InFlight(host); got != 0 { |
| + t.Fatalf("InFlight = %d after both released, want 0", got) |
| + } |
| + l.mu.Lock() |
| + for k, e := range l.entries { |
| + if e.holders != 0 { |
| + t.Errorf("entry %+v holds %d after every release, want 0", k, e.holders) |
| + } |
| + } |
| + l.mu.Unlock() |
| +} |
| + |
| func TestHoldSlotFalseTakesNoSlot(t *testing.T) { |
| clk := newClock() |
| l := NewLimiter(clk.now, 100) |
diff --git a/internal/ratelimit/settings.go b/internal/ratelimit/settings.go| index acd970f07..4ae20f3b3 100644 |
| --- a/internal/ratelimit/settings.go |
| +++ b/internal/ratelimit/settings.go |
| @@ -196,13 +196,14 @@ func capAt(v, ceiling int) int { |
| } |
| // RuleFor returns the rule for a request path: the path rule with the longest |
| -// matching prefix, and its index in e.Paths, or the base rule and -1. Matching |
| -// is a case-sensitive prefix test on the path as given. |
| -func (e Effective) RuleFor(path string) (idx int, r Rule) { |
| - for i, p := range e.Paths { |
| +// matching prefix, and that prefix, or the base rule and "". Validate refuses |
| +// an empty prefix, so "" names only the base rule. Matching is a |
| +// case-sensitive prefix test on the path as given. |
| +func (e Effective) RuleFor(path string) (prefix string, r Rule) { |
| + for _, p := range e.Paths { |
| if strings.HasPrefix(path, p.Prefix) { |
| - return i, p.Rule |
| + return p.Prefix, p.Rule |
| } |
| } |
| - return -1, e.Base |
| + return "", e.Base |
| } |
diff --git a/internal/ratelimit/settings_test.go b/internal/ratelimit/settings_test.go| index 97dd467d6..02305aef5 100644 |
| --- a/internal/ratelimit/settings_test.go |
| +++ b/internal/ratelimit/settings_test.go |
| @@ -154,25 +154,22 @@ func TestRuleForLongestPrefix(t *testing.T) { |
| }} |
| e := Resolve(s, Ceiling{}) |
| for _, c := range []struct { |
| - path string |
| - rate int |
| - neg bool // expect Base |
| + path string |
| + rate int |
| + prefix string // "" expects Base |
| }{ |
| - {"/api/v2/x", 20, false}, |
| - {"/api/x", 10, false}, |
| - {"/api", 300, true}, |
| - {"/API/x", 300, true}, |
| - {"/", 300, true}, |
| + {"/api/v2/x", 20, "/api/v2/"}, |
| + {"/api/x", 10, "/api/"}, |
| + {"/api", 300, ""}, |
| + {"/API/x", 300, ""}, |
| + {"/", 300, ""}, |
| } { |
| - idx, r := e.RuleFor(c.path) |
| + prefix, r := e.RuleFor(c.path) |
| if r.Rate != c.rate { |
| t.Errorf("%s: rate %d, want %d", c.path, r.Rate, c.rate) |
| } |
| - if (idx == -1) != c.neg { |
| - t.Errorf("%s: idx %d, base expected %v", c.path, idx, c.neg) |
| - } |
| - if idx >= 0 && e.Paths[idx].Rule != r { |
| - t.Errorf("%s: idx %d does not index the returned rule", c.path, idx) |
| + if prefix != c.prefix { |
| + t.Errorf("%s: prefix %q, want %q", c.path, prefix, c.prefix) |
| } |
| } |
| } |