江苏龙网

  1. 主页 > 生活百科 > >

漏桶、令牌桶限流算法的Go语言实现( 二 )

上面的代码根据记录每次请求的间隔时间和上一次请求的时刻来计算当次请求需要阻塞的时间——sleepFor，这里需要留意的是sleepFor的值可能为负，在经过间隔时间长的两次访问之后会导致随后大量的请求被放行，所以代码中针对这个场景有专门的优化处理 。maxSlack默认值可以通过创建限制器的New函数看到 。
func New(rate int, opts ...Option) Limiter { l := &limiter{  perRequest: time.Second / time.Duration(rate),  maxSlack:   -10 * time.Second / time.Duration(rate), } for _, opt := range opts {  opt(l) } if l.clock == nil {  l.clock = clock.New() } return l}令牌桶令牌桶其实和漏桶的原理类似，令牌桶按固定的速率往桶里放入令牌，并且只要能从桶里取出令牌就能通过，令牌桶支持突发流量的快速处理 。

漏桶、令牌桶限流算法的Go语言实现

文章插图
 
令牌桶原理
对于从桶里取不到令牌的场景，我们可以选择等待也可以直接拒绝并返回 。
对于令牌桶的 Go 语言实现，大家可以参照https://github.com/juju/ratelimit 。
这个库支持多种令牌桶模式，并且使用起来也比较简单 。
创建令牌桶的方法：
// 创建指定填充速率和容量大小的令牌桶func NewBucket(fillInterval time.Duration, capacity int64) *Bucket// 创建指定填充速率、容量大小和每次填充的令牌数的令牌桶func NewBucketWithQuantum(fillInterval time.Duration, capacity, quantum int64) *Bucket// 创建填充速度为指定速率和容量大小的令牌桶// NewBucketWithRate(0.1, 200) 表示每秒填充20个令牌func NewBucketWithRate(rate float64, capacity int64) *Bucket取出令牌的方法：
// 非阻塞的取tokenfunc (tb *Bucket) Take(count int64) time.Durationfunc (tb *Bucket) TakeAvailable(count int64) int64// 最多等maxWait时间取tokenfunc (tb *Bucket) TakeMaxDuration(count int64, maxWait time.Duration) (time.Duration, bool)// 阻塞的取tokenfunc (tb *Bucket) Wait(count int64)func (tb *Bucket) WaitMaxDuration(count int64, maxWait time.Duration) bool虽说是令牌桶，但是我们没有必要真的去生成令牌放到桶里，我们只需要每次来取令牌的时候计算一下，当前是否有足够的令牌可以使用就可以了，具体的计算公式如下 。
当前令牌数 = 上一次剩余的令牌数 + (本次取令牌的时刻-上一次取令牌的时刻)/放置令牌的时间间隔 * 每次放置的令牌数github.com/juju/ratelimit这个库中关于令牌数计算的具体实现如下：
func (tb *Bucket) adjustavailableTokens(tick int64) { if tb.availableTokens >= tb.capacity {  return } tb.availableTokens += (tick - tb.latestTick) * tb.quantum if tb.availableTokens > tb.capacity {  tb.availableTokens = tb.capacity } tb.latestTick = tick return}获取令牌的TakeAvailable函数关键部分的源码如下：
func (tb *Bucket) takeAvailable(now time.Time, count int64) int64 { if count <= 0 {  return 0 } tb.adjustavailableTokens(tb.currentTick(now)) if tb.availableTokens <= 0 {  return 0 } if count > tb.availableTokens {  count = tb.availableTokens } tb.availableTokens -= count return count}大家从代码中也可以看到其实令牌桶的实现并没有很复杂 。
gin 框架中使用限流中间件在 gin 框架构建的项目中，我们可以将限流组件定义成中间件 。
这里使用令牌桶作为限流策略，编写一个限流中间件如下：
func RateLimitMiddleware(fillInterval time.Duration, cap int64) func(c *gin.Context) { bucket := ratelimit.NewBucket(fillInterval, cap) return func(c *gin.Context) {  // 如果取不到令牌就返回响应  if bucket.TakeAvailable(1) == 0 {   c.String(http.StatusOK, "rate limit...")   c.Abort()   return  }  c.Next() }}


推荐阅读


上一篇：理解卷积神经网络中的自注意力机制

下一篇：微软提交补丁，用Linux替代Hyper-V根分区的Win