GO--基于令牌桶和漏桶的限流策略

至于为什么要限流，字面意思已经很清楚了，就是为了减轻服务器的压力

下面我们将介绍两个限流策略----漏桶和令牌桶。

漏桶

原理介绍

漏桶，顾名思义就是一个漏斗，漏斗嘴的大小是固定的，所以不管漏斗现容量多大，都不会影响漏斗出水的速度。类比到我们的web服务中，我们可以为web服务准备一个固定速率的请求处理器（漏桶），对我们的请求，如果此时桶内请求量超过了桶的最大容量，那么就执行特定的抛弃策略（直接丢弃或者阻塞请求）。与此同时，对我们的请求，进行固定速率的处理。这样就会形成一种限流的效果。

使用

基于这个原理，有很多个人或者组织都自己开发了相应的算法实现，我们今天介绍的是github.com/uber-go/ratelimithttps://github.com/uber-go/ratelimit

这个算法比较简单，库的使用也比较清晰，有需要的可以去看一看源代码，下面我将介绍

如何在项目中使用这个库，下面我将以gin框架为例：

导包

go.uber.org/ratelimit

编写中间件

Go 复制代码

package ratelimit

// 基于gin 框架 编写限流中间件 并 使用
// 限流中间件：漏桶 和 令牌桶

import (
	"time"

	"github.com/gin-gonic/gin"
	loutong "go.uber.org/ratelimit" // 漏桶中间件(具名导入)
)

// 漏桶策略限流器
func RateLimitMiddleware(rate int) func(c *gin.Context) {
    // 使用闭包写中间件，可以在使用中间件的时候，传入指定参数( rate 是请求数，每秒处理请求数 )
    // 创建(漏桶)限流器，指定处理速率（不限制速率：NewUnlimited()）
    rl := loutong.New(rate)
    // 返回一个函数
    return func(c *gin.Context) {
        // 获得当前的时间
        now := time.Now()
        // 尝试从漏桶中获得一个请求去处理
        rl.Take()
        // 获得处理请求花费的时间
        cost := time.Since(now)
        if cost > 0 {
            // 如果有等待，那么花费时间大于0
            c.Abort()       // 丢弃请求
            // 对待等待的请求进行处理：阻塞 或者 丢弃
        }
        c.Next()
    }
}
/*
使用Take()方法的作用：
如果当前请求的流入速度太快，超过了设定的每秒请求数（RPS）的标准，那么这个 "Take" 操作就需要阻塞（暂停、等待），
不让过多的请求一下子涌入后续的处理流程，直到请求流入的速度符合设定好的每秒请求数这个要求，
以此来实现精准的限流控制，保证系统按照既定的请求处理速率稳定运行。
Take()方法返回一个time值，反映的是 请求等待的时间 
*/

测试

Go 复制代码

package main

