深入Redis（六）简单限流

简单限流

限流算法是一个在分布式领域经常被提起的话题，当系统处理能力有限时，如何阻止计划外的请求继续对系统施压，这是一个需要重视的问题。

除了控制流量，限流还有一个应用目的是用于控制用户行为，避免垃圾请求，比如严格限定在规定时间内某行为的允许次数。

如何使用Redis来实现简单限流策略？

系统要限定用户的某个行为在指定时间内只能允许发生N次，如何利用Redis的数据结构来实现这个限流的功能？

先定义这个接口

def is_action_allowed(user_id, action_key, period, max_count):
    return True

can_reply = is_action_allowed('tom', 'reply', 60, 5)
if can_reply:
    do_reply()
else:
    raise ActionThresholdOverflow()

解决方案

利用zset数据结构的score值来作为时间窗口，value保证唯一性即可，用uuid比较浪费空间，用毫秒时间戳即可。

用一个zset结构记录用户的历史行为，每一个行为都会作为zset中的一个key保存下来，同一个用户的同一种行为用一个zset记录。

为节省内存，只保留时间窗口内的行为记录，如果用户是冷用户，窗口内的行为是空记录，则这个zset可以从内存中移除。

通过统计窗口内的行为数量与阈值进行比较就可以得出当前行为是否允许。

import time
import redis

pool = redis.ConnectionPool(host='localhost', port=6379)
client = redis.StrictRedis(connection_pool=pool)

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

for i in range(20):
    print(is_action_allowed('tom', 'reply', 60, 5))

这几个连续的Redis操作都是针对同一个key的，因此可以用pipeline来显著提升Redis存取效率，但这个方案也有缺点，它要记录窗口内的所有记录，如果这个量特别大，就会消耗大量空间，因此是不适合简单限流的。