当前位置：首页 > article >正文

【造个轮子】使用Golang实现简易令牌桶算法

article 2025/3/4 15:19:05

在这里插入图片描述

本文目录

1. 令牌桶算法
2. 调用第三方库实现令牌桶
3. 手撕令牌桶

前言：之前在Bluebell社区项目中，我们使用了开源的库来实现令牌桶限流，这次我们试着使用Go来手撕实现下令牌桶算法。

1. 令牌桶算法

为了防止网络拥塞，需要限制流出或者流入网络的流量，使流量以比较均匀的速度向外发送。令牌桶算法就实现了这个功能，可控制发送到网络上数据的数目，并允许突发数据的发送。

令牌桶算法是网络流量整形和速率限制中最常使用的一种算法。大小固定的令牌桶可自行以恒定的速率源源不断地产生令牌。如果令牌不被消耗，或者被消耗的速度小于产生的速度，令牌就会不断地增多，直到把桶填满。后面再产生的令牌就会从桶中溢出。最后桶中可以保存的最大令牌数永远不会超过桶的大小。

在这里插入图片描述

传送到令牌桶的数据包需要消耗令牌。不同大小的数据包，消耗的令牌数量不一样。

令牌桶中的每一个令牌都代表一次请求或者一个字节。如果令牌桶中存在令牌，则允许发送流量；而如果令牌桶中不存在令牌，则不允许发送流量。因此，如果突发门限被合理地配置并且令牌桶中有足够的令牌，那么流量就可以以峰值速率发送。

2. 调用第三方库实现令牌桶

上次我们调用了封装好的令牌桶算法，代码如下。

回顾一下关键点，fillInterval time.Duration：表示令牌桶的填充间隔时间（例如，每秒填充一次）。cap int64：表示令牌桶的最大容量（即桶中最多可以存储的令牌数量）。

它返回一个函数，该函数的签名是 func(c *gin.Context)，这符合 gin 中间件的标准形式。也就是接收一个 *gin.Context 的参数，用于处理每个 HTTP 请求。

这里的 func(c *gin.Context) { ... } 是一个匿名函数。它捕获 bucket 变量（闭包特性），因此可以在每次 HTTP 请求时使用同一个 bucket 实例。

func RateLimitMiddleware(fillInterval time.Duration, cap int64) func(c *gin.Context) {
	bucket := ratelimit.NewBucket(fillInterval, cap)
	return func(c *gin.Context) {
		// 如果取不到令牌就中断本次请求返回 rate limit...
		if bucket.TakeAvailable(1) == 0 {
			c.String(http.StatusOK, "rate limit...")
			c.Abort()
			return
		}
		// 取到令牌就放行
		c.Next()
	}
}

我们在路由组中使用Use调用中间件，代码如下。

在这里插入图片描述

r.Use(Logger.GinLogger(), Logger.GinRecovery(stack: true), middlewares.RateLimitMiddleware(2*time.Second, cap: 1))

Logger.GinLogger()：添加了一个日志记录中间件，用于记录 HTTP 请求的日志信息。

Logger.GinRecovery(stack: true)：添加了一个错误恢复中间件，用于捕获和处理在请求处理过程中发生的任何 panic，并可选地记录堆栈跟踪。

middlewares.RateLimitMiddleware(2*time.Second, cap: 1)：添加了速率限制中间件，配置为每两秒钟向令牌桶中添加一个令牌，桶的容量为 1。这意味着在任何给定的两秒时间内，最多只能处理一个请求。

3. 手撕令牌桶

直接上代码吧，逻辑比较简单。咱们来看看总体代码，再进行部分讲解。

package awesomeProject

import (
	"sync"
	"time"
)

// 定义令牌桶结构
type tokenBucket struct {
	limitRate int           // 限制频率，即每分钟加入多少个令牌
	tokenChan chan struct{} // 令牌通道，可以理解为桶
	cap       int           // 令牌桶的容量
	muLock    *sync.Mutex   // 令牌桶锁，保证线程安全
	stop      bool          // 停止标记，结束令牌桶
}

// NewTokenBucket 创建令牌桶
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
}

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,
	}
}

tokenChan chan struct{}：tokenChan 是一个类型为 chan struct{} 的通道，用于存储令牌。

make(chan struct{}, cap) 创建了一个容量为 cap 的缓冲通道。这里 struct{} 表示通道中存储的是空结构体，仅用于占位，因为我们关心的是通道中元素的数量，而不是元素的值。这个通道模拟了令牌桶中令牌的数量，通道中的每个空结构体代表桶中的一个令牌。

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)
	}
}