feat: 实现 Redis 令牌桶限流器
- RedisLimiter 基于 Lua 脚本保证原子性 - Lua 脚本实现完整令牌桶算法(填充、消耗、TTL) - fail-open 策略:Redis 故障时允许请求通过 - FormatKey 辅助函数格式化限流 key - 完整单元测试(10 个测试用例,使用 miniredis) - 测试覆盖:首次请求、耗尽、不同用户、补充、容量上限、零速率、TTL、故障降级
This commit is contained in:
128
backend/internal/ratelimit/redis_bucket.go
Normal file
128
backend/internal/ratelimit/redis_bucket.go
Normal file
@@ -0,0 +1,128 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user