From b74fb3564d833cae14d2b7360033f4716e6c8fc9 Mon Sep 17 00:00:00 2001 From: hhs <386998068@qq.com> Date: Sat, 20 Jun 2026 23:55:37 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=AE=9E=E7=8E=B0=20Redis=20=E4=BB=A4?= =?UTF-8?q?=E7=89=8C=E6=A1=B6=E9=99=90=E6=B5=81=E5=99=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - RedisLimiter 基于 Lua 脚本保证原子性 - Lua 脚本实现完整令牌桶算法（填充、消耗、TTL） - fail-open 策略：Redis 故障时允许请求通过 - FormatKey 辅助函数格式化限流 key - 完整单元测试（10 个测试用例，使用 miniredis） - 测试覆盖：首次请求、耗尽、不同用户、补充、容量上限、零速率、TTL、故障降级 --- backend/go.mod | 2 + backend/go.sum | 4 + backend/internal/ratelimit/redis_bucket.go | 128 ++++++++++ .../internal/ratelimit/redis_bucket_test.go | 228 ++++++++++++++++++ 4 files changed, 362 insertions(+) create mode 100644 backend/internal/ratelimit/redis_bucket.go create mode 100644 backend/internal/ratelimit/redis_bucket_test.go diff --git a/backend/go.mod b/backend/go.mod index 6f652f0..e9b487e 100644 --- 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 diff --git a/backend/go.sum b/backend/go.sum index 9d58137..93d4a44 100644 --- 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= diff --git a/backend/internal/ratelimit/redis_bucket.go b/backend/internal/ratelimit/redis_bucket.go new file mode 100644 index 0000000..eb1e606 --- /dev/null +++ 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) diff --git a/backend/internal/ratelimit/redis_bucket_test.go b/backend/internal/ratelimit/redis_bucket_test.go new file mode 100644 index 0000000..becc606 --- /dev/null +++ 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) +}