fa4d89d0b9
- new module cooler: Install, Handle, OpenRedis (no URL in errors), RateLimitStore (atomic fixed-window Lua script with MemoryStore semantics, per-instance fallback) - serve installs cooler and publishes the Redis surf.Store before BuildRouter; BuildRouter looks up a published Store, route:list keeps the memory store - config keys http.rate_limit.store and redis.url; go-redis v9.23.0 - cooler README, root table row, surf README and rate-limiting, configuration and scaling docs
122 lines
3.7 KiB
Go
122 lines
3.7 KiB
Go
package cooler
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log/slog"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
)
|
|
|
|
// attemptScript is one atomic fixed-window admission with the semantics of
|
|
// surf.MemoryStore: a missing or expired key opens a window of decay
|
|
// milliseconds, a count at max denies without incrementing, otherwise the
|
|
// count is incremented. INCR keeps the TTL, so the window never extends.
|
|
// It returns {allowed (0|1), count, remaining window in ms}.
|
|
//
|
|
// KEYS[1] = <prefix>rl:<key>; ARGV[1] = max; ARGV[2] = decay in ms (>= 1).
|
|
var attemptScript = redis.NewScript(`
|
|
local max = tonumber(ARGV[1])
|
|
local decay = tonumber(ARGV[2])
|
|
local count = tonumber(redis.call('GET', KEYS[1]) or '0')
|
|
local ttl = redis.call('PTTL', KEYS[1])
|
|
if ttl < 0 then
|
|
count = 0
|
|
redis.call('SET', KEYS[1], 0, 'PX', decay)
|
|
ttl = decay
|
|
end
|
|
if count >= max then
|
|
return {0, count, ttl}
|
|
end
|
|
count = redis.call('INCR', KEYS[1])
|
|
return {1, count, ttl}
|
|
`)
|
|
|
|
// warnEvery bounds how often a Redis failure is logged.
|
|
const warnEvery = time.Minute
|
|
|
|
// RateLimitStore is a fixed-window rate-limit store on Redis, shared by every
|
|
// application instance that points at the same Redis. Its Attempt has the
|
|
// exact signature of surf.Store, so surf uses it without cooler importing
|
|
// surf. When Redis does not answer, Attempt counts in the fallback store
|
|
// (per instance) and logs a warning at most once per minute.
|
|
type RateLimitStore struct {
|
|
client *redis.Client
|
|
prefix string
|
|
fallback Attempter
|
|
log *slog.Logger
|
|
|
|
mu sync.Mutex
|
|
lastWarn time.Time
|
|
}
|
|
|
|
// NewRateLimitStore returns a store that keeps its counters under
|
|
// <prefix>rl:<key> on client. fallback answers while Redis fails; surf passes
|
|
// a surf.MemoryStore. A nil fallback admits every request during an outage.
|
|
func NewRateLimitStore(client *redis.Client, prefix string, fallback Attempter) *RateLimitStore {
|
|
return &RateLimitStore{client: client, prefix: prefix, fallback: fallback, log: slog.Default()}
|
|
}
|
|
|
|
// Attempt admits and counts one request when key is below max, in one
|
|
// atomic Redis script bounded by AttemptTimeout. The first attempt opens the
|
|
// window, a denied attempt does not increment, and retryAfter is the time
|
|
// left in the window.
|
|
func (s *RateLimitStore) Attempt(key string, max int, decay time.Duration) (allowed bool, attempts int, retryAfter time.Duration) {
|
|
ms := decay.Milliseconds()
|
|
if ms < 1 {
|
|
ms = 1
|
|
}
|
|
ctx, cancel := context.WithTimeout(context.Background(), AttemptTimeout)
|
|
defer cancel()
|
|
res, err := s.run(ctx, s.prefix+"rl:"+key, max, ms)
|
|
if err != nil {
|
|
s.warn(err)
|
|
if s.fallback == nil {
|
|
return true, 0, 0
|
|
}
|
|
return s.fallback.Attempt(key, max, decay)
|
|
}
|
|
ttl := time.Duration(res[2]) * time.Millisecond
|
|
if ttl < 0 {
|
|
ttl = 0
|
|
}
|
|
return res[0] == 1, int(res[1]), ttl
|
|
}
|
|
|
|
func (s *RateLimitStore) run(ctx context.Context, key string, max int, ms int64) ([]int64, error) {
|
|
if s == nil || s.client == nil {
|
|
return nil, fmt.Errorf("cooler: rate-limit store has no Redis client")
|
|
}
|
|
res, err := attemptScript.Run(ctx, s.client, []string{key}, max, ms).Int64Slice()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(res) != 3 {
|
|
return nil, fmt.Errorf("cooler: rate-limit script returned %d values, want 3", len(res))
|
|
}
|
|
return res, nil
|
|
}
|
|
|
|
// warn logs a Redis failure at most once per warnEvery.
|
|
func (s *RateLimitStore) warn(err error) {
|
|
if s == nil {
|
|
return
|
|
}
|
|
s.mu.Lock()
|
|
now := time.Now()
|
|
if !s.lastWarn.IsZero() && now.Sub(s.lastWarn) < warnEvery {
|
|
s.mu.Unlock()
|
|
return
|
|
}
|
|
s.lastWarn = now
|
|
s.mu.Unlock()
|
|
log := s.log
|
|
if log == nil {
|
|
log = slog.Default()
|
|
}
|
|
log.Warn("cooler: redis rate-limit store failed; counting per instance until it answers",
|
|
"endpoint", redisEndpoint(s.client), "error", err)
|
|
}
|