用redis的zset实现简单的限流

Posted 2021-04-07 梅沙技术团队

当系统处理能力有限时，限制计划外的请求对系统造成的压力显得特别重要，接口限流有点断尾求生的意思，牺牲一部分异常的请求，保证大部分的请求正常。除了控制流量，限流还有一个应用就是控制用户的行为，避免垃圾请求，比如限制用户在一定时间内的发帖、回复、点赞等数量，现在给某个手机号发送验证码的频率，超过一定频率就可以判断为非法请求。

常见的限流手段

1. 漏桶是指由一个固定容量的桶，以固定的速率流出水滴。当然也需要流入水桶里才有水。如果桶是空的，则流不出水滴。如果流入的水滴超过了桶的容量，水就溢出了。

2. 令牌桶则是匀速往桶里添加令牌，如果有请求要处理，则向桶里取出一个令牌，当桶里的令牌被取完了继续等待或者拒绝服务。

这两种方法看起来很像，都有一个桶，都是往桶里流入和取出，但还是有区别的。漏桶流出的速率固定，而令牌桶只有桶里有令牌就可以拿，可以看出漏桶不支持瞬间并发，处理的服务会匀速进行；而令牌桶则支持一定程度的并发。例如同一时刻，有100个请求，只要令牌桶有100个令牌，那这100个请求就能被全部处理。令牌桶在没有令牌的时候也可以成为漏桶。

实际应用中令牌桶应用较为广泛，有许多流行的限流器都是基于令牌桶思想实现的。

从功能上看，令牌桶模型实际上就是对一个全局的计数器做加减法操作的过程，使用计数器需要加读写锁。用GO语言可以使用带有缓冲通道来实现简单的加令牌、取令牌的操作。

 
   
   
 

  
    
    
  
func main() {
  
    
    
  
 fillInterval := 10 * time.Microsecond //添加令牌的时间间隔
  
    
    
  
 capacity := 100 //令牌桶的容量
  
    
    
  
 tokenBucket := make(chan struct{}, capacity) //初始化一个令牌桶
  
    
    
  


  
    
    
  
 //每隔一秒钟往令牌桶里添加令牌，如果桶已经满了，则直接放弃
  
    
    
  
 fillToken := func() {
  
    
    
  
 ticker := time.NewTicker(fillInterval)
  
    
    
  
 for {
  
    
    
  
 select {
  
    
    
  
 case <-ticker.C:
  
    
    
  
 select {
  
    
    
  
 case tokenBucket <- struct{}{}:
  
    
    
  
 default:
  
    
    
  
 }
  
    
    
  


  
    
    
  
 }
  
    
    
  
 }
  
    
    
  
 }
  
    
    
  


  
    
    
  
 go fillToken()
  
    
    
  
 select {}
  
    
    
  
}
 
   
   
 

用Redis的zset限流

假设现在我们需要实现一个简单的限流功能：

系统要在限定用户的某个行为在指定的时间里只允许发送N次，比如1分钟内只能给某个手机号发送3次短信，用redis如何实现？

答案是可以用redis的有序集合实现，在实现之前我们先了解一下有序集合。

zset（有序集合）一方面它是一个set（集合），保证了内部value的唯一性，另一方面它可以给每个value赋予一个score，从而对value的排序。它的内部实现是用跳跃列表的数据结构。

跳跃列表是一种层级的结构，最下面一层所有的元素都会串起来，然后每个几个元素挑选出一个代表，再将这个几个代表使用另一级指针串起来，再在里边挑出二级代表。

这个限流需求中存在一个滑动的时间窗口（定宽，1分钟），zset的score用来记录每次行为的时间戳，然后每次我们判断的时候就圈出一个时间窗口（这题是最近1分钟），窗口之外的数据都可以丢弃了。那value存什么值呢，value在zset内，就要求是唯一的，我们可以给一次行为一个value，但要这个value也唯一可以也用时间戳吧。这样value、score都用时间戳，时间戳的精度可以定为毫秒。

伪代码如下，当然也有需要优化的地方，比如可以用redis的管道去执行。

 
   
   
 

  
    
    
  
func IsActionAllow(userId, actionKey string, period, maxCount int) bool {
  
    
    
  
 //用userId和actionKey拼接成zset的key
  
    
    
  
 key := "limiter:" + userId + ":" + actionKey
  
    
    
  
 //现在毫秒时间戳
  
    
    
  
 nowTime := time.Now().UnixNano() / 1e6
  
    
    
  
 //设置一个值，score和value都为毫秒时间戳
  
    
    
  
 redis.zadd(key, nowTime, nowTime)
  
    
    
  
 //移除时间窗口之前的值，剩下的就是当前的时间窗口
  
    
    
  
 redis.zremrangeByScore(key, 0, nowTime-period*1000)
  
    
    
  
 //获取时间窗口内的数量
  
    
    
  
 count := redis.zcard(key)
  
    
    
  
 return count <= maxCount
  
    
    
  
}
 
   
   
 

整体的一个思路是，每次一个行为过来都会调用这个函数去判断，每次执行这个函数都维护一个时间窗口，将时间窗口外的记录都清理掉，只保留窗口内的记录。zset中的score值非常重要，value值没有特别的意义，只需要保证唯一就可以了。

这个方案也有缺点，因为它要记录时间窗口内的所有行为记录，如果这个量很大，就会消耗比较多的内存，比如“限定1小时内100万次”之类的。这只是一个简单的方案，使用的场景也比较有限，本文只是拿来扩展思路，加深对zset的理解。