golang中怎么利用go-rate實現一個令牌桶算法

golang中怎么利用go-rate實現一個令牌桶算法，很多新手對此不是很清楚，為了幫助大家解決這個難題，下面小編將為大家詳細講解，有這方面需求的人可以來學習下，希望你能有所收獲。

構造限流器

我們首先構造一個限流器對象：

limiter := NewLimiter(10, 1);

這里有兩個參數：

1. 第一個參數是 r Limit。代表每秒可以向 Token 桶中產生多少 token。Limit 實際上是 float64 的別名。

2. 第二個參數是 b int。b 代表 Token 桶的容量大小。

上述的限流器的含義是：擁有一個容量為1的令牌桶，以每鈔10個的速度向桶中放令牌。

除了直接指定每秒產生的 Token 個數外，還可以用 Every 方法來指定向 Token 桶中放置 Token 的間隔，例如：

limiter := NewLimiter(Every(100 * time.Millisecond), 1);

以上就表示每 100ms 往桶中放一個 Token。本質上也就是一秒鐘產生 10 個。

消費令牌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 滿足條件。如果充足則直接返回。

這里可以看到，Wait 方法有一個 context 參數。我們可以設置 context 的 Deadline 或者 Timeout，來決定此次 Wait 的最長時間。

Allow/AllowN

Allow 實際上就是 AllowN(time.Now(),1)。

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

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

Reserve/ReserveN

Reserve 相當于 ReserveN(time.Now(), 1)。

ReserveN 的用法就相對來說復雜一些，當調用完成后，無論 Token 是否充足，都會返回一個 Reservation * 對象。

你可以調用該對象的 Delay() 方法，該方法返回了需要等待的時間。如果等待時間為 0，則說明不用等待。必須等到等待時間之后，才能進行接下來的工作。

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

使用一個偽代碼來舉例，我們可以如何使用 Reserve 方法。

r := lim.Reserve()
//是否愿意等待
f !r.OK() {
    //不愿意等待直接退出
    return
}

//如果愿意等待，將等待時間拋給用戶 time.Sleep代表用戶需要等待的時間。
time.Sleep(r.Delay())
Act() // 一段時間后生成生成新的令牌，開始執行相關邏輯

動態調整速率

Limiter 支持可以調整速率和桶大小：

1. SetLimit(Limit) 改變放入 Token 的速率

2. SetBurst(int) 改變 Token 桶大小

有了這兩個方法，可以根據現有環境和條件以及我們的需求，動態地改變 Token 桶大小和速率。

案例1-單位時間只允許一次郵件發送操作

客戶端軟件客戶點擊發送郵件，如果客戶一秒鐘內點擊10次，就會發送10次，這明顯是不合適的。如果使用速率限制，我們就可以限制一秒內只能發送一次，實現方法為：

(令牌桶)容量為1，速度為每一秒生成一個令牌，這樣可以保證一秒鐘只會被執行一次，偽代碼實現如下

//初始化 limiter 每秒生成1個令牌，令牌桶容量為20
limiter := rate.NewLimiter(rate.Every(time.Second), 1)
//模擬單位時間執行多次操作
for i := 0; i < 5; i++ {
	if limiter.Allow() {
		fmt.Println("發送郵件")
	} else {
		fmt.Println("請求多次，過濾")
	}
}
if limiter.Allow() {
		fmt.Println("發送郵件")
}

執行結果

發送郵件 請求多次，過濾 請求多次，過濾 請求多次，過濾 請求多次，過濾 發送郵件

我們發現，第一次執行是可以被允許的因為第一次的令牌被允許，之后的請求失敗是因為還沒有生成新的令牌，所以需要等待1秒，之后又可以進行發送郵件操作。

通過這樣一個案例，相信大家對令牌桶的實現場景有了一個基本的了解。

案例2——令牌取出單個和多個

初始化令牌桶容量為20，設置每100毫秒生成一個令牌，即1秒生產10個令牌。編碼測試功能

//初始化 limiter 每秒10個令牌，令牌桶容量為20
limiter := rate.NewLimiter(rate.Every(time.Millisecond*100), 20)
for i := 0; i < 25; i++ {
	if limiter.Allow() {
		fmt.Println("success") //do something
	} else {
		fmt.Println("busy")
	}
}

//阻塞直到獲取足夠的令牌或者上下文取消
ctx, _ := context.WithTimeout(context.Background(), time.Second*2)
fmt.Println("start get token", time.Now())
err := limiter.WaitN(ctx, 20)
if err != nil {
	fmt.Println("error", err)
	return
}
fmt.Println("success get token", time.Now())

看完上述內容是否對您有幫助呢？如果還想對相關知識有進一步的了解或閱讀更多相關文章，請關注億速云行業資訊頻道，感謝您對億速云的支持。