feat(19.3-01): share rate limits across instances through cooler's Redis store
- 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
This commit is contained in:
121
modules/cooler/ratelimit.go
Normal file
121
modules/cooler/ratelimit.go
Normal file
@@ -0,0 +1,121 @@
|
||||
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)
|
||||
}
|
||||
Reference in New Issue
Block a user