SpringCloud、Soul网关中的令牌桶限流算法-lua

local tokens_key = KEYS[1] --桶的key
local timestamp_key = KEYS[2] --计时器的key
--redis.log(redis.LOG_WARNING, "tokens_key " .. tokens_key)

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) --过期时间

--redis.log(redis.LOG_WARNING, "rate " .. ARGV[1])
--redis.log(redis.LOG_WARNING, "capacity " .. ARGV[2])
--redis.log(redis.LOG_WARNING, "now " .. ARGV[3])
--redis.log(redis.LOG_WARNING, "requested " .. ARGV[4])
--redis.log(redis.LOG_WARNING, "filltime " .. fill_time)
--redis.log(redis.LOG_WARNING, "ttl " .. ttl)

local last_tokens = tonumber(redis.call("get", tokens_key)) --获取当前桶中令牌的个数
if last_tokens == nil then
  last_tokens = capacity --如果不存在则初始化令牌为当前的桶的容量
end
--redis.log(redis.LOG_WARNING, "last_tokens " .. last_tokens)

local last_refreshed = tonumber(redis.call("get", timestamp_key)) --获取上次获取时间
if last_refreshed == nil then
  last_refreshed = 0 --没有则初始化
end
--redis.log(redis.LOG_WARNING, "last_refreshed " .. last_refreshed)

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 --最新的令牌个数
local allowed_num = 0
if allowed then
  new_tokens = filled_tokens - requested --最新的令牌个数-当前请求个数
  allowed_num = 1 --设置允许请求
end

--redis.log(redis.LOG_WARNING, "delta " .. delta)
--redis.log(redis.LOG_WARNING, "filled_tokens " .. filled_tokens)
--redis.log(redis.LOG_WARNING, "allowed_num " .. allowed_num)
--redis.log(redis.LOG_WARNING, "new_tokens " .. new_tokens)

redis.call("setex", tokens_key, ttl, new_tokens) --重置桶中的令牌个数
redis.call("setex", timestamp_key, ttl, now) -- 重置请求时间

return { allowed_num, new_tokens } --返回结果