Skip to content

Commit dee6aae

Browse files
authored
fix(ratelimit): bound InProcessLimiter.keys with LRU cap (#344)
Replace the unbounded map[string]*rate.Limiter with an LRU cache backed by container/list (stdlib). Evicts the least-recently-used bucket when the cap is reached, preventing unbounded memory growth from clients cycling subject values in metadata-auth mode. - Default cap: 100_000 (configurable via WithMaxBuckets option) - Eviction preserves atomicity: rl is captured under the lock before release - Uses only stdlib (container/list) — no new dependency - Race-clean: go test -race -count=5 ./internal/ratelimit/ passes Closes #287
1 parent 106fcc2 commit dee6aae

3 files changed

Lines changed: 84 additions & 15 deletions

File tree

.agents/context/security-review.md

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,9 @@ Implemented in `internal/ratelimit/`. Per-tenant + per-method in-process
129129
token-bucket limiter via `golang.org/x/time/rate`. Health check exempt.
130130
Returns `codes.ResourceExhausted` + `RetryInfo` detail. `Limiter` interface
131131
allows future Redis-backed replacement. `OTEL_METRICS_RATE_LIMIT` counter
132-
for observability.
132+
for observability. Bucket map is LRU-capped (`defaultMaxBuckets=100_000`,
133+
configurable via `WithMaxBuckets`) — prevents unbounded growth from cycling
134+
subjects (#287).
133135

134136
### 6. Schema-complexity bounds missing — High
135137

internal/ratelimit/limiter.go

Lines changed: 50 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,42 +1,78 @@
11
package ratelimit
22

33
import (
4+
"container/list"
45
"sync"
56

67
"golang.org/x/time/rate"
78
)
89

10+
const defaultMaxBuckets = 100_000
11+
912
// Limiter decides whether a request for a given key should proceed.
1013
// Implementations may be in-process or backed by an external store (e.g. Redis).
1114
type Limiter interface {
1215
Allow(key string) bool
1316
}
1417

1518
// InProcessLimiter is a per-key token-bucket limiter backed by golang.org/x/time/rate.
16-
// Each unique key gets its own bucket; buckets are allocated on first use.
19+
// Each unique key gets its own bucket, allocated on first use. Buckets are evicted
20+
// via LRU when the cache reaches capacity, preventing unbounded memory growth.
1721
type InProcessLimiter struct {
18-
limit rate.Limit
19-
burst int
20-
mu sync.Mutex
21-
keys map[string]*rate.Limiter
22+
limit rate.Limit
23+
burst int
24+
maxBuckets int
25+
mu sync.Mutex
26+
ll *list.List
27+
keys map[string]*list.Element
28+
}
29+
30+
type lruEntry struct {
31+
key string
32+
rl *rate.Limiter
33+
}
34+
35+
// InProcessOption configures an InProcessLimiter.
36+
type InProcessOption func(*InProcessLimiter)
37+
38+
// WithMaxBuckets sets the maximum number of key buckets held in memory.
39+
// When the cap is reached the least-recently-used bucket is evicted; the next
40+
// request for that key starts a fresh full bucket. Defaults to 100_000.
41+
func WithMaxBuckets(n int) InProcessOption {
42+
return func(l *InProcessLimiter) { l.maxBuckets = n }
2243
}
2344

2445
// NewInProcess returns an InProcessLimiter with the given rate (events/sec) and burst size.
25-
func NewInProcess(limit rate.Limit, burst int) *InProcessLimiter {
26-
return &InProcessLimiter{
27-
limit: limit,
28-
burst: burst,
29-
keys: make(map[string]*rate.Limiter),
46+
func NewInProcess(limit rate.Limit, burst int, opts ...InProcessOption) *InProcessLimiter {
47+
l := &InProcessLimiter{
48+
limit: limit,
49+
burst: burst,
50+
maxBuckets: defaultMaxBuckets,
51+
ll: list.New(),
52+
keys: make(map[string]*list.Element),
3053
}
54+
for _, o := range opts {
55+
o(l)
56+
}
57+
return l
3158
}
3259

3360
// Allow reports whether a request for key should be allowed under the rate limit.
3461
func (l *InProcessLimiter) Allow(key string) bool {
3562
l.mu.Lock()
36-
rl, ok := l.keys[key]
37-
if !ok {
38-
rl = rate.NewLimiter(l.limit, l.burst)
39-
l.keys[key] = rl
63+
if el, ok := l.keys[key]; ok {
64+
l.ll.MoveToFront(el)
65+
rl := el.Value.(*lruEntry).rl
66+
l.mu.Unlock()
67+
return rl.Allow()
68+
}
69+
rl := rate.NewLimiter(l.limit, l.burst)
70+
el := l.ll.PushFront(&lruEntry{key: key, rl: rl})
71+
l.keys[key] = el
72+
if l.ll.Len() > l.maxBuckets {
73+
back := l.ll.Back()
74+
l.ll.Remove(back)
75+
delete(l.keys, back.Value.(*lruEntry).key)
4076
}
4177
l.mu.Unlock()
4278
return rl.Allow()

internal/ratelimit/limiter_test.go

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
package ratelimit_test
2+
3+
import (
4+
"testing"
5+
6+
"github.com/stretchr/testify/assert"
7+
"golang.org/x/time/rate"
8+
9+
"github.com/opendecree/decree/internal/ratelimit"
10+
)
11+
12+
// TestInProcessLimiter_LRUEviction verifies that inserting cap+1 unique keys evicts
13+
// the least-recently-used bucket, and the evicted key gets a fresh full bucket on reuse.
14+
func TestInProcessLimiter_LRUEviction(t *testing.T) {
15+
const cap = 3
16+
// rate=0: no token replenishment; burst=1: one token per fresh bucket.
17+
lim := ratelimit.NewInProcess(rate.Limit(0), 1, ratelimit.WithMaxBuckets(cap))
18+
19+
// Fill to capacity; each call consumes the sole token.
20+
// After three inserts, LRU order (back→front): key-0, key-1, key-2.
21+
lim.Allow("key-0")
22+
lim.Allow("key-1")
23+
lim.Allow("key-2")
24+
25+
// Inserting key-3 evicts key-0 (LRU).
26+
lim.Allow("key-3")
27+
28+
// key-0 was evicted; requesting it yields a fresh bucket with 1 token.
29+
assert.True(t, lim.Allow("key-0"), "evicted key should get a fresh bucket")
30+
assert.False(t, lim.Allow("key-0"), "fresh bucket has only one token")
31+
}

0 commit comments

Comments
 (0)