-- 定義返回值，是個數(shù)組，包含：是否觸發(fā)限流（1限流 0通過）、當(dāng)前桶中的令牌數(shù)
local ret={}
ret[1]=0
-- Redis集群分片Key，KEYS[1]是限流目標
local cl_key = '{' .. KEYS[1] .. '}'

-- 獲取限流懲罰的當(dāng)前設(shè)置，觸發(fā)限流懲罰時會寫一個有過期時間的KV
-- 如果存在限流懲罰，則返回結(jié)果[1,-1]
local lock_key=cl_key .. '-lock'
local lock_val=redis.call('get',lock_key)
if lock_val == '1' then
    ret[1]=1
    ret[2]=-1
    return ret;
end

-- 這里省略部分代碼

-- 獲取[上次向桶中投放令牌的時間]，如果沒有設(shè)置過這個投放時間，則令牌桶也不存在，此時：
-- 一種情況是：首次執(zhí)行，此時定義令牌桶就是滿的。
-- 另一種情況是：較長時間沒有執(zhí)行過限流處理，導(dǎo)致承載這個時間的KV被釋放了，
-- 這個過期時間會超過自然投放令牌到桶中直到桶滿的時間，所以令牌桶也應(yīng)該是滿的。
local last_time=redis.call('get',st_key)
if(last_time==false)
then
 -- 本次執(zhí)行后剩余令牌數(shù)量：桶的容量- 本次執(zhí)行消耗的令牌數(shù)量
    bucket_amount = capacity - amount;
    -- 將這個令牌數(shù)量更新到令牌桶中，同時這里有個過期時間，如果長時間不執(zhí)行這個程序，令牌桶KV會被回收
    redis.call('set',KEYS[1],bucket_amount,'PX',key_expire_time)
    -- 設(shè)置[上次向桶中放入令牌的時間]，后邊計算應(yīng)放入桶中的令牌數(shù)量時會用到
    redis.call('set',st_key,start_time,'PX',key_expire_time)
    -- 返回值[當(dāng)前桶中的令牌數(shù)]
    ret[2]=bucket_amount
    -- 無需其它處理
    return ret
end

-- 令牌桶存在，獲取令牌桶中的當(dāng)前令牌數(shù)
local current_value = redis.call('get',KEYS[1])
current_value = tonumber(current_value)

-- 判斷是不是該放入新令牌到桶中了：當(dāng)前時間-上次投放的時間 >= 投放的時間間隔
last_time=tonumber(last_time)
local last_time_changed=0
local past_time=current_time-last_time
if(past_time<inflow_unit)
then
 -- 不到投放的時候，直接從令牌桶中取走令牌
    bucket_amount=current_value-amount
else
 -- 需要放入一些令牌, 預(yù)計投放數(shù)量 = (距上次投放過去的時間/投放的時間間隔)*每單位時間投放的數(shù)量
    local past_inflow_unit_quantity = past_time/inflow_unit
    past_inflow_unit_quantity=math.floor(past_inflow_unit_quantity)
    last_time=last_time+past_inflow_unit_quantity*inflow_unit
    last_time_changed=1
    local past_inflow_quantity=past_inflow_unit_quantity*inflow_quantity_per_unit
    bucket_amount=current_value+past_inflow_quantity-amount
end

-- 這里省略部分代碼

ret[2]=bucket_amount

-- 如果桶中剩余數(shù)量小于0，則看看是否需要限流懲罰，如果需要則寫入一個懲罰KV，過期時間為懲罰的秒數(shù)
if(bucket_amount<0)
then
    if lock_seconds>0 then
        redis.call('set',lock_key,'1','EX',lock_seconds,'NX')
    end
    ret[1]=1
    return ret
end

-- 來到這里，代表可以成功扣減令牌，則需要更新令牌桶KV
if last_time_changed==1 then
    redis.call('set',KEYS[1],bucket_amount,'PX',key_expire_time)
 -- 有新投放，更新[上次投放時間]為本次投放時間
    redis.call('set',st_key,last_time,'PX',key_expire_time)
else
    redis.call('set',KEYS[1],bucket_amount,'PX',key_expire_time)
end
return ret