import (
	"component/ratelimit"

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

func main() {
	r := gin.Default()

	r.Use(ratelimit.RateLimitMiddleware(5))

	r.GET("/ping", func(c *gin.Context) {
		c.JSON(200, gin.H{
			"message": "pong",
		})
	})

	r.NoRoute(func(c *gin.Context){
		c.JSON(404, gin.H{
			"msg": "404,您的也页面好像找不到了~",
		})
	})

	r.Run("127.0.0.1:8080")
}

优点
- 流量整形效果好
  - 漏桶算法能够平滑突发流量。例如，在网络服务中，当有大量用户在短时间内发起请求（如某热门商品限时抢购活动导致大量用户同时访问电商网站），这些请求会先进入 "漏桶" 缓存起来，然后以固定的速率流出进行处理。这就使得后端服务接收到的请求是平稳的，避免了后端服务因突发的高流量而崩溃。
- 易于理解和实现
  - 其原理简单直观，类似于生活中的漏斗。从编程角度看，实现一个基本的漏桶算法不需要复杂的代码结构。比如，可以用一个简单的队列来模拟 "漏桶"，记录请求的到达时间，按照固定的时间间隔从队列头部取出请求进行处理。
- 对资源的消耗可预测
  - 因为处理请求的速率是固定的，所以系统在单位时间内消耗的资源（如 CPU、内存等）是可以预估的。例如，一个按照每秒 10 个请求速率处理的服务，只要提前评估每个请求处理所需的资源量，就可以大致计算出系统在一段时间内需要的资源总量。
缺点
- 可能导致响应延迟增加
  - 当请求流量突发且超过桶的容量时，新的请求可能会被阻塞或丢弃。对于被阻塞的请求，会导致用户等待时间变长。例如，在一个在线支付系统中，如果因为流量限制而导致支付请求被阻塞，用户可能会等待很长时间才能完成支付，这会影响用户体验。
- 不能充分利用系统资源
  - 由于漏桶算法是按照固定速率处理请求，即使系统有足够的资源来处理更多的请求，也不会加快处理速度。比如，在深夜时段，服务器资源利用率很低，但是因为漏桶算法限制了请求处理速率，使得服务器不能更快地处理请求，从而导致资源闲置。

令牌桶

原理介绍

令牌桶的原理其实和漏桶差不多，漏桶相当于在单位时间（比如1秒内处理多少个请求）。那么令牌桶就是表现的间接了点。他是通过指定每秒的令牌数，从而指定每秒处理请求的个数（每个请求在被处理的时候都会拿取一个令牌）

使用

同样，这个算法也有很多人去实现，我们这里以下面这个为例：

github.com/juju/ratelimithttps://github.com/juju/ratelimit这个库支持多种令牌桶模式，并且使用起来也比较简单。

导包

github.com/juju/ratelimit

编写中间件

Go 复制代码

package ratelimit

// 基于gin 框架 编写限流中间件 并 使用
// 限流中间件：漏桶 和 令牌桶

import (
	"net/http"
	"time"
	"github.com/gin-gonic/gin"
	lingpai "github.com/juju/ratelimit"
)

// 令牌桶中间件
// 限流中间件: 每 fillInterval 的时间，可最多处理 1 个请求(相当于对处理请求的 打点器)
// 特点： 如果一段时间内不请求，token会存储，直到桶的最大容量cap，应对突发多个请求
func RateLimitMiddleware(fillInterval time.Duration, cap int64) func(c *gin.Context) {
    // 创建令牌桶 ： 参数： 1. 填充时间（单位时间） 2. 容量
    bucket := lingpai.NewBucket(fillInterval, cap)
    return func(c *gin.Context) {
        // 判断是否限流(如果取不到 token 就限流)
        if bucket.TakeAvailable(1) <= 0 {
            c.String(http.StatusOK, "你点击的太快了，请慢点点击~")
            // 终止执行
            c.Abort()
            return
        }
        // 取到令牌继续执行
        c.Next()
    }
}

创建令牌桶的方法：

Go 复制代码

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

取出令牌的方法如下：

Go 复制代码

// 取token（非阻塞）
func (tb *Bucket) Take(count int64) time.Duration
func (tb *Bucket) TakeAvailable(count int64) int64

// 最多等maxWait时间取token
func (tb *Bucket) TakeMaxDuration(count int64, maxWait time.Duration) (time.Duration, bool)

// 取token（阻塞）
func (tb *Bucket) Wait(count int64)
func (tb *Bucket) WaitMaxDuration(count int64, maxWait time.Duration) bool

测试

Go 复制代码

package main

import (
	"component/ratelimit"
	"time"

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

func main() {
	r := gin.Default()

	r.Use(ratelimit.RateLimitMiddleware(1*time.Second, 10))

	r.GET("/ping", func(c *gin.Context) {
		c.JSON(200, gin.H{
			"message": "pong",
		})
	})

	r.NoRoute(func(c *gin.Context){
		c.JSON(404, gin.H{
			"msg": "404,您的也页面好像找不到了~",
		})
	})

	r.Run("127.0.0.1:8080")
}

优点
- 支持突发流量处理
  - 令牌桶算法允许流量在一定范围内突发。因为令牌是以固定速率生成并积累在桶中的，所以在短时间内，如果桶中有足够的令牌，就可以允许高于平均速率的突发流量通过。例如，一个令牌桶每秒产生 10 个令牌，桶的容量为 100 个令牌。在某一时刻，桶内积累了 90 个令牌，此时如果有 50 个请求同时到达，只要令牌足够，这 50 个请求就可以立即通过，这对于应对突发的高流量情况（如限时抢购活动）非常有效。
- 有效利用系统资源
  - 相比于漏桶算法，令牌桶算法能够更好地利用系统资源。当系统资源充足且有足够的令牌时，请求可以快速通过，不会像漏桶算法那样限制请求只能以固定的较低速率通过。例如，在服务器负载较低的时段，大量请求可以利用积累的令牌快速处理，充分发挥系统的处理能力。
- 流量控制灵活
  - 可以通过调整令牌生成速率和桶的容量来灵活地控制流量。比如，对于不同重要性的服务或者不同的用户级别，可以设置不同的令牌桶参数。对于高级别的用户或者重要的服务接口，可以设置较大的令牌生成速率和桶容量，以保证其流量优先通过。
- 平均速率限制准确
  - 虽然允许突发流量，但令牌桶算法依然能够保证在较长的时间范围内，请求通过的平均速率不会超过令牌生成的速率。这就使得系统可以在允许一定程度的灵活性的同时，维持一个稳定的整体流量水平。
缺点
- 实现相对复杂
  - 与漏桶算法相比，令牌桶算法的实现较为复杂。它需要维护令牌的生成、存储和消耗的逻辑。在代码层面，需要考虑如何准确地按照固定速率生成令牌，如何有效地存储令牌（可能涉及到数据结构的选择，如队列、计数器等），以及如何正确地在请求到达时消耗令牌。
- 参数调整难度较大
  - 令牌桶算法的性能高度依赖于令牌生成速率和桶容量这两个参数的设置。如果参数设置不合理，可能会导致流量控制效果不佳。例如，若令牌生成速率设置过高，可能无法有效地限制流量，导致系统过载；若桶容量设置过小，可能无法充分利用系统允许的突发流量特性，频繁地拒绝请求。而且在实际的复杂系统中，要准确地确定这两个参数的合适值需要进行大量的测试和性能评估。