From b98b7faafaaa1693df9609116c259fb607de9b28 Mon Sep 17 00:00:00 2001 From: zijiren233 Date: Mon, 20 Jul 2026 00:02:12 +0800 Subject: [PATCH 1/2] fix: prevent Redis rate limiter scripts from blocking --- core/common/reqlimit/mem.go | 5 +- core/common/reqlimit/mem_internal_test.go | 26 ++++ core/common/reqlimit/redis.go | 91 ++++++------ .../common/reqlimit/redis_integration_test.go | 130 ++++++++++++++++++ 4 files changed, 204 insertions(+), 48 deletions(-) diff --git a/core/common/reqlimit/mem.go b/core/common/reqlimit/mem.go index e20118c9..37279585 100644 --- a/core/common/reqlimit/mem.go +++ b/core/common/reqlimit/mem.go @@ -82,7 +82,10 @@ func (m *InMemoryRecord) rebuildAggregateLocked(e *entry, windowSeconds, cutoff func (m *InMemoryRecord) refreshAggregateLocked(e *entry, nowSecond, windowSeconds int64) { cutoff := nowSecond - windowSeconds - if !e.aggregateInitialized || e.windowSeconds != windowSeconds { + if !e.aggregateInitialized || + e.windowSeconds != windowSeconds || + e.lastCleanedCutoff > cutoff || + cutoff-e.lastCleanedCutoff > windowSeconds { m.rebuildAggregateLocked(e, windowSeconds, cutoff) return } diff --git a/core/common/reqlimit/mem_internal_test.go b/core/common/reqlimit/mem_internal_test.go index 89000b63..5677d843 100644 --- a/core/common/reqlimit/mem_internal_test.go +++ b/core/common/reqlimit/mem_internal_test.go @@ -32,6 +32,32 @@ func TestInMemoryRecordDropsExpiredMinuteWindowWithoutRealWait(t *testing.T) { require.True(t, recentExists) } +func TestInMemoryRecordRebuildsAggregateAfterTimeDiscontinuity(t *testing.T) { + rl := &InMemoryRecord{} + e := rl.getEntry([]string{"group1", "model1"}) + nowSecond := time.Now().Unix() + + e.Lock() + e.windows[nowSecond-61] = &windowCounts{normal: 3} + e.windows[nowSecond] = &windowCounts{normal: 2} + e.windowSeconds = 60 + e.totalNormal = 5 + e.lastCleanedCutoff = nowSecond + 60 + e.aggregateInitialized = true + e.Unlock() + + totalCount, secondCount := rl.GetRequest(time.Minute, "group1", "model1") + require.Equal(t, int64(2), totalCount) + require.Equal(t, int64(2), secondCount) + + e.Lock() + _, expiredExists := e.windows[nowSecond-61] + lastCleanedCutoff := e.lastCleanedCutoff + e.Unlock() + require.False(t, expiredExists) + require.Equal(t, nowSecond-60, lastCleanedCutoff) +} + func TestInMemoryRecordCleanupInactiveEntries(t *testing.T) { rl := &InMemoryRecord{} diff --git a/core/common/reqlimit/redis.go b/core/common/reqlimit/redis.go index c3a5f792..9d077628 100644 --- a/core/common/reqlimit/redis.go +++ b/core/common/reqlimit/redis.go @@ -52,9 +52,9 @@ const pushRequestLuaScript = ` local bucket_key = KEYS[1] local meta_key = KEYS[2] local window_seconds = tonumber(ARGV[1]) -local current_time = tonumber(ARGV[2]) -local max_requests = tonumber(ARGV[3]) -local n = tonumber(ARGV[4]) +local current_time = tonumber(redis.call('TIME')[1]) +local max_requests = tonumber(ARGV[2]) +local n = tonumber(ARGV[3]) local cutoff_slice = current_time - window_seconds local function parse_count(value) @@ -81,7 +81,8 @@ local count = tonumber(redis.call('HGET', meta_key, 'total_normal')) or 0 local over_count = tonumber(redis.call('HGET', meta_key, 'total_over')) or 0 local last_cleaned = tonumber(redis.call('HGET', meta_key, 'last_cleaned_second')) -if not last_cleaned then +if not last_cleaned or last_cleaned > cutoff_slice or + cutoff_slice - last_cleaned > window_seconds then last_cleaned = cutoff_slice local all_fields = redis.call('HGETALL', bucket_key) count = 0 @@ -131,9 +132,8 @@ return string.format("%d:%d:%d", count, over_count, current_second_count) const getRequestCountLuaScript = ` local exact_meta_key = KEYS[1] -local pattern = KEYS[2] local window_seconds = tonumber(ARGV[1]) -local current_time = tonumber(ARGV[2]) +local current_time = tonumber(redis.call('TIME')[1]) local cutoff_slice = current_time - window_seconds local function parse_count(value) @@ -143,11 +143,18 @@ local function parse_count(value) end local function cleanup_meta(bucket_key, meta_key) + local bucket_ttl = redis.call('PTTL', bucket_key) + if bucket_ttl == -2 then + redis.call('DEL', meta_key) + return 0, 0 + end + local count = tonumber(redis.call('HGET', meta_key, 'total_normal')) or 0 local over = tonumber(redis.call('HGET', meta_key, 'total_over')) or 0 local last_cleaned = tonumber(redis.call('HGET', meta_key, 'last_cleaned_second')) - if not last_cleaned then + if not last_cleaned or last_cleaned > cutoff_slice or + cutoff_slice - last_cleaned > window_seconds then last_cleaned = cutoff_slice local all_fields = redis.call('HGETALL', bucket_key) count = 0 @@ -185,32 +192,20 @@ local function cleanup_meta(bucket_key, meta_key) cutoff_slice ) - return count, over -end - -if exact_meta_key ~= '' then - local exact_bucket_key = string.gsub(exact_meta_key, ':meta$', ':buckets') - local count, over = cleanup_meta(exact_bucket_key, exact_meta_key) - local current_value = redis.call('HGET', exact_bucket_key, tostring(current_time)) - local current_c, current_oc = parse_count(current_value) - return string.format("%d:%d", count + over, current_c + current_oc) -end - -local total = 0 -local current_second_count = 0 - -local keys = redis.call('KEYS', pattern) -for _, meta_key in ipairs(keys) do - local bucket_key = string.gsub(meta_key, ':meta$', ':buckets') - local count, over = cleanup_meta(bucket_key, meta_key) - total = total + count + over + if bucket_ttl == -1 then + bucket_ttl = window_seconds * 1000 + redis.call('PEXPIRE', bucket_key, bucket_ttl) + end + redis.call('PEXPIRE', meta_key, bucket_ttl) - local current_value = redis.call('HGET', bucket_key, tostring(current_time)) - local current_c, current_oc = parse_count(current_value) - current_second_count = current_second_count + current_c + current_oc + return count, over end -return string.format("%d:%d", total, current_second_count) +local exact_bucket_key = string.gsub(exact_meta_key, ':meta$', ':buckets') +local count, over = cleanup_meta(exact_bucket_key, exact_meta_key) +local current_value = redis.call('HGET', exact_bucket_key, tostring(current_time)) +local current_c, current_oc = parse_count(current_value) +return string.format("%d:%d", count + over, current_c + current_oc) ` var ( @@ -240,20 +235,25 @@ func (r *redisRateRecord) GetRequest( return 0, 0, errors.New("redis client is nil") } - exactMetaKey := "" + if hasWildcard(keys) { + snapshots, err := r.SnapshotByPattern(ctx, duration, keys...) + if err != nil { + return 0, 0, err + } - pattern := r.buildMetaKey(keys...) - if !hasWildcard(keys) { - exactMetaKey = pattern - pattern = "" + for _, snapshot := range snapshots { + totalCount += snapshot.TotalCount + secondCount += snapshot.SecondCount + } + + return totalCount, secondCount, nil } result, err := getRequestCountScript.Run( ctx, rdb, - []string{exactMetaKey, pattern}, + []string{r.buildMetaKey(keys...)}, duration.Seconds(), - time.Now().Unix(), ).Text() if err != nil { return 0, 0, err @@ -297,7 +297,6 @@ func (r *redisRateRecord) PushRequest( rdb, []string{bucketKey, metaKey}, duration.Seconds(), - time.Now().Unix(), overed, n, ).Text() @@ -345,26 +344,24 @@ func (r *redisRateRecord) SnapshotByPattern( return nil, errors.New("redis client is nil") } - metaPattern := r.buildMetaKey(keys...) - if !hasWildcard(keys) { - metaPattern = r.buildMetaKey(keys...) - } - - pattern := metaPattern + pattern := r.buildMetaKey(keys...) iter := rdb.Scan(ctx, 0, pattern, 0).Iterator() - nowUnix := time.Now().Unix() windowSeconds := duration.Seconds() snapshots := make([]recordSnapshot, 0) + seenMetaKeys := make(map[string]struct{}) for iter.Next(ctx) { metaKey := iter.Val() + if _, seen := seenMetaKeys[metaKey]; seen { + continue + } + seenMetaKeys[metaKey] = struct{}{} result, err := getRequestCountScript.Run( ctx, rdb, - []string{metaKey, ""}, + []string{metaKey}, windowSeconds, - nowUnix, ).Text() if err != nil { return nil, err diff --git a/core/common/reqlimit/redis_integration_test.go b/core/common/reqlimit/redis_integration_test.go index e1211fae..e3513411 100644 --- a/core/common/reqlimit/redis_integration_test.go +++ b/core/common/reqlimit/redis_integration_test.go @@ -321,6 +321,136 @@ func TestRedisRateRecordGetDoesNotRefreshTTL(t *testing.T) { }, 4*time.Second, 100*time.Millisecond) } +func TestRedisRateRecordGetMissingDoesNotCreateMeta(t *testing.T) { + ctx := context.Background() + + redisClient, cleanup := setupRedisForReqLimitTest(t, ctx) + defer cleanup() + + record := newRedisChannelModelRecord(func() *redis.Client { return redisClient }) + bucketKey := record.buildBucketKey("missing-channel", "missing-model") + metaKey := record.buildMetaKey("missing-channel", "missing-model") + + totalCount, secondCount, err := record.GetRequest( + ctx, + time.Minute, + "missing-channel", + "missing-model", + ) + require.NoError(t, err) + require.Zero(t, totalCount) + require.Zero(t, secondCount) + + exists, err := redisClient.Exists(ctx, bucketKey, metaKey).Result() + require.NoError(t, err) + require.Zero(t, exists) +} + +func TestRedisRateRecordGetRemovesOrphanedMeta(t *testing.T) { + ctx := context.Background() + + redisClient, cleanup := setupRedisForReqLimitTest(t, ctx) + defer cleanup() + + record := newRedisChannelModelRecord(func() *redis.Client { return redisClient }) + metaKey := record.buildMetaKey("orphan-channel", "orphan-model") + require.NoError(t, redisClient.HSet( + ctx, + metaKey, + "total_normal", 0, + "total_over", 0, + "last_cleaned_second", 1, + ).Err()) + + totalCount, secondCount, err := record.GetRequest( + ctx, + time.Minute, + "orphan-channel", + "orphan-model", + ) + require.NoError(t, err) + require.Zero(t, totalCount) + require.Zero(t, secondCount) + + exists, err := redisClient.Exists(ctx, metaKey).Result() + require.NoError(t, err) + require.Zero(t, exists) +} + +func TestRedisRateRecordGetRebuildsStaleMeta(t *testing.T) { + ctx := context.Background() + + redisClient, cleanup := setupRedisForReqLimitTest(t, ctx) + defer cleanup() + + record := newRedisChannelModelRecord(func() *redis.Client { return redisClient }) + bucketKey := record.buildBucketKey("stale-channel", "stale-model") + metaKey := record.buildMetaKey("stale-channel", "stale-model") + now := time.Now().Unix() + + pipe := redisClient.Pipeline() + pipe.HSet(ctx, bucketKey, strconv.FormatInt(now, 10), "7:2") + pipe.Expire(ctx, bucketKey, time.Minute) + pipe.HSet( + ctx, + metaKey, + "total_normal", 100, + "total_over", 50, + "last_cleaned_second", 1, + ) + pipe.Expire(ctx, metaKey, time.Minute) + _, err := pipe.Exec(ctx) + require.NoError(t, err) + + totalCount, secondCount, err := record.GetRequest( + ctx, + time.Minute, + "stale-channel", + "stale-model", + ) + require.NoError(t, err) + require.Equal(t, int64(9), totalCount) + require.Equal(t, int64(9), secondCount) + + metaTTL, err := redisClient.TTL(ctx, metaKey).Result() + require.NoError(t, err) + require.Positive(t, metaTTL) +} + +func TestRedisRateRecordPushRebuildsStaleMeta(t *testing.T) { + ctx := context.Background() + + redisClient, cleanup := setupRedisForReqLimitTest(t, ctx) + defer cleanup() + + record := newRedisChannelModelRecord(func() *redis.Client { return redisClient }) + metaKey := record.buildMetaKey("stale-push-channel", "stale-push-model") + require.NoError(t, redisClient.HSet( + ctx, + metaKey, + "total_normal", 100, + "total_over", 50, + "last_cleaned_second", 1, + ).Err()) + + normalCount, overCount, secondCount, err := record.PushRequest( + ctx, + 0, + time.Minute, + 1, + "stale-push-channel", + "stale-push-model", + ) + require.NoError(t, err) + require.Equal(t, int64(1), normalCount) + require.Zero(t, overCount) + require.Equal(t, int64(1), secondCount) + + metaTTL, err := redisClient.TTL(ctx, metaKey).Result() + require.NoError(t, err) + require.Positive(t, metaTTL) +} + func setupRedisForReqLimitTest(t *testing.T, ctx context.Context) (*redis.Client, func()) { t.Helper() From d97da9833f43de61c2f182cb9f2c80bc33ab5306 Mon Sep 17 00:00:00 2001 From: zijiren233 Date: Mon, 20 Jul 2026 00:06:58 +0800 Subject: [PATCH 2/2] fix: ci lint --- core/common/reqlimit/redis.go | 1 + 1 file changed, 1 insertion(+) diff --git a/core/common/reqlimit/redis.go b/core/common/reqlimit/redis.go index 9d077628..8837fa52 100644 --- a/core/common/reqlimit/redis.go +++ b/core/common/reqlimit/redis.go @@ -355,6 +355,7 @@ func (r *redisRateRecord) SnapshotByPattern( if _, seen := seenMetaKeys[metaKey]; seen { continue } + seenMetaKeys[metaKey] = struct{}{} result, err := getRequestCountScript.Run(