基于令牌桶算法的流量控制

发表于 | 分类于 ，

问题起源

作为一款IM软件，负责收发消息的通讯服务器需要仔细处理QPS的压力。为了应对潜在的恶意流量攻击，比如第三方短时间发送大流量到服务器，引起通讯服务器瘫痪，最近以每个会话(例如a<->b的会话)为基本单位加上了流量控制。

限流算法

常用算法比较

常用限流算法包含计时器，漏桶，令牌桶算法。具体的比较与分析，在谈谈高并发系统的限流这篇文章中有比较详细的介绍。

需要注意的是漏桶算法和令牌桶算法的区别：

  • 漏桶算法强制限制了数据的通过速率，并不负责处理突发的大流量；
  • 令牌桶算法则既平均了数据的通过速率，又能够处理突发的流量；

令牌桶

利用令牌桶的算法进行限流是目前比较通用的方法，所以本文主要聚焦于令牌桶算法的介绍。

根据wiki上的介绍令牌桶算法，基本定义是：

一个最大容量为b个令牌的令牌桶，按照一定速率r往桶中添加令牌。当需要处理的数据包通过时，先向该令牌桶索要一定数量n个的令牌(数据大小<->令牌个数 的对应关系自定义)，如果这时候桶内有>=n个的令牌，那么就从桶中移除n个令牌，然后该数据包通过；如果桶内的令牌数不足，那么该数据包就被表示为目前不能通过，不通过的数据包通常有3种处理方式：

  • 被丢弃
  • 阻塞等待，直到桶内产生了足够多的令牌
  • 可以继续传输，但是需要被特殊标示，当网络拥堵的时候该特殊表示的数据包可能会被优先丢弃

总结来说的话，令牌桶包括三个关键点：

  • 每过1/r秒，桶内增加一个令牌(r是增加速率 单位：个/s)
  • 桶内最多包含b个令牌，桶满之后，再过来的令牌会被丢弃
  • 当一个包含n个bytes的数据包到来的时候(此时假设 1个byte需要1个令牌)：
    • 假如现在桶内有>=n个令牌，那么就从桶内移除n个令牌，该数据包通过验证；
    • 假如桶内现在没有n个令牌，就不从桶内移除任何令牌，该数据包被标示为不通过，不通过的数据包再做特殊处理

所以说，令牌桶算法通过设置通过率r，限制通过数据的最大平均速率不会超过r；通过设置最大容量b，又可以处理突发性的数据流量。

令牌桶的实现

Google的开源guava项目中有令牌桶限流算法的java实现rateLimiter，有兴趣可以看看。

因为我们的服务端主要采用Go语言，所以本文主要介绍Go语言实现的令牌桶算法，Go语言官方已经提供了基于令牌桶令牌桶算法的速率限制器rateLimiter

该rateLimiter的基本结构如下

1
2
3
4
5
6
7
8
9
10
11
12

type Limiter struct {
    limit Limit
    burst int

    mu     sync.Mutex
    tokens float64
    // last is the last time the limiter's tokens field was updated
    last time.Time
    // lastEvent is the latest time of a rate-limited event (past or future)
    lastEvent time.Time
}

针对获取不到足够令牌的数据包进行的不同处理，该限制器主要提供了三种方法。首先如果令牌充足的情况下，三种方法都消耗n个令牌；但是如果令牌不足，那么每个方法有各自的处理：

  • AllowN：令牌不足时，直接返回false，适合不通过的数据包需要被直接丢弃的情况。

  • ReserveN：令牌不足时，提前预定还需要的令牌的差值个数(比如本次总共需要5个，目前只有2个，那么就预定3个），以及调用者在能够使用这些令牌前必须等待的时间。

  • WaitN：令牌不足时，阻塞等待直到桶内有对应的令牌个数；或者如果传入了取消信号，那么就取消等待，返回错误。

当回到我们目前通讯服务器的限流情况时，因为是以会话为单位限制通讯流量，预防恶意流量攻击，所以设定了一个用户发送消息的上限，该上限是正常用户输入发送消息所不能达到的，所以主要采用了AllowN的方法，超出限制的部分直接丢弃。

利用LRUcache限制内存中的令牌桶个数

因为是以会话为单位进行限流的，每个会话一个限流器，所以如果全部保存在内存中的话，数量非常的大。所以采用cache的方式缓存热点数据，达到控制内存中限流器个数的目的。cache的操作主要是基于Go语言的LRUcache里的LRUcache来实现的。

LRUcache利用LRU算法实现cache的操作，LRU即为最近最少使用(Least recently used)，核心思想就是最新的添加/查看操作，使用频率最高。按照这个条件，当容量满了之后，淘汰频率最低的数据(即距离上次使用时间最长的数据)。

利用定时器主动淘汰过期数据

如果只使用LRUcache，那么还有一个弊端就是只有当cache满了之后才会开始淘汰数据，因为限流是以会话id做key，如果在加入该限制器的时候，同时增加一个过期销毁限制器的操作，就能够更加主动灵活的淘汰掉一些老的限制器，使得内存占用更小。

具体定时器的实现详见之前的文章基于最小堆的定时器的实现

上面就是为了应对潜在的恶意流量攻击，而引入限流优化的详细过程。