使用Golang如何實現簡易的令牌桶算法

更新時間：2023年07月10日 15:59:39 作者：CG國斌

這篇文章主要介紹了使用Golang如何實現簡易的令牌桶算法問題，具有很好的參考價值，希望對大家有所幫助。如有錯誤或未考慮完全的地方，望不吝賜教

簡介

在網絡中傳輸數據的時候時，為了防止網絡擁塞，需限制流出網絡的流量，使流量以比較均勻的速度向外發(fā)送。

令牌桶算法就實現了這個功能，可控制發(fā)送到網絡上數據的數目，并允許突發(fā)數據的發(fā)送。

令牌桶算法是網絡流量整形和速率限制中最常使用的一種算法。

大小固定的令牌桶可自行以恒定的速率源源不斷地產生令牌。

如果令牌不被消耗，或者被消耗的速度小于產生的速度，令牌就會不斷地增多，直到把桶填滿。

后面再產生的令牌就會從桶中溢出。最后桶中可以保存的最大令牌數永遠不會超過桶的大小。

ipt-1

傳送到令牌桶的數據包需要消耗令牌。不同大小的數據包，消耗的令牌數量不一樣。

令牌桶這種控制機制基于令牌桶中是否存在令牌來指示什么時候可以發(fā)送流量。令牌桶中的每一個令牌都代表一個字節(jié)。

如果令牌桶中存在令牌，則允許發(fā)送流量；而如果令牌桶中不存在令牌，則不允許發(fā)送流量。

因此，如果突發(fā)門限被合理地配置并且令牌桶中有足夠的令牌，那么流量就可以以峰值速率發(fā)送。

lpt-2

與“令牌桶算法”類似的算法還有“漏桶算法”，這兩種算法的主要區(qū)別在于“漏桶算法”能夠強行限制數據的傳輸速率，而“令牌桶算法”在能夠限制數據的平均傳輸速率外，還允許某種程度的突發(fā)傳輸。

在“令牌桶算法”中，只要令牌桶中存在令牌，那么就允許突發(fā)地傳輸數據直到達到用戶配置的門限，因此它適合于具有突發(fā)特性的流量。

在本文中，我們使用 Golong 語言實現一個簡單的“令牌桶算法”，或者說是“漏桶算法”更為合適。

實現

首先，我們假設令牌桶的放入令牌的速率是恒定的，不考慮流量速率突變的情況。

package awesomeProject
import (
	"sync"
	"time"
)
// 定義令牌桶結構
type tokenBucket struct {
	limitRate int           // 限制頻率，即每分鐘加入多少個令牌
	tokenChan chan struct{} // 令牌通道，可以理解為桶
	cap       int           // 令牌桶的容量
	muLock    *sync.Mutex   // 令牌桶鎖，保證線程安全
	stop      bool          // 停止標記，結束令牌桶
}
// NewTokenBucket 創(chuàng)建令牌桶
func NewTokenBucket(limitRate, cap int) *tokenBucket {
	if cap < 1 {
		panic("token bucket cap must be large 1")
	}
	return &tokenBucket{
		tokenChan: make(chan struct{}, cap),
		limitRate: limitRate,
		muLock:    new(sync.Mutex),
		cap:       cap,
	}
}
// Start 開啟令牌桶
func (b *tokenBucket) Start() {
	go b.produce()
}
// 生產令牌
func (b *tokenBucket) produce() {
	for {
		b.muLock.Lock()
		if b.stop {
			close(b.tokenChan)
			b.muLock.Unlock()
			return
		}
		b.tokenChan <- struct{}{}
		d := time.Minute / time.Duration(b.limitRate)
		b.muLock.Unlock()
		time.Sleep(d)
	}
}
// Consume 消費令牌
func (b *tokenBucket) Consume() {
	<-b.tokenChan
}
// Stop 停止令牌桶
func (b *tokenBucket) Stop() {
	b.muLock.Lock()
	defer b.muLock.Unlock()
	b.stop = true
}

其中，

tokenBucket為令牌桶的結構，包括限制頻率、令牌桶容量和通道等；
NewTokenBucket為對外提供的創(chuàng)建令牌桶的方法；
Start為開啟令牌桶的方法；
produce為以恒定速率生成令牌的方法，以協(xié)程的方式啟動；
Consume為消費令牌的方法；
Stop為停止令牌桶的方法。

