Reason: The original steps 1 and 3 in the redisRateLimitHandler method were not atomic, leading to poor precision under high concurrent requests. For example, with a rate limit set to 60, sending 200 concurrent requests would result in none being blocked, whereas theoretically around 140 should be intercepted. Solution: I chose not to merge steps 1 and 3 into a single Lua script because a single atomic operation involving read, write, and delete operations could suffer from performance issues under high concurrency. Instead, I implemented a token bucket algorithm to optimize this, reducing the atomic operation to just read and write steps while significantly decreasing the memory footprint.
44 lines
1.2 KiB
Lua
44 lines
1.2 KiB
Lua
-- 令牌桶限流器
|
||
-- KEYS[1]: 限流器唯一标识
|
||
-- ARGV[1]: 请求令牌数 (通常为1)
|
||
-- ARGV[2]: 令牌生成速率 (每秒)
|
||
-- ARGV[3]: 桶容量
|
||
|
||
local key = KEYS[1]
|
||
local requested = tonumber(ARGV[1])
|
||
local rate = tonumber(ARGV[2])
|
||
local capacity = tonumber(ARGV[3])
|
||
|
||
-- 获取当前时间(Redis服务器时间)
|
||
local now = redis.call('TIME')
|
||
local nowInSeconds = tonumber(now[1])
|
||
|
||
-- 获取桶状态
|
||
local bucket = redis.call('HMGET', key, 'tokens', 'last_time')
|
||
local tokens = tonumber(bucket[1])
|
||
local last_time = tonumber(bucket[2])
|
||
|
||
-- 初始化桶(首次请求或过期)
|
||
if not tokens or not last_time then
|
||
tokens = capacity
|
||
last_time = nowInSeconds
|
||
else
|
||
-- 计算新增令牌
|
||
local elapsed = nowInSeconds - last_time
|
||
local add_tokens = elapsed * rate
|
||
tokens = math.min(capacity, tokens + add_tokens)
|
||
last_time = nowInSeconds
|
||
end
|
||
|
||
-- 判断是否允许请求
|
||
local allowed = false
|
||
if tokens >= requested then
|
||
tokens = tokens - requested
|
||
allowed = true
|
||
end
|
||
|
||
---- 更新桶状态并设置过期时间
|
||
redis.call('HMSET', key, 'tokens', tokens, 'last_time', last_time)
|
||
--redis.call('EXPIRE', key, math.ceil(capacity / rate) + 60) -- 适当延长过期时间
|
||
|
||
return allowed and 1 or 0 |