feat: 实现 Redis 令牌桶限流器

- RedisLimiter 基于 Lua 脚本保证原子性 - Lua 脚本实现完整令牌桶算法（填充、消耗、TTL） - fail-open 策略：Redis 故障时允许请求通过 - FormatKey 辅助函数格式化限流 key - 完整单元测试（10 个测试用例，使用 miniredis） - 测试覆盖：首次请求、耗尽、不同用户、补充、容量上限、零速率、TTL、故障降级
2026-06-20 23:55:37 +08:00
parent 3b6226394b
commit b74fb3564d
4 changed files with 362 additions and 0 deletions
--- a/backend/go.mod
+++ b/backend/go.mod
@@ -19,6 +19,7 @@ require (
 )

 require (
+	github.com/alicebob/miniredis/v2 v2.38.0 // indirect
 	github.com/bahlo/generic-list-go v0.2.0 // indirect
 	github.com/buger/jsonparser v1.1.1 // indirect
 	github.com/bytedance/gopkg v0.1.3 // indirect
@@ -68,6 +69,7 @@ require (
 	github.com/ugorji/go/codec v1.2.12 // indirect
 	github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
 	github.com/yargevad/filepathx v1.0.0 // indirect
+	github.com/yuin/gopher-lua v1.1.1 // indirect
 	go.uber.org/atomic v1.11.0 // indirect
 	go.uber.org/multierr v1.10.0 // indirect
 	go.yaml.in/yaml/v3 v3.0.4 // indirect
--- a/backend/go.sum
+++ b/backend/go.sum
@@ -1,4 +1,6 @@
 github.com/airbrake/gobrake v3.6.1+incompatible/go.mod h1:wM4gu3Cn0W0K7GUuVWnlXZU11AGBXMILnrdOU8Kn00o=
+github.com/alicebob/miniredis/v2 v2.38.0 h1:nZAzCR+Lj+Vxk4ZXzm2NuKq2O33RXj1XxJ2e2uP9jiw=
+github.com/alicebob/miniredis/v2 v2.38.0/go.mod h1:TcL7YfarKPGDAthEtl5NBeHZfeUQj6OXMm/+iu5cLMM=
 github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk=
 github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg=
 github.com/bitly/go-simplejson v0.5.0/go.mod h1:cXHtHw4XUPsvGaxgjIAn8PhEWG9NfngEKAMDJEczWVA=
@@ -190,6 +192,8 @@ github.com/x-cray/logrus-prefixed-formatter v0.5.2 h1:00txxvfBM9muc0jiLIEAkAcIMJ
 github.com/x-cray/logrus-prefixed-formatter v0.5.2/go.mod h1:2duySbKsL6M18s5GU7VPsoEPHyzalCE06qoARUCeBBE=
 github.com/yargevad/filepathx v1.0.0 h1:SYcT+N3tYGi+NvazubCNlvgIPbzAk7i7y2dwg3I5FYc=
 github.com/yargevad/filepathx v1.0.0/go.mod h1:BprfX/gpYNJHJfc35GjRRpVcwWXS89gGulUIU5tK3tA=
+github.com/yuin/gopher-lua v1.1.1 h1:kYKnWBjvbNP4XLT3+bPEwAXJx262OhaHDWDVOPjL46M=
+github.com/yuin/gopher-lua v1.1.1/go.mod h1:GBR0iDaNXjAgGg9zfCvksxSRnQx76gclCIb7kdAd1Pw=
 github.com/zeebo/xxh3 v1.1.0 h1:s7DLGDK45Dyfg7++yxI0khrfwq9661w9EN78eP/UZVs=
 github.com/zeebo/xxh3 v1.1.0/go.mod h1:IisAie1LELR4xhVinxWS5+zf1lA4p0MW4T+w+W07F5s=
 go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE=
--- a/backend/internal/ratelimit/redis_bucket.go
+++ b/backend/internal/ratelimit/redis_bucket.go
@@ -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)
--- a/backend/internal/ratelimit/redis_bucket_test.go
+++ b/backend/internal/ratelimit/redis_bucket_test.go
@@ -0,0 +1,228 @@
+package ratelimit
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/alicebob/miniredis/v2"
+	"github.com/hhs/camtalk/internal/config"
+	"github.com/redis/go-redis/v9"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// setupMiniRedis 创建一个内存 Redis 实例用于测试。
+func setupMiniRedis(t *testing.T) (*miniredis.Miniredis, *redis.Client) {
+	mr, err := miniredis.Run()
+	require.NoError(t, err)
+
+	client := redis.NewClient(&redis.Options{
+		Addr: mr.Addr(),
+	})
+
+	t.Cleanup(func() {
+		client.Close()
+		mr.Close()
+	})
+
+	return mr, client
+}
+
+func TestRedisLimiter_Allow_FirstRequest(t *testing.T) {
+	_, client := setupMiniRedis(t)
+
+	cfg := config.RateLimitConfig{
+		Enabled: true,
+		Query:   config.BucketConfig{Capacity: 5, Rate: 0.2},
+	}
+	limiter := NewRedisLimiter(client, cfg)
+
+	ctx := context.Background()
+	allowed, retryAfter := limiter.Allow(ctx, "user1:query")
+
+	assert.True(t, allowed)
+	assert.Equal(t, time.Duration(0), retryAfter)
+}
+
+func TestRedisLimiter_Allow_ConsumeUntilEmpty(t *testing.T) {
+	_, client := setupMiniRedis(t)
+
+	cfg := config.RateLimitConfig{
+		Enabled: true,
+		Query:   config.BucketConfig{Capacity: 3, Rate: 0.2},
+	}
+	limiter := NewRedisLimiter(client, cfg)
+
+	ctx := context.Background()
+	key := "user1:query"
+
+	// 连续消耗 3 个令牌
+	for i := 0; i < 3; i++ {
+		allowed, _ := limiter.Allow(ctx, key)
+		assert.True(t, allowed, "request %d should be allowed", i+1)
+	}
+
+	// 第 4 个请求应被拒绝
+	allowed, retryAfter := limiter.Allow(ctx, key)
+	assert.False(t, allowed)
+	assert.Greater(t, retryAfter, time.Duration(0))
+}
+
+func TestRedisLimiter_Allow_DifferentKeys(t *testing.T) {
+	_, client := setupMiniRedis(t)
+
+	cfg := config.RateLimitConfig{
+		Enabled: true,
+		Query:   config.BucketConfig{Capacity: 2, Rate: 1.0},
+	}
+	limiter := NewRedisLimiter(client, cfg)
+
+	ctx := context.Background()
+
+	// user1 消耗 2 个令牌
+	allowed, _ := limiter.Allow(ctx, "user1:query")
+	assert.True(t, allowed)
+	allowed, _ = limiter.Allow(ctx, "user1:query")
+	assert.True(t, allowed)
+
+	// user1 第 3 个被拒绝
+	allowed, _ = limiter.Allow(ctx, "user1:query")
+	assert.False(t, allowed)
+
+	// user2 应该不受影响
+	allowed, _ = limiter.Allow(ctx, "user2:query")
+	assert.True(t, allowed)
+}
+
+func TestRedisLimiter_Allow_RefillAfterWait(t *testing.T) {
+	_, client := setupMiniRedis(t)
+
+	cfg := config.RateLimitConfig{
+		Enabled: true,
+		Query:   config.BucketConfig{Capacity: 2, Rate: 10.0}, // 每秒 10 个令牌
+	}
+	limiter := NewRedisLimiter(client, cfg)
+
+	ctx := context.Background()
+	key := "user1:query"
+
+	// 消耗 2 个令牌
+	limiter.Allow(ctx, key)
+	limiter.Allow(ctx, key)
+
+	// 真实等待 150ms（Lua 脚本使用系统时间）
+	time.Sleep(150 * time.Millisecond)
+
+	// 应该补充了至少 1 个令牌
+	allowed, _ := limiter.Allow(ctx, key)
+	assert.True(t, allowed)
+}
+
+func TestRedisLimiter_Allow_CapacityLimit(t *testing.T) {
+	_, client := setupMiniRedis(t)
+
+	cfg := config.RateLimitConfig{
+		Enabled: true,
+		Query:   config.BucketConfig{Capacity: 3, Rate: 1.0},
+	}
+	limiter := NewRedisLimiter(client, cfg)
+
+	ctx := context.Background()
+	key := "user1:query"
+
+	// 真实等待让桶"溢出"
+	time.Sleep(100 * time.Millisecond)
+
+	// 但最多只能消耗 capacity 个令牌
+	for i := 0; i < 3; i++ {
+		allowed, _ := limiter.Allow(ctx, key)
+		assert.True(t, allowed, "request %d should be allowed", i+1)
+	}
+
+	// 第 4 个应被拒绝
+	allowed, _ := limiter.Allow(ctx, key)
+	assert.False(t, allowed)
+}
+
+func TestRedisLimiter_Allow_ZeroRate(t *testing.T) {
+	_, client := setupMiniRedis(t)
+
+	cfg := config.RateLimitConfig{
+		Enabled: true,
+		Query:   config.BucketConfig{Capacity: 1, Rate: 0.0},
+	}
+	limiter := NewRedisLimiter(client, cfg)
+
+	ctx := context.Background()
+	key := "user1:query"
+
+	// 第一个通过
+	allowed, _ := limiter.Allow(ctx, key)
+	assert.True(t, allowed)
+
+	// 第二个被拒绝，retryAfter 应该很大
+	allowed, retryAfter := limiter.Allow(ctx, key)
+	assert.False(t, allowed)
+	assert.Greater(t, retryAfter, 1*time.Hour)
+}
+
+func TestRedisLimiter_Allow_KeyTTL(t *testing.T) {
+	mr, client := setupMiniRedis(t)
+
+	cfg := config.RateLimitConfig{
+		Enabled: true,
+		Query:   config.BucketConfig{Capacity: 5, Rate: 1.0},
+	}
+	limiter := NewRedisLimiter(client, cfg)
+
+	ctx := context.Background()
+	key := "user1:query"
+
+	// 第一次请求
+	limiter.Allow(ctx, key)
+
+	// 验证 key 已设置 TTL
+	ttl := mr.TTL(key)
+	assert.Greater(t, ttl, time.Duration(0))
+	assert.LessOrEqual(t, ttl, 600*time.Second)
+}
+
+func TestRedisLimiter_Allow_FailOpen(t *testing.T) {
+	mr, client := setupMiniRedis(t)
+
+	cfg := config.RateLimitConfig{
+		Enabled: true,
+		Query:   config.BucketConfig{Capacity: 1, Rate: 1.0},
+	}
+	limiter := NewRedisLimiter(client, cfg)
+
+	ctx := context.Background()
+
+	// 关闭 Redis 模拟故障
+	mr.Close()
+
+	// 应该 fail-open（允许请求）
+	allowed, retryAfter := limiter.Allow(ctx, "user1:query")
+	assert.True(t, allowed)
+	assert.Equal(t, time.Duration(0), retryAfter)
+}
+
+func TestRedisLimiter_Stop(t *testing.T) {
+	_, client := setupMiniRedis(t)
+
+	cfg := config.RateLimitConfig{
+		Enabled: true,
+		Query:   config.BucketConfig{Capacity: 1, Rate: 1.0},
+	}
+	limiter := NewRedisLimiter(client, cfg)
+
+	// Stop 应该不会 panic（即使多次调用）
+	limiter.Stop()
+	limiter.Stop()
+}
+
+func TestFormatKey(t *testing.T) {
+	key := FormatKey("user123", "query")
+	assert.Equal(t, "ratelimit:user123:query", key)
+}