Files
summercms/modules/cooler/ratelimit.go
Jakub Zych fd7f204979 feat(19.3-01): add cooler.Cache with otter memory and Redis drivers
- Cache: Get, Set, Has, Forget, Remember (degrades on cache errors), Increment, Flush
  with per-entry TTLs under <prefix>cache:
- memory driver on otter v2.3.0 bounded by cache.memory.max_entries; Redis driver
  with SET PX, INCRBY and a Flush limited to SCAN MATCH <prefix>cache:* plus UNLINK
- Install publishes the Cache (sharing one Redis client with the rate limiter);
  FromApp serves jobs and commands without a dial
- new docs/services/caching.md, configuration rows, cooler README
2026-10-15 17:40:23 +02:00

95 lines
3.2 KiB
Go

package cooler
import (
"context"
"fmt"
"log/slog"
"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}
`)
// 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
warn *warner
}
// 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, warn: newWarner(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.warn("cooler: redis rate-limit store failed; counting per instance until it answers",
"endpoint", redisEndpoint(s.client), "error", 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
}