限速神器RateLimiter原始碼解析

京東雲開發者發表於2023-05-16

作者:京東科技 李玉亮

目錄指引

限流場景

軟體系統中一般有兩種場景會用到限流:

場景一、高併發的使用者端場景。 尤其是C端系統,經常面對海量使用者請求,如不做限流,遇到瞬間高併發的場景,則可能壓垮系統。

場景二、內部交易處理場景。 如某類交易任務處理時有速率要求,再如上下游呼叫時下游對上游有速率要求。

•無論哪種場景,都需要對請求處理的速率進行限制,或者單個請求處理的速率相對固定,或者批次請求的處理速率相對固定,見下圖:

常用的限流演算法有如下幾種:

演算法一、訊號量演算法。 維護最大的併發請求數(如連線數),當併發請求數達到閾值時報錯或等待,如執行緒池。

演算法二、漏桶演算法。 模擬一個按固定速率漏出的桶,當流入的請求量大於桶的容量時溢位。

演算法三、令牌桶演算法。 以固定速率向桶內發放令牌。請求處理時,先從桶裡獲取令牌,只服務有令牌的請求。

本次要介紹的RateLimiter使用的是令牌桶演算法。RateLimiter是google的guava包中的一個輕巧限流元件,它主要有兩個java類檔案,RateLimiter.java和SmoothRateLimiter.java。兩個類檔案共有java程式碼301行、註釋420行,註釋比java程式碼還要多,寫的非常詳細,後面的介紹也有相關內容是翻譯自其註釋,有些描述英文原版更加準確清晰,有興趣的也可以結合原版註釋進行更詳細的瞭解。

使用介紹

RateLimiter使用時只需引入guava jar便可,最新的版本是31.1-jre, 本文介紹的原始碼也是此版本。

            <dependency>
                <groupId>com.google.guava</groupId>
                <artifactId>guava</artifactId>
                <version>31.1-jre</version>
            </dependency>

原始碼中提供了兩個直觀的使用示例。

示例一、有一系列任務列表要提交執行,控制提交速率不超過每秒2個。

 final RateLimiter rateLimiter = RateLimiter.create(2.0); // 建立一個每秒2個許可的RateLimiter物件.
 void submitTasks(List<Runnable> tasks, Executor executor) {
   for (Runnable task : tasks) {
     rateLimiter.acquire(); // 此處可能有等待
     executor.execute(task);
   }
 }

示例二、以不超過5kb/s的速率產生資料流。

 final RateLimiter rateLimiter = RateLimiter.create(5000.0); // 建立一個每秒5k個許可的RateLimiter物件
 void submitPacket(byte[] packet) {
   rateLimiter.acquire(packet.length);
   networkService.send(packet);
 }

可以看出RateLimiter的使用非常簡單,只需要構造限速器,呼叫獲取許可方法便可,不需要釋放許可.

演算法介紹

在介紹之前,先說一下RateLimiter中的幾個名詞:

許可( permit ): 代表一個令牌,獲取到許可的請求才能放行。

資源利用不足( underunilization ): 許可的發放一般是勻速的,但請求未必是勻速的,有時會有無請求(資源利用不足)的場景,令牌桶會有貯存機制。

貯存許可( storedPermit ): 令牌桶支援對空閒資源進行許可貯存,許可請求時優先使用貯存許可。

新鮮許可( freshPermit ): 當貯存許可為空時,採用透支方式,下發新鮮許可,同時設定下次許可生效時間為本次新鮮許可的結束時間。

•如下為一個許可發放示例,矩形代表整個令牌桶,許可產生速度為1個/秒,令牌桶裡有一個貯存桶,容量為2。

以上示例中,在T1貯存容量為0,許可請求時直接返回1個新鮮許可,貯存容量隨著時間推移,增長至最大值2,在T2時收到3個許可的請求,此時會先從貯存桶中取出2個,然後再產生1個新鮮許可,0.5s後在T3時刻又來了1個許可請求,由於最近的許可0.5s後才會下發,因此先sleep0.5s再下發。

