前言
本文是记录的是" Go初级之http限流令牌桶的基本实现 "
此文是个人学习归纳的记录，腾讯云独家发布，未经允许，严禁转载，如有不对, 还望斧正, 感谢!
关于令牌桶
令牌桶是一种常用的流量控制技术，其本身没有丢弃和优先级策略。令牌桶的工作原理如下：
1. 令牌以一定的速率放入桶中。 2. 每个令牌允许源发送一定数量的比特。 3. 发送一个包，流量调节器就要从桶中删除与包大小相等的令牌数。 4. 如果没有足够的令牌发送包，这个包就会等待直到有足够的令牌（在整形器的情况下）或者包被丢弃，也有可能被标记更低的DSCP（在策略者的情况下）。 5. 桶有特定的容量，如果桶已经满了，新加入的令牌就会被丢弃。因此，在任何时候，源发送到网络上的最大突发数据量与桶的大小成比例。令牌桶允许突发，但是不能超过限制。
令牌桶算法是网络流量整形（Traffic Shaping）和速率限制（Rate Limiting）中最常使用的一种算法。典型情况下，令牌桶算法用来控制发送到网络上的数据的数目，并允许突发数据的发送。
常见的限流算法
除了令牌桶，还有其他几种常用的限流算法，包括滑动窗口、计数器和漏桶。
- 滑动窗口：滑动窗口算法通过将时间分为固定大小的窗口来限流。每个窗口内的请求数量被计数，当请求数量超过限制时，请求将被拒绝。当窗口滑动到下一个时间段时，计数器被重置，新的请求计数开始。滑动窗口算法可以处理突发流量，但可能导致请求在窗口之间的边界处被拒绝。
- 计数器：计数器算法通过在固定的时间间隔内计算请求数量来限流。例如，每秒钟计算请求数量，并在达到限制时拒绝请求。计数器算法简单易实现，但可能导致突发流量问题。
- 漏桶：漏桶算法通过将请求放入一个有限容量的桶中来限流。请求以固定的速率进入桶，而桶以固定的速率消耗请求。当桶满时，新的请求将被丢弃。漏桶算法可以处理突发流量，并在限制请求速率的同时，允许短时间的突发流量通过。
这些限流算法各有优缺点，可以根据具体场景选择合适的算法。例如，如果需要精确控制请求速率，可以使用令牌桶或漏桶算法；如果需要简单的速率限制，可以使用计数器或滑动窗口算法。
简单地用go语言的代码实现一个限流的令牌桶
上面我已经解释很清楚了，我们通过控制令牌桶中的令牌的使用和生成来对http请求之类的流量进行控制，所以我们主要关心的就是桶的容积，桶中令牌的数量。
然后我们简单定义一个桶
代码语言：go复制type Bucket struct {
cap int // 容积
value int // 当前令牌数
lastTime int // 上次添加的时间
lock sync.Mutex // 来个互斥锁,保证原子性
}上次添加的时间是什么意思呢，就是我们需要动态控制桶中的令牌的数量，可以通过这个字段和当前时间进行比较，如果时间差距很大，说明请求很稀疏，然后我们可以通过自定义的方法，动态控制令牌的数量。
定义好了上面的Bucket桶结构之后，我们再写入以下代码，
代码语言：go复制// 创建一个桶，传入桶的容积和当前的令牌数
func NewBucket(cap int, value int) *Bucket {
return &Bucket{cap: cap, value: value, lastTime: int(time.Now().Unix())}
}
// 定义一个add函数，用来动态控制令牌的数量，比较简陋
func (B *Bucket) add() {
now := int(time.Now().Unix()) // 得到当前的时间
lastTime := B.lastTime // 得到上一次请求的时间
B.lastTime = now
// 当前令牌数大于容量时，不添加令牌
if B.value >= l.cap {
return
}
l.value = (now - lastTime) / 1000 // 此处就简单用时间来控制了。
if B.value > l.cap {
l.value = l.cap
}
}
// 判断当前令牌桶的令牌还有没有，返回布尔值
func (B *Bucket) IsOk() bool {
B.lock.Lock()
defer B.lock.Unlock()
B.add() // 调用令牌控制函数动态控制令牌数量
if B.value > 0 {
l.value--
return true
}
return false
}上面的是一个简陋的，只实现了基本原理。
下面我们写一个gin的中间件，基于第三方包的实现
先要安装 "github.com/juju/ratelimit" 这个包
代码语言：go复制package middleware
import (
"github.com/gin-gonic/gin"
"github.com/juju/ratelimit"
"strings"
"time"
)
// SpeedLimit 限速路由结构体, 用来指定要限速的路由
type SpeedLimit struct {
Limit map[string]*ratelimit.Bucket // 键值对，前面放路由，后面是令牌桶
}
// LimiterBucketRule 是一个结构体，用于定义令牌桶规则
type LimiterBucketRule struct {
// Key 是键也是路由
Key string
// FillInterval 是令牌桶的填充间隔
FillInterval time.Duration
// Capacity 是令牌桶的容量
Capacity int64
// Quantum 是令牌桶的单次请求所需的令牌数量
Quantum int64
}
// NewSpeedLimit 新建限速结构体
func NewSpeedLimit() *SpeedLimit {
return &SpeedLimit{
Limit: make(map[string]*ratelimit.Bucket),
}
}
// NewLimiterBucketRule 新建限速规则
func NewLimiterBucketRule(key string, fillInterval time.Duration, capacity int64, quantum int64) LimiterBucketRule {
return LimiterBucketRule{
Key: key,
FillInterval: fillInterval,
Capacity: capacity,
Quantum: quantum,
}
}
// GetKey 获取限速的键，也就是限速的路由
func (sl *SpeedLimit) GetKey(c *gin.Context) string {
url := c.Request.URL.Path
// 切分url
index := strings.Index(url, "?")
if index != -1 {
url = url[:index]
}
return url
}
// GetBucket 获取限速的令牌桶
func (sl *SpeedLimit) GetBucket(key string) (*ratelimit.Bucket, bool) {
bucket, flag := sl.Limit[key]
return bucket, flag
}
// AddToLimit 添加限速规则
func (sl *SpeedLimit) AddToLimit(rules ...LimiterBucketRule) *SpeedLimit {
for _, rule := range rules {
// 检查令牌桶是否已经存在
if _, flag := sl.Limit[rule.Key]; flag {
sl.Limit[rule.Key] = ratelimit.NewBucketWithQuantum(rule.FillInterval, rule.Capacity, rule.Quantum)
}
}
return sl
}
// SpeedLimitMiddleware gin路由限速路由中间件
func SpeedLimitMiddleware(sl *SpeedLimit) gin.HandlerFunc {
return func(c *gin.Context) {
// 获取限速的键，也就是请求的路由
key := sl.GetKey(c)
// 获取该条路由限速的令牌桶
bucket, flag := sl.GetBucket(key)
// 判断是否要限速
if flag {
// 获取令牌数量
count := bucket.TakeAvailable(1)
if count == 0 {
c.AbortWithStatusJSON(403, gin.H{
"code": 403,
"msg": "访问过于频繁",
})
return
}
}
c.Next() // 放行
}
}到这里，我们也就差不多了，前面实现的简陋的和这个实际上也差不多，有区别的地方，就是在于判断是否限速，这个需要多方考虑。
简陋的里面只是简单的用一定时间的请求量来判断了，还可以优化。
我正在参与2024腾讯技术创作特训营最新征文，快来和我瓜分大奖！
