基于RateLimiter的分布式限流实现

限速控制

1. 令牌桶模型

首先定义令牌桶模型，与RateLimiter中类似，包括几个关键属性与关键方法。其中关键属性定义如下，

@Data
public class RedisPermits {
   

    /**
     * 最大存储令牌数
     */
    private double maxPermits;
    /**
     * 当前存储令牌数
     */
    private double storedPermits;
    /**
     * 添加令牌的时间间隔/毫秒
     */
    private double intervalMillis;
    /**
     * 下次请求可以获取令牌的时间，可以是过去（令牌积累）也可以是将来的时间（令牌预消费）
     */
    private long nextFreeTicketMillis;

    //...

关键方法定义与RateLimiter也大同小异，方法注释基本已描述各方法用途，不再赘述。

  /**
     * 构建Redis令牌数据模型
     *
     * @param permitsPerSecond     每秒放入的令牌数
     * @param maxBurstSeconds      maxPermits由此字段计算，最大存储maxBurstSeconds秒生成的令牌
     * @param nextFreeTicketMillis 下次请求可以获取令牌的起始时间，默认当前系统时间
     */
    public RedisPermits(double permitsPerSecond, double maxBurstSeconds, Long nextFreeTicketMillis) {
   
        this.maxPermits = permitsPerSecond * maxBurstSeconds;
        this.storedPermits = maxPermits;
        this.intervalMillis = TimeUnit.SECONDS.toMillis(1) / permitsPerSecond;
        this.nextFreeTicketMillis = nextFreeTicketMillis;
    }

    /**
     * 基于当前时间，若当前时间晚于nextFreeTicketMicros，则计算该段时间内可以生成多少令牌，将生成的令牌加入令牌桶中并更新数据
     */
    public void resync(long nowMillis) {
   
        if (nowMillis > nextFreeTicketMillis) {
   
            double newPermits = (nowMillis - nextFreeTicketMillis) / intervalMillis;
            storedPermits = Math.min(maxPermits, storedPermits + newPermits);
            nextFreeTicketMillis = nowMillis;
        }
    }

    /**
    * 保留指定数量令牌，并返回需要等待的时间
    */
    public long reserveAndGetWaitLength(long nowMillis, int permits) {
   
        resync(nowMillis);
        double storedPermitsToSpend = Math.min(permits, storedPermits); // 可以消耗的令牌数
        double freshPermits = permits - storedPermitsToSpend; // 需要等待的令牌数
        long waitMillis = (long) (freshPermits * intervalMillis); // 需要等待的时间

        nextFreeTicketMillis = LongMath.saturatedAdd(nextFreeTicketMillis, waitMillis);
        storedPermits -= storedPermitsToSpend;
        return waitMillis;
    }

    /**
    * 在超时时间内，是否有指定数量的令牌可用
    */
    public boolean canAcquire(long nowMillis, int permits, long timeoutMillis) {
   
        return queryEarliestAvailable(nowMillis, permits) <= timeoutMillis;
    }

    /**
     * 指定数量令牌数可用需等待的时间
     *
     * @param permits 需保留的令牌数
     * @return 指定数量令牌可用的等待时间，如果为0或负数，表示当前可用
     */
    private long queryEarliestAvailable(long nowMillis, int permits) {
   
        resync(nowMillis);
        double storedPermitsToSpend = Math.min(permits, storedPermits); // 可以消耗的令牌数
        double freshPermits = permits - storedPermitsToSpend; // 需要等待的令牌数
        long waitMillis = (long) (freshPermits * intervalMillis); // 需要等待的时间

        return LongMath.saturatedAdd(nextFreeTicketMillis - nowMillis, waitMillis);
    }

2. 令牌桶控制类

Guava RateLimiter中的控制都在RateLimiter及其子类中（如SmoothBursty），本处涉及到分布式环境下的同步，因此将其解耦，令牌桶模型存储于Redis中，对其同步操作的控制放置在如下控制类，其中同步控制使用到了前面介绍的分布式锁

@Slf4j
public class RedisRateLimiter {
   

    /**
     * 获取一个令牌，阻塞一直到获取令牌，返回阻塞等待时间
     *
     * @return time 阻塞等待时间/毫秒
     */
    public long acquire(String key) throws IllegalArgumentException {
   
        return acquire(key, 1);
    }

    /**
     * 获取指定数量的令牌，如果令牌数不够，则一直阻塞，返回阻塞等待的时间
     *
     * @param permits 需要获取的令牌数
     * @return time 等待的时间/毫秒
     * @throws IllegalArgumentException tokens值不能为负数或零
     */
    public long acquire(String key, int permits) throws IllegalArgumentException {
   
        long millisToWait = reserve(key, permits);
        log.info("acquire {} permits for key[{}], waiting for {}ms", permits, key, millisToWait);
        try {
   
            Thread.sleep(millisToWait);
        } catch (InterruptedException e) {
   
            log.error("Interrupted when trying to acquire {} permits for key[{}]", permits, key, e);
        }
        return millisToWait;
    }

    /**
     * 在指定时间内获取一个令牌，如果获取不到则一直阻塞，直到超时
     *
     * @param timeout 最大等待时间（超时时间），为0则不等待立即返回
     * @param unit    时间单元
     * @return 获取到令牌则true，否则false
     * @throws IllegalArgumentException
     */
    public boolean tryAcquire(String key, long timeout, TimeUnit unit