CamTalk/backend/internal/ratelimit/redis_bucket.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)