利用Redis实现时间滑动窗口限流

背景

redisson开源框架已经提供了限流的功能，但由于项目较旧，没引入redisson，为了降低风险，用spring-redis的RedisTemplate来自行实现一个限流工具类。

实现

• 用sortSet来存储时间窗口数据，命令是ZADD，score的值为当前时间戳
• 利用sortSet的ZREMRANGEBYSCORE命令删除时间窗口外的数据
• 用sortSet的ZCARD来判断是否超过限流阈值

代码实现

@Component
@Slf4j
public class RateLimiter {

    @Autowired
    RedisTemplate<String, Long> redisTemplate;

    private static final String redisPrefix = "wechat-data:limiter:";

    /**
     * @param source 资源
     * @param limit 限制值
     * @param period 间隔
     * @param unit 时间单位
     * @return 是否允许
     */
    public boolean isAllow(String source, int limit, long period, TimeUnit unit){
        if (StringUtils.isEmpty(source)) {
            log.warn("参数为空??");
            return false;
        }
        String key = redisPrefix + source;
        long now = System.currentTimeMillis();
        final long ms = TimeUnit.MILLISECONDS.convert(period, unit);
        final Boolean result = redisTemplate.execute(RateLimiterScript.instance, Collections.singletonList(key), limit, now,  ms);
        return result != null && result;
    }


    static class RateLimiterScript implements RedisScript<Boolean> {
        static RateLimiterScript instance = new RateLimiterScript();
        static String shaCache;

        @Override
        public String getSha1() {
            if (shaCache != null)
                return shaCache;
            try {
                MessageDigest mdigest = MessageDigest.getInstance("SHA-1");
                byte[] s = mdigest.digest(getScriptAsString().getBytes());
                shaCache = ByteBufUtil.hexDump(s);
            } catch (Exception e) {
                throw new IllegalStateException(e);
            }
            return shaCache;
        }

        @Override
        public Class<Boolean> getResultType() {
            return Boolean.class;
        }

        @Override
        public String getScriptAsString() {
            return "redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, tonumber(ARGV[2]) - tonumber(ARGV[3])); " +
                    "if (redis.call('ZCARD', KEYS[1]) >= tonumber(ARGV[1])) then " +
                    "return nil;" +
                    "end; " +
                    "redis.call('ZADD', KEYS[1], ARGV[2], ARGV[2]); " +
                    "redis.call('pexpire', KEYS[1], ARGV[3]); " +
                    "return true;";
        }
    }
}

说明：

• 考虑到原子性，这里使用了lua进行操作
• 利用RedisTemplate的execute方法执行脚本
• redis.call(‘ZREMRANGEBYSCORE’, KEYS[1], 0, tonumber(ARGV[2]) - tonumber(ARGV[3]))：删除时间窗口外的数据
• if (redis.call(‘ZCARD’, KEYS[1]) >= tonumber(ARGV[1])) ：判断是否超过限流的阈值
• redis.call(‘ZADD’, KEYS[1], ARGV[2], ARGV[2])：向时间滑动窗口加入数据
• redis.call(‘pexpire’, KEYS[1], ARGV[3])：设置过期时间，防止无效key堆积