面試官：來談談限流-RateLimiter源碼分析

RateLimiter有兩個實現類：SmoothBursty和SmoothWarmingUp，其都是令牌桶算法的變種實現，區別在於SmoothBursty加令牌的速度是恆定的，而SmoothWarmingUp會有個預熱期，在預熱期內加令牌的速度是慢慢增加的，直到達到固定速度爲止。其適用場景是，對於有的系統而言剛啓動時能承受的QPS較小，需要預熱一段時間後才能達到最佳狀態。

基本使用

RateLimiter的使用很簡單：

//create方法傳入的是每秒生成令牌的個數
RateLimiter rateLimiter= RateLimiter.create(1);
for (int i = 0; i < 5; i++) {
    //acquire方法傳入的是需要的令牌個數，當令牌不足時會進行等待，該方法返回的是等待的時間
    double waitTime=rateLimiter.acquire(1);
    System.out.println(System.currentTimeMillis()/1000+" , "+waitTime);
}

輸出如下：

1548070953 , 0.0
1548070954 , 0.998356
1548070955 , 0.998136
1548070956 , 0.99982

需要注意的是，當令牌不足時，acquire方法並不會阻塞本次調用，而是會算在下次調用的頭上。比如第一次調用時，令牌桶中並沒有令牌，但是第一次調用也沒有阻塞，而是在第二次調用的時候阻塞了1秒。也就是說，每次調用欠的令牌（如果桶中令牌不足）都是讓下一次調用買單。

RateLimiter rateLimiter= RateLimiter.create(1);
double waitTime=rateLimiter.acquire(1000);
System.out.println(System.currentTimeMillis()/1000+" , "+waitTime);
waitTime=rateLimiter.acquire(1);
System.out.println(System.currentTimeMillis()/1000+" , "+waitTime);

輸出如下：

1548072250 , 0.0
1548073250 , 999.998773

這樣設計的目的是：

Last, but not least: consider a RateLimiter with rate of 1 permit per second, currently completely unused, and an expensive acquire(100) request comes. It would be nonsensical to just wait for 100 seconds, and /then/ start the actual task. Why wait without doing anything? A much better approach is to /allow/ the request right away (as if it was an acquire(1) request instead), and postpone /subsequent/ requests as needed. In this version, we allow starting the task immediately, and postpone by 100 seconds future requests, thus we allow for work to get done in the meantime instead of waiting idly.

簡單的說就是，如果每次請求都爲本次買單會有不必要的等待。比如說令牌增加的速度爲每秒1個，初始時桶中沒有令牌，這時來了個請求需要100個令牌，那需要等待100s後才能開始這個任務。所以更好的辦法是先放行這個請求，然後延遲之後的請求。

另外，RateLimiter還有個tryAcquire方法，如果令牌夠會立即返回true，否則立即返回false。

源碼分析

本文主要分析SmoothBursty的實現。

首先看SmoothBursty中的幾個關鍵字段:

// 桶中最多存放多少秒的令牌數
final double maxBurstSeconds;
//桶中的令牌個數
double storedPermits;
//桶中最多能存放多少個令牌，=maxBurstSeconds*每秒生成令牌個數
double maxPermits;
//加入令牌的平均間隔，單位爲微秒，如果加入令牌速度爲每秒5個，則該值爲1000*1000/5
double stableIntervalMicros;
//下一個請求需要等待的時間
private long nextFreeTicketMicros = 0L;

RateLimiter的創建

先看創建RateLimiter的create方法。

// permitsPerSecond爲每秒生成的令牌數
public static RateLimiter create(double permitsPerSecond) {
    return create(permitsPerSecond, SleepingStopwatch.createFromSystemTimer());
}

//SleepingStopwatch主要用於計時和休眠
static RateLimiter create(double permitsPerSecond, SleepingStopwatch stopwatch) {
    //創建一個SmoothBursty
    RateLimiter rateLimiter = new SmoothBursty(stopwatch, 1.0 /* maxBurstSeconds */);
    rateLimiter.setRate(permitsPerSecond);
    return rateLimiter;
}

