【重复造轮子】基于Redis的RateLimiter

【重复造轮子】基于Redis的RateLimiter

造轮子的过程也是学习的过程。

如果公司的 Redis 不支持发布订阅指令的话，是没法用 Redisson 的，因为 Redisson 的大部分功能都依赖于 Redis 的发布订阅指令。

下面这段 lua 代码是基于 Redis 的 RateLimiter 的实现：

local maxToken = tonumber(ARGV[1])      -- 单位时间内的最大 token 数
local intervalTime = tonumber(ARGV[2])  -- 生成一个 token 的间隔毫秒数
local curTime = tonumber(ARGV[3])       -- 当前时间的毫秒数
local remainToken                       -- 当前可用的 token 数

local latestTime = redis.call("HGET", KEYS[1], "latestTime")    -- 上次生成 token 的时间

-- 将当前时间对齐到 intervalTime
-- 比如 intervalTime 是 200，当前时间是 1662444699797，则对齐后的时间就是 1662444699600
local alignTime = curTime - (curTime % intervalTime)

if latestTime then
    -- 获取剩余可以用 token 数
    remainToken = tonumber(redis.call("HGET", KEYS[1], "remainToken"))

    -- 计算与上次生成 token 的时间差，是否允许再次生成一个 token
    local diffTime = alignTime - tonumber(latestTime)

    -- 与上次生成时间相差不到 1秒
    if diffTime < 1000 then
        -- 计算可以生成几个 token
        local newTokenNum = math.floor(diffTime / intervalTime)
        -- 将新生成的 token 累加到剩余 token 中，然后更新生成 token 的时间
        if newTokenNum > 0 then
            remainToken = remainToken + newTokenNum
            redis.call("HSET", KEYS[1], "latestTime", alignTime)
        end
    elseif diffTime == 1000 then
        -- 新的一秒的开始
        remainToken = 1
        redis.call("HSET", KEYS[1], "latestTime", alignTime)
    else
        -- 计算到整秒时间戳的差
        diffTime = alignTime - (curTime - (curTime % 1000))
        local newTokenNum = math.floor(diffTime / intervalTime)
        if newTokenNum > 0 then
            remainToken = remainToken + newTokenNum
            redis.call("HSET", KEYS[1], "latestTime", alignTime)
        end
     end
else
    -- 第一次获取，先放入一个 token，然后初始化刷新时间
    remainToken = 1
    redis.call("HSET", KEYS[1], "latestTime", alignTime)
end
if remainToken <= 0 then
    redis.call("HSET", KEYS[1], "remainToken", 0)
    return 0
else
    redis.call("HSET", KEYS[1], "remainToken", remainToken - 1)
    return 1
end

有任何问题欢迎提出。