72 lines
2.0 KiB
Go
72 lines
2.0 KiB
Go
|
package mw
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"golang.org/x/time/rate"
|
||
|
|
)
|
||
|
|
|
||
|
|
var _ RateLimiterStore = (*InMemoryRateLimiterStore)(nil)
|
||
|
|
|
||
|
|
const inMemoryCleanupInterval = 5 * time.Minute
|
||
|
|
|
||
|
|
type limiterEntry struct {
|
||
|
|
limiter *rate.Limiter
|
||
|
|
lastSeen time.Time
|
||
|
|
}
|
||
|
|
|
||
|
|
// InMemoryRateLimiterStore is a per-key token-bucket rate limiter backed by
|
||
|
|
// an in-memory map. Suitable for single-instance deployments or development.
|
||
|
|
// For distributed rate limiting implement [RateLimiterStore] with a Valkey or
|
||
|
|
// Redis backend and pass it to [IPRateLimit] or [UserRateLimit] instead.
|
||
|
|
//
|
||
|
|
// Stale entries are evicted every 5 minutes by a background goroutine.
|
||
|
|
// [Allow] always returns a nil error — in-memory never has infrastructure failures.
|
||
|
|
type InMemoryRateLimiterStore struct {
|
||
|
|
mu sync.Map
|
||
|
|
rps rate.Limit
|
||
|
|
burst int
|
||
|
|
}
|
||
|
|
|
||
|
|
// NewInMemoryRateLimiterStore creates a per-key token bucket store.
|
||
|
|
// rps is the sustained request rate per second per key.
|
||
|
|
// burst is the maximum instantaneous burst per key.
|
||
|
|
func NewInMemoryRateLimiterStore(rps float64, burst int) *InMemoryRateLimiterStore {
|
||
|
|
s := &InMemoryRateLimiterStore{
|
||
|
|
rps: rate.Limit(rps),
|
||
|
|
burst: burst,
|
||
|
|
}
|
||
|
|
go s.cleanupLoop()
|
||
|
|
return s
|
||
|
|
}
|
||
|
|
|
||
|
|
// Allow returns true when the request for key is within the configured rate limit.
|
||
|
|
func (s *InMemoryRateLimiterStore) Allow(_ context.Context, key string) (bool, error) {
|
||
|
|
now := time.Now()
|
||
|
|
v, _ := s.mu.LoadOrStore(key, &limiterEntry{
|
||
|
|
limiter: rate.NewLimiter(s.rps, s.burst),
|
||
|
|
lastSeen: now,
|
||
|
|
})
|
||
|
|
e := v.(*limiterEntry)
|
||
|
|
e.lastSeen = now
|
||
|
|
return e.limiter.Allow(), nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// cleanupLoop runs for the lifetime of the process; it stops only when the process exits.
|
||
|
|
// This is intentional: InMemoryRateLimiterStore is stateless from a lifecycle perspective.
|
||
|
|
func (s *InMemoryRateLimiterStore) cleanupLoop() {
|
||
|
|
ticker := time.NewTicker(inMemoryCleanupInterval)
|
||
|
|
defer ticker.Stop()
|
||
|
|
for range ticker.C {
|
||
|
|
cutoff := time.Now().Add(-inMemoryCleanupInterval)
|
||
|
|
s.mu.Range(func(k, v any) bool {
|
||
|
|
if v.(*limiterEntry).lastSeen.Before(cutoff) {
|
||
|
|
s.mu.Delete(k)
|
||
|
|
}
|
||
|
|
return true
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|