上篇文章提到固定時(shí)間窗口限流無(wú)法處理突然請(qǐng)求洪峰情況，本文講述的令牌桶線路算法則可以比較好的處理此場(chǎng)景。

工作原理

單位時(shí)間按照一定速率勻速的生產(chǎn) token 放入桶內(nèi)，直到達(dá)到桶容量上限。

處理請(qǐng)求，每次嘗試獲取一個(gè)或多個(gè)令牌，如果拿到則處理請(qǐng)求，失敗則拒絕請(qǐng)求。

優(yōu)缺點(diǎn)

優(yōu)點(diǎn)

可以有效處理瞬間的突發(fā)流量，桶內(nèi)存量 token 即可作為流量緩沖區(qū)平滑處理突發(fā)流量。

缺點(diǎn)

實(shí)現(xiàn)較為復(fù)雜。

代碼實(shí)現(xiàn)

1. core/limit/tokenlimit.go

分布式環(huán)境下考慮使用 redis 作為桶和令牌的存儲(chǔ)容器，采用 lua 腳本實(shí)現(xiàn)整個(gè)算法流程。

redis lua 腳本

1. -- 每秒生成token數(shù)量即token生成速度
2. local rate = tonumber(ARGV[1])
3. -- 桶容量
4. local capacity = tonumber(ARGV[2])
5. -- 當(dāng)前時(shí)間戳
6. local now = tonumber(ARGV[3])
7. -- 當(dāng)前請(qǐng)求token數(shù)量
8. local requested = tonumber(ARGV[4])
9. -- 需要多少秒才能填滿桶
10. local fill_time = capacity/rate
11. -- 向下取整,ttl為填滿時(shí)間的2倍
12. local ttl = math.floor(fill_time*2)
13. -- 當(dāng)前時(shí)間桶容量
14. local last_tokens = tonumber(redis.call("get", KEYS[1]))
15. -- 如果當(dāng)前桶容量為0,說明是第一次進(jìn)入,則默認(rèn)容量為桶的最大容量
16. if last_tokens == nil then
17. last_tokens = capacity
18. end
19. -- 上一次刷新的時(shí)間
20. local last_refreshed = tonumber(redis.call("get", KEYS[2]))
21. -- 第一次進(jìn)入則設(shè)置刷新時(shí)間為0
22. if last_refreshed == nil then
23. last_refreshed = 0
24. end
25. -- 距離上次請(qǐng)求的時(shí)間跨度
26. local delta = math.max(0, now-last_refreshed)
27. -- 距離上次請(qǐng)求的時(shí)間跨度,總共能生產(chǎn)token的數(shù)量,如果超多最大容量則丟棄多余的token
28. local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
29. -- 本次請(qǐng)求token數(shù)量是否足夠
30. local allowed = filled_tokens >= requested
31. -- 桶剩余數(shù)量
32. local new_tokens = filled_tokens
33. -- 允許本次token申請(qǐng),計(jì)算剩余數(shù)量
34. if allowed then
35. new_tokens = filled_tokens - requested
36. end
37. -- 設(shè)置剩余token數(shù)量
38. redis.call("setex", KEYS[1], ttl, new_tokens)
39. -- 設(shè)置刷新時(shí)間
40. redis.call("setex", KEYS[2], ttl, now)
42. return allowed

令牌桶限流器定義

1. type TokenLimiter struct {
2. // 每秒生產(chǎn)速率
3. rate int
4. // 桶容量
5. burst int
6. // 存儲(chǔ)容器
7. store *redis.Redis
8. // redis key
9. tokenKey string
10. // 桶刷新時(shí)間key
11. timestampKey string
12. // lock
13. rescueLock sync.Mutex
14. // redis健康標(biāo)識(shí)
15. redisAlive uint32
16. // redis故障時(shí)采用進(jìn)程內(nèi) 令牌桶限流器
17. rescueLimiter *xrate.Limiter
18. // redis監(jiān)控探測(cè)任務(wù)標(biāo)識(shí)
19. monitorStarted bool
20. }
22. func NewTokenLimiter(rate, burst int, store *redis.Redis, key string) *TokenLimiter {
23. tokenKey := fmt.Sprintf(tokenFormat, key)
24. timestampKey := fmt.Sprintf(timestampFormat, key)
26. return &TokenLimiter{
27. rate: rate,
28. burst: burst,
29. store: store,
30. tokenKey: tokenKey,
31. timestampKey: timestampKey,
32. redisAlive: 1,
33. rescueLimiter: xrate.NewLimiter(xrate.Every(time.Second/time.Duration(rate)), burst),
34. }
35. }

