mirror of
https://github.com/zeromicro/go-zero.git
synced 2025-01-23 09:00:20 +08:00
31 lines
932 B
Lua
31 lines
932 B
Lua
-- to be compatible with aliyun redis, we cannot use `local key = KEYS[1]` to reuse the key
|
|
-- KEYS[1] as tokens_key
|
|
-- KEYS[2] as timestamp_key
|
|
local rate = tonumber(ARGV[1])
|
|
local capacity = tonumber(ARGV[2])
|
|
local now = tonumber(ARGV[3])
|
|
local requested = tonumber(ARGV[4])
|
|
local fill_time = capacity/rate
|
|
local ttl = math.floor(fill_time*2)
|
|
local last_tokens = tonumber(redis.call("get", KEYS[1]))
|
|
if last_tokens == nil then
|
|
last_tokens = capacity
|
|
end
|
|
|
|
local last_refreshed = tonumber(redis.call("get", KEYS[2]))
|
|
if last_refreshed == nil then
|
|
last_refreshed = 0
|
|
end
|
|
|
|
local delta = math.max(0, now-last_refreshed)
|
|
local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
|
|
local allowed = filled_tokens >= requested
|
|
local new_tokens = filled_tokens
|
|
if allowed then
|
|
new_tokens = filled_tokens - requested
|
|
end
|
|
|
|
redis.call("setex", KEYS[1], ttl, new_tokens)
|
|
redis.call("setex", KEYS[2], ttl, now)
|
|
|
|
return allowed |