RateLimiter原始碼解析

納蘭小依發表於2021-03-14
MIT原始碼

RateLimiter是Guava包提供的限流器，採用了令牌桶演算法，特定是均勻地向桶中新增令牌，每次消費時也必須持有令牌，否則就需要等待。應用場景之一是限制訊息消費的速度，避免訊息消費過快而對下游的資料庫造成較大的壓力。
本文主要介紹RateLimiter的原始碼，包括他的基本限流器SmoothBursty，以及帶預熱效果的SmoothWarmingUp。

RateLimiter作為限流器的頂層類，只有兩個屬性：

  private final SleepingStopwatch stopwatch;
  private volatile Object mutexDoNotUseDirectly;

stopwatch用來計算時間間隔，以及實現了當拿不到令牌時將執行緒阻塞的功能；mutexDoNotUseDirectly主要用來進行執行緒同步。
RateLimiter作為一個抽象類，本身不能直接例項化，可以使用靜態工廠方法來建立：

 public static RateLimiter create(double permitsPerSecond);  //①
 public static RateLimiter create(double permitsPerSecond, Duration warmupPeriod);  //②
 public static RateLimiter create(double permitsPerSecond, long warmupPeriod, TimeUnit unit)  //③

RateLimiter對外提供了3個構造器，分成兩類，構造器①是第一類，底層會建立基本限流器SmoothBursty；構造器②和③是第二類，底層會建立帶warm up效果的SmoothWarmingUp。引數permitsPerSecond表示每秒產生多少個令牌，引數warmupPeriod是限流器warm up階段的時間，即限流器產生令牌從最慢到最快所需要的時間，引數unit是warm up的時間單位。
SmoothRateLimiter新增了4個屬性：

  //桶中儲存的令牌數
  double storedPermits;
  //桶中允許的最大令牌數
  double maxPermits;
  //穩定狀態下產生令牌是速度，其值為1/permitsPerSecond
  double stableIntervalMicros;
  //下一次請求需要等待的時間
  private long nextFreeTicketMicros = 0L; // could be either in the past or future

這其中比較有意思的是nextFreeTicketMicros欄位，它表示下一次獲取令牌的請求到來時需要等待的時間，該欄位可以實現上一次獲取令牌的請求預支的等待時間由下一次請求來兌現。
接下來先介紹下SmoothBursty的構造過程：

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

static RateLimiter create(double permitsPerSecond, SleepingStopwatch stopwatch) {
	RateLimiter rateLimiter = new SmoothBursty(stopwatch, 1.0 /* maxBurstSeconds */);
	rateLimiter.setRate(permitsPerSecond);
	return rateLimiter;
}

構造SmoothBursty時出傳入了兩個引數，stopwatch好理解，第二個引數意思是當限流器長時間沒用時，令牌桶內最多儲存多少秒的令牌，這裡限定了最多隻儲存1秒鐘的令牌，也就是permitsPerSecond個。
我們繼續分析setRate方法的實現：

  public final void setRate(double permitsPerSecond) {
    checkArgument(
        permitsPerSecond > 0.0 && !Double.isNaN(permitsPerSecond), "rate must be positive");
    synchronized (mutex()) {
      doSetRate(permitsPerSecond, stopwatch.readMicros());
    }
  }

setRate方法先校驗permitsPerSecond必須為整數，然後在同步塊中執行doSetRate方法。mutex方法通過雙重檢測的方式例項化mutexDoNotUseDirectly欄位，詳細程式碼略去，doSetRate是抽象方法，其具體的實現在抽象子類SmoothRateLimiter中：

  final void doSetRate(double permitsPerSecond, long nowMicros) {
    resync(nowMicros);
    double stableIntervalMicros = SECONDS.toMicros(1L) / permitsPerSecond;
    this.stableIntervalMicros = stableIntervalMicros;
    doSetRate(permitsPerSecond, stableIntervalMicros);
  }

doSetRate方法主要是設定了stableIntervalMicros欄位，呼叫的兩個方法resync和過載方法doSetRate我們接著分析。resync方法主要用來設定storedPermits和nextFreeTicketMicros這倆欄位，程式碼如下：

  void resync(long nowMicros) {
    // if nextFreeTicket is in the past, resync to now
    if (nowMicros > nextFreeTicketMicros) {
      //計算超過的這些時間裡產生了多少新的令牌
      double newPermits = (nowMicros - nextFreeTicketMicros) / coolDownIntervalMicros();
      //重新計算當前令牌桶內持有的令牌數
      storedPermits = min(maxPermits, storedPermits + newPermits);
      //更新下次准許獲取令牌的時間為當前時間
      nextFreeTicketMicros = nowMicros;
    }
  }

