chore: import clean Agent Desk snapshot (2/7)
This commit is contained in:
101
internal/pkg/ratelimit/ratelimit.go
Normal file
101
internal/pkg/ratelimit/ratelimit.go
Normal file
@@ -0,0 +1,101 @@
|
||||
// Package ratelimit provides a small in-process fixed-window counter for the
|
||||
// public endpoints that an unauthenticated caller can drive.
|
||||
//
|
||||
// It is deliberately not distributed. A deployment running several replicas gets
|
||||
// the configured limit per replica rather than per cluster, which weakens the
|
||||
// bound by a factor of the replica count. That is the right trade here because
|
||||
// these limits exist to make flooding expensive, not to meter a quota, and
|
||||
// because reaching for a shared store would put a network round trip on the
|
||||
// hottest unauthenticated paths in the application. If this ever needs to be
|
||||
// exact across replicas, the Limiter interface is small enough to swap.
|
||||
package ratelimit
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type bucket struct {
|
||||
count int
|
||||
resetAt time.Time
|
||||
}
|
||||
|
||||
// Limiter counts requests per key inside a fixed window. The zero value is not
|
||||
// usable; call New. A nil *Limiter allows everything, so a caller can represent
|
||||
// "disabled" by holding nil instead of branching at every use site.
|
||||
type Limiter struct {
|
||||
limit int
|
||||
window time.Duration
|
||||
|
||||
mu sync.Mutex
|
||||
buckets map[string]bucket
|
||||
sweptAt time.Time
|
||||
}
|
||||
|
||||
// New returns a Limiter that allows limit requests per window for each distinct
|
||||
// key. A limit of zero or less, or a non-positive window, disables the limiter.
|
||||
func New(limit int, window time.Duration) *Limiter {
|
||||
if limit <= 0 || window <= 0 {
|
||||
return &Limiter{limit: limit, window: window, buckets: make(map[string]bucket)}
|
||||
}
|
||||
return &Limiter{limit: limit, window: window, buckets: make(map[string]bucket)}
|
||||
}
|
||||
|
||||
// Allow records one request for key and reports whether it may proceed. When it
|
||||
// returns false, retryAfter is how long the caller should wait before the window
|
||||
// resets, and is meant for a Retry-After header.
|
||||
func (l *Limiter) Allow(key string) (bool, time.Duration) {
|
||||
if l == nil || l.limit <= 0 || l.window <= 0 {
|
||||
return true, 0
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
l.sweep(now)
|
||||
|
||||
current, ok := l.buckets[key]
|
||||
if !ok || !now.Before(current.resetAt) {
|
||||
l.buckets[key] = bucket{count: 1, resetAt: now.Add(l.window)}
|
||||
return true, 0
|
||||
}
|
||||
|
||||
current.count++
|
||||
l.buckets[key] = current
|
||||
if current.count > l.limit {
|
||||
retryAfter := current.resetAt.Sub(now)
|
||||
if retryAfter < 0 {
|
||||
retryAfter = 0
|
||||
}
|
||||
return false, retryAfter
|
||||
}
|
||||
return true, 0
|
||||
}
|
||||
|
||||
// Len reports how many keys are currently tracked. It exists for tests.
|
||||
func (l *Limiter) Len() int {
|
||||
if l == nil {
|
||||
return 0
|
||||
}
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
return len(l.buckets)
|
||||
}
|
||||
|
||||
// sweep drops buckets whose window has elapsed, so a long-running process does
|
||||
// not accumulate one entry per address it has ever seen. It runs at most once per
|
||||
// window and under the same lock as Allow, which keeps the cost amortised and
|
||||
// avoids a background goroutine whose lifetime somebody would have to own.
|
||||
func (l *Limiter) sweep(now time.Time) {
|
||||
if !l.sweptAt.IsZero() && now.Sub(l.sweptAt) < l.window {
|
||||
return
|
||||
}
|
||||
for key, current := range l.buckets {
|
||||
if !now.Before(current.resetAt) {
|
||||
delete(l.buckets, key)
|
||||
}
|
||||
}
|
||||
l.sweptAt = now
|
||||
}
|
||||
132
internal/pkg/ratelimit/ratelimit_test.go
Normal file
132
internal/pkg/ratelimit/ratelimit_test.go
Normal file
@@ -0,0 +1,132 @@
|
||||
package ratelimit
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestLimiterAllowsUpToTheLimitThenRefuses(t *testing.T) {
|
||||
l := New(3, time.Minute)
|
||||
for i := 1; i <= 3; i++ {
|
||||
if ok, retryAfter := l.Allow("203.0.113.7"); !ok {
|
||||
t.Fatalf("request %d was refused although the limit is 3 (retryAfter=%v)", i, retryAfter)
|
||||
}
|
||||
}
|
||||
|
||||
ok, retryAfter := l.Allow("203.0.113.7")
|
||||
if ok {
|
||||
t.Fatal("the fourth request was allowed although the limit is 3")
|
||||
}
|
||||
if retryAfter <= 0 || retryAfter > time.Minute {
|
||||
t.Fatalf("retryAfter = %v, want a positive duration within the window", retryAfter)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimiterKeysAreIndependent(t *testing.T) {
|
||||
l := New(1, time.Minute)
|
||||
if ok, _ := l.Allow("203.0.113.7"); !ok {
|
||||
t.Fatal("the first request for a key was refused")
|
||||
}
|
||||
if ok, _ := l.Allow("203.0.113.7"); ok {
|
||||
t.Fatal("the second request for the same key was allowed")
|
||||
}
|
||||
// A different caller must not inherit somebody else's exhaustion. Without key
|
||||
// isolation one blocked address would lock out every user behind it.
|
||||
if ok, _ := l.Allow("198.51.100.9"); !ok {
|
||||
t.Fatal("a different key was refused because another key was exhausted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimiterWindowResets(t *testing.T) {
|
||||
l := New(1, 20*time.Millisecond)
|
||||
if ok, _ := l.Allow("k"); !ok {
|
||||
t.Fatal("the first request was refused")
|
||||
}
|
||||
if ok, _ := l.Allow("k"); ok {
|
||||
t.Fatal("the second request inside the window was allowed")
|
||||
}
|
||||
time.Sleep(40 * time.Millisecond)
|
||||
if ok, _ := l.Allow("k"); !ok {
|
||||
t.Fatal("the window did not reset")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNilLimiterAllowsEverything(t *testing.T) {
|
||||
var l *Limiter
|
||||
for i := 0; i < 100; i++ {
|
||||
if ok, retryAfter := l.Allow("k"); !ok {
|
||||
t.Fatalf("a nil limiter refused request %d (retryAfter=%v)", i, retryAfter)
|
||||
}
|
||||
}
|
||||
if l.Len() != 0 {
|
||||
t.Fatalf("Len() on a nil limiter = %d want 0", l.Len())
|
||||
}
|
||||
}
|
||||
|
||||
func TestDisabledLimiterAllowsEverything(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
limiter *Limiter
|
||||
}{
|
||||
{"zero limit", New(0, time.Minute)},
|
||||
{"negative limit", New(-1, time.Minute)},
|
||||
{"zero window", New(5, 0)},
|
||||
{"negative window", New(5, -time.Second)},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
for i := 0; i < 50; i++ {
|
||||
if ok, _ := tc.limiter.Allow("k"); !ok {
|
||||
t.Fatalf("%s: request %d was refused although the limiter is disabled", tc.name, i)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSweepReleasesExpiredKeys(t *testing.T) {
|
||||
l := New(10, 20*time.Millisecond)
|
||||
for i := 0; i < 200; i++ {
|
||||
l.Allow(fmt.Sprintf("key-%d", i))
|
||||
}
|
||||
if got := l.Len(); got != 200 {
|
||||
t.Fatalf("Len() = %d before the sweep, want 200", got)
|
||||
}
|
||||
|
||||
time.Sleep(40 * time.Millisecond)
|
||||
// The sweep is driven by Allow rather than by a goroutine, so one more call
|
||||
// has to reclaim everything that expired.
|
||||
l.Allow("fresh")
|
||||
if got := l.Len(); got != 1 {
|
||||
t.Fatalf("Len() = %d after the sweep, want only the fresh key", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestLimiterCountsExactlyUnderConcurrency is the reason Allow holds a single
|
||||
// mutex across the read-modify-write rather than checking and then recording.
|
||||
func TestLimiterCountsExactlyUnderConcurrency(t *testing.T) {
|
||||
const goroutines = 16
|
||||
const perGoroutine = 200
|
||||
const limit = 1000
|
||||
|
||||
l := New(limit, time.Minute)
|
||||
var allowed atomic.Int64
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < goroutines; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := 0; j < perGoroutine; j++ {
|
||||
if ok, _ := l.Allow("shared"); ok {
|
||||
allowed.Add(1)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
if got := allowed.Load(); got != limit {
|
||||
t.Fatalf("allowed %d of %d requests, want exactly the limit of %d", got, goroutines*perGoroutine, limit)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user