Go 分佈式令牌桶限流 - 兜底保障

上篇文章提到固定時間窗口限流無法處理突然請求洪峯情況,本文講述的令牌桶線路算法則可以比較好的處理此場景。

工作原理

  1. 單位時間按照一定速率勻速的生產 token 放入桶內,直到達到桶容量上限。

  2. 處理請求,每次嘗試獲取一個或多個令牌,如果拿到則處理請求,失敗則拒絕請求。

優缺點

優點

可以有效處理瞬間的突發流量,桶內存量 token 即可作爲流量緩衝區平滑處理突發流量。

缺點

實現較爲複雜。

代碼實現

core/limit/tokenlimit.go

分佈式環境下考慮使用 redis 作爲桶和令牌的存儲容器,採用 lua 腳本實現整個算法流程。

redis lua 腳本

-- 每秒生成token數量即token生成速度
local rate = tonumber(ARGV[1])
-- 桶容量
local capacity = tonumber(ARGV[2])
-- 當前時間戳
local now = tonumber(ARGV[3])
-- 當前請求token數量
local requested = tonumber(ARGV[4])
-- 需要多少秒才能填滿桶
local fill_time = capacity/rate
-- 向下取整,ttl爲填滿時間的2倍
local ttl = math.floor(fill_time*2)
-- 當前時間桶容量
local last_tokens = tonumber(redis.call("get", KEYS[1]))
-- 如果當前桶容量爲0,說明是第一次進入,則默認容量爲桶的最大容量
if last_tokens == nil then
last_tokens = capacity
end
-- 上一次刷新的時間
local last_refreshed = tonumber(redis.call("get", KEYS[2]))
-- 第一次進入則設置刷新時間爲0
if last_refreshed == nil then
last_refreshed = 0
end
-- 距離上次請求的時間跨度
local delta = math.max(0, now-last_refreshed)
-- 距離上次請求的時間跨度,總共能生產token的數量,如果超多最大容量則丟棄多餘的token
local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
-- 本次請求token數量是否足夠
local allowed = filled_tokens >= requested
-- 桶剩餘數量
local new_tokens = filled_tokens
-- 允許本次token申請,計算剩餘數量
if allowed then
new_tokens = filled_tokens - requested
end
-- 設置剩餘token數量
redis.call("setex", KEYS[1], ttl, new_tokens)
-- 設置刷新時間
redis.call("setex", KEYS[2], ttl, now)

return allowed

令牌桶限流器定義

type TokenLimiter struct {
    // 每秒生產速率
    rate int
    // 桶容量
    burst int
    // 存儲容器
    store *redis.Redis
    // redis key
    tokenKey       string
    // 桶刷新時間key
    timestampKey   string
    // lock
    rescueLock     sync.Mutex
    // redis健康標識
    redisAlive     uint32
    // redis故障時採用進程內 令牌桶限流器
    rescueLimiter  *xrate.Limiter
    // redis監控探測任務標識
    monitorStarted bool
}

func NewTokenLimiter(rate, burst int, store *redis.Redis, key string) *TokenLimiter {
    tokenKey := fmt.Sprintf(tokenFormat, key)
    timestampKey := fmt.Sprintf(timestampFormat, key)

    return &TokenLimiter{
        rate:          rate,
        burst:         burst,
        store:         store,
        tokenKey:      tokenKey,
        timestampKey:  timestampKey,
        redisAlive:    1,
        rescueLimiter: xrate.NewLimiter(xrate.Every(time.Second/time.Duration(rate)), burst),
    }
}

獲取令牌

func (lim *TokenLimiter) reserveN(now time.Time, n int) bool {
    // 判斷redis是否健康
    // redis故障時採用進程內限流器
    // 兜底保障
    if atomic.LoadUint32(&lim.redisAlive) == 0 {
        return lim.rescueLimiter.AllowN(now, n)
    }
    // 執行腳本獲取令牌
    resp, err := lim.store.Eval(
        script,
        []string{
            lim.tokenKey,
            lim.timestampKey,
        },
        []string{
            strconv.Itoa(lim.rate),
            strconv.Itoa(lim.burst),
            strconv.FormatInt(now.Unix(), 10),
            strconv.Itoa(n),
        })
    // redis allowed == false
    // Lua boolean false -> r Nil bulk reply
    // 特殊處理key不存在的情況
    if err == redis.Nil {
        return false
    } else if err != nil {
        logx.Errorf("fail to use rate limiter: %s, use in-process limiter for rescue", err)
        // 執行異常,開啓redis健康探測任務
        // 同時採用進程內限流器作爲兜底
        lim.startMonitor()
        return lim.rescueLimiter.AllowN(now, n)
    }

    code, ok := resp.(int64)
    if !ok {
        logx.Errorf("fail to eval redis script: %v, use in-process limiter for rescue", resp)
        lim.startMonitor()
        return lim.rescueLimiter.AllowN(now, n)
    }

    // redis allowed == true
    // Lua boolean true -> r integer reply with value of 1
    return code == 1
}

redis 故障時兜底策略

兜底策略的設計考慮得非常細節,當 redis 不可用的時候,啓動單機版的 ratelimit 做備用限流,確保基本的限流可用,服務不會被沖垮。

// 開啓redis健康探測
func (lim *TokenLimiter) startMonitor() {
    lim.rescueLock.Lock()
    defer lim.rescueLock.Unlock()
    // 防止重複開啓
    if lim.monitorStarted {
        return
    }

    // 設置任務和健康標識
    lim.monitorStarted = true
    atomic.StoreUint32(&lim.redisAlive, 0)
    // 健康探測
    go lim.waitForRedis()
}

// redis健康探測定時任務
func (lim *TokenLimiter) waitForRedis() {
    ticker := time.NewTicker(pingInterval)
    // 健康探測成功時回調此函數
    defer func() {
        ticker.Stop()
        lim.rescueLock.Lock()
        lim.monitorStarted = false
        lim.rescueLock.Unlock()
    }()

    for range ticker.C {
        // ping屬於redis內置健康探測命令
        if lim.store.Ping() {
            // 健康探測成功,設置健康標識
            atomic.StoreUint32(&lim.redisAlive, 1)
            return
        }
    }
}

項目地址

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

歡迎使用 go-zerostar 支持我們!

微信交流羣

關注『微服務實踐』公衆號並點擊 交流羣 獲取社區羣二維碼。

本文由 Readfog 進行 AMP 轉碼,版權歸原作者所有。
來源https://mp.weixin.qq.com/s/ulGRw4qkWbGKdF83VaIb7A