Golang官方限流器庫(kù)實(shí)現(xiàn)限流示例詳解

更新時(shí)間：2022年08月18日 11:17:56 作者：jiaxwu

這篇文章主要為大家介紹了Golang官方限流器庫(kù)使用示例詳解，有需要的朋友可以借鑒參考下，希望能夠有所幫助，祝大家多多進(jìn)步，早日升職加薪

前言

在翻Golang官方庫(kù)的過(guò)程中，發(fā)現(xiàn)一個(gè)有趣的庫(kù)golang.org/x/time ，里面只有一個(gè)類rate，研究了一下發(fā)現(xiàn)它是一個(gè)限流器，實(shí)現(xiàn)了很多的功能，當(dāng)然它的核心原理并不復(fù)雜，也就是令牌桶算法。

令牌桶算法的原理是：令牌桶會(huì)不斷地把令牌添加到桶里，而請(qǐng)求會(huì)從桶中獲取令牌，只有擁有令牌地請(qǐng)求才能被接受。因?yàn)橥爸锌梢蕴崆氨Ａ粢恍┝钆?，所以它允許一定地突發(fā)流量通過(guò)。

例子

下面是限流算法常見的寫法，首先判斷是否有令牌，如果有就通過(guò)，否則直接失敗。

package main
import (
	"fmt"
	"time"
	"golang.org/x/time/rate"
)
func main() {
        // 每0.1秒生成一個(gè)令牌，也就是一秒10個(gè)令牌，最大保留令牌上限10
	l := rate.NewLimiter(rate.Every(time.Second/10), 10)
	for i := 0; i < 10; i++ {
		go func(i int) {
			for {
                                // 判斷是否有令牌，如果有就輸出
				if l.Allow() {
					fmt.Printf("allow %d\n", i)
				}
                                // 每0.5秒請(qǐng)求一次
				time.Sleep(time.Second / 2)
			}
		}(i)
	}
	time.Sleep(time.Second * 10)
}

上面的rate.Every(time.Second/10)會(huì)返回一個(gè)Limit類型，代表每秒生成多少個(gè)令牌。

這個(gè)庫(kù)還提供了另外一種寫法，等待直到有令牌為止（或超時(shí)）：

func main() {
	l := rate.NewLimiter(rate.Every(time.Second/10), 100)
	for i := 0; i < 10; i++ {
		go func(i int) {
			for {
                                // 等待直到有令牌
				if err := l.Wait(context.TODO()); err != nil {
				} else {
					fmt.Printf("allow %d\n", i)
				}
				time.Sleep(time.Second / 2)
			}
		}(i)
	}
	time.Sleep(time.Second * 10)
}

這樣在某些場(chǎng)景下我們可以讓請(qǐng)求等待一會(huì)，而不是直接失敗。

還有一個(gè)更加特殊的請(qǐng)求令牌方式，也就是先預(yù)留令牌，到指定時(shí)間不再需要去獲取令牌，直接執(zhí)行操作即可：

func main() {
	l := rate.NewLimiter(rate.Every(time.Second/10), 10)
	for i := 0; i < 10; i++ {
		go func(i int) {
			for {
                                // 先預(yù)留令牌
				if r := l.Reserve(); r.OK() {
                                        // 休眠直到令牌生效
					time.Sleep(r.Delay())
					fmt.Printf("allow %d\n", i)
				}
				time.Sleep(time.Second / 2)
			}
		}(i)
	}
	time.Sleep(time.Second * 10)
}

當(dāng)然，如果預(yù)留的令牌不想使用了，也可以使用r.Cancel()歸還已預(yù)留的令牌。

上面的Allow()、Wait()、Reserve()都是一次消耗一個(gè)令牌，其實(shí)都有對(duì)應(yīng)的AllowN()、WaitN()、ReserveN()方法，一次消耗N個(gè)令牌，這樣就可以根據(jù)任務(wù)消耗的資源靈活的消耗令牌。

實(shí)現(xiàn)

不管我們從Allow()、Wait()還是Reserve()進(jìn)去，最終都會(huì)進(jìn)入到reserveN(now time.Time, n int, maxFutureReserve time.Duration) Reservation 方法：

首先在進(jìn)入方法的時(shí)候，會(huì)先處理兩種特殊情況：

        // 如果無(wú)限生成令牌，則直接返回
        if lim.limit == Inf {
		return Reservation{
			ok:        true,
			lim:       lim,
			tokens:    n,
			timeToAct: now,
		}
        // 如果不會(huì)生成令牌，則在初始令牌里面拿，直到拿完為止
	} else if lim.limit == 0 {
		var ok bool
		if lim.burst >= n {
			ok = true
			lim.burst -= n
		}
		return Reservation{
			ok:        ok,
			lim:       lim,
			tokens:    lim.burst,
			timeToAct: now,
		}
	}

然后重新計(jì)算當(dāng)前有多少令牌，減去要消耗的令牌：

        // 計(jì)算當(dāng)前有多少令牌
	now, last, tokens := lim.advance(now)
	// 減去要消耗的N個(gè)令牌
	tokens -= float64(n)
        // 如果剩余令牌為負(fù)數(shù)，那么計(jì)算一下要等待多久才能拿到令牌
        var waitDuration time.Duration
	if tokens < 0 {
		waitDuration = lim.limit.durationFromTokens(-tokens)
	}
        // 判斷請(qǐng)求是否成功，maxFutureReserve代表最大可以等待的時(shí)間，也就是請(qǐng)求能否接收拿到令牌需要等待的時(shí)間
        ok := n <= lim.burst && waitDuration <= maxFutureReserve

余下的代碼就是更新限流器的狀態(tài)。