hhs
edc66625ba
添加 trace-aware 日志： - Error: Redis 限流检查失败（fail-open 降级） - Warn: 限流触发，记录 key 和 retry_after_sec
133 lines
3.3 KiB
Go
133 lines
3.3 KiB
Go
package ratelimit
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"strconv"
|
||
"time"
|
||
|
||
"github.com/hhs/camtalk/internal/config"
|
||
"github.com/hhs/camtalk/internal/trace"
|
||
"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) {
|
||
log := trace.FromContext(ctx)
|
||
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 {
|
||
log.Errorw("rate limit check failed", "key", key, "error", err)
|
||
// 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
|
||
log.Warnw("rate limit triggered", "key", key, "retry_after_sec", retryAfterSec)
|
||
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)
|