此方法會根據當前的時間決定是否進行欄位賦值，如果當前時間已經超過了nextFreeTicketMicros的值，那麼就重新計算storedPermits和nextFreeTicketMicros欄位，其中計算storedPermits的程式碼雖然容易理解，但是思路挺巧妙。一般來說，令牌桶演算法的令牌需要以固定的速率進行新增，那麼很自然想到可以起一個任務，按照一定的速度產生令牌，但是起一個新任務會佔用一定的資源，從而加重系統的負擔，此處的實現是根據利用時間差來計算這段時間產生的令牌數，以簡單的計算完成了新任務需要做的事情，開銷大大減少了。coolDownIntervalMicros方法是抽象方法，在SmoothBursty和SmoothWarmingUp有不同的實現，在SmoothBursty的實現是直接返回stableIntervalMicros欄位，這個欄位目前還沒設定過值，取預設值0.0，這裡double的除零操作並不會拋異常，而是會返回無窮大。
我們接著看一下doSetRate方法，這也是個抽象方法，在SmoothBursty的實現如下：

    void doSetRate(double permitsPerSecond, double stableIntervalMicros) {
      double oldMaxPermits = this.maxPermits;
      maxPermits = maxBurstSeconds * permitsPerSecond;
      if (oldMaxPermits == Double.POSITIVE_INFINITY) {
        // if we don't special-case this, we would get storedPermits == NaN, below
        storedPermits = maxPermits;
      } else {
        storedPermits =
            (oldMaxPermits == 0.0)
                ? 0.0 // initial state
                : storedPermits * maxPermits / oldMaxPermits;
      }
    }

maxPermits在此之前並沒有設定過值，因此預設是0.0，這裡只是將storedPermits初始化成了0。不過這裡的程式碼也說明，在執行期間maxPermits是可以在其他地方被修改的，如果出現了更改，就會等比例修改storedPermits的值。
到這裡SmoothBursty的初始化過程就結束了，大體上是將內部的欄位賦予了初始值。我們接下來看看SmoothBursty的使用：

  public double acquire() {
    return acquire(1);
  }

  public double acquire(int permits) {
    long microsToWait = reserve(permits);
    stopwatch.sleepMicrosUninterruptibly(microsToWait);
    return 1.0 * microsToWait / SECONDS.toMicros(1L);
  }

acquire方法用於從令牌桶中獲取令牌，引數permits表示需要獲取的令牌數量，如果當前沒辦法拿到需要的令牌，執行緒會阻塞一段時間，該方法返回等待的時間，reserve的實現如下：

  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);
    //返回等待時間，如果不需要等待，返回0
    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;
    long waitMicros =
        storedPermitsToWaitTime(this.storedPermits, storedPermitsToSpend)
            + (long) (freshPermits * stableIntervalMicros);

    this.nextFreeTicketMicros = LongMath.saturatedAdd(nextFreeTicketMicros, waitMicros);
    //扣減令牌桶庫存
    this.storedPermits -= storedPermitsToSpend;
    return returnValue;
  }

reserve的核心邏輯在reserveEarliestAvailable方法中，該方法的主要思想是檢查當前令牌桶內令牌數是否滿足需求，如果滿足則不需要額外的等待時間，否則需要將額外等待時間追加到nextFreeTicketMicros，需要注意的是方法返回的不是更新過後的nextFreeTicketMicros，而是上一次請求更新的時間，這個時間就是當前執行緒需要阻塞的時間，也就是說，當前請求所需要等待的時間是由下次請求完成的，下次請求需要的等待時間由下下次請求完成，以此類推。當前請求的令牌數超過令牌桶中的令牌數越多，下次請求需要等待的時間就越長。並且這裡並沒有對requiredPermits的上限做檢查，這就允許預支令牌，即假設桶的上限是100個令牌，一次請求可以允許超過100個令牌，只是生成多餘令牌的時間需要算到下一個請求上。同時這裡的邏輯也說明，獲取令牌是直接成功的，只是獲取完令牌後才需要一小段等待時間。
到這裡SmoothBursty的初始化以及獲取令牌的所有邏輯就介紹完了，接下來看看另一個類SmoothWarmingUp的原始碼。

  static RateLimiter create(
      double permitsPerSecond,
      long warmupPeriod,
      TimeUnit unit,
      double coldFactor,
      SleepingStopwatch stopwatch) {
    RateLimiter rateLimiter = new SmoothWarmingUp(stopwatch, warmupPeriod, unit, coldFactor);
    rateLimiter.setRate(permitsPerSecond);
    return rateLimiter;
  }

