Redis和Lua實(shí)現(xiàn)分布式限流器的方法詳解

-- 計(jì)數(shù)器限流
-- 此處支持的最小單位時(shí)間是秒, 若將 expire 改成 pexpire 則可支持毫秒粒度.
-- KEYS[1] string 限流的key
-- ARGV[1] int  限流數(shù)
-- ARGV[2] int  單位時(shí)間(秒)

local cnt = tonumber(redis.call("incr", KEYS[1]))

if (cnt == 1) then
 -- cnt 值為1說(shuō)明之前不存在該值, 因此需要設(shè)置其過(guò)期時(shí)間
 redis.call("expire", KEYS[1], tonumber(ARGV[2]))
elseif (cnt > tonumber(ARGV[1])) then
 return -1
end 

return cnt

-- 令牌桶限流: 不支持預(yù)消費(fèi), 初始桶是滿的
-- KEYS[1] string 限流的key

-- ARGV[1] int  桶最大容量
-- ARGV[2] int  每次添加令牌數(shù)
-- ARGV[3] int  令牌添加間隔(秒)
-- ARGV[4] int  當(dāng)前時(shí)間戳

local bucket_capacity = tonumber(ARGV[1])
local add_token = tonumber(ARGV[2])
local add_interval = tonumber(ARGV[3])
local now = tonumber(ARGV[4])

-- 保存上一次更新桶的時(shí)間的key
local LAST_TIME_KEY = KEYS[1].."_time";   
-- 獲取當(dāng)前桶中令牌數(shù)
local token_cnt = redis.call("get", KEYS[1]) 
-- 桶完全恢復(fù)需要的最大時(shí)長(zhǎng)
local reset_time = math.ceil(bucket_capacity / add_token) * add_interval;

if token_cnt then -- 令牌桶存在
 -- 上一次更新桶的時(shí)間
 local last_time = redis.call('get', LAST_TIME_KEY)
 -- 恢復(fù)倍數(shù)
 local multiple = math.floor((now - last_time) / add_interval)
 -- 恢復(fù)令牌數(shù)
 local recovery_cnt = multiple * add_token
 -- 確保不超過(guò)桶容量
 local token_cnt = math.min(bucket_capacity, token_cnt + recovery_cnt) - 1
 
 if token_cnt  0 then
  return -1;
 end
 
 -- 重新設(shè)置過(guò)期時(shí)間, 避免key過(guò)期
 redis.call('set', KEYS[1], token_cnt, 'EX', reset_time)      
 redis.call('set', LAST_TIME_KEY, last_time + multiple * add_interval, 'EX', reset_time)
 return token_cnt
 
else -- 令牌桶不存在
 token_cnt = bucket_capacity - 1
 -- 設(shè)置過(guò)期時(shí)間避免key一直存在
 redis.call('set', KEYS[1], token_cnt, 'EX', reset_time);
 redis.call('set', LAST_TIME_KEY, now, 'EX', reset_time + 1); 
 return token_cnt 
end