Redis+Lua脚本三步实现分布式系统限流

  在分布式系统中，说到限流方案我们一般会使用redis结合限流算法来做，一般的限流算法有令牌桶算法、漏桶算法、固定窗口、滑动窗口等，这几种算法优缺点大家可以自行百度，不过实际都需要我们自己代理里实现，虽然也可以使用一些第三方包如guava，但不一定适配自身项目。
   今天我打算介绍我们公司一直使用的接口限流方案，无需代码实现复杂算法，14行lua脚本即可实现固定窗口算法实现限流。lua脚本类似于mysql中的存储过程，可以在redis中一次性执行几个命令。
   大体思路：实现一个限流拦截器，在拦截器中执行lua脚本,在脚本中redis设置一个带过期时间的值，每次+1，达到窗口阈值就返回0，限制接口访问。
   核心代码如下：

/**
 * 限流拦截器
 */
public class RateLimitInterceptor  implements HandlerInterceptor {

    private final Logger logger = LoggerFactory.getLogger(getClass());

    /**
     * 已过期标志
     */
    private final Long EXPIRED = 0L;

    /**
     * 滑动窗口算法使用ttl实现，脚本语言很容易理解，KEYS与ARGV都是脚本 
     * 执行参数，KEYS[1]是redis中的key,ARGV为参数数组
     */
    private final byte[] script = SafeEncoder
            .encode("local values = redis.call('get',KEYS[1])\n"
                    + "local defaultValues = ARGV[1]\n"
                    + "local ttl = ARGV[2]\n"
                    + "local maxValues = tonumber(ARGV[3])\n"
                    + "if values == false then\n"
                    + "        redis.call('set',KEYS[1],defaultValues,'EX',ttl)\n"
                    + "        return 1\n"
                    + "end\n"
                    + "values = values + 1\n"
                    + "if values > maxValues then\n"
                    + "        return 0\n"
                    + "end\n"
                    + "redis.call('incr',KEYS[1])\n"
                    + "return 1");

    @Autowired
    private RedisService redisService;

    @Override
    public boolean preHandle(final HttpServletRequest request, final HttpServletResponse response, final Object handler) {
        // 以ip或者customerId做key
        String key = SecurityUtils.getIp();
        // 配置为10秒100个请求
        RateLimitConfig config = new RateLimitConfig(10, 100);
        List<byte[]> argvList = Arrays.asList(String.valueOf(1).getBytes(),
                String.valueOf(config.getUnitTime()).getBytes(),
                String.valueOf(config.getLimit()).getBytes());;
        try {
            Response<Object> rsp = redisService.execScript(key, argvList, script);
            if (!rsp.isSuccess()) {
                logger.error(rsp.getErrorMsg());
                return false;
            }
            Long result = (Long) rsp.getResult();
            if (EXPIRED.equals(result)) {
                return false;
            }
        } catch (Exception e) {
            logger.error(e.getLocalizedMessage(), e);
            return false;
        }
        return true;
    }


    @Data
    static class RateLimitConfig {

        public static RateLimitConfig DEFAULT = new RateLimitConfig(10, 200);

        // 单位时间(秒)
        int unitTime;

        // 限流次数
        int limit;

        public RateLimitConfig() {
        }

        public RateLimitConfig(final int unitTime, final int limit) {
            this.unitTime = unitTime;
            this.limit = limit;
        }
    }
}