From e2a0c4b16aa6c2e365f0845cf073c52b933385a2 Mon Sep 17 00:00:00 2001 From: BradleyZ <2111976953@qq.com> Date: Tue, 20 Oct 2026 22:28:33 +0800 Subject: [PATCH] fix(limiter): recover redis script after flush (#6486) - Retry the Lua rate limiter with EVAL when Redis returns NOSCRIPT - Add regression coverage for Redis script cache flush recovery --- common/limiter/limiter.go | 15 +++++++ common/limiter/limiter_test.go | 80 ++++++++++++++++++++++++++++++++++ 2 files changed, 95 insertions(+) create mode 100644 common/limiter/limiter_test.go diff --git a/common/limiter/limiter.go b/common/limiter/limiter.go index 46cd8c614..30a187727 100644 --- a/common/limiter/limiter.go +++ b/common/limiter/limiter.go @@ -4,6 +4,7 @@ import ( "context" _ "embed" "fmt" + "strings" "sync" "github.com/QuantumNous/new-api/common" @@ -63,6 +64,20 @@ func (rl *RedisLimiter) Allow(ctx context.Context, key string, opts ...Option) ( ).Int() if err != nil { + if strings.HasPrefix(err.Error(), "NOSCRIPT ") { + result, err := rl.client.Eval( + ctx, + rateLimitScript, + []string{key}, + config.Requested, + config.Rate, + config.Capacity, + ).Int() + if err != nil { + return false, fmt.Errorf("rate limit failed: %w", err) + } + return result == 1, nil + } return false, fmt.Errorf("rate limit failed: %w", err) } return result == 1, nil diff --git a/common/limiter/limiter_test.go b/common/limiter/limiter_test.go new file mode 100644 index 000000000..8850d31cc --- /dev/null +++ b/common/limiter/limiter_test.go @@ -0,0 +1,80 @@ +package limiter + +import ( + "context" + "sync" + "testing" + + "github.com/alicebob/miniredis/v2" + "github.com/go-redis/redis/v8" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func useLimiterMiniRedis(t *testing.T) (*miniredis.Miniredis, *redis.Client) { + t.Helper() + + previousInstance := instance + + redisServer := miniredis.RunT(t) + redisClient := redis.NewClient(&redis.Options{Addr: redisServer.Addr()}) + require.NoError(t, redisClient.Ping(context.Background()).Err()) + + instance = nil + once = sync.Once{} + + t.Cleanup(func() { + _ = redisClient.Close() + instance = previousInstance + once = sync.Once{} + if previousInstance != nil { + once.Do(func() {}) + } + }) + + return redisServer, redisClient +} + +func TestAllowFallsBackToEvalAfterScriptFlush(t *testing.T) { + ctx := context.Background() + _, redisClient := useLimiterMiniRedis(t) + + redisLimiter := New(ctx, redisClient) + + allowed, err := redisLimiter.Allow( + ctx, + "rateLimit:test-before-flush", + WithCapacity(60), + WithRate(60), + WithRequested(60), + ) + require.NoError(t, err) + assert.True(t, allowed) + + exists, err := redisClient.ScriptExists(ctx, redisLimiter.limitScriptSHA).Result() + require.NoError(t, err) + require.Len(t, exists, 1) + assert.True(t, exists[0]) + + require.NoError(t, redisClient.ScriptFlush(ctx).Err()) + + exists, err = redisClient.ScriptExists(ctx, redisLimiter.limitScriptSHA).Result() + require.NoError(t, err) + require.Len(t, exists, 1) + assert.False(t, exists[0]) + + allowed, err = redisLimiter.Allow( + ctx, + "rateLimit:test-after-flush", + WithCapacity(60), + WithRate(60), + WithRequested(60), + ) + require.NoError(t, err) + assert.True(t, allowed) + + exists, err = redisClient.ScriptExists(ctx, redisLimiter.limitScriptSHA).Result() + require.NoError(t, err) + require.Len(t, exists, 1) + assert.True(t, exists[0]) +}