獲取令牌

1. func (lim *TokenLimiter) reserveN(now time.Time, n int) bool {
2. // 判斷redis是否健康
3. // redis故障時(shí)采用進(jìn)程內(nèi)限流器
4. // 兜底保障
5. if atomic.LoadUint32(&lim.redisAlive) == 0 {
6. return lim.rescueLimiter.AllowN(now, n)
7. }
8. // 執(zhí)行腳本獲取令牌
9. resp, err := lim.store.Eval(
10. script,
11. []string{
12. lim.tokenKey,
13. lim.timestampKey,
14. },
15. []string{
16. strconv.Itoa(lim.rate),
17. strconv.Itoa(lim.burst),
18. strconv.FormatInt(now.Unix(), 10),
19. strconv.Itoa(n),
20. })
21. // redis allowed == false
22. // Lua boolean false -> r Nil bulk reply
23. // 特殊處理key不存在的情況
24. if err == redis.Nil {
25. return false
26. } else if err != nil {
27. logx.Errorf("fail to use rate limiter: %s, use in-process limiter for rescue", err)
28. // 執(zhí)行異常，開啟redis健康探測(cè)任務(wù)
29. // 同時(shí)采用進(jìn)程內(nèi)限流器作為兜底
30. lim.startMonitor()
31. return lim.rescueLimiter.AllowN(now, n)
32. }
34. code, ok := resp.(int64)
35. if !ok {
36. logx.Errorf("fail to eval redis script: %v, use in-process limiter for rescue", resp)
37. lim.startMonitor()
38. return lim.rescueLimiter.AllowN(now, n)
39. }
41. // redis allowed == true
42. // Lua boolean true -> r integer reply with value of 1
43. return code == 1
44. }

redis 故障時(shí)兜底策略

兜底策略的設(shè)計(jì)考慮得非常細(xì)節(jié)，當(dāng) redis 不可用的時(shí)候，啟動(dòng)單機(jī)版的 ratelimit 做備用限流，確保基本的限流可用，服務(wù)不會(huì)被沖垮。

1. // 開啟redis健康探測(cè)
2. func (lim *TokenLimiter) startMonitor() {
3. lim.rescueLock.Lock()
4. defer lim.rescueLock.Unlock()
5. // 防止重復(fù)開啟
6. if lim.monitorStarted {
7. return
8. }
10. // 設(shè)置任務(wù)和健康標(biāo)識(shí)
11. lim.monitorStarted = true
12. atomic.StoreUint32(&lim.redisAlive, 0)
13. // 健康探測(cè)
14. go lim.waitForRedis()
15. }
17. // redis健康探測(cè)定時(shí)任務(wù)
18. func (lim *TokenLimiter) waitForRedis() {
19. ticker := time.NewTicker(pingInterval)
20. // 健康探測(cè)成功時(shí)回調(diào)此函數(shù)
21. defer func() {
22. ticker.Stop()
23. lim.rescueLock.Lock()
24. lim.monitorStarted = false
25. lim.rescueLock.Unlock()
26. }()
28. for range ticker.C {
29. // ping屬于redis內(nèi)置健康探測(cè)命令
30. if lim.store.Ping() {
31. // 健康探測(cè)成功，設(shè)置健康標(biāo)識(shí)
32. atomic.StoreUint32(&lim.redisAlive, 1)
33. return
34. }
35. }
36. }

項(xiàng)目地址

https://github.com/zeromicro/go-zero

歡迎使用 go-zero 并 star 支持我們!

原文鏈接：https://mp.weixin.qq.com/s/ulGRw4qkWbGKdF83VaIb7A