使用Redis实现一个简单的限流

​ 有时候我们想对我们的服务进行某些请求访问限制，即控制流量。还有就是控制用户行为，避免垃圾请求。这个时候我们需要实现一个限流策略，来满足我们的需求。

下面我们实现一个简单的限流策略：系统要限定用户某个行为在指定的时间里只能允许发生N次。

实验环境

python3.7
redis==3.3.8

实现方式

使用zset数据结构，我们使用时间戳来当值，还有就是时间戳当score来进行值比较，这样我们就能通过时间戳来比较最新的访问记录。

在这个实现方式中，有一个滑动时间窗口（定宽：即指定的时间范围），我们通过score 的差值来圈出这个时间窗口。我们只需要保留这个时间窗口之内的数据，之外的数据都可以删掉。

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

为了节省内存，我们只需要记录保留窗口内的行为记录，不在窗口内的数据，我们直接删除。

通过统计滑动窗口内的行为数量与阈值max_count之间进行比较就知道当前行为是否被允许。

实现代码

# -*- coding: utf-8 -*-

import time
import redis

client = redis.StrictRedis()


def is_action_allowed(user_id, action_key, period, max_count):
    key = f"history:{user_id}:{action_key}"
    now_ts = int(time.time() * 10000)

    with client.pipeline() as pipe:
        # 记录行为 value和score都使用时间戳
        pipe.zadd(key, {now_ts: now_ts})
        # 移除时间窗口之前的 记录 ，剩下的都是时间窗口之内的记录
        pipe.zremrangebyscore(key, 0, now_ts - period * 10000)
        # 获取窗口内的行为数量
        pipe.zcard(key)

        # 设置 zset 过期时间 避免冷用户占用内存 过期时间应该等于窗口时间的长度 再多宽限 1秒
        pipe.expire(key, period + 1)

        # 批量执行
        add_result, rem_result, current_count, expire_result = pipe.execute()
        print(add_result, rem_result, current_count, expire_result)

    # 比较数量是否超标
    return current_count <= max_count


if __name__ == '__main__':
    for i in range(10):
        print(is_action_allowed("primary_coder", "coding", 60, 5))

# 输出结果
1 0 1 True
True
1 0 2 True
True
1 0 3 True
True
1 0 4 True
True
1 0 5 True
True
1 0 6 True
False
1 0 7 True
False
1 0 8 True
False
1 0 9 True
False
1 0 10 True
False

我们看到只有5次执行成功，之后的不再成功。这样我们就实现了一个简单的限流策略。

文章为《Redis深度历险 核心原理与应用实践》阅读笔记。