wangjia
3c687e5e0b
新包 internal/ratelimit：内存/redis(GCRA+Lua) 双实现 + 出错逐调用降级内存 （fail-open 到内存不 fail-closed）。REDIS_ADDR 空=内存模式，行为与既往一致； 配置后跨重启保状态、支持多实例。miniredis 全覆盖测试，零真实外部依赖。 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
132 lines
4.1 KiB
Go
132 lines
4.1 KiB
Go
package ratelimit
|
||
|
||
import (
|
||
"context"
|
||
"time"
|
||
|
||
redis_rate "github.com/go-redis/redis_rate/v10"
|
||
"github.com/redis/go-redis/v9"
|
||
"golang.org/x/time/rate"
|
||
)
|
||
|
||
// redis 实现。键统一前缀 jiu:rl:（同 redis 将来可能被同宿主其他服务复用）。
|
||
// 所有方法返回 error 供 fallback 包装器降级判断；本文件不直接暴露给调用方。
|
||
|
||
const (
|
||
redisKeyPrefix = "jiu:rl:"
|
||
redisOpTimeout = 500 * time.Millisecond
|
||
// 失败计数累计窗口（对齐内存版 janitor 的 30min 空闲清理语义）。
|
||
redisFailTTL = 30 * time.Minute
|
||
)
|
||
|
||
type redisStore struct {
|
||
rdb *redis.Client
|
||
rl *redis_rate.Limiter
|
||
}
|
||
|
||
func newRedisStore(rdb *redis.Client) *redisStore {
|
||
return &redisStore{rdb: rdb, rl: redis_rate.NewLimiter(rdb)}
|
||
}
|
||
|
||
func opCtx() (context.Context, context.CancelFunc) {
|
||
return context.WithTimeout(context.Background(), redisOpTimeout)
|
||
}
|
||
|
||
// ── Limiter：GCRA ────────────────────────────────────────────────────────────
|
||
|
||
type redisLimiter struct {
|
||
rl *redis_rate.Limiter
|
||
limit redis_rate.Limit
|
||
}
|
||
|
||
// newLimiter 把 x/time/rate 的「令牌/秒」换算成 GCRA 的每事件间隔：
|
||
// Limit{Rate:1, Period: 1/r} 与 rate.Limiter(r) 的稳态速率等价，burst 语义一致。
|
||
func (s *redisStore) newLimiter(r rate.Limit, burst int) *redisLimiter {
|
||
if r <= 0 {
|
||
r = rate.Limit(1.0 / 60.0) // 防御：非法速率按每分钟 1 次
|
||
}
|
||
period := time.Duration(float64(time.Second) / float64(r))
|
||
return &redisLimiter{
|
||
rl: s.rl,
|
||
limit: redis_rate.Limit{Rate: 1, Period: period, Burst: burst},
|
||
}
|
||
}
|
||
|
||
func (l *redisLimiter) allow(key string) (Result, error) {
|
||
ctx, cancel := opCtx()
|
||
defer cancel()
|
||
res, err := l.rl.Allow(ctx, redisKeyPrefix+"lim:"+key, l.limit)
|
||
if err != nil {
|
||
return Result{}, err
|
||
}
|
||
ra := res.RetryAfter
|
||
if ra <= 0 {
|
||
ra = memRetryAfter
|
||
}
|
||
return Result{Allowed: res.Allowed > 0, RetryAfter: ra}, nil
|
||
}
|
||
|
||
// ── Counter ──────────────────────────────────────────────────────────────────
|
||
|
||
type redisCounter struct{ rdb *redis.Client }
|
||
|
||
func (c *redisCounter) Incr(key string, ttl time.Duration) (int64, error) {
|
||
ctx, cancel := opCtx()
|
||
defer cancel()
|
||
k := redisKeyPrefix + "cnt:" + key
|
||
n, err := c.rdb.Incr(ctx, k).Result()
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
if n == 1 {
|
||
// 首次自增设 TTL；失败不致命（键会在下个周期被重建）
|
||
_ = c.rdb.Expire(ctx, k, ttl).Err()
|
||
}
|
||
return n, nil
|
||
}
|
||
|
||
// ── FailLocker ───────────────────────────────────────────────────────────────
|
||
|
||
// recordFailScript：INCR 失败计数 + 续窗口 TTL；达阈值则置锁并清计数（原子）。
|
||
// KEYS[1]=fail 键, KEYS[2]=lock 键；ARGV[1]=failTTL 秒, ARGV[2]=max, ARGV[3]=lock 秒。
|
||
var recordFailScript = redis.NewScript(`
|
||
local f = redis.call('INCR', KEYS[1])
|
||
redis.call('EXPIRE', KEYS[1], ARGV[1])
|
||
if tonumber(ARGV[2]) > 0 and f >= tonumber(ARGV[2]) then
|
||
redis.call('SET', KEYS[2], 1, 'EX', ARGV[3])
|
||
redis.call('DEL', KEYS[1])
|
||
end
|
||
return f
|
||
`)
|
||
|
||
type redisFailLocker struct{ rdb *redis.Client }
|
||
|
||
func failKeys(key string) (string, string) {
|
||
return redisKeyPrefix + "fail:" + key, redisKeyPrefix + "lock:" + key
|
||
}
|
||
|
||
func (l *redisFailLocker) locked(key string) (bool, error) {
|
||
ctx, cancel := opCtx()
|
||
defer cancel()
|
||
_, lockKey := failKeys(key)
|
||
n, err := l.rdb.Exists(ctx, lockKey).Result()
|
||
return n > 0, err
|
||
}
|
||
|
||
func (l *redisFailLocker) recordFailure(key string, max int, lockFor time.Duration) error {
|
||
ctx, cancel := opCtx()
|
||
defer cancel()
|
||
failKey, lockKey := failKeys(key)
|
||
return recordFailScript.Run(ctx, l.rdb,
|
||
[]string{failKey, lockKey},
|
||
int(redisFailTTL.Seconds()), max, int(lockFor.Seconds()),
|
||
).Err()
|
||
}
|
||
|
||
func (l *redisFailLocker) reset(key string) error {
|
||
ctx, cancel := opCtx()
|
||
defer cancel()
|
||
failKey, lockKey := failKeys(key)
|
||
return l.rdb.Del(ctx, failKey, lockKey).Err()
|
||
}
|