基于Redis实现一套支持排队等待的限流器

一、背景

由于项目中调用了一个政府官方系统，前段时间得到政府通知说我们调用频率太高，目前给我们开放的接口调用频率是每秒一次。然后还发过来一个我们请求通过与超频的比例报告，显示失败率高达80%，也就是说我们百分之80的请求都被拦截了，这里应该会有有伙伴疑问80%的异常请求你们系统怎么开展业务的。实际上我们系统内部有做缓存，然后后台有主动和被动两种方式去刷新缓存，所以这80%失败请求中绝大多数都是后台刷新缓存的请求，并非客户端用户的请求所以呢对我们的业务也没有实质性的影响。基于此我方也需要做请求限制，不然政府方面会考虑以请求失败率过高而把我们的接口权限下掉。

二、调研

关于限流的经典算法漏斗和令牌通这里就不多说了，这类算法介绍网上已经很多内容了。我这里整理下目前市面上开源的限流工具以及为什么我们没选择使用开源工具而要自己造轮子的原因。

Google Guava

首先就是谷歌的Guava工具类库，该类库提供很多比较好用的工具类，其中就包括基于令牌通算法实现的限流器RateLimiter

// 1、声明一个qps最大为1的限流器
RateLimiter limiter = RateLimiter.create(1);
// 2、尝试阻塞获取令牌
limiter.acquire();

Alibaba Sentinel

然后就是阿里巴巴的Sentinel，这个就比较强大了，应该是目前市面上限流方面做的最全面的开源项目了。不仅支持流量控制，同时还支持分布式限流，熔断降级，系统监控等，还有比较灵活的限流策略配置支持。这里我也没用过，可能需要花些时间才能掌握吧。

Redisson RRateLimiter

最后要介绍的是基于令牌通算法实现的RRateLimiter, 它是Redisson类库的限流工具，支持分布式限流，使用起来也相当的方便

// 1、 声明一个限流器
RRateLimiter rateLimiter = redissonClient.getRateLimiter(key);
 
// 2、 设置速率，5秒中产生3个令牌
rateLimiter.trySetRate(RateType.OVERALL, 3, 5, RateIntervalUnit.SECONDS);
 
// 3、试图获取一个令牌，获取到返回true
rateLimiter.tryAcquire(1)

选型

1. 首先我的需求是限流器必须要支持分布式，那Guava首先可以排除了。
2. 然后Sentinel对于我们的需求来说有些笨重，太过于重量所以也排除了。
3. 最后RRateLimiter虽然支持分布式，使用也比较简单，但是好像它不支持公平排队（不太确定）。

三、造轮子

基于以上我决定自己手撸一个支持公平排队的分布式限流器。实现方案是基于Redis Lua脚本然后配合业务层代码支持，直接上菜

限流器的主体代码

public class RedisRateLimiter {

    public static final GenericToStringSerializer argsSerializer = new GenericToStringSerializer<>(Object.class);
    public static final GenericToStringSerializer resultSerializer = new GenericToStringSerializer<>(Long.class);

    @Resource
    private RedisTemplate<String, Object> redisTemplate;

    @Resource
    private RedisUtil redisUtil;

    public static final int DEFAULT_MAX_PERMIT_COUNT = 1;
    public static final float DEFAULT_INTERVAL_SECONDS = 1.3f;
    public static final int DEFAULT_TIMEOUT_SECONDS = 5;

    // TODO 目前不支持自定义该值
    /**
     * 一个周期内的最大许可数量
     */
    private int maxPermitCount;

    public RedisRateLimiter() {
        this.maxPermitCount = DEFAULT_MAX_PERMIT_COUNT;
    }

    public static DefaultRedisScript<Long> redisScript;

    static {
        redisScript = new DefaultRedisScript<>();
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("rateLimiter.lua")));
        redisScript.setResultType(Long.class);
    }

    /**
     *
     * @param redisKey
     * @param intervalSeconds 间隔几秒创建一个许可
     * @param timeoutSeconds 获取许可超时时间
     * @return
     * @throws InterruptedException
     */
    public boolean tryAcquire(String redisKey, float intervalSeconds, long timeoutSeconds) throws InterruptedException {
        try {
            if (redisKey == null) {
                throw new BusinessException(BusinessExceptionCode.REQUEST_PARAM_ERROR, "redisKey不能为空！");
            }
            Preconditions.checkArgument(intervalSeconds > 0.0 && !Double.isNaN(intervalSeconds), "rate must be positive");

            long intervalMillis = (long) (intervalSeconds * 1000);
            long timeoutMillis = Math.max(TimeUnit.SECONDS.toMillis(timeoutSeconds), 0);

            long pttl = redisTemplate.execute(redisScript, argsSerializer, resultSerializer, Arrays.asList(redisKey), maxPermitCount, intervalMillis, timeoutMillis);

            if (pttl == 0) {
                log.info("----------------无需排队，直接通过, 当前许可数量={}", redisUtil.get(redisKey));
                return true;
            }else if(pttl < timeoutMillis) {
                Thread.sleep(pttl);
                log.info("----------------排队{}毫秒后，通过, 当前许可数量={}", pttl, redisUtil.get(redisKey));
                return true;
            }else {
                // 直接超时
                log.info("----------------需排队{}毫秒，直接超时, 当前许可数量={}", pttl, redisUtil.get(redisKey));
                return false;
            }
        }catch (Exception e) {
            log.error("限流异常", e);
        }
        return true;
    }
}

核心Lua脚本文件: rateLimiter.lua

local limit = tonumber(ARGV[1])
local interval = tonumber(ARGV[2])
local timeout = tonumber(ARGV[3])
local count = tonumber(redis.call('get', KEYS[1]) or "0")
local pttl = tonumber(redis.call('pttl', KEYS[1]) or "0")
if pttl < 0 then
 pttl = 0
end
-- 这个代表已被预定的令牌数
local currentCount = count - math.max(math.floor((count*interval - pttl)/interval), 0)
-- 新增一个令牌
local newCount = currentCount + 1
-- 所有令牌总的失效毫秒
local newPTTL = pttl + interval

if newCount <= limit then

 --无需排队直接通过
 redis.call("PSETEX", KEYS[1], newPTTL, newCount)
 return 0
elseif pttl < timeout then

 --排队pttl毫秒后可通过
 redis.call("PSETEX", KEYS[1], newPTTL, newCount)
else

 -- 超时
 redis.call("PSETEX", KEYS[1], pttl, currentCount)
end
-- 返回需等待毫秒数
return pttl