我們之前介紹的另外領個構造器的底層呼叫的是這個包級的create方法，該方法的5個引數中，只有coldFactor是新出現的，字面意思是冷啟動因子，原始碼寫死了是3.0，該值表示指在warm up階段開始時，以多大的速率產生令牌，速率是穩定速率的三分之一，冷啟動階段結束後恢復到正常速率。
setRate方法底層會呼叫如下的doSetRate方法：

  final void doSetRate(double permitsPerSecond, long nowMicros) {
    resync(nowMicros);
    double stableIntervalMicros = SECONDS.toMicros(1L) / permitsPerSecond;
    this.stableIntervalMicros = stableIntervalMicros;
    doSetRate(permitsPerSecond, stableIntervalMicros);
  }

   void doSetRate(double permitsPerSecond, double stableIntervalMicros) {
	  double oldMaxPermits = maxPermits;
	  //設定冷啟動生成令牌的間隔是正常值的3倍(codeFactor固定為3)
	  double coldIntervalMicros = stableIntervalMicros * coldFactor;
	  thresholdPermits = 0.5 * warmupPeriodMicros / stableIntervalMicros;
	  maxPermits = thresholdPermits + 2.0 * warmupPeriodMicros / (stableIntervalMicros + coldIntervalMicros);
	  //slope是梯形部分斜線的斜率
	  slope = (coldIntervalMicros - stableIntervalMicros) / (maxPermits - thresholdPermits);
	  if (oldMaxPermits == Double.POSITIVE_INFINITY) {
	    // if we don't special-case this, we would get storedPermits == NaN, below
	    storedPermits = 0.0;
	  } else {
	    storedPermits =
	        (oldMaxPermits == 0.0)
	            ? maxPermits // initial state is cold
	            : storedPermits * maxPermits / oldMaxPermits;
	  }
	}


doSetRate的程式碼不容易理解，原始碼中利用圖示介紹了幾個變數之間的關係（但是本人仍然不是很理解，因此只能將結論放在這裡，無法進行更多解釋），如圖所示，原始碼註釋中說明了如下的兩個等式：

  • 梯形的面積等於預熱時間warmupPeriodMicros
warmupPeriodMicros = 0.5 * (coldIntervalMicros + stableIntervalMicros) * (maxPermits - thresholdPermits)

由此可以得到maxPermits的值：

maxPermits = thresholdPermits + 2.0 * warmupPeriodMicros / (stableIntervalMicros + coldIntervalMicros);
  • 左邊矩形的面積是梯形面積的一半，由此可知：
warmupPeriodMicros * 0.5 = thresholdPermits * stableIntervalMicros

計算出thresholdPermits的值為：

thresholdPermits = 0.5 * warmupPeriodMicros / stableIntervalMicros

SmoothWarmingUp的初始化邏輯到這裡就結束了，接下來介紹下它獲取令牌的流程，acquire方法的其他部分上文已經結束過，此處重點介紹storedPermitsToWaitTime方法：

    long storedPermitsToWaitTime(double storedPermits, double permitsToTake) {
      //儲存的令牌數量超出thresholdPermits的部分，這部分反應在梯形區域
      double availablePermitsAboveThreshold = storedPermits - thresholdPermits;
      long micros = 0;
      // measuring the integral on the right part of the function (the climbing line)
      if (availablePermitsAboveThreshold > 0.0) {
      	//permitsAboveThresholdToTake表示梯形區域的高
        double permitsAboveThresholdToTake = min(availablePermitsAboveThreshold, permitsToTake);
        //length計算的是梯形的上底+下底
        double length =
            permitsToTime(availablePermitsAboveThreshold)
                + permitsToTime(availablePermitsAboveThreshold - permitsAboveThresholdToTake);
        //梯形區域的面積，即生產梯形區域的令牌數所需要的時間
        micros = (long) (permitsAboveThresholdToTake * length / 2.0);
        //扣除掉需要消耗的梯形區域的令牌數，表示還需要從左側矩形區域取得的令牌數量
        permitsToTake -= permitsAboveThresholdToTake;
      }
      // measuring the integral on the left part of the function (the horizontal line)
      //等待時間=梯形區域的時間+矩形區域的時間
      micros += (long) (stableIntervalMicros * permitsToTake);
      return micros;
    }

    //由前文可知，slope = =y/x = 產生令牌間隔/令牌數，permits * slope表示產生令牌間隔的增量，加上stableIntervalMicros表示梯形的底
    private double permitsToTime(double permits) {
      return stableIntervalMicros + permits * slope;
    }


此處的storedPermitsToWaitTime與SmoothBursty中的實現大不相同，SmoothBursty由於不需要預熱，可以直接獲取桶中的令牌，因此直接返回了0，而此處存在預熱階段，不能直接獲取到令牌，因此計算邏輯稍微複雜些，總體來說，就是求圖中陰影部分的面積。

相關文章