限速神器RateLimiter原始碼解析
作者：京東科技 李玉亮
目錄指引
限流場景
軟體系統中一般有兩種場景會用到限流：
•場景一、高並行的使用者端場景。 尤其是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進行重寫，貯存許可的生產和消費演演算法是容易變化和重寫的點，將整個原始碼拷貝出來進行二次修改也是一種方案。