RateLimiter的核心功能是限速,我們首先想到的限速方案是記住最後一次下發令牌許可(permit)時間,下次許可請求時,如果與最後一次下發許可時間的間隔小於1/QPS,則進行sleep至1/QPS,否則直接發放,但該方法不能感知到資源利用不足的場景。一方面,隔了很長一段再來請求許可,則可能系統此時相對空閒,可下發更多的許可以充分利用資源;另一方面,隔了很長一段時間再來請求許可,也可能意味著處理請求的資源變冷(如快取失效),處理效率會下降。因此在RateLimiter中,增加了資源利用不足(underutilization)的管理,在程式碼中體現為貯存許可(storedPermits),貯存許可值最開始為0,隨著時間的增加,一直增長為最大貯存許可數。許可獲取時,首先從貯存許可中獲取,然後再根據下次新鮮許可獲取時間來進行新鮮許可獲取。這裡要說的是RateLimiter是記住了下次令牌發放的時間,類似於透支的功能,當前許可獲取時立刻返回,同時記錄下次獲取許可的時間。

程式碼結構和主體流程

程式碼結構

整體類圖如下:

RateLimiter類

RateLimiter類是頂級類,也是唯一暴露給使用者的類,它提供了工廠方法來建立RateLimiter方法。 create(double permitsPerSecond) 方法建立的是突發限速器,create(double permitsPerSecond, Duration warmupPeriod)方法建立的是預熱限速器。同時它提供了acquire方法用於獲取令牌,提供了tryAcquire方法用於嘗試獲取令牌。該類的內部實現上,一方面有一個SleepingStopWatch 用於sleep操作,另一方面有一個mutexDoNotUseDirectly變數和mutex()方法進行互斥加鎖。

SmoothRateLimiter類

該類繼承了RateLimiter類,是一個抽象類,含義為平滑限速器,限制速率是平滑的,maxPermits和storedPermits維護了最大儲存許可數量和當前儲存許可數量;stableIntervalMicros指規定的穩定許可發放間隔,nextFreeTicketMicros指下一個空閒許可時間。

SmoothBursty類

平滑突發限速器,該類繼承了SmoothRateLimiter,它儲存許可的發放頻率同設定的stableIntervalMicros,有一個成員變數maxBurstSeconds,代表最多儲存多長時間的令牌許可。

SmoothWarmingUp類

平滑預熱限速器,繼承了SmoothRateLimiter,與SmoothBursty平級,它的預熱演算法需要一定的理解成本。

主體流程

獲取許可的主體流程如下:

主體流程主要是對貯存許可數量和新鮮許可數量進行計算和更新,得到當前許可請求的等待時間。SmoothBursty演算法和SmoothWarmingUp演算法共用這一套主體流程,差異主要是貯存許可的管理策略,兩種演算法的不同策略在兩個子類中各自實現,SmoothBursty演算法相對簡單一些,下面先介紹該演算法,然後再介紹SmoothWarmingUp演算法。

SmoothBursty演算法

限速器建立

採用的是工廠模式建立,原始碼如下:

  public static RateLimiter create(double permitsPerSecond) {
    // permitsPerSecond指每秒允許的許可數. 該方法呼叫了下面的方法
    return create(permitsPerSecond, SleepingStopwatch.createFromSystemTimer());
  }
  // 建立SmoothBursty(固定貯存1s的貯存許可), 然後設定速率
  static RateLimiter create(double permitsPerSecond, SleepingStopwatch stopwatch) {
    RateLimiter rateLimiter = new SmoothBursty(stopwatch, 1.0 /* maxBurstSeconds */);
    rateLimiter.setRate(permitsPerSecond);
    return rateLimiter;
  }

1、SmoothBursty的構造方法相對簡單:

    SmoothBursty(SleepingStopwatch stopwatch, double maxBurstSeconds) {
      super(stopwatch);
      this.maxBurstSeconds = maxBurstSeconds;
    }

2、rateLimiter.setRate的定義在父類RateLimiter中

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

該方法使用synchronized(mutex())方法對互斥鎖進行同步,以保證多執行緒呼叫的安全,然後呼叫子類的doSetRate方法。 第二個引數nowMicros傳的值是呼叫了stopwatch的方法,將限速器建立的時間定義為0,然後計算了當前時間和建立時間的時間差,因此採用的是相對時間。

