Codevault initial commit
This commit is contained in:
145
internal/ratelimit/limiter.go
Normal file
145
internal/ratelimit/limiter.go
Normal file
@@ -0,0 +1,145 @@
|
||||
// Package ratelimit implements a per-IP sliding-window failure counter used to
|
||||
// hide the database behind a cap on failed lookups.
|
||||
package ratelimit
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type bucket struct {
|
||||
index int64
|
||||
count int
|
||||
}
|
||||
|
||||
type ipState struct {
|
||||
current bucket
|
||||
previous bucket
|
||||
}
|
||||
|
||||
// Limiter is a two-bucket sliding-window counter per IP. Allowed is a pure read
|
||||
// (RLock) that estimates the failure count in the trailing window from whatever
|
||||
// is already stored, without mutating it. RecordFailure takes the write lock and
|
||||
// performs the actual window roll and increment. Keeping the previous window's
|
||||
// count means a client blocked near the end of one window stays throttled into
|
||||
// the next, instead of getting an instant reset at the window boundary.
|
||||
type Limiter struct {
|
||||
mu sync.RWMutex
|
||||
state map[string]ipState
|
||||
window time.Duration
|
||||
max int
|
||||
}
|
||||
|
||||
func NewLimiter(max int, window time.Duration) *Limiter {
|
||||
return &Limiter{
|
||||
state: make(map[string]ipState),
|
||||
window: window,
|
||||
max: max,
|
||||
}
|
||||
}
|
||||
|
||||
func (l *Limiter) windowIndex(t time.Time) int64 {
|
||||
return t.UnixNano() / int64(l.window)
|
||||
}
|
||||
|
||||
// Allowed reports whether ip is currently under the failure threshold. It never
|
||||
// mutates the limiter's state.
|
||||
func (l *Limiter) Allowed(ip string) bool {
|
||||
now := time.Now()
|
||||
idx := l.windowIndex(now)
|
||||
|
||||
l.mu.RLock()
|
||||
state, ok := l.state[ip]
|
||||
l.mu.RUnlock()
|
||||
if !ok {
|
||||
return true
|
||||
}
|
||||
|
||||
current, previous := effectiveBuckets(state, idx)
|
||||
|
||||
overlap := overlapFraction(now, l.window, previous.index)
|
||||
estimate := float64(current.count) + float64(previous.count)*overlap
|
||||
|
||||
return estimate < float64(l.max)
|
||||
}
|
||||
|
||||
// RecordFailure increments the failure count for ip in the current window,
|
||||
// rolling the previous stored window forward if it has become stale.
|
||||
func (l *Limiter) RecordFailure(ip string) {
|
||||
now := time.Now()
|
||||
idx := l.windowIndex(now)
|
||||
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
state := l.state[ip]
|
||||
current, previous := effectiveBuckets(state, idx)
|
||||
current.count++
|
||||
l.state[ip] = ipState{current: current, previous: previous}
|
||||
}
|
||||
|
||||
// Cleanup removes IPs whose state is more than 2 windows stale, bounding memory.
|
||||
func (l *Limiter) Cleanup() {
|
||||
idx := l.windowIndex(time.Now())
|
||||
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
for ip, state := range l.state {
|
||||
if idx-state.current.index > 2 {
|
||||
delete(l.state, ip)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Run periodically calls Cleanup until ctx is done.
|
||||
func (l *Limiter) Run(done <-chan struct{}, interval time.Duration) {
|
||||
ticker := time.NewTicker(interval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-done:
|
||||
return
|
||||
case <-ticker.C:
|
||||
l.Cleanup()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// effectiveBuckets derives the current/previous buckets as seen from window idx,
|
||||
// without mutating the stored state: if the stored current bucket is stale by
|
||||
// exactly one window it becomes the effective previous bucket; if it is older
|
||||
// than that, both buckets are treated as empty.
|
||||
func effectiveBuckets(state ipState, idx int64) (current, previous bucket) {
|
||||
switch {
|
||||
case state.current.index == idx:
|
||||
current = state.current
|
||||
if state.previous.index == idx-1 {
|
||||
previous = state.previous
|
||||
} else {
|
||||
previous = bucket{index: idx - 1}
|
||||
}
|
||||
case state.current.index == idx-1:
|
||||
current = bucket{index: idx}
|
||||
previous = state.current
|
||||
default:
|
||||
current = bucket{index: idx}
|
||||
previous = bucket{index: idx - 1}
|
||||
}
|
||||
return current, previous
|
||||
}
|
||||
|
||||
// overlapFraction returns how much of the previous window (identified by
|
||||
// previousIdx) still falls within the trailing lookback window of length
|
||||
// windowSize measured back from now.
|
||||
func overlapFraction(now time.Time, windowSize time.Duration, previousIdx int64) float64 {
|
||||
previousWindowEnd := time.Unix(0, (previousIdx+1)*int64(windowSize))
|
||||
elapsedSincePreviousEnd := now.Sub(previousWindowEnd)
|
||||
if elapsedSincePreviousEnd >= windowSize {
|
||||
return 0
|
||||
}
|
||||
if elapsedSincePreviousEnd <= 0 {
|
||||
return 1
|
||||
}
|
||||
return 1 - float64(elapsedSincePreviousEnd)/float64(windowSize)
|
||||
}
|
||||
101
internal/ratelimit/limiter_test.go
Normal file
101
internal/ratelimit/limiter_test.go
Normal file
@@ -0,0 +1,101 @@
|
||||
package ratelimit
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestLimiter_AllowsUnderThreshold(t *testing.T) {
|
||||
l := NewLimiter(10, time.Minute)
|
||||
for i := range 9 {
|
||||
if !l.Allowed("1.2.3.4") {
|
||||
t.Fatalf("expected allowed at failure %d", i)
|
||||
}
|
||||
l.RecordFailure("1.2.3.4")
|
||||
}
|
||||
if !l.Allowed("1.2.3.4") {
|
||||
t.Fatalf("expected allowed after 9 failures")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimiter_BlocksAtThreshold(t *testing.T) {
|
||||
l := NewLimiter(10, time.Minute)
|
||||
for range 10 {
|
||||
l.RecordFailure("1.2.3.4")
|
||||
}
|
||||
if l.Allowed("1.2.3.4") {
|
||||
t.Fatalf("expected blocked after 10 failures")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimiter_BlockedClientStaysThrottledAcrossWindowBoundary(t *testing.T) {
|
||||
window := 300 * time.Millisecond
|
||||
l := NewLimiter(10, window)
|
||||
|
||||
for range 20 {
|
||||
l.RecordFailure("1.2.3.4")
|
||||
}
|
||||
if l.Allowed("1.2.3.4") {
|
||||
t.Fatalf("expected blocked immediately after 20 failures")
|
||||
}
|
||||
|
||||
idx := l.windowIndex(time.Now())
|
||||
windowEnd := time.Unix(0, (idx+1)*int64(window))
|
||||
// Sleep just past the window boundary so the failures roll into the previous
|
||||
// bucket, then check that their weighted contribution still blocks the client.
|
||||
time.Sleep(time.Until(windowEnd) + window/10)
|
||||
|
||||
if l.Allowed("1.2.3.4") {
|
||||
t.Fatalf("expected still blocked shortly after the window boundary")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimiter_ResetsAfterIdlePeriod(t *testing.T) {
|
||||
window := 50 * time.Millisecond
|
||||
l := NewLimiter(10, window)
|
||||
|
||||
for range 10 {
|
||||
l.RecordFailure("1.2.3.4")
|
||||
}
|
||||
if l.Allowed("1.2.3.4") {
|
||||
t.Fatalf("expected blocked immediately after 10 failures")
|
||||
}
|
||||
|
||||
time.Sleep(window * 3)
|
||||
|
||||
if !l.Allowed("1.2.3.4") {
|
||||
t.Fatalf("expected allowed after idling for 2+ windows")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimiter_Cleanup(t *testing.T) {
|
||||
window := 20 * time.Millisecond
|
||||
l := NewLimiter(10, window)
|
||||
l.RecordFailure("1.2.3.4")
|
||||
|
||||
time.Sleep(window * 4)
|
||||
l.Cleanup()
|
||||
|
||||
l.mu.RLock()
|
||||
_, ok := l.state["1.2.3.4"]
|
||||
l.mu.RUnlock()
|
||||
if ok {
|
||||
t.Fatalf("expected stale IP to be cleaned up")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOverlapFraction(t *testing.T) {
|
||||
window := 100 * time.Millisecond
|
||||
previousIdx := int64(1000)
|
||||
previousWindowEnd := time.Unix(0, (previousIdx+1)*int64(window))
|
||||
|
||||
if got := overlapFraction(previousWindowEnd.Add(-10*time.Millisecond), window, previousIdx); got != 1 {
|
||||
t.Fatalf("expected overlap 1 before window end, got %v", got)
|
||||
}
|
||||
if got := overlapFraction(previousWindowEnd.Add(20*time.Millisecond), window, previousIdx); got < 0.79 || got > 0.81 {
|
||||
t.Fatalf("expected overlap ~0.8 at 20ms into next window, got %v", got)
|
||||
}
|
||||
if got := overlapFraction(previousWindowEnd.Add(200*time.Millisecond), window, previousIdx); got != 0 {
|
||||
t.Fatalf("expected overlap 0 once fully elapsed, got %v", got)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user