mirror of
https://github.com/zeromicro/go-zero.git
synced 2026-05-10 08:29:58 +08:00
feat: migrate lua script to lua file (#4069)
This commit is contained in:
31
core/limit/tokenscript.lua
Normal file
31
core/limit/tokenscript.lua
Normal file
@@ -0,0 +1,31 @@
|
||||
-- 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
|
||||
Reference in New Issue
Block a user