RateLimit 限流中间件

前言

为什么需要限流中间件

在大数据量高并发访问时经常会出现服务或接口面对大量的请求而导致数据库崩溃的情况甚至引发连锁反映导致整个系统崩溃。或者有人恶意攻击网站大量的无用请求出现会导致缓存穿透的情况出现。使用限流中间件可以在短时间内对请求进行限制数量起到降级的作用从而保障了网站的安全性。

应对大量并发请求的策略

1. 使用消息中间件进行统一限制降速
2. 使用限流方案将多余请求返回限流
3. 升级服务器
4. 负载均衡升级
5. 等等

可以看出在代码已经无法提升的情况下只能去提升硬件水平。或者改动架构再加一层也可以使用消息中间件统一处理。而结合看来限流方案是一种既不需要大幅改动也不需要高额开销的策略。

常见的限流方案

1. 令牌桶算法
2. 漏桶算法
3. 滑动窗口算法
4. 等等

这里主要根据golang的库介绍令牌桶和漏桶的实现原理以及在实际项目中如何应用。

漏桶

引入ratelimit库

go get -u go.uber.org/ratelimit

库函数源代码

// New returns a Limiter that will limit to the given RPS.
func New(rate int, opts ...Option) Limiter {
	return newAtomicBased(rate, opts...)
}

// newAtomicBased returns a new atomic based limiter.
func newAtomicBased(rate int, opts ...Option) *atomicLimiter {
	// TODO consider moving config building to the implementation
	// independent code.
	config := buildConfig(opts)
	perRequest := config.per / time.Duration(rate)
	l := &atomicLimiter{
		perRequest: perRequest,
		maxSlack:   -1 * time.Duration(config.slack) * perRequest,
		clock:      config.clock,
	}

	initialState := state{
		last:     time.Time{},
		sleepFor: 0,
	}
	atomic.StorePointer(&l.state, unsafe.Pointer(&initialState))
	return l
}

该函数使用了函数选项模式对多个结构体对象进行初始化

首先根据传入的值来初始化一个桶结构体 rate 为int传参 (time.Duration(rate)单位为纳秒 = 1/1e9秒)

初始化过程中包括了

• 每一滴水需要的时间 perquest = config.per / time.Duration(rate)

• maxSlack 宽松度宽松度为负值-1 * time.Duration(config.slack) * perRequest 松紧度是用来规范等待时间的

// Clock is the minimum necessary interface to instantiate a rate limiter with
// a clock or mock clock, compatible with clocks created using
// github.com/andres-erbsen/clock.
type Clock interface {
   Now() time.Time
   Sleep(time.Duration)
}

同时还需要结构体Clock来记录当前请求的时间now和此刻的请求所需要花费等待的时间sleep

type state struct {
   last     time.Time
   sleepFor time.Duration
}

state 主要用来记录上次执行的时间以及当前执行请求需要花费等待的时间作为中间状态记录

最重要的Take逻辑

// Take blocks to ensure that the time spent between multiple
// Take calls is on average time.Second/rate.
func (t *atomicLimiter) Take() time.Time {
   var (
      newState state
      taken    bool
      interval time.Duration
   )
   for !taken {
      now := t.clock.Now()

      previousStatePointer := atomic.LoadPointer(&t.state)
      oldState := (*state)(previousStatePointer)

      newState = state{
         last:     now,
         sleepFor: oldState.sleepFor,
      }

      // If this is our first request, then we allow it.
      if oldState.last.IsZero() {
         taken = atomic.CompareAndSwapPointer(&t.state, previousStatePointer, unsafe.Pointer(&newState))
         continue
      }
      // 计算是否需要进行等待取水操作
      newState.sleepFor += t.perRequest(每两滴水之间的间隔时间) - now.Sub(oldState.last)(当前时间与上次取水时间的间隔)
       
       // 如果等待取水时间特别小就需要松紧度进行维护
      if newState.sleepFor < t.maxSlack {
         newState.sleepFor = t.maxSlack
      }
       // 如果等待时间大于0就进行更新
      if newState.sleepFor > 0 {
         newState.last = newState.last.Add(newState.sleepFor)
         interval, newState.sleepFor = newState.sleepFor, 0
      }
      taken = atomic.CompareAndSwapPointer(&t.state, previousStatePointer, unsafe.Pointer(&newState))
   }
   t.clock.Sleep(interval)
   // 最后返回需要等待的时间
    return newState.last
}

实现一个Take方法

