Skip to content

Commit 67c0506

Browse files
committed
fix(billing): add window expiration check to Redis rate limit Lua script
The updateRateLimitUsageScript Lua script previously performed unconditional HINCRBYFLOAT on all usage counters without checking whether the rate limit window had expired. This caused usage to accumulate across window boundaries in Redis while the DB correctly reset on expiration, leading to incorrect 429 rate limiting that could persist for up to 24 hours. The Lua script now checks each window timestamp before incrementing: - If the window has expired, usage is reset to the current cost and the window timestamp is updated (matching DB-side semantics) - If the window is still valid, usage is accumulated normally This also resolves the async race condition where stale HINCRBYFLOAT tasks from the worker queue could pollute a freshly rebuilt cache after invalidation, since the script now self-corrects expired windows. Closes Wei-Shaw#1049
1 parent 6447be4 commit 67c0506

1 file changed

Lines changed: 42 additions & 6 deletions

File tree

backend/internal/repository/billing_cache.go

Lines changed: 42 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,11 @@ const (
2020
billingCacheTTL = 5 * time.Minute
2121
billingCacheJitter = 30 * time.Second
2222
rateLimitCacheTTL = 7 * 24 * time.Hour // 7 days matches the longest window
23+
24+
// Rate limit window durations — must match service.RateLimitWindow* constants.
25+
rateLimitWindow5h = 5 * time.Hour
26+
rateLimitWindow1d = 24 * time.Hour
27+
rateLimitWindow7d = 7 * 24 * time.Hour
2328
)
2429

2530
// jitteredTTL 返回带随机抖动的 TTL,防止缓存雪崩
@@ -90,17 +95,40 @@ var (
9095
return 1
9196
`)
9297

93-
// updateRateLimitUsageScript atomically increments all three rate limit usage counters.
94-
// Returns 0 if the key doesn't exist (cache miss), 1 on success.
98+
// updateRateLimitUsageScript atomically increments all three rate limit usage counters
99+
// with window expiration checking. If a window has expired, its usage is reset to cost
100+
// (instead of accumulated) and the window timestamp is updated, matching the DB-side
101+
// IncrementRateLimitUsage semantics.
102+
//
103+
// ARGV: [1]=cost, [2]=ttl_seconds, [3]=now_unix, [4]=window_5h_seconds, [5]=window_1d_seconds, [6]=window_7d_seconds
95104
updateRateLimitUsageScript = redis.NewScript(`
96105
local exists = redis.call('EXISTS', KEYS[1])
97106
if exists == 0 then
98107
return 0
99108
end
100109
local cost = tonumber(ARGV[1])
101-
redis.call('HINCRBYFLOAT', KEYS[1], 'usage_5h', cost)
102-
redis.call('HINCRBYFLOAT', KEYS[1], 'usage_1d', cost)
103-
redis.call('HINCRBYFLOAT', KEYS[1], 'usage_7d', cost)
110+
local now = tonumber(ARGV[3])
111+
local win5h = tonumber(ARGV[4])
112+
local win1d = tonumber(ARGV[5])
113+
local win7d = tonumber(ARGV[6])
114+
115+
-- Helper: check if window is expired and update usage + window accordingly
116+
-- Returns nothing, modifies the hash in-place.
117+
local function update_window(usage_field, window_field, window_duration)
118+
local w = tonumber(redis.call('HGET', KEYS[1], window_field) or 0)
119+
if w == 0 or (now - w) >= window_duration then
120+
-- Window expired or never started: reset usage to cost, start new window
121+
redis.call('HSET', KEYS[1], usage_field, tostring(cost))
122+
redis.call('HSET', KEYS[1], window_field, tostring(now))
123+
else
124+
-- Window still valid: accumulate
125+
redis.call('HINCRBYFLOAT', KEYS[1], usage_field, cost)
126+
end
127+
end
128+
129+
update_window('usage_5h', 'window_5h', win5h)
130+
update_window('usage_1d', 'window_1d', win1d)
131+
update_window('usage_7d', 'window_7d', win7d)
104132
redis.call('EXPIRE', KEYS[1], ARGV[2])
105133
return 1
106134
`)
@@ -280,7 +308,15 @@ func (c *billingCache) SetAPIKeyRateLimit(ctx context.Context, keyID int64, data
280308

281309
func (c *billingCache) UpdateAPIKeyRateLimitUsage(ctx context.Context, keyID int64, cost float64) error {
282310
key := billingRateLimitKey(keyID)
283-
_, err := updateRateLimitUsageScript.Run(ctx, c.rdb, []string{key}, cost, int(rateLimitCacheTTL.Seconds())).Result()
311+
now := time.Now().Unix()
312+
_, err := updateRateLimitUsageScript.Run(ctx, c.rdb, []string{key},
313+
cost,
314+
int(rateLimitCacheTTL.Seconds()),
315+
now,
316+
int(rateLimitWindow5h.Seconds()),
317+
int(rateLimitWindow1d.Seconds()),
318+
int(rateLimitWindow7d.Seconds()),
319+
).Result()
284320
if err != nil && !errors.Is(err, redis.Nil) {
285321
log.Printf("Warning: update rate limit usage cache failed for api key %d: %v", keyID, err)
286322
return err

0 commit comments

Comments
 (0)