前言
本文是记录的是" Go初级之http限流令牌桶的基本实现 "
此文是个人学习归纳的记录,腾讯云独家发布,未经允许,严禁转载,如有不对, 还望斧正, 感谢!
令牌桶是一种常用的流量控制技术,其本身没有丢弃和优先级策略。令牌桶的工作原理如下:
1. 令牌以一定的速率放入桶中。 2. 每个令牌允许源发送一定数量的比特。 3. 发送一个包,流量调节器就要从桶中删除与包大小相等的令牌数。 4. 如果没有足够的令牌发送包,这个包就会等待直到有足够的令牌(在整形器的情况下)或者包被丢弃,也有可能被标记更低的DSCP(在策略者的情况下)。 5. 桶有特定的容量,如果桶已经满了,新加入的令牌就会被丢弃。因此,在任何时候,源发送到网络上的最大突发数据量与桶的大小成比例。令牌桶允许突发,但是不能超过限制。
令牌桶算法是网络流量整形(Traffic Shaping)和速率限制(Rate Limiting)中最常使用的一种算法。典型情况下,令牌桶算法用来控制发送到网络上的数据的数目,并允许突发数据的发送。
除了令牌桶,还有其他几种常用的限流算法,包括滑动窗口、计数器和漏桶。
这些限流算法各有优缺点,可以根据具体场景选择合适的算法。例如,如果需要精确控制请求速率,可以使用令牌桶或漏桶算法;如果需要简单的速率限制,可以使用计数器或滑动窗口算法。
上面我已经解释很清楚了,我们通过控制令牌桶中的令牌的使用和生成来对http请求之类的流量进行控制,所以我们主要关心的就是桶的容积,桶中令牌的数量。
然后我们简单定义一个桶
type Bucket struct {
cap int // 容积
value int // 当前令牌数
lastTime int // 上次添加的时间
lock sync.Mutex // 来个互斥锁,保证原子性
}
上次添加的时间是什么意思呢,就是我们需要动态控制桶中的令牌的数量,可以通过这个字段和当前时间进行比较,如果时间差距很大,说明请求很稀疏,然后我们可以通过自定义的方法,动态控制令牌的数量。
定义好了上面的Bucket桶结构之后,我们再写入以下代码,
// 创建一个桶,传入桶的容积和当前的令牌数
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
}
上面的是一个简陋的,只实现了基本原理。
先要安装 "github.com/juju/ratelimit" 这个包
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() // 放行
}
}
到这里,我们也就差不多了,前面实现的简陋的和这个实际上也差不多,有区别的地方,就是在于判断是否限速,这个需要多方考虑。
简陋的里面只是简单的用一定时间的请求量来判断了,还可以优化。
原创声明:本文系作者授权腾讯云开发者社区发表,未经许可,不得转载。
如有侵权,请联系 cloudcommunity@tencent.com 删除。
原创声明:本文系作者授权腾讯云开发者社区发表,未经许可,不得转载。
如有侵权,请联系 cloudcommunity@tencent.com 删除。