快捷導(dǎo)航

使用Redis實(shí)現(xiàn)令牌桶算法原理解析

更新時(shí)間：2021年12月08日 08:51:33 作者：波斯馬

這篇文章主要介紹了使用Redis實(shí)現(xiàn)令牌桶算法,該算法可以應(yīng)對(duì)短暫的突發(fā)流量，這對(duì)于現(xiàn)實(shí)環(huán)境中流量不怎么均勻的情況特別有用，不會(huì)頻繁的觸發(fā)限流，對(duì)調(diào)用方比較友好，需要的朋友可以參考下

在限流算法中有一種令牌桶算法，該算法可以應(yīng)對(duì)短暫的突發(fā)流量，這對(duì)于現(xiàn)實(shí)環(huán)境中流量不怎么均勻的情況特別有用，不會(huì)頻繁的觸發(fā)限流，對(duì)調(diào)用方比較友好。

例如，當(dāng)前限制10qps，大多數(shù)情況下不會(huì)超過(guò)此數(shù)量，但偶爾會(huì)達(dá)到30qps，然后很快就會(huì)恢復(fù)正常，假設(shè)這種突發(fā)流量不會(huì)對(duì)系統(tǒng)穩(wěn)定性產(chǎn)生影響，我們可以在一定程度上允許這種瞬時(shí)突發(fā)流量，從而為用戶(hù)帶來(lái)更好的可用性體驗(yàn)。這就是使用令牌桶算法的地方。

令牌桶算法原理

如下圖所示，該算法的基本原理是：有一個(gè)容量為X的令牌桶，每Y單位時(shí)間內(nèi)將Z個(gè)令牌放入該桶。如果桶中的令牌數(shù)量超過(guò)X，那么它將被丟棄。處理請(qǐng)求時(shí)，需要先從令牌桶中取出令牌，如果拿到了令牌，則繼續(xù)處理；如果拿不到令牌，則拒絕請(qǐng)求。

可以看出，在令牌桶算法中設(shè)置X，Y和Z的數(shù)量尤為重要。Z應(yīng)該比每Y單位時(shí)間內(nèi)的請(qǐng)求數(shù)稍大，系統(tǒng)將長(zhǎng)時(shí)間處于此狀態(tài)；X是系統(tǒng)允許的瞬時(shí)最大請(qǐng)求數(shù)，并且系統(tǒng)不應(yīng)該長(zhǎng)時(shí)間處于此狀態(tài)，否則就會(huì)頻繁觸發(fā)限流，此時(shí)表明流量出現(xiàn)了超預(yù)期的情況，需要及時(shí)調(diào)查原因并采取相應(yīng)措施。

Redis實(shí)現(xiàn)令牌桶算法

之前看過(guò)有些程序?qū)崿F(xiàn)的令牌桶，其向桶中放入令牌的方法是啟動(dòng)一個(gè)線程，每隔Y單位時(shí)間增加一次令牌數(shù)量，或者在Timer中定時(shí)執(zhí)行這一過(guò)程。我不太滿(mǎn)意這種方法，原因有二，一是浪費(fèi)線程資源，二是因?yàn)檎{(diào)度的問(wèn)題執(zhí)行時(shí)間不精確。

這里確定令牌桶中令牌數(shù)量的方法是通過(guò)計(jì)算得出，首先算出從上次請(qǐng)求到這次請(qǐng)求經(jīng)過(guò)了多長(zhǎng)時(shí)間，是否達(dá)到發(fā)令牌的時(shí)間閾值，然后增加的令牌數(shù)是多少，這些令牌能夠放到桶中的是多少。

Talk is cheap!

下邊就來(lái)看看Redis中怎么實(shí)現(xiàn)的，因?yàn)樯婕暗蕉啻闻cRedis的交互，這里為了提高限流處理的吞吐量，減少程序與Redis的交互次數(shù)，采用了Redis支持的Lua script，Lua script的執(zhí)行是原子的，所以也不用擔(dān)心出現(xiàn)臟數(shù)據(jù)的問(wèn)題。

代碼節(jié)選自 FireflySoft.RateLimit ，它不僅支持普通主從部署Redis，還支持集群Redis，所以吞吐量可以通過(guò)水平擴(kuò)展的方式進(jìn)行提升。為了方便閱讀，這里增加一些注釋?zhuān)瑢?shí)際是沒(méi)有的。

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

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

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

-- 判斷是不是該放入新令牌到桶中了：當(dāng)前時(shí)間-上次投放的時(shí)間 >= 投放的時(shí)間間隔
last_time=tonumber(last_time)
local last_time_changed=0
local past_time=current_time-last_time
if(past_time<inflow_unit)
then
 -- 不到投放的時(shí)候，直接從令牌桶中取走令牌
    bucket_amount=current_value-amount
else
 -- 需要放入一些令牌, 預(yù)計(jì)投放數(shù)量 = (距上次投放過(guò)去的時(shí)間/投放的時(shí)間間隔)*每單位時(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，則看看是否需要限流懲罰，如果需要?jiǎng)t寫(xiě)入一個(gè)懲罰KV，過(guò)期時(shí)間為懲罰的秒數(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

-- 來(lái)到這里，代表可以成功扣減令牌，則需要更新令牌桶KV
if last_time_changed==1 then
    redis.call('set',KEYS[1],bucket_amount,'PX',key_expire_time)
 -- 有新投放，更新[上次投放時(shí)間]為本次投放時(shí)間
    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

通過(guò)以上代碼，可以看出，其主要處理過(guò)程是：

1、判斷有沒(méi)有被限流懲罰，有則直接返回，無(wú)則進(jìn)入下一步。

2、判斷令牌桶是否存在，不存在則先創(chuàng)建令牌桶，然后扣減令牌返回，存在則進(jìn)入下一步。

3、判斷是否需要投放令牌，不需要?jiǎng)t直接扣減令牌，需要?jiǎng)t先投放令牌再扣減令牌。

4、判斷扣減后的令牌數(shù)，如果小于0則返回限流，同時(shí)設(shè)置限流懲罰，如果大于等于0則進(jìn)入下一步。

5、更新桶中的令牌數(shù)到Redis。

你可以在任何一種開(kāi)發(fā)語(yǔ)言的Redis庫(kù)中提交并運(yùn)行這段Lua script腳本，如果你使用的是.NET平臺(tái)，可以參考這篇文章：ASP.NET Core中使用令牌桶限流。