令牌桶算法实现限流

服务稳定性对于软件开发来说是十分重要的，在有些场景下可以使用服务限流来保障服务的稳定性。假设服务器的 QPS 已经达到限速阈值了, 但是并不想将所有的流量都拒之门外, 仍然让部分流量能够正常通过限流器. 这样我们的服务器在面对突发流量时能够有一定的伸缩空间, 而不是一直处于不可用状态。要实现流量的控制，必须有一种机制可以对通过设备的流量进行度量。令牌桶（Token-Bucket）是目前最常采用的一种流量测量和控制的方法，用来评估流量速率是否超过了规定值。

令牌桶

令牌桶能够在限制请求平均速率的同时还允许一定程度的突发调用。在令牌桶算法中，存在一个桶，用来存放固定数量的令牌。该算法以一定的速率往桶中放入令牌。每次请求需要先获取到桶中的令牌才能继续执行，否则等待可用的令牌，或者直接拒绝。
  放令牌的动作是持续不断进行的，如果桶中令牌数达到上限，则丢弃令牌，因此桶中可能一直持有大量的可用令牌。此时请求进来可以直接拿到令牌执行。比如设置qps为100，那么限流器初始化完成1秒后，桶中就已经有100个令牌了，如果此前还没有请求过来，这时突然来了100个请求，该限流器可以抵挡瞬时的100个请求。由此可见，只有桶中没有令牌时，请求才会进行等待，最终表现的效果即为以一定的速率执行。令牌桶的示意图如下：

当请求到达服务端时首先会根据请求数据的大小从令牌桶中取出与数据大小相当的令牌数量用来传输数据。也就是说要使请求数据被传输必须保证令牌桶里有足够多的令牌，如果令牌数量不够，则数据就不会正常被服务器处理。这就可以限制报文的流量只能小于等于令牌生成的速度，达到限制流量的目的。

ratelimit源码

ratelimit 是大部分项目都在使用的 golang 令牌桶的实现方案。RateLimiter的主要特性是它的“稳定速率”，表示正常情况下允许的最大速率。这是根据“节流”输入请求来强制执行的，例如，对于一个请求而言，计算合适的节流时间，然后让该线程等待相应的时间。下面会结合其用法, 源码剖析令牌桶的实现的方案.

话不多说,直接上代码。首先定义令牌桶数据结构Bucket：

// Bucket represents a token bucket that fills at a predetermined rate.
// Methods on Bucket may be called concurrently.
type Bucket struct {
	clock Clock

	// startTime holds the moment when the bucket was
	// first created and ticks began.
	startTime time.Time

	// capacity holds the overall capacity of the bucket.
	capacity int64

	// quantum holds how many tokens are added on each tick.
	quantum int64

	// fillInterval holds the interval between each tick.
	fillInterval time.Duration

	// mu guards the fields below it.
	mu sync.Mutex

	// availableTokens holds the number of available
	// tokens as of the associated latestTick.
	// It will be negative when there are consumers
	// waiting for tokens.
	availableTokens int64

	// latestTick holds the latest tick for which
	// we know the number of tokens in the bucket.
	latestTick int64
}

创建令牌桶的方法有以下几种：

// 创建指定填充速率和容量大小的令牌桶
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

Bucket中的availableTokens就是可用的令牌个数，fillInterval是已经走过的时间tick。限速就是要控制每秒的速度，RateLimiter速度是这样设定的：（10的九次方（纳秒）*每次填充令牌数）除以时间间隔：

// Rate returns the fill rate of the bucket, in tokens per second.
func (tb *Bucket) Rate() float64 {
	return 1e9 * float64(tb.quantum) / float64(tb.fillInterval)
}

当我们需要使用限速的时候就去查询可用的令牌，通过方法Available查询当前可用令牌数：

func (tb *Bucket) Available() int64 {
	return tb.available(tb.clock.Now())
}

// available is the internal version of available - it takes the current time as
// an argument to enable easy testing.
func (tb *Bucket) available(now time.Time) int64 {
	tb.mu.Lock()
	defer tb.mu.Unlock()
	tb.adjustavailableTokens(tb.currentTick(now))
	return tb.availableTokens
}

查询令牌数Available方法中通过adjustavailableTokens调整令牌桶当前令牌的数量：