如上述所示，即為令牌桶的簡易實現。

輪子

實際上，在 Go 語言中已經提供了對令牌桶的支持了，因此不需要我們重復造輪子。

lpt-3

令牌桶，go語言創(chuàng)建和使用令牌桶

什么是令牌桶

百度百科

令牌桶算法是網絡流量整形（Traffic Shaping）和速率限制（Rate Limiting）中最常使用的一種算法。

典型情況下，令牌桶算法用來控制發(fā)送到網絡上的數據的數目，并允許突發(fā)數據的發(fā)送。

更詳細的自行搜索理解，這里只提供一下代碼思路

基本使用

代碼

package tokenBucket
import (
? ?"log"
? ?"sync"
? ?"time"
)
type TokensBucket struct {
? ?limiter float64 ? ?//速率
? ?burst ? int ? ? ? ?//桶大小
? ?mu ? ? ?sync.Mutex //鎖
? ?tokens ?float64 ? ?//桶里面的令牌數量
? ?last ? ?time.Time ?//最后一次消耗令牌的時間
}
// NewTokensBucket 創(chuàng)建令牌桶
func NewTokensBucket(limiter float64, burst int) *TokensBucket {
? ?return &TokensBucket{limiter: limiter, burst: burst}
}
// Allow 使用，每次消耗一個令牌
func (t *TokensBucket) Allow() bool {
? ?return t.AllowN(time.Now(), 1)
}
// AllowN 當前時間，一次消耗的令牌
func (t *TokensBucket) AllowN(now time.Time, i int) bool {
? ?t.mu.Lock()
? ?defer t.mu.Unlock()
? ?//當前時間-最后一次添加令牌的時間 * 桶速率 = 應該補充的令牌
? ?delta := now.Sub(t.last).Seconds() * t.limiter
? ?t.tokens += delta
? ?//桶內令牌 > 桶總大小 ?= ?只補充最大令牌數
? ?if t.tokens > float64(t.burst) {
? ? ? t.tokens = float64(t.burst)
? ?}
? ?//桶內令牌 < 需要的令牌 = 返回false
? ?if t.tokens < float64(i) {
? ? ? return false
? ?}
? ?//否則返回true，并用桶的剩余令牌 - 消耗令牌
? ?t.tokens -= float64(i)
? ?//桶最后一次補充時間重置為當前時間
? ?t.last = now
? ?//返回true
? ?return true
}

測試

func main() {
? ?bucket := NewTokensBucket(3, 5)
? ?for true {
? ? ? n := 4
? ? ? for i := 0; i < n; i++ {
? ? ? ? ?go func(i int) {
? ? ? ? ? ? if bucket.Allow() {
? ? ? ? ? ? ? ?log.Printf("allow [%d]", i)
? ? ? ? ? ? } else {
? ? ? ? ? ? ? ?log.Printf("forbid [%d]", i)
? ? ? ? ? ? }
? ? ? ? ?}(i)
? ? ? }
? ? ? time.Sleep(time.Second)
? ? ? log.Println("========================================")
? ?}
}

在開發(fā)中使用

最基本的使用，實際開發(fā)肯定不是這樣的要考慮到更多的情況，這里只是一個小的演示而已

func main() {
?? ?app := gin.Default()
?? ?bucket := tokenBucket.NewTokensBucket(1, 2)
?? ?app.Use(func(context *gin.Context) {
? ?//拿到令牌就給放行
?? ??? ?if bucket.Allow() {
?? ??? ??? ?context.Next()
? ? //拿不到就不給過
?? ??? ?} else {
?? ??? ??? ?context.JSON(500, gin.H{
?? ??? ??? ??? ?"msg": "false",
?? ??? ??? ?})
?? ??? ??? ?context.Abort()
?? ??? ?}
?? ?})
?? ?app.GET("/", func(context *gin.Context) {
?? ??? ?context.JSON(200, gin.H{
?? ??? ??? ?"msg": "success",
?? ??? ?})
?? ?})
?? ?app.Run(":80")
}