1. 该Take方法会进行原子性操作(可以理解为加锁和解锁)在大量并发请求下仍可以保证正常使用。
2. 记录下当前的时间 now := t.clock.Now()
3. oldState.last.IsZero() 判断是不是第一次取水如果是就直接将state结构体中的值进行返回。而这个结构体中初始化了上次执行时间如果是第一次取水就作为当前时间直接传参。
4. 如果 newState.sleepFor 非常小就会出现问题因此需要借助宽松度一旦这个最小值比宽松度小就用宽松度对取水时间进行维护。
5. 如果newState.sleepFor > 0 就直接更新结构体中上次执行时间newState.last = newState.last.Add(newState.sleepFor)并记录需要等待的时间interval, newState.sleepFor = newState.sleepFor, 0。
6. 如果允许取水和等待操作那就说明没有发生并发竞争的情况就模拟睡眠时间t.clock.Sleep(interval)。然后将取水的目标时间进行返回由服务端代码来判断是否打回响应或者等待该时间后继续响应。

t.clock.Sleep(interval)

func (c *clock) Sleep(d time.Duration) { time.Sleep(d) }

实际上在一个请求来的时候限流器就会进行睡眠对应的时间并在睡眠后将最新取水时间返回。

实际应用使用Gin框架

func ratelimit1() func(ctx *gin.Context) {
	r1 := rate1.New(100)
	return func(ctx *gin.Context) {
		now := time.Now()
		//  Take 返回的是一个 time.Duration的时间
		if r1.Take().Sub(now) > 0 {
			// 返回的时间比当前的时间还大说明需要进行等待
			// 如果需要等待, 就 time.Sleep(r1.Take().Sub(now())) 然后放行
			// 如果不需要等待请求时间就直接进行Abort 然后返回
			response(ctx, http.StatusRequestTimeout, "rate1 limit...")
			fmt.Println("rate1 limit...")
			ctx.Abort()
			return
		}
		// 放行
		ctx.Next()
	}
}

这里你可以进行选择是否返回。因为Take一定会执行sleep函数所以当执行take结束后表示当前请求已经接到了水。当前演示使用第一种情况。

• 如果你的业务要求响应不允许进行等待。那么可以在该请求接完水之后然后如上例。

• 如果你的业务允许响应等待那么该请求等待对应的接水时间后进行下一步。具体代码就是将if中的内容直接忽略。建议使用

测试代码

这里定义了一个响应函数和一个handler函数方便测试

func response(c *gin.Context, code int, info any) {
   c.JSON(code, info)
}

func pingHandler(c *gin.Context) {
   response(c, 200, "ping ok~")
}

执行go test -run=Run -v先开启一个web服务

func TestRun(t *testing.T) {
   r := gin.Default()

   r.GET("/ping1", ratelimit1(), pingHandler)
   r.GET("/ping2", ratelimit2(), helloHandler)

   _ = r.Run(":4399")
}

使用接口压力测试工具go-wrk进行测试->tsliwowicz/go-wrk: go-wrk)

在golang install版本可以直接通过go install github.com/tsliwowicz/go-wrk@latest下载

