package ratelimit import ( "context" "fmt" "strconv" "time" "github.com/hhs/camtalk/internal/config" "github.com/redis/go-redis/v9" ) // luaScript 是 Redis 令牌桶算法的 Lua 脚本。 // 保证原子性:读取-计算-回写在一个事务中完成。 const luaScript = ` -- KEYS[1] = 限流 key -- ARGV[1] = capacity(桶容量) -- ARGV[2] = rate(每秒填充数) -- ARGV[3] = now(当前时间戳,秒,浮点) -- ARGV[4] = ttl(key 过期时间,秒) local key = KEYS[1] local capacity = tonumber(ARGV[1]) local rate = tonumber(ARGV[2]) local now = tonumber(ARGV[3]) local ttl = tonumber(ARGV[4]) local data = redis.call('HMGET', key, 'tokens', 'last_refill') local tokens = tonumber(data[1]) or capacity local last_refill = tonumber(data[2]) or now -- 计算新令牌 local elapsed = math.max(0, now - last_refill) tokens = math.min(capacity, tokens + elapsed * rate) local allowed = 0 local retry_after = 0 if tokens >= 1 then tokens = tokens - 1 allowed = 1 else if rate == 0 then retry_after = 86400 -- 24小时 else retry_after = (1 - tokens) / rate end end -- 回写状态 redis.call('HMSET', key, 'tokens', tokens, 'last_refill', now) redis.call('EXPIRE', key, ttl) return {allowed, tostring(retry_after)} ` // RedisLimiter Redis 令牌桶限流器。 type RedisLimiter struct { client *redis.Client config config.RateLimitConfig script *redis.Script } // NewRedisLimiter 创建 Redis 限流器。 func NewRedisLimiter(client *redis.Client, cfg config.RateLimitConfig) *RedisLimiter { return &RedisLimiter{ client: client, config: cfg, script: redis.NewScript(luaScript), } } // Allow 实现 Limiter 接口。 func (l *RedisLimiter) Allow(ctx context.Context, key string) (bool, time.Duration) { cfg := l.getBucketConfig(key) now := float64(time.Now().UnixNano()) / 1e9 // 秒,浮点 ttl := 600 // key 过期时间 10 分钟 result, err := l.script.Run(ctx, l.client, []string{key}, cfg.Capacity, cfg.Rate, now, ttl).Result() if err != nil { // Redis 错误时降级:允许请求(fail-open 策略) return true, 0 } // 解析返回值 vals, ok := result.([]interface{}) if !ok || len(vals) != 2 { return true, 0 } allowed, _ := vals[0].(int64) retryAfterStr, _ := vals[1].(string) retryAfterSec, _ := strconv.ParseFloat(retryAfterStr, 64) if allowed == 1 { return true, 0 } retryAfter := time.Duration(retryAfterSec*1000) * time.Millisecond return false, retryAfter } // Stop 实现 Limiter 接口(Redis 不需要清理资源)。 func (l *RedisLimiter) Stop() { // Redis 客户端由外部管理,这里不需要操作 } // getBucketConfig 根据 key 获取桶配置。 func (l *RedisLimiter) getBucketConfig(key string) config.BucketConfig { // 简化实现:默认使用 query 配置 return l.config.Query } // KeyPrefix 返回限流 key 的前缀。 func KeyPrefix() string { return "ratelimit:" } // FormatKey 格式化限流 key。 func FormatKey(userID, action string) string { return fmt.Sprintf("%s%s:%s", KeyPrefix(), userID, action) } // 编译期接口检查 var _ Limiter = (*RedisLimiter)(nil)