go-zero 是如何實現計數器限流的？

原文鏈接： 如何實現計數器限流？

上一篇文章 go-zero 是如何做路由管理的？ 介紹了路由管理，這篇文章來說說限流，主要介紹計數器限流演算法，具體的代碼實現，我們還是來分析微服務框架 go-zero 的源碼。

在微服務架構中，一個服務可能需要頻繁地與其他服務交互，而過多的請求可能導致性能下降或系統崩潰。為了確保系統的穩定性和高可用性，限流演算法應運而生。

限流演算法允許在給定時間段內，對服務的請求流量進行控制和調整，以防止資源耗盡和服務過載。

計數器限流演算法主要有兩種實現方式，分別是：

1. 固定視窗計數器
2. 滑動視窗計數器

下麵分別來介紹。

固定視窗計數器

演算法概念如下：

• 將時間劃分為多個視窗；
• 在每個視窗內每有一次請求就將計數器加一；
• 如果計數器超過了限制數量，則本視窗內所有的請求都被丟棄當時間到達下一個視窗時，計數器重置。

固定視窗計數器是最為簡單的演算法，但這個演算法有時會讓通過請求量允許為限制的兩倍。

考慮如下情況：限制 1 秒內最多通過 5 個請求，在第一個視窗的最後半秒內通過了 5 個請求，第二個視窗的前半秒內又通過了 5 個請求。這樣看來就是在 1 秒內通過了 10 個請求。

滑動視窗計數器

演算法概念如下：

• 將時間劃分為多個區間；
• 在每個區間內每有一次請求就將計數器加一維持一個時間視窗，占據多個區間；
• 每經過一個區間的時間，則拋棄最老的一個區間，並納入最新的一個區間；
• 如果當前視窗內區間的請求計數總和超過了限制數量，則本視窗內所有的請求都被丟棄。

滑動視窗計數器是通過將視窗再細分，並且按照時間滑動，這種演算法避免了固定視窗計數器帶來的雙倍突發請求，但時間區間的精度越高，演算法所需的空間容量就越大。

go-zero 實現

go-zero 實現的是固定視窗的方式，計算一段時間內對同一個資源的訪問次數，如果超過指定的 limit，則拒絕訪問。當然如果在一段時間內訪問不同的資源，每一個資源訪問量都不超過 limit，此種情況是不會拒絕的。

而在一個分散式系統中，存在多個微服務提供服務。所以當瞬間的流量同時訪問同一個資源，如何讓計數器在分散式系統中正常計數？

這裡要解決的一個主要問題就是計算的原子性，保證多個計算都能得到正確結果。

通過以下兩個方面來解決：

• 使用 redis 的 incrby 做資源訪問計數
• 採用 lua script 做整個視窗計算，保證計算的原子性

接下來先看一下 lua script 的源碼：

// core/limit/periodlimit.go

const periodScript = `local limit = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local current = redis.call("INCRBY", KEYS[1], 1)
if current == 1 then
    redis.call("expire", KEYS[1], window)
end
if current < limit then
    return 1
elseif current == limit then
    return 2
else
    return 0
end`

主要就是使用 INCRBY 命令來實現，第一次請求需要給 key 加上一個過期時間，到達過期時間之後，key 過期被清楚，重新計數。

限流器初始化：

type (
    // PeriodOption defines the method to customize a PeriodLimit.
    PeriodOption func(l *PeriodLimit)

    // A PeriodLimit is used to limit requests during a period of time.
    PeriodLimit struct {
        period     int  // 視窗大小，單位 s
        quota      int  // 請求上限
        limitStore *redis.Redis
        keyPrefix  string   // key 首碼
        align      bool
    }
)

// NewPeriodLimit returns a PeriodLimit with given parameters.
func NewPeriodLimit(period, quota int, limitStore *redis.Redis, keyPrefix string,
    opts ...PeriodOption) *PeriodLimit {
    limiter := &PeriodLimit{
        period:     period,
        quota:      quota,
        limitStore: limitStore,
        keyPrefix:  keyPrefix,
    }

    for _, opt := range opts {
        opt(limiter)
    }

    return limiter
}

調用限流：

// key 就是需要被限制的資源標識
func (h *PeriodLimit) Take(key string) (int, error) {
    return h.TakeCtx(context.Background(), key)
}

// TakeCtx requests a permit with context, it returns the permit state.
func (h *PeriodLimit) TakeCtx(ctx context.Context, key string) (int, error) {
    resp, err := h.limitStore.EvalCtx(ctx, periodScript, []string{h.keyPrefix + key}, []string{
        strconv.Itoa(h.quota),
        strconv.Itoa(h.calcExpireSeconds()),
    })
    if err != nil {
        return Unknown, err
    }

    code, ok := resp.(int64)
    if !ok {
        return Unknown, ErrUnknownCode
    }

    switch code {
    case internalOverQuota: // 超過上限
        return OverQuota, nil
    case internalAllowed:   // 未超過，允許訪問
        return Allowed, nil
    case internalHitQuota:  // 正好達到限流上限
        return HitQuota, nil
    default:
        return Unknown, ErrUnknownCode
    }
}

上文已經介紹了，固定時間視窗會有臨界突發問題，並不是那麼嚴謹，下篇文章我們來介紹令牌桶限流。

以上就是本文的全部內容，如果覺得還不錯的話歡迎點贊，轉發和關註，感謝支持。

參考文章：

• https://juejin.cn/post/6895928148521648141
• https://juejin.cn/post/7051406419823689765
• https://www.infoq.cn/article/Qg2tX8fyw5Vt-f3HH673

推薦閱讀：

• go-zero 是如何做路由管理的？