使用Golang如何實現(xiàn)簡易的令牌桶算法

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

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

簡介

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

令牌桶算法就實現(xiàn)了這個功能，可控制發(fā)送到網(wǎng)絡(luò)上數(shù)據(jù)的數(shù)目，并允許突發(fā)數(shù)據(jù)的發(fā)送。

令牌桶算法是網(wǎng)絡(luò)流量整形和速率限制中最常使用的一種算法。

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

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

后面再產(chǎn)生的令牌就會從桶中溢出。最后桶中可以保存的最大令牌數(shù)永遠(yuǎn)不會超過桶的大小。

ipt-1

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

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

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

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

lpt-2

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

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

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

實現(xiàn)

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

package awesomeProject
import (
	"sync"
	"time"
)
// 定義令牌桶結(jié)構(gòu)
type tokenBucket struct {
	limitRate int           // 限制頻率，即每分鐘加入多少個令牌
	tokenChan chan struct{} // 令牌通道，可以理解為桶
	cap       int           // 令牌桶的容量
	muLock    *sync.Mutex   // 令牌桶鎖，保證線程安全
	stop      bool          // 停止標(biāo)記，結(jié)束令牌桶
}
// 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()
}
// 生產(chǎn)令牌
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為令牌桶的結(jié)構(gòu)，包括限制頻率、令牌桶容量和通道等；
NewTokenBucket為對外提供的創(chuàng)建令牌桶的方法；
Start為開啟令牌桶的方法；
produce為以恒定速率生成令牌的方法，以協(xié)程的方式啟動；
Consume為消費令牌的方法；
Stop為停止令牌桶的方法。

如上述所示，即為令牌桶的簡易實現(xiàn)。

輪子

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

lpt-3

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

什么是令牌桶

百度百科

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

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

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

基本使用

代碼

package tokenBucket
import (
? ?"log"
? ?"sync"
? ?"time"
)
type TokensBucket struct {
? ?limiter float64 ? ?//速率
? ?burst ? int ? ? ? ?//桶大小
? ?mu ? ? ?sync.Mutex //鎖
? ?tokens ?float64 ? ?//桶里面的令牌數(shù)量
? ?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 當(dāng)前時間，一次消耗的令牌
func (t *TokensBucket) AllowN(now time.Time, i int) bool {
? ?t.mu.Lock()
? ?defer t.mu.Unlock()
? ?//當(dāng)前時間-最后一次添加令牌的時間 * 桶速率 = 應(yīng)該補充的令牌
? ?delta := now.Sub(t.last).Seconds() * t.limiter
? ?t.tokens += delta
? ?//桶內(nèi)令牌 > 桶總大小 ?= ?只補充最大令牌數(shù)
? ?if t.tokens > float64(t.burst) {
? ? ? t.tokens = float64(t.burst)
? ?}
? ?//桶內(nèi)令牌 < 需要的令牌 = 返回false
? ?if t.tokens < float64(i) {
? ? ? return false
? ?}
? ?//否則返回true，并用桶的剩余令牌 - 消耗令牌
? ?t.tokens -= float64(i)
? ?//桶最后一次補充時間重置為當(dāng)前時間
? ?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")
}