2.1 mutex方法的實現如下:

  // Can't be initialized in the constructor because mocks don't call the constructor.
  // 從上行註釋可看出,這是因為mock才用了懶載入, 實際上即時載入程式碼更簡潔
  @CheckForNull private volatile Object mutexDoNotUseDirectly;
  // 雙重檢查鎖的懶載入模式
  private Object mutex() {
    Object mutex = mutexDoNotUseDirectly;
    if (mutex == null) {
      synchronized (this) {
        mutex = mutexDoNotUseDirectly;
        if (mutex == null) {
          mutexDoNotUseDirectly = mutex = new Object();
        }
      }
    }
    return mutex;
  }

該方法使用了雙重檢查鎖來對鎖物件mutexDoNotUseDirectly進行懶載入,另外該方法透過mutex臨時變數來解決了雙重檢查鎖失效的問題。

2.2 doSetRate方法的主體實現在SmoothRateLimiter類中:

  final void doSetRate(double permitsPerSecond, long nowMicros) {
    // 同步貯存許可和時間
    resync(nowMicros);
    double stableIntervalMicros = SECONDS.toMicros(1L) / permitsPerSecond;
    this.stableIntervalMicros = stableIntervalMicros;
    doSetRate(permitsPerSecond, stableIntervalMicros);
  }

該方法在限速器建立時會呼叫,建立後呼叫限速器的setRate重置速率時也會呼叫。

2.2.1 resync方法用於基於當前時間重新整理計算最新的storedPermis和nextFreeTicketMicros.

  /** Updates {@code storedPermits} and {@code nextFreeTicketMicros} based on the current time. */
  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;
    }
  }

該方法從現實場景上講,代表的是隨著時間的流逝,貯存許可不斷增加,但從技術實現的角度,並不是真正的持續重新整理,而是僅在需要時呼叫重新整理。該方法如果當前時間小於等於下次許可時間,則貯存許可數量和下次許可時間不需要重新整理;否則透過 (當前時間-下次許可時間)/貯存許可的發放間隔計算出的值域最大貯存數量取小,則為已貯存的許可數量,需要注意的是貯存許可數量是double型別的。

限速器使用

限速器常用的方法主要有accquire和tryAccquire。

先說一下accquire方法, 共有兩個共有方法,一個是無參的,每次獲取1個許可,再一個是整數引數的,每次呼叫獲取多個許可。

  // 獲取1個許可
  public double acquire() {
    return acquire(1);
  }
 
  // 獲取多個許可
  public double acquire(int permits) {
    // 留出permits個許可,得到需要sleep的微秒數.
    long microsToWait = reserve(permits);
    // 該方法如果小於等於零則直接返回,否則sleep
    stopwatch.sleepMicrosUninterruptibly(microsToWait);
    // 返回休眠的秒數.
    return 1.0 * microsToWait / SECONDS.toMicros(1L);
  }

從以上原始碼可看出,獲取許可的邏輯很簡單:留出permits個許可,根據返回值決定是否sleep等待。留出許可的方法實現如下:

 // 預留出permits個許可
  final long reserve(int permits) {
    checkPermits(permits);
    synchronized (mutex()) {
      return reserveAndGetWaitLength(permits, stopwatch.readMicros());
    }
  }
   
  // 預留出permits個需求,得到需要等待的時間
  final long reserveAndGetWaitLength(int permits, long nowMicros) {
    long momentAvailable = reserveEarliestAvailable(permits, nowMicros);
    return max(momentAvailable - nowMicros, 0);
  }
  abstract long reserveEarliestAvailable(int permits, long nowMicros);

reserveEarliestAvailable為抽象方法,實現在SmoothRateLimiter類中,該方法是核心主鏈路方法,該方法先從貯存許可中獲取,如果數量足夠則直接返回,否則先將全部貯存許可取出,再計算還需要的等待時間,邏輯如下:

  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);
    // 透支後的下次許可可用時間=當前時間(nextFreeTicketMicros)+等待時間(waitMicros)
    this.nextFreeTicketMicros = LongMath.saturatedAdd(nextFreeTicketMicros, waitMicros);
    // 貯存許可數量減少
    this.storedPermits -= storedPermitsToSpend;
    return returnValue;
  }