create方法主要就是創建了一個SmoothBursty實例，並調用了其setRate方法。注意這裏的maxBurstSeconds寫死爲1.0。

@Override
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大，說明上一個請求欠的令牌已經補充好了，本次請求不用等待
    if (nowMicros > nextFreeTicketMicros) {
      // 計算這段時間內需要補充的令牌，coolDownIntervalMicros返回的是stableIntervalMicros
      double newPermits = (nowMicros - nextFreeTicketMicros) / coolDownIntervalMicros();
     // 更新桶中的令牌，不能超過maxPermits
      storedPermits = min(maxPermits, storedPermits + newPermits);
      // 這裏先設置爲nowMicros
      nextFreeTicketMicros = nowMicros;
    }
}

@Override
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 {
        //第一次調用oldMaxPermits爲0，所以storedPermits（桶中令牌個數）也爲0
        storedPermits =
                (oldMaxPermits == 0.0)
                        ? 0.0 // initial state
                        : storedPermits * maxPermits / oldMaxPermits;
    }
}

setRate方法中設置了maxPermits=maxBurstSeconds * permitsPerSecond;而maxBurstSeconds爲1，所以maxBurstSeconds 只會保存1秒中的令牌數。

需要注意的是SmoothBursty是非public的類，也就是說只能通過RateLimiter.create方法創建，而該方法中的maxBurstSeconds 是寫死1.0的，也就是說我們只能創建桶大小爲permitsPerSecond*1的SmoothBursty對象（當然反射的方式不在討論範圍），在guava的github倉庫裏有好幾條issue（issue1,issue2,issue3,issue4）希望能由外部設置maxBurstSeconds ，但是並沒有看到官方人員的回覆。而在唯品會的開源項目vjtools中，有人提出了這個問題，唯品會的同學對guava的RateLimiter進行了拓展。

對於guava的這樣設計我很不理解，有清楚的朋友可以說下~

到此爲止一個SmoothBursty對象就創建好了，接下來我們分析其acquire方法。

acquire方法

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方法，可能會更新桶中的令牌值和nextFreeTicketMicros
    resync(nowMicros);
    // 如果上次請求花費的令牌還沒有補齊，這裏returnValue爲上一次請求後需要等待的時間，否則爲nowMicros
    long returnValue = nextFreeTicketMicros;
    double storedPermitsToSpend = min(requiredPermits, this.storedPermits);
    // 缺少的令牌數
    double freshPermits = requiredPermits - storedPermitsToSpend;
    // waitMicros爲下一次請求需要等待的時間；SmoothBursty的storedPermitsToWaitTime返回0
    long waitMicros =
        storedPermitsToWaitTime(this.storedPermits, storedPermitsToSpend)
            + (long) (freshPermits * stableIntervalMicros);
    // 更新nextFreeTicketMicros
    this.nextFreeTicketMicros = LongMath.saturatedAdd(nextFreeTicketMicros, waitMicros);
    // 減少令牌
    this.storedPermits -= storedPermitsToSpend;
    return returnValue;
}

acquire中會調用reserve方法獲得當前請求需要等待的時間，然後進行休眠。reserve方法最終會調用到reserveEarliestAvailable，在該方法中會先調用上文提到的resync方法對桶中的令牌進行補充（如果需要的話），然後減少桶中的令牌，以及計算這次請求欠的令牌數及需要等待的時間（由下次請求負責等待）。

如果上一次請求沒有欠令牌或欠的令牌已經還清則返回值爲nowMicros，否則返回值爲上一次請求缺少的令牌個數*生成一個令牌所需要的時間。

End

本文講解了RateLimiter子類SmoothBursty的源碼，對於另一個子類SmoothWarmingUp的原理大家可以自行分析。相對於傳統意義上的令牌桶，RateLimiter的實現還是略有不同，主要體現在一次請求的花費由下一次請求來承擔這一點上。

本人免費整理了Java高級資料，涵蓋了Java、Redis、MongoDB、MySQL、Zookeeper、Spring Cloud、Dubbo高併發分佈式等教程，一共30G，需要自己領取。
傳送門：https://mp.weixin.qq.com/s/JzddfH-7yNudmkjT0IRL8Q