Golang 標準庫 限流器 time/rate 設計與實現

友情提示：此篇文章大約需要閱讀 12分鐘7秒，不足之處請多指教，感謝你的閱讀。 訂閱本站

限流器是後臺服務中十分重要的元件，在實際的業務場景中使用居多，其設計在微服務、閘道器、和一些後臺服務中會經常遇到。限流器的作用是用來限制其請求的速率，保護後臺響應服務，以免服務過載導致服務不可用現象出現。

限流器的實現方法有很多種，例如 Token Bucket、滑動視窗法、Leaky Bucket等。

在 Golang 庫中官方給我們提供了限流器的實現golang.org/x/time/rate，它是基於令牌桶演算法（Token Bucket）設計實現的。

令牌桶演算法

令牌桶設計比較簡單，可以簡單的理解成一個只能存放固定數量雪糕的一個冰箱，每個請求可以理解成來拿雪糕的人，有且只能每一次請求拿一塊，那雪糕拿完了會怎麼樣呢？這裡會有一個固定放雪糕的工人，並且他往冰箱裡放雪糕的頻率都是一致的，例如他 1s 中只能往冰箱裡放 10 塊雪糕，這裡就可以看出請求響應的頻率了。

令牌桶設計概念：

• 令牌：每次請求只有拿到 Token 令牌後，才可以繼續訪問；
• 桶：具有固定數量的桶，每個桶中最多隻能放設計好的固定數量的令牌；
• 入桶頻率：按照固定的頻率往桶中放入令牌，放入令牌不能超過桶的容量。

也就是說，基於令牌桶設計演算法就限制了請求的速率，達到請求響應可控的目的，特別是針對於高併發場景中突發流量請求的現象，後臺就可以輕鬆應對請求了，因為到後端具體服務的時候突發流量請求已經經過了限流了。

具體設計

限流器定義

type Limiter struct {
    mu        sync.Mutex // 互斥鎖（排他鎖）
    limit     Limit      // 放入桶的頻率  float64 型別
    burst     int        // 桶的大小
    tokens    float64    // 令牌 token 當前剩餘的數量
    last      time.Time  // 最近取走 token 的時間
    lastEvent time.Time  // 最近限流事件的時間
}

limit、burst 和 token 是這個限流器中核心的引數，請求併發的大小在這裡實現的。

在令牌發放之後，會儲存在 Reservation 預約物件中：

type Reservation struct {
    ok        bool      // 是否滿足條件分配了 token
    lim       *Limiter  // 傳送令牌的限流器
    tokens    int       // 傳送 token 令牌的數量
    timeToAct time.Time // 滿足令牌發放的時間
    limit     Limit     // 令牌發放速度
}

消費 Token

Limiter 提供了三類方法供使用者消費 Token，使用者可以每次消費一個 Token，也可以一次性消費多個 Token。而每種方法代表了當 Token 不足時，各自不同的對應手段。

Wait、WaitN

func (lim *Limiter) Wait(ctx context.Context) (err error)
func (lim *Limiter) WaitN(ctx context.Context, n int) (err error)

其中，Wait 就是 WaitN(ctx, 1)，在下面的方法介紹實現也是一樣的。

使用 Wait 方法消費 Token 時，如果此時桶內 Token 陣列不足 ( 小於 n )，那麼 Wait 方法將會阻塞一段時間，直至 Token 滿足條件。如果充足則直接返回。

Allow、AllowN

func (lim *Limiter) Allow() bool
func (lim *Limiter) AllowN(now time.Time, n int) bool 

AllowN 方法表示，截止到當前某一時刻，目前桶中數目是否至少為 n 個，滿足則返回 true，同時從桶中消費 n 個 token。
反之返回不消費 Token，false。

通常對應這樣的線上場景，如果請求速率過快，就直接丟到某些請求。

Reserve、ReserveN

官方提供的限流器有阻塞等待式的 Wait，也有直接判斷方式的 Allow，還有提供了自己維護預留式的，但核心的實現都是下面的 reserveN 方法。

func (lim *Limiter) Reserve() *Reservation
func (lim *Limiter) ReserveN(now time.Time, n int) *Reservation

當呼叫完成後，無論 Token 是否充足，都會返回一個Reservation *物件。

你可以呼叫該物件的 Delay() 方法，該方法返回了需要等待的時間。如果等待時間為 0，則說明不用等待。
必須等到等待時間結束之後，才能進行接下來的工作。

或者，如果不想等待，可以呼叫 Cancel()方法，該方法會將 Token 歸還。

