RateLimiter源码解读

最近使用了Guava提供的限流器RateLimiter觉得十分好用，遍想一探究竟。

RateLimiter使用的限流算法为令牌桶，这里对令牌桶算法不做过多的解释，不懂的同学可以搜索一番。

限流器创建

public static RateLimiter create(double permitsPerSecond) {
    return create(permitsPerSecond, SleepingStopwatch.createFromSystemTimer());
}


static RateLimiter create(double permitsPerSecond, SleepingStopwatch stopwatch) {
    
    // 1.new SmoothBursty对象
    RateLimiter rateLimiter = new SmoothBursty(stopwatch, 1.0);
    // 2.设置 限流值
    rateLimiter.setRate(permitsPerSecond);
    return rateLimiter;
}

SmoothBursty是RateLimiter的子类，它有一个参数名为maxBurstSeconds。这个参数的作用为，当限流器没有被使用时最多可以存放多少秒的令牌桶。默认的参数为1，即只可以保存1s产生的令牌。 下面来具体看一下限流值是如何被设置的。

public final void setRate(double permitsPerSecond) {
    
    // 校验参数必须大于0
    checkArgument(
        permitsPerSecond > 0.0 && !Double.isNaN(permitsPerSecond), "rate must be positive");
    
    // 加锁设置限流值
    synchronized (mutex()) {
      doSetRate(permitsPerSecond, stopwatch.readMicros());
    }
}

final void doSetRate(double permitsPerSecond, long nowMicros) {
    
   	// 基于现在的时间更新令牌数和下一次可以获取令牌桶的时间
    resync(nowMicros);
    
    // 计算时间间隔
    double stableIntervalMicros = SECONDS.toMicros(1L) / permitsPerSecond;
    this.stableIntervalMicros = stableIntervalMicros;
    
    // 设置时间间隔
    doSetRate(permitsPerSecond, stableIntervalMicros);
}

void resync(long nowMicros) {
    
    // 如果当前的时间大于下一次可以获取令牌桶的时间 （初始化nextFreeTicketMicros为0）
    if (nowMicros > nextFreeTicketMicros) {
        
      // 计算令牌数         (当前时间 - 下一次获取令牌的时间) / 冷启动时间
      double newPermits = (nowMicros - nextFreeTicketMicros) / coolDownIntervalMicros();
      
      // 计算可存储的令牌数 取最大的令牌数与存储的令牌数和新生成的令牌数之间较小的值
      storedPermits = min(maxPermits, storedPermits + newPermits);
        
      // 将下一次获取令牌数的时间设置为当前时间
      nextFreeTicketMicros = nowMicros;
    }
  }


void doSetRate(double permitsPerSecond, double stableIntervalMicros) {
      double oldMaxPermits = this.maxPermits;
    
      // 最大令牌数 = 最多可以储存的秒数 * 每秒产生的令牌数
      maxPermits = maxBurstSeconds * permitsPerSecond;
      if (oldMaxPermits == Double.POSITIVE_INFINITY) {
        storedPermits = maxPermits;
      } else {
        storedPermits =
            (oldMaxPermits == 0.0)
                ? 0.0 // initial state
                : storedPermits * maxPermits / oldMaxPermits;
      }
    }

在设置限流值的代码中，比较重要的方法是resync，该方法重新设置了下一次获取令牌数的时间nextFreeTicketMicros和可以获取的令牌数。在该方法内有一个coolDownIntervalMicros()方法，该方法是一个抽象方法具体的逻辑有子类实现SmoothBursty中的实现使用的是使用固定的时间间隔，即1s/限流值。参看和另一个类SmoothWarmingUp实现之差别可以得知SmoothWarmingUp提供的是带慢启动的限流器。

限流器的使用

public double acquire(int permits) {
    
  // 使用指定数量的令牌数
  long microsToWait = reserve(permits);
  
  // 休息一下
  stopwatch.sleepMicrosUninterruptibly(microsToWait);
    
  // 计算等待时间
  return 1.0 * microsToWait / SECONDS.toMicros(1L);
}


final long reserve(int permits) {
    checkPermits(permits);
    synchronized (mutex()) {
      // 使用并且获取等待时间
      return reserveAndGetWaitLength(permits, stopwatch.readMicros());
    }
}

final long reserveAndGetWaitLength(int permits, long nowMicros) {
    
    // 使用最先可以使用的令牌
    long momentAvailable = reserveEarliestAvailable(permits, nowMicros);
    return max(momentAvailable - nowMicros, 0);
}

final long reserveEarliestAvailable(int requiredPermits, long nowMicros) {
    
    // 基于现在的时间更新令牌数和下一次可以获取令牌桶的时间
    resync(nowMicros);
    long returnValue = nextFreeTicketMicros;
    double storedPermitsToSpend = min(requiredPermits, this.storedPermits);
    
    // 计算是否还需要新产生令牌
    double freshPermits = requiredPermits - storedPermitsToSpend;
    
    // 计算等待时间  如果 requiredPermits > storedPermits 则freshPermits 不为0 则需要等待
    //             如果 requiredPermits < storedPermits 则freshPermits 为0   则不需要等待
    long waitMicros =
        storedPermitsToWaitTime(this.storedPermits, storedPermitsToSpend)
            + (long) (freshPermits * stableIntervalMicros);

    // 计算下一次的获取令牌的时间
    this.nextFreeTicketMicros = LongMath.saturatedAdd(nextFreeTicketMicros, waitMicros);
    this.storedPermits -= storedPermitsToSpend;
    return returnValue;
}

获取令牌的具体逻辑已经写到注释里面了，具体可以参看相关注释。这里说一下获取令牌的具体算法，短短几行代码就可以实现核心逻辑，真是代码的艺术。

1. 当当前时间 > 下一次获取令牌的时间时 则产生新的令牌 否则则不产生新的令牌
2. 需要的令牌数和存储的令牌数两者之间取小值和需要的令牌数做差，如果差值 > 0则说明当前存储的令牌数不足，需要等待时间产生新的令牌。如若 = 0 则说明存储的令牌数大于需要的令牌数，则不需要等待。