使用帮助

   Usage: go-wrk <options> <url>
   Options:
    -H       Header to add to each request (you can define multiple -H flags) (Default )
    -M       HTTP method (Default GET)
    -T       Socket/request timeout in ms (Default 1000)
    -body    request body string or @filename (Default )
    -c       Number of goroutines to use (concurrent connections) (Default 10)
    -ca      CA file to verify peer against (SSL/TLS) (Default )
    -cert    CA certificate file to verify peer against (SSL/TLS) (Default )
    -d       Duration of test in seconds (Default 10)
    -f       Playback file name (Default <empty>)
    -help    Print help (Default false)
    -host    Host Header (Default )
    -http    Use HTTP/2 (Default true)
    -key     Private key file name (SSL/TLS (Default )
    -no-c    Disable Compression - Prevents sending the "Accept-Encoding: gzip" header (Default false)
    -no-ka   Disable KeepAlive - prevents re-use of TCP connections between different HTTP requests (Default false)
    -no-vr   Skip verifying SSL certificate of the server (Default false)
    -redir   Allow Redirects (Default false)
    -v       Print version details (Default false)

-t 8个线程 -c 400个连接 -n 模拟1k次请求 -d 替换-n 表示连接时间

输入go-wrk -t=8 -c=400 -n=1000 http://127.0.0.1:4399/ping1

可以稍微等待一下水流积攒否则一个请求也不一定能够响应。

可以看出89个请求全部返回。也就是说在一段请求高峰期不会有请求进行响应。因此我认为既然内部已经睡眠那么就应该对请求放行处理。限流器实现的比较纯粹

令牌桶

引入ratelimit库

go get -u github.com/juju/ratelimit

初始化

// NewBucket returns a new token bucket that fills at the
// rate of one token every fillInterval, up to the given
// maximum capacity. Both arguments must be
// positive. The bucket is initially full.
func NewBucket(fillInterval time.Duration, capacity int64) *Bucket {
   return NewBucketWithClock(fillInterval, capacity, nil)
}

// NewBucketWithClock is identical to NewBucket but injects a testable clock
// interface.
func NewBucketWithClock(fillInterval time.Duration, capacity int64, clock Clock) *Bucket {
   return NewBucketWithQuantumAndClock(fillInterval, capacity, 1, clock)
}

进行Bucket桶的初始化。

/ NewBucketWithQuantumAndClock is like NewBucketWithQuantum, but
// also has a clock argument that allows clients to fake the passing
// of time. If clock is nil, the system clock will be used.
func NewBucketWithQuantumAndClock(fillInterval time.Duration, capacity, quantum int64, clock Clock) *Bucket {
   if clock == nil {
      clock = realClock{}
   }
    // 填充速率
   if fillInterval <= 0 {
      panic("token bucket fill interval is not > 0")
   }
    // 最大令牌容量
   if capacity <= 0 {
      panic("token bucket capacity is not > 0")
   }
    // 单次令牌生成量
   if quantum <= 0 {
      panic("token bucket quantum is not > 0")
   }
   return &Bucket{
      clock:           clock,
      startTime:       clock.Now(),
      latestTick:      0,
      fillInterval:    fillInterval,
      capacity:        capacity,
      quantum:         quantum,
      availableTokens: capacity,
   }
}

令牌桶初始化过程初始化结构体 fillInterval(填充速率) cap(最大令牌量) quannum(每次令牌生成量)。

如果三个变量有一个小于或者等于0的话直接进行报错返回。在最开始就将当前令牌数初始化为最大容量。

调用

// 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函数传入参数为需要取出的令牌数量返回参数是实际能够取出的令牌数量。

内部实现

// 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 {
   // 如果需要取出的令牌数小于等于零那么就返回0个令牌
    if count <= 0 {
      return 0
   }
    // 根据时间对当前桶中令牌数进行计算
   tb.adjustavailableTokens(tb.currentTick(now))
    // 计算之后的令牌总数小于等于0说明当前令牌不足取出那么就直接返回0个令牌
   if tb.availableTokens <= 0 {
      return 0
   }
    // 如果当前存储的令牌数量多于请求数量那么就返回取出令牌数
   if count > tb.availableTokens {
      count = tb.availableTokens
   }
    // 调整令牌数
   tb.availableTokens -= count
   return count
}

调整令牌

// 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
   }
    // 当前令牌数 += (当前时间 - 上次取出令牌数的时间) * quannum(每次生成令牌量)
   tb.availableTokens += (tick - lastTick) * tb.quantum
    // 如果当前令牌数大于最大等于容量, 将当前令牌数 = 最大容量 然后返回 当前令牌数
   if tb.availableTokens > tb.capacity {
      tb.availableTokens = tb.capacity
   }
   return
}

实现原理

1. 加锁 defer 解锁
2. 判断count(想要取出的令牌数) 是否小于等于 0如果是直接返回 0
3. 调用函数adjustTokens 获取可用的令牌数量该函数实现原理
   1. 如果当前令牌数大于最大等于容量直接返回最大容量
   2. 当前令牌数 += (当前时间 - 上次取出令牌数的时间) * quannum(每次生成令牌量)
   3. 如果当前令牌数大于最大等于容量, 将当前令牌数 = 最大容量 然后返回 当前令牌数
4. 如果当前可以取出的令牌数小于等于0 直接返回 0
5. 如果当前可以取出的令牌数小于当前想要取出的令牌数(count) count = 当前可以取出的令牌数
6. 当前的令牌数 -= 取出的令牌数(count)
7. 返回 count

额外介绍

take函数能够返回等待时间和布尔值允许欠账没有令牌也可以取出。

func (tb *Bucket) Take(count int64) time.Duration

takeMaxDuration函数可以根据最大等待时间来进行判断。

func (tb *Bucket) TakeMaxDuration(count int64, maxWait time.Duration) (time.Duration, bool)

因为他们内部的实现都基于令牌调整我这里不做过多介绍如果感兴趣可以自行研究一下。

测试

func ratelimit2() func(ctx *gin.Context) {
	// 生成速率 最大容量
	r2 := rate2.NewBucket(time.Second, 200)
	return func(ctx *gin.Context) {
		//r2.Take() // 允许欠账令牌不够也可以接收请求
		if r2.TakeAvailable(1) == 1 {
			// 如果想要取出1个令牌并且能够取出就放行
			ctx.Next()
			return
		}
		response(ctx, http.StatusRequestTimeout, "rate2 limit...")
		ctx.Abort()
		return
	}
}

由于压测速度过于快速在实际过程中可以根据调整令牌生成速率来进行具体限流

小结

令牌桶可以允许自己判断请求是否继续不用进行睡眠。而漏桶需要进行睡眠并没有提供方法让程序员进行判断是否放行。
个人用令牌桶还是多的也可能是我对漏桶源码的解析有误没有看到相关的点。