refactor: extract shared rate limiter types into ratelimit_shared.go
Move RateLimiter, ProgressiveRateLimiter, and cleanup functions to ratelimit_shared.go. Add dev and prod rate limiter tests. Remove inline cleanup goroutines. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
This commit is contained in:
@@ -0,0 +1,104 @@
|
||||
// Shared types, registration, and cleanup for RateLimiter + ProgressiveRateLimiter.
|
||||
// Used by both production (!dev) and dev (dev) builds — keep tag-free.
|
||||
|
||||
package mw
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crussell/clock"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// RateLimiter implements a simple in-memory rate limiter
|
||||
type RateLimiter struct {
|
||||
requests map[string][]time.Time
|
||||
mu sync.RWMutex
|
||||
limit int
|
||||
window time.Duration
|
||||
}
|
||||
|
||||
// ProgressiveRateLimiter implements per-IP rate limiting with increasing backoff.
|
||||
// Designed for bot-spam prevention across accounts (not account-specific lockout).
|
||||
type ProgressiveRateLimiter struct {
|
||||
requests map[string]*ipProgressiveState
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
type ipProgressiveState struct {
|
||||
// Timestamps of all requests within the tracking window
|
||||
timestamps []time.Time
|
||||
}
|
||||
|
||||
var (
|
||||
registeredLimiters []*RateLimiter
|
||||
registeredLimitersMu sync.Mutex
|
||||
)
|
||||
|
||||
func registerLimiter(rl *RateLimiter) {
|
||||
registeredLimitersMu.Lock()
|
||||
registeredLimiters = append(registeredLimiters, rl)
|
||||
registeredLimitersMu.Unlock()
|
||||
}
|
||||
|
||||
// CleanupAllRateLimiters runs Cleanup on every registered RateLimiter.
|
||||
// Called by the centralised jobs scheduler.
|
||||
func CleanupAllRateLimiters(ctx context.Context) (int, error) {
|
||||
registeredLimitersMu.Lock()
|
||||
limiters := make([]*RateLimiter, len(registeredLimiters))
|
||||
copy(limiters, registeredLimiters)
|
||||
registeredLimitersMu.Unlock()
|
||||
for _, l := range limiters {
|
||||
l.Cleanup()
|
||||
}
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// CleanupProgressiveRateLimiter runs Cleanup on the global progressive rate limiter.
|
||||
// Called by the centralised jobs scheduler.
|
||||
func CleanupProgressiveRateLimiter(ctx context.Context) (int, error) {
|
||||
globalProgressiveLimiter.Cleanup()
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// Cleanup removes expired entries from the rate limiter map.
|
||||
func (rl *RateLimiter) Cleanup() {
|
||||
rl.mu.Lock()
|
||||
defer rl.mu.Unlock()
|
||||
now := clock.Now()
|
||||
for key, times := range rl.requests {
|
||||
var valid []time.Time
|
||||
for _, t := range times {
|
||||
if now.Sub(t) < rl.window {
|
||||
valid = append(valid, t)
|
||||
}
|
||||
}
|
||||
if len(valid) == 0 {
|
||||
delete(rl.requests, key)
|
||||
} else {
|
||||
rl.requests[key] = valid
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Cleanup removes expired entries from the progressive rate limiter map.
|
||||
func (prl *ProgressiveRateLimiter) Cleanup() {
|
||||
prl.mu.Lock()
|
||||
defer prl.mu.Unlock()
|
||||
cutoff := clock.Now().Add(-60 * time.Second)
|
||||
for ip, state := range prl.requests {
|
||||
var valid []time.Time
|
||||
for _, t := range state.timestamps {
|
||||
if t.After(cutoff) {
|
||||
valid = append(valid, t)
|
||||
}
|
||||
}
|
||||
if len(valid) == 0 {
|
||||
delete(prl.requests, ip)
|
||||
} else {
|
||||
state.timestamps = valid
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var globalProgressiveLimiter = NewProgressiveRateLimiter()
|
||||
Reference in New Issue
Block a user