Redis zset 实现简单限流

文章目录

1. 1. 如何使用 Redis 来实现简单限流策略？
   1. 1.1. 错误方案
      1. 1.1.1. 问题模拟分析
   2. 1.2. 正确方案
      1. 1.2.1. 缺点
   3. 1.3. 参考

除了控制流量，限流还有一个应用目的是用于控制用户行为，避免垃圾请求。比如在 UGC 社区，用户的发帖、回复、点赞等行为都要严格受控，一般要严格限定某行为在规定时间内允许的次数，超过了次数那就是非法行为。对非法行为，业务必须规定适当的惩处策略。

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

首先我们来看一个常见的简单的限流策略。系统要限定用户的某个行为在指定的时间里只能允许发生 N 次，如何使用 Redis 的数据结构来实现这个限流的功能？

我们先定义这个接口，理解了这个接口的定义，读者就应该能明白我们期望达到的功能。

// 指定用户 user_id 的某个行为 action_key，在特定的时间内 period，只允许发生一定的次数 max_count
public boolean isActionAllowed(String userId, String actionKey, int period, int maxCount) {}

// 调用这个接口，5 分钟之内只能发帖 10 次
limiter.isActionAllowed("xiaoming", "publish", 5 * 60, 10)

错误方案

将 userId 和 actionKey 拼成 key，在第一次请求时设置 value 为 1，设置过期时间 expire 为特定的时间 period

每次请求的时候获取 value 值，若存在则 incr 自增 1，超过 maxCount 则做限制

问题模拟分析

如上图：5 分钟之内只能发帖 10 次。

11:01 用户发帖 1 次，此时 redis 中存放数据 key 为 userId:actionKey，vlaue 为 1，过期时间 5 分钟；

11:05 用户发帖 8 次，发帖成功；

11:05 之后，key 过期时间到，被移除；

11:06 用户发帖，此时 redis 中 key 不存在，重新存放 key，发帖 8 次，发帖成功；

那么 11:05 -> 11:06 时间段 2 分钟发帖 16 次，没有达到期望的功能；

正确方案

这个限流需求中存在一个滑动时间窗口，想想 zset 数据结构的 score 值，是不是可以通过 score 来圈出这个时间窗口来。而且我们只需要保留这个时间窗口，窗口之外的数据都可以砍掉。那这个 zset 的 value 填什么比较合适呢？它只需要保证唯一性即可，用 uuid 会比较浪费空间，那就改用毫秒时间戳吧。

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

为节省内存，我们只需要保留时间窗口内的行为记录，同时如果用户是冷用户，滑动时间窗口内的行为是空记录，那么这个 zset 就可以从内存中移除，不再占用空间。

通过统计滑动窗口内的行为数量与阈值 max_count 进行比较就可以得出当前的行为是否允许。用代码表示如下：

public class SimpleRateLimiter {
    private Jedis jedis;

    public SimpleRateLimiter(Jedis jedis) {
        this.jedis = jedis;
    }

    /**
     * @param userId    用户 user_id
     * @param actionKey 某个行为
     * @param period    特定的时间内，单位秒
     * @param maxCount  最大允许的次数
     * @return
     */
    public boolean isActionAllowed(String userId, String actionKey, int period, int maxCount) {
        String key = String.format("hist:%s:%s", userId, actionKey);
        // 毫秒时间戳
        long nowTs = System.currentTimeMillis();

        Pipeline pipe = jedis.pipelined();

        // 用了multi，也就是事务，能保证一系列指令的原子顺序执行
        pipe.multi();

        // 存放数据，value 和 score 都使用毫秒时间戳
        pipe.zadd(key, nowTs, "" + nowTs);

        // zremrangebyscore key min max 命令用于移除有序集中，指定分数（score）区间内的所有成员
        // 移除时间窗口之前的数据，剩下的都是时间窗口之内的
        Response<Long> longResponse = pipe.zremrangeByScore(key, 0, nowTs - period * 1000);

        // 相当于 count()，获取时间窗口内的行为数量
        Response<Long> count = pipe.zcard(key);

        // 设置 zset 过期时间，避免冷用户持续占用内存
        // 过期时间应该等于时间窗口的长度，再多宽限 1s
        pipe.expire(key, period + 1);
        pipe.exec();
        pipe.close();

        // 比较数量是否超标
        return count.get() <= maxCount;
    }

    public static void main(String[] args) {
        Jedis jedis = new Jedis();
        SimpleRateLimiter limiter = new SimpleRateLimiter(jedis);
        for (int i = 0; i < 20; i++) {
            System.out.println(limiter.isActionAllowed("xiaoming", "publish", 5 * 60, 10));
        }
    }
}

缺点

因为它要记录时间窗口内所有的行为记录，如果这个量很大，比如限定 60s 内操作不得超过 100w 次这样的参数，它是不适合做这样的限流的，因为会消耗大量的存储空间。

参考

《Redis 深度历险：核心原理与应用实践》 作者：钱文品