該方法有兩點說明:1、returnValue為之前計算的下次空閒時間(前面有說RateLimiter採用預支的模式,本次直接返回,同時計算下次的最早空閒時間) 2、貯存許可的等待時間不同的實現方邏輯不同,SmoothBursty演算法認為貯存許可直接可用,所以返回0, 後面的SmoothWarmingUp演算法認為貯存許可需要消耗比正常速率更多的預熱時間,有一定演算法邏輯.

至此整個accquire方法的呼叫鏈路分析結束,下面再看tryAccquire方法就比較簡單了,tryAccquire比accquire差異的邏輯在於tryAccquire方法會判斷下次許可時間-當前時間是否大於超時時間,如果是則直接返回false,否則進行sleep並返回true. 方法原始碼如下:

  public boolean tryAcquire(Duration timeout) {
    return tryAcquire(1, toNanosSaturated(timeout), TimeUnit.NANOSECONDS);
  }

  public boolean tryAcquire(long timeout, TimeUnit unit) {
    return tryAcquire(1, timeout, unit);
  }

  public boolean tryAcquire(int permits) {
    return tryAcquire(permits, 0, MICROSECONDS);
  }

  public boolean tryAcquire() {
    return tryAcquire(1, 0, MICROSECONDS);
  }

  public boolean tryAcquire(int permits, Duration timeout) {
    return tryAcquire(permits, toNanosSaturated(timeout), TimeUnit.NANOSECONDS);
  }

  public boolean tryAcquire(int permits, long timeout, TimeUnit unit) {
    long timeoutMicros = max(unit.toMicros(timeout), 0);
    checkPermits(permits);
    long microsToWait;
    synchronized (mutex()) {
      long nowMicros = stopwatch.readMicros();
      // 判斷超時微秒數是否可等到下個許可時間
      if (!canAcquire(nowMicros, timeoutMicros)) {
        return false;
      } else {
        microsToWait = reserveAndGetWaitLength(permits, nowMicros);
      }
    }
    // 休眠等待
    stopwatch.sleepMicrosUninterruptibly(microsToWait);
    return true;
  }
  
  // 下次許可時間-超時時間<=當前時間
  private boolean canAcquire(long nowMicros, long timeoutMicros) {
    return queryEarliestAvailable(nowMicros) - timeoutMicros <= nowMicros;
  }

SmoothWarmingUp演算法

SmoothWarmingUp演算法的主體處理流程同SmoothBurstry演算法,主要在貯存許可時間計算上的兩個方法進行了新實現,該演算法不像SmoothBurstry演算法那麼直觀好理解,需要先了解演算法邏輯,再看原始碼。

演算法說明

該演算法在原始碼註釋中已經描述的比較清晰了,主要思想是限流器的初始貯存許可數量便是最大貯存許可值, 貯存許可執行時按一定演算法由慢到快的產生,直至設定的固定速率,以此來達到預熱過程。該演算法涉及到一些數學知識,如果不是很感興趣,則瞭解其主要思想便可。下面詳細說一下該演算法。

說到該演算法前,我們再回頭看一下SmoothRateLimiter的貯存許可,貯存許可有當前數量和最大數量,另外還有兩個演算法邏輯,一個是貯存許可生產的速率控制,再一個是貯存許可消費速率的控制,在Bursty演算法中,生產的速率同設定的固定速率,而消費的速率為無窮大(立刻消費,不佔用時間);在WarmingUp演算法中,需對照下圖進行分析:

該圖可這樣理解,每個貯存許可的消費耗時為右側梯形面積,梯形面積=(上邊長+下邊長)/2 * 高. 可以看到每個貯存許可的面積越來越小,直到固定速率的長方形面積。

在限速器初始化時,輸入的變數有固定速率和預熱時間,另外冷卻因子是固定值3;在作者演算法中,首先計算的是閾值許可數 = 0.5 * 預熱週期 / 固定速率. 然後計算的是最大許可數,我們知道了梯形的面積、上邊(大速率)、下邊(小速率),便能推到出高,最大許可=閥值許可數 + 高。

void doSetRate(double permitsPerSecond, double stableIntervalMicros) {
      double oldMaxPermits = maxPermits;
      double coldIntervalMicros = stableIntervalMicros * coldFactor;
      thresholdPermits = 0.5 * warmupPeriodMicros / stableIntervalMicros;
      maxPermits =
          thresholdPermits + 2.0 * warmupPeriodMicros / (stableIntervalMicros + coldIntervalMicros);
      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;
      }
    }