func (lim *Limiter) reserveN(now time.Time, n int, maxFutureReserve time.Duration) Reservation {
    lim.mu.Lock()

    // 首先判斷是否放入頻率是否為無窮大
    // 如果為無窮大，說明暫時不限流
    if lim.limit == Inf {
        lim.mu.Unlock()
        return Reservation{
            ok:        true,
            lim:       lim,
            tokens:    n,
            timeToAct: now,
        }
    }

    // 拿到截至 now 時間時
    // 可以獲取的令牌 tokens 數量及上一次拿走令牌的時間 last
    now, last, tokens := lim.advance(now)

    // 更新 tokens 數量
    tokens -= float64(n)

    // 如果 tokens 為負數，代表當前沒有 token 放入桶中
    // 說明需要等待，計算等待的時間
    var waitDuration time.Duration
    if tokens < 0 {
        waitDuration = lim.limit.durationFromTokens(-tokens)
    }

    // 計算是否滿足分配條件
    // 1、需要分配的大小不超過桶的大小
    // 2、等待時間不超過設定的等待時長
    ok := n <= lim.burst && waitDuration <= maxFutureReserve

    // 預處理 reservation
    r := Reservation{
        ok:    ok,
        lim:   lim,
        limit: lim.limit,
    }
    // 若當前滿足分配條件
    // 1、設定分配大小
    // 2、滿足令牌發放的時間 = 當前時間 + 等待時長
    if ok {
        r.tokens = n
        r.timeToAct = now.Add(waitDuration)
    }

    // 更新 limiter 的值，並返回
    if ok {
        lim.last = now
        lim.tokens = tokens
        lim.lastEvent = r.timeToAct
    } else {
        lim.last = last
    }

    lim.mu.Unlock()
    return r
}

具體使用

rate 包中提供了對限流器的使用，只需要指定 limit（放入桶中的頻率）、burst（桶的大小）。

func NewLimiter(r Limit, b int) *Limiter {
    return &Limiter{
        limit: r, // 放入桶的頻率
        burst: b, // 桶的大小
    }
}

在這裡，使用一個 http api 來簡單的驗證一下 time/rate 的強大：

func main() {
    r := rate.Every(1 * time.Millisecond)
    limit := rate.NewLimiter(r, 10)
    http.HandleFunc("/", func(writer http.ResponseWriter, request *http.Request) {
        if limit.Allow() {
            fmt.Printf("請求成功，當前時間：%s\n", time.Now().Format("2006-01-02 15:04:05"))
        } else {
            fmt.Printf("請求成功，但是被限流了。。。\n")
        }
    })

    _ = http.ListenAndServe(":8081", nil)
}

在這裡，我把桶設定成了每一毫秒投放一次令牌，桶容量大小為 10，起一個 http 的服務，模擬後臺 API。

接下來做一個壓力測試，看看效果如何：

func GetApi() {
    api := "http://localhost:8081/"
    res, err := http.Get(api)
    if err != nil {
        panic(err)
    }
    defer res.Body.Close()

    if res.StatusCode == http.StatusOK {
        fmt.Printf("get api success\n")
    }
}

func Benchmark_Main(b *testing.B) {
    for i := 0; i < b.N; i++ {
        GetApi()
    }
}

效果如下：

......
請求成功，當前時間：2020-08-24 14:26:52
請求成功，但是被限流了。。。
請求成功，但是被限流了。。。
請求成功，但是被限流了。。。
請求成功，但是被限流了。。。
請求成功，但是被限流了。。。
請求成功，當前時間：2020-08-24 14:26:52
請求成功，但是被限流了。。。
請求成功，但是被限流了。。。
請求成功，但是被限流了。。。
請求成功，但是被限流了。。。
......

在這裡，可以看到，當使用 AllowN 方法中，只有當令牌 Token 生產出來，才可以消費令牌，繼續請求，剩餘的則是將其請求拋棄，當然在實際的業務處理中，可以用比較友好的方式反饋給前端。

在這裡，先有的幾次請求都會成功，是因為服務啟動後，令牌桶會初始化，將令牌放入到桶中，但是隨著突發流量的請求，令牌按照預定的速率生產令牌，就會出現明顯的令牌供不應求的現象。

開源文化

目前 time/rate 是一個獨立的限流器開源解決方案，感興趣的小夥伴可以給此專案一個 Star，謝謝。

GitHub
golang/time

參考文章

• 限流器系列(2) — Token Bucket 令牌桶
• Golang 限流器的使用和實現
• Golang 標準庫限流器 time/rate 使用介紹
• https://github.com/golang/time/rate.go

本作品採用《CC 協議》，轉載必須註明作者和本文連結