// adjustavailableTokens adjusts the current number of tokens
// available in the bucket at the given time, which must
// be in the future (positive) with respect to tb.latestTick.
func (tb *Bucket) adjustavailableTokens(tick int64) {
	lastTick := tb.latestTick
	tb.latestTick = tick
	if tb.availableTokens >= tb.capacity {
		return
	}
	tb.availableTokens += (tick - lastTick) * tb.quantum
	if tb.availableTokens > tb.capacity {
		tb.availableTokens = tb.capacity
	}
	return
}

adjustavailableTokens调整令牌桶中令牌数量主要是通过：

1. 拿当前时间减去初始时间除以时间间隔获取tick个数，那这个tick数减去之前已经分配tick个数，就是增量的时间Tick
2. 再拿这个增量的Tick*因子就算出要增加的令牌个数，这样就调整了availableTokens 的个数了，并且更新latestTick 个数
3. 因为不能超过容量，所以才有了tb.availableTokens = tb.capacity。
4. 设计的简单可靠，不需要开启单独协程定时的往这个bucket里面加入令牌。

查到可用的令牌，就可以获取令牌了，通过Take方法获取count数量的令牌，并从令牌桶中减去count数量的令牌：

func (tb *Bucket) Take(count int64) time.Duration {
	tb.mu.Lock()
	defer tb.mu.Unlock()
	d, _ := tb.take(tb.clock.Now(), count, infinityDuration)
	return d
}

// take is the internal version of Take - it takes the current time as
// an argument to enable easy testing.
func (tb *Bucket) take(now time.Time, count int64, maxWait time.Duration) (time.Duration, bool) {
	if count <= 0 {
		return 0, true
	}

	tick := tb.currentTick(now)
	tb.adjustavailableTokens(tick)
	avail := tb.availableTokens - count
	if avail >= 0 {
		tb.availableTokens = avail
		return 0, true
	}
	// Round up the missing tokens to the nearest multiple
	// of quantum - the tokens won't be available until
	// that tick.

	// endTick holds the tick when all the requested tokens will
	// become available.
	endTick := tick + (-avail+tb.quantum-1)/tb.quantum
	endTime := tb.startTime.Add(time.Duration(endTick) * tb.fillInterval)
	waitTime := endTime.Sub(now)
	if waitTime > maxWait {
		return 0, false
	}
	tb.availableTokens = avail
	return waitTime, true
}

Take方法里面不仅涉及到令牌获取，获取后tb.availableTokens - count，而且还可以估算等待时间，如果超过现有令牌，可以预估等待时间waitTime，这样不仅可以获取，还可以通过Wait方法等待，如果成功返回0, true。

Wait等待的方式获取令牌，直到令牌数符合要求：

// Wait takes count tokens from the bucket, waiting until they are
// available.
func (tb *Bucket) Wait(count int64) {
	if d := tb.Take(count); d > 0 {
		tb.clock.Sleep(d)
	}
}

TakeAvailable方法会安全的拿到可用的令牌，如果获取令牌超过的话，会获取并返回当前现有的令牌。当然，如果你已经知道速度，就可以创建一个已知rate的令牌桶了

// TakeAvailable takes up to count immediately available tokens from the
// bucket. It returns the number of tokens removed, or zero if there are
// no available tokens. It does not block.
func (tb *Bucket) TakeAvailable(count int64) int64 {
	tb.mu.Lock()
	defer tb.mu.Unlock()
	return tb.takeAvailable(tb.clock.Now(), count)
}

// takeAvailable is the internal version of TakeAvailable - it takes the
// current time as an argument to enable easy testing.
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
}

以上就是ratelimit源码的核心内容，了解了这些，就可以去按照需求实现自己的限流器了。

限流中间件

在gin框架中结合ratelimit加入简易版限流中间件：

package main

import (
	"log"
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/juju/ratelimit"
)

// 定义令牌桶初始化参数
var limiter = ratelimit.NewBucketWithQuantum(time.Second, 10, 10)

// 限流中间件
func tokenRateLimiter() gin.HandlerFunc {
	return func(context *gin.Context) {
		// 每次请求拿出一个令牌
		if limiter.TakeAvailable(1) == 0 {
			// 没有可用令牌，则请求失败
			log.Printf("available token :%d", limiter.Available())
			context.AbortWithStatusJSON(http.StatusTooManyRequests, "Too Many Request")
		} else {
			// 有可用令牌，请求成功
			context.Next()
		}
	}
}

func main() {
	e := gin.Default()
	// 请求中加入限流中间件
	e.GET("/test", tokenRateLimiter(), func(context *gin.Context) {
		context.JSON(200, true)
	})
	e.Run(":8080")
}