使用Redis实现简单限流和漏洞限流

目录

1. 简单限流

2. 漏斗限流

2.1 Redis-Cell

---

限流在分布式领域是一个经常被提起的话题，当系统的处理能力有限的时候，需要阻止计划外的请求继续对系统施压。除了流量控制，限流还有一个应用的目的是用于控制用户的行为，避免垃圾请求。比如在BBS中，用户的发帖、回复、点赞等行为要严格限定在规定时间内允许的次数，超过了次数就是非法行为。

1. 简单限流

我们先来看一个简单的限流策略，系统要限定用户的某个行为在指定的时间里只能允许发生N次。这个限流需求中存在一个滑动时间窗口，我们可以使用zset数据结构的score值，来圈出这个时间窗口，毫秒时间戳作为score值。

如图所示，用一个zset结构记录用户的历史行为，每一个行为都作为zset中的一个key保存下来，同一个用户同一个行为用zset记录。为了节省内存，我们只需要保留时间窗口内的行为记录，如果用户是冷用户，滑动时间窗口内的行为是空记录，那么这个zset就可以直接从内存中移除，不再占用空间。

实现代码如下：

# simple limit
import time
import redis
import my_tools as tools

client = tools.get_redis()


def is_action_allowed(user_id, action_key, period, max_count):
    key = 'hist:%s:%s' % (user_id, action_key)
    now_ts = int(time.time() * 1000)
    with client.pipeline() as pipe:
        # 记录行为，value和score都是使用毫秒时间戳
        mapping = {now_ts: now_ts}
        pipe.zadd(key, mapping)
        # 移除时间窗口之前的行为记录，剩下的都是时间窗口内的
        pipe.zremrangebyscore(key, 0, now_ts - period * 1000)
        # 获得窗口内的行为数量
        pipe.zcard(key)
        # 设置zset过期时间，避免冷用户持续占用内存
        # 过期时间应该等于时间窗口的长度，再多1s
        pipe.expire(key, period + 1)
        # 批量执行
        _, _, current_count, _ = pipe.execute()

    print(current_count)
    return current_count <= max_count


for i in range(20):
    limit_flag = is_action_allowed('martin', 'reply', 5, 5)
    if not limit_flag:
        print('被限流')

上述代码的整体思路是在某一个行为到来时，都维护一次时间窗口。将时间窗口外的记录全部清理掉，只保留窗口内的记录。zset集合中只有score值非常重要，value的值没有特别的意义，但是需要保证它是唯一的。因为这个几个连续的Redis操作都是针对同一个key的，使用pipeline可以显著的提升Redis存取效率。但是这种方案也有缺点，因为它要记录时间窗口内的所有的行为记录，如果量非常大，比如限定60s内操作不得超过100W次这样的参数，它是不适合做这样的限流，因为会消耗大量的存储空间。

2.漏斗限流

漏斗限流是最常用的限流方法之一，该算法的灵感来源于漏斗的结构。

漏斗的容量是有限的，如果将漏斗堵住，然后一直往里面灌水，它就会变满，甚至再也装不进去。如果将漏嘴放开，水就会往下流，流走一部分之后，就可以继续往里面灌水。如果漏嘴流水的速率大于灌水的速率，那么漏斗永远都装不满。如果漏嘴流水速率小于灌水的速率，那么一旦漏斗满了，灌水就需要暂停并等待漏斗腾空。漏斗算法实现的代码如下：

import time


class Funnel(object):
    def __init__(self, capacity, leaking_rate):
        self.capacity = capacity  # 漏斗容量
        self.leaking_rate = leaking_rate  # 漏斗流水速度
        self.left_quota = capacity  # 漏斗剩余空间
        self.leaking_ts = time.time()  # 上一次漏水时间

    def make_space(self):
        now_ts = time.time()
        delta_ts = now_ts - self.leaking_ts  # 距离上一次漏水过去了多少时间
        delta_quota = delta_ts * self.leaking_rate  # 可以腾出的空间
        if delta_quota < 1:  # 腾出的空间不足，直接返回
            return
        self.left_quota += delta_quota  # 增加剩余空间
        self.leaking_ts = now_ts  # 记录漏水时间
        if self.left_quota > self.capacity:  # 剩余空间不得高于容量
            self.left_quota = self.capacity

    def watering(self, quota):
        self.make_space()
        if self.left_quota >= quota:
            self.left_quota -= quota
            return True

        return False


funnels = {}  # 所有的漏斗


# capacity 漏斗容量
# leaking_rate 漏嘴流水速率 quota/s
def is_action_allowed(user_id, action_key, capacity, leaking_rate):
    key = '%s:%s' % (user_id, action_key)
    funnel = funnels.get(key)
    if not funnel:
        funnel = Funnel(capacity, leaking_rate)
        funnels[key] = funnel

    return funnel.watering(1)


for i in range(20):
    print(is_action_allowed('martin', 'reply', 15, 0.5))

运行结果：
True
True
True
True
True
True
True
True
True
True
True
True
True
True
True
False
False
False
False

Funnel对象的make_space方法是漏斗算法的核心，其在每次灌水前都会被调用以触发漏水，给漏斗腾出空间来。能腾出多少空间取决于过去了多久以及流水的速率。

上面的实现代码仅仅适用于单机系统，分布式的漏斗算法如何实现呢？是否可以使用Redis的基础数据结构来实现呢？

我们可以将Funnel对象的几个字段，存储到Redis的hash结构中，灌水的时候将hash结构的字段取出来进行逻辑运算后，再将新值回填到hash结构中就完成了一次行为频度的检测。但是这里有个问题，我们无法保证整个过程的原子性，从hash结构中取值，然后在内存里运算，再回填到hash结构，这三个过程无法原子化，意味着需要进行适当的加锁控制。而一旦加锁，就意味着会有加锁失败，加锁失败就需要选择重试或者放弃，这将会影响用户的体验。同时，代码的复杂度也跟着升高很多。那我们该如何解决这个问题呢？答案是Redis-Cell。

2.1 Redis-Cell

Redis4.0提供了一个限流Redis模块，叫作redis-cell。该模块使用了漏斗算法，并提供了原子的限流指令。安装步骤参考：https://github.com/brandur/redis-cell

该模块只有1条指令cl.throttle，它的参数和返回值都有些复杂，指令具体如下：

cl.throttle martin:reply 15 30 60 1

上面的指令表示是允许“martin回复行为”的频率为每60s最多30次（漏水速率），漏斗的初始容量为15，也就是说一开始可以连续回复15个帖子，然后才会受到漏水速率的影响。1表示该指令执行占用了1个单位空间，可以不写，默认值为1。

cl.throttle martin:reply 15 30 60
1) (integer) 0 # 0 表示允许， 1 表示拒绝
2) (integer) 15 # 漏斗容量 capacity
3) (integer) 14 # 漏斗剩余空间 left_quota
4) (integer) -1 # 如果拒绝了，需要多长时间后再试(漏斗有空间了，单位秒)
5) (integer) 2 # 多长时间后，漏斗完全空出来(left_quota==capacity，单位秒)

在执行限流指令时，如果被拒绝了，就需要丢弃或者重试。cl.throttle指令返回的第四个值就是重试时间，我们拿到该值进行sleep即可，如果不想阻塞线程，可以使用异步定时任务来进行重试。