redis-lua 限流脚本

参考：https://gist.github.com/ptarjan/e38f45f2dfe601419ca3af937fff574d

令牌桶算法(token bucket)

local tokens_key = KEYS[1] -- token键
local timestamp_key = KEYS[2] -- 上次更新时间戳

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) -- 填充时间的两倍，避免过期

-- 获取tokens，过期则相当于充满
local last_tokens = tonumber(redis.call("GET", tokens_key))
if last_tokens == nil then
  last_tokens = capacity
end

local last_refreshed = tonumber(redis.call("GET", timestamp_key))
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", tokens_key, ttl, new_tokens)
redis.call("SETEX", timestamp_key, ttl, now)

-- new_token是消耗令牌后的tokens
return { allowed, new_tokens }

并行限制 ZSET实现

并行限制可以控制同时并行的请求数，使用zset实现：

1. 请求来到时，ZCARD查看count数，如果小于则ZADD添加，否则拒绝。
2. 请求结束时，需要调用ZREM+id删除这个member

local key = KEYS[1]

local capacity = tonumber(ARGV[1]) -- 最大请求数
local timestamp = tonumber(ARGV[2]) -- 当前时间戳
local id = ARGV[3] -- 请求标识

local count = redis.call("ZCARD", key)
local allowed = count < capacity

if allowed then
  redis.call("ZADD", key, timestamp, id)
end

return { allowed, count }

定点刷新

还可以实现定点刷新脚本：假定每天凌晨四点刷新

利用redis过期机制刷新，先计算现在与下一个刷新点的时间差，检查key是否存在，如果不存在则SET

存在则DECRBY这个key，那么每天到点后，key过期就会自动刷新值。