Files
stat_center/common/redis/limit.lua
T
2026-05-20 16:13:58 +08:00

90 lines
2.5 KiB
Lua

local CAPACITY = "ca"
local RATE = "ra"
local SIZE = "si"
local GENERATE_AT = "ge_at"
local CREATED_AT = "cr_at"
local UPDATED_AT = "up_at"
local PASSED_COUNT = "pa_co"
local REJECT_COUNT = "re_co"
local PASSED_HOURS = "pa_ho"
local REJECT_HOURS = "re_ho"
local SUCCESS = "success"
local FAILED = "failed"
local sKey = tostring(KEYS[1] or "")
local sName = tostring(KEYS[2] or "")
local iCapacity = tonumber(KEYS[3] or 0)
local iRate = tonumber(KEYS[4] or 0)
local iCount = tonumber(KEYS[5] or 0)
local iNow = tonumber(KEYS[6] or 0)
local sHour = tostring(KEYS[7] or "")
if sKey == "" or sName == "" or sHour == "" then
return FAILED .. " param error1"
end
if iCapacity <= 0 or iRate <= 0 or iCount <= 0 or iNow <= 0 then
return FAILED .. "param error2"
end
local bucket = {}
local data = redis.pcall("HGET", sKey, sName)
local init = false
if data == nil or data == false or data == "" then
init = true
else
bucket = cjson.decode(data)
if bucket == nil then
init = true
end
end
if init then
bucket[CAPACITY] = iCapacity
bucket[RATE] = iRate
bucket[SIZE] = iCapacity
bucket[GENERATE_AT] = iNow
bucket[CREATED_AT] = iNow
bucket[UPDATED_AT] = iNow
bucket[PASSED_COUNT] = 0
bucket[REJECT_COUNT] = 0
bucket[PASSED_HOURS] = {}
bucket[REJECT_HOURS] = {}
else
bucket[CAPACITY] = iCapacity
bucket[RATE] = iRate
end
local bk_size = tonumber(bucket[SIZE] or 0)
local bk_capacity = tonumber(bucket[CAPACITY] or 0)
if bk_size < bk_capacity then
local bk_last = tonumber(bucket[GENERATE_AT] or 0)
local dura = iNow - bk_last
local incr = math.floor(dura * iRate / 1000)
if incr >= 1 then
bk_size = math.min(bk_size + incr, bk_capacity)
bucket[SIZE] = bk_size
bucket[GENERATE_AT] = iNow
end
end
local message = FAILED
if iCount <= bk_size then
bucket[SIZE] = bk_size - iCount
bucket[PASSED_COUNT] = tonumber(bucket[PASSED_COUNT] or 0) + 1
local hours = bucket[PASSED_HOURS]
if hours == nil then
hours = {}
end
hours[sHour] = tonumber(hours[sHour] or 0) + 1
bucket[PASSED_HOURS] = hours
message = SUCCESS
else
bucket[REJECT_COUNT] = tonumber(bucket[REJECT_COUNT] or 0) + 1
local hours = bucket[REJECT_HOURS]
if hours == nil then
hours = {}
end
hours[sHour] = tonumber(hours[sHour] or 0) + 1
bucket[REJECT_HOURS] = hours
message = FAILED
end
bucket[UPDATED_AT] = iNow
data = cjson.encode(bucket)
redis.pcall("HSET", sKey, sName, data)
return message