在具體使用中,一個是生產的速率,固定為預熱時間/最大許可數,原始碼如下:

  double coolDownIntervalMicros() {
      return warmupPeriodMicros / maxPermits;
    }

再一個是消費的速率,按如上曲線從右至左的面積=梯形面積+長方形面積,梯形面積=(上邊+下邊) /2 * 高 ,原始碼如下:

    long storedPermitsToWaitTime(double storedPermits, double permitsToTake) {
      double availablePermitsAboveThreshold = storedPermits - thresholdPermits;
      long micros = 0;
      // measuring the integral on the right part of the function (the climbing line)
      if (availablePermitsAboveThreshold > 0.0) {
        double permitsAboveThresholdToTake = min(availablePermitsAboveThreshold, permitsToTake);
        // TODO(cpovirk): Figure out a good name for this variable.
        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;
    }

原始碼分析

瞭解了以上演算法後,再看下面的原始碼就相對簡單了。

  static final class SmoothWarmingUp extends SmoothRateLimiter {
    // 預熱時間
    private final long warmupPeriodMicros;
    //斜率
    private double slope;
    //閾值許可
    private double thresholdPermits;
    //冷卻因子
    private double coldFactor;

    SmoothWarmingUp(
        SleepingStopwatch stopwatch, long warmupPeriod, TimeUnit timeUnit, double coldFactor) {
      super(stopwatch);
      this.warmupPeriodMicros = timeUnit.toMicros(warmupPeriod);
      this.coldFactor = coldFactor;
    }
    
    // 引數初始化
    @Override
    void doSetRate(double permitsPerSecond, double stableIntervalMicros) {
      double oldMaxPermits = maxPermits;
      double coldIntervalMicros = stableIntervalMicros * coldFactor;
      thresholdPermits = 0.5 * warmupPeriodMicros / stableIntervalMicros;
      maxPermits =
          thresholdPermits + 2.0 * warmupPeriodMicros / (stableIntervalMicros + coldIntervalMicros);
      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;
      }
    }

    // 有storedPermits個貯存許可,要使用permitsToTake個時的等待時間計算
    @Override
    long storedPermitsToWaitTime(double storedPermits, double permitsToTake) {
      double availablePermitsAboveThreshold = storedPermits - thresholdPermits;
      long micros = 0;
      // measuring the integral on the right part of the function (the climbing line)
      if (availablePermitsAboveThreshold > 0.0) {
        double permitsAboveThresholdToTake = min(availablePermitsAboveThreshold, permitsToTake);
        // TODO(cpovirk): Figure out a good name for this variable.
        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;
    }
    // 許可耗時=固定速率+許可值*斜率
    private double permitsToTime(double permits) {
      return stableIntervalMicros + permits * slope;
    }
    // 冷卻間隔固定為預熱時間/最大許可數.
    @Override
    double coolDownIntervalMicros() {
      return warmupPeriodMicros / maxPermits;
    }
  }

思考總結

sleep說明和相對時間

RateLimiter內部使用類StopWatch進行了一個相對時間的度量,RateLimiter建立時,時間為0,然後向後累計,sleep時不受interrupt異常影響。

double浮點數

RateLimiter暴露的API的許可數量入參為整數型別,但內部計算時實際是浮點double型別,支援小數許可數量,一方面浮點存在丟失精度,另一方面也不便於理解;是否可以使用整數值得考慮。

只支援單機

RateLimiter的這幾種演算法只支援單機限流,如要支援叢集限流,一種方式是先根據負載均衡的權重計算出單機的限速值,再進行單節點限速;另一種方式是參考該元件使用redis等中心化數量管理的中介軟體,但效能和穩定性會降低一些。

擴充套件性

RateLimiter提供了有限的擴充套件能力,自帶的SmoothBursty和SmoothWarmingUp類不是公開類,不能直接建立或調整引數,如關閉貯存功能或調整預熱係數等。這種場景需要繼承SmoothRateLimiter進行重寫,貯存許可的生產和消費演算法是容易變化和重寫的點,將整個原始碼複製出來進行二次修改也是一種方案。

相關文章