【Redis Lua脚本 分布式限流 窗口 令牌桶 漏桶】

前言

本文提供几个可在生产上使用的Redis+Lua分布式限流实现方案。

• 固定窗口限流
• 滑动串口限流
• 令牌桶限流

---

一、固定窗口限流

使用Redis String对象自增实现固定窗口算法限流。

Lua脚本

local key = KEYS[1] -- 限流资源
local limitCount = ARGV[1] -- 限流请求数
local limitTime = ARGV[2] -- 限流时间
local currentCount = redis.call('get', key) -- 当前请求数
-- 限流存在并且超过限流大小，则返回剩余可用请求数=0
if (currentCount and tonumber(currentCount) >= tonumber(limitCount)) then
    return 0
end
-- 请求数自增
currentCount = redis.call('incr', key)
-- 第一次请求，则设置过期时间
if (tonumber(currentCount) == 1) then
    redis.call('expire', key, limitTime)
end
-- 返回剩余可用请求数
return tonumber(limitCount) - tonumber(currentCount)

演示效果

能正常限流

二、滑动窗口限流

使用Redis ZSet对象，score控制时间窗口，实现活动窗口自算法限流。

Lua脚本

local key = KEYS[1]
local limitCount = ARGV[1] -- 限流请求数
local startTime = ARGV[2] -- 限流开始时间戳
local endTime = ARGV[3] -- 限流结束时间戳
local timeout = ARGV[4] -- 限流超时时间-用于清除内存-毫秒
local holdMember= endTime -- 占位member
local currentCount = redis.call('zcount', key, startTime, endTime) -- 当前请求数
-- 限流存在并且超过限流大小，则返回剩余可用请求数=0
if (currentCount and tonumber(currentCount) >= tonumber(limitCount)) then
    return 0
end
-- 记录本次请求
redis.call('zadd', key, endTime, holdMember)
-- 设置超时时间
redis.call('expire', key, timeout)
-- 返回剩余可用请求数
return tonumber(limitCount) - tonumber(currentCount)

演示效果

能正常限流

三、令牌桶限流

使用Redis Hash对象可以实现令牌桶。
Hash对象需要存储维护，‘最近更新时间’、‘桶内令牌数’。节省内存，Lua脚本稍微复杂。

Lua脚本

local key = KEYS[1] -- 限流资源
local intervalPerTokenTime = ARGV[1] -- 生成单个令牌的间隔-毫秒
local currentTime = ARGV[2]
local initTokenCount = ARGV[3] -- 令牌桶初始令牌数-可选
local limitTokenCount = ARGV[4] -- 令牌桶上限令牌数
local timeout = ARGV[5] -- 令牌桶过期时间

local durationTime -- 距离上次更新令牌桶时间
local durationTimeAvailableTokenCount -- durationTime可分配令牌数

-- 当前key的value=令牌桶对象
local bucketExists = redis.call('exists', key) == 1
local bucketTokenCount_field_name = 'bucketTokenCount' -- field name
local lastUpdateTime_field_name = 'lastUpdateTime' -- field name
local bucketTokenCount = redis.call('hget', key, bucketTokenCount_field_name)
local lastUpdateTime = redis.call('hget', key, lastUpdateTime_field_name)

-- 获取当前令牌桶对象
bucket = redis.call('hgetall', key)
-- 令牌桶初始化
if (table.maxn(bucket) == 0) then
    bucketTokenCount = initTokenCount
    redis.call('hset', key, lastUpdateTime, currentTime)
    redis.call('hset', key, bucketTokenCount_field_name, currentTokenCount)
    redis.call('expire', key, timeout)
    -- 返回令牌数
    return math.max(1, bucketTokenCount)
-- 令牌桶已存在，则处理放置令牌
else
    -- 计算该生成多少令牌
    durationTime = currentTime - lastUpdateTime
    durationTimeAvailableTokenCount = math.floor(durationTime / intervalPerTokenTime) -- 向下取整
    if (durationTimeAvailableTokenCount > 0) then
        -- 计算此次的lastUpdateTime，不能取currentTime，因为lastUpdateTime要对齐，下次计算生成令牌才准确
        lastUpdateTime = lastUpdateTime + (durationTimeAvailableTokenCount * intervalPerTokenTime)
        redis.call('hset', key, lastUpdateTime_field_name, lastUpdateTime)
    end
    -- 放入令牌
    bucketTokenCount = bucketTokenCount + durationTimeAvailableTokenCount
    local currentTokenCount = bucketTokenCount - 1
    -- 令牌桶不能超过容量，不能少于0
    currentTokenCount = math.min(currentTokenCount, limitTokenCount)
    currentTokenCount = math.max(currentTokenCount, 0)
    redis.call('hset', key, bucketTokenCount_field_name, currentTokenCount)
    -- 返回可用令牌数
    return bucketTokenCount
end

演示效果

能正常限流

---

总结

以上都是生产中可以直接使用的Redis Lua分布式限流实现。