令牌桶算法与Guava的实现RateLimiter源码分析

本文介绍: 令牌桶算法是一种限流算法。令牌桶算法的原理就是以一个恒定的速度往桶里放入令牌，每一个请求的处理都需要从桶里先获取一个令牌，当桶里没有令牌时，则请求不会被处理，要么排队等待，要么降级处理，要么直接拒绝服务。当桶里令牌满时，新添加的令牌会被丢弃或拒绝。令牌桶算法主要是可以控制请求的平均处理速率，它允许预消费，即可以提前消费令牌，以应对突发请求，但是后面的请求需要为预消费买单（等待更长的时间），以满足请求处理的平均速率是一定的。

令牌桶算法与Guava的实现RateLimiter源码分析

令牌桶RateLimiter

令牌桶RateLimiter

简介

令牌桶算法是一种限流算法。

令牌桶算法的原理就是以一个恒定的速度往桶里放入令牌，每一个请求的处理都需要从桶里先获取一个令牌，当桶里没有令牌时，则请求不会被处理，要么排队等待，要么降级处理，要么直接拒绝服务。当桶里令牌满时，新添加的令牌会被丢弃或拒绝。

令牌桶算法主要是可以控制请求的平均处理速率，它允许预消费，即可以提前消费令牌，以应对突发请求，但是后面的请求需要为预消费买单（等待更长的时间），以满足请求处理的平均速率是一定的。

RateLimiter使用示例

导入maven依赖

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

编写测试代码

public class RateLimiterTest {
    public void limit() {
        // 创建一个限流器，设置每秒放置的令牌数为1个
        RateLimiter rateLimiter = RateLimiter.create(1);
        IntStream.range(1, 10).forEach(i -> {
            // 一次获取i个令牌
            double waitTime = rateLimiter.acquire(i);
            System.out.println("acquire：" + i + " waitTime:" + waitTime);
        });
    }

    public static void main(String[] args) {
        RateLimiterTest rateLimiterTest = new RateLimiterTest();
        rateLimiterTest.limit();
    }
}

这段代码创建一个限流器，设置每秒放置的令牌数为1个，并循环获取令牌，每次获取i个。
执行结果：

第一次获取一个令牌时，等待0s立即可获取到（这里之所以不需要等待是因为令牌桶的预消费特性），第二次获取两个令牌，等待时间1s，这个1s就是前面获取一个令牌时因为预消费没有等待延到这次来等待的时间，这次获取两个又是预消费，所以下一次获取（取3个时）就要等待这次预消费需要的2s了，依此类推。可见预消费不需要等待的时间都由下一次来买单，以保障一定的平均处理速率（上例为1s一次）。

RateLimiter的实现

RateLimiter类在guava里是一个抽象类，其有两个具体实现：

SmoothBursty（平滑突发）：以恒定的速率生成令牌。
SmoothWarmingUp（顺利热身）：令牌生成的速度逐渐提升，直到达到一个稳定的值。

其类图如下：

他们的关系与作用：

RateLimiter是顶层封装，提供新建令牌桶的方法
SleepingStopwatch是guava实现的一个时钟类，提供时钟功能
SmoothRateLimiter是令牌桶抽象，提供操作令牌桶的抽象方法，其有两个内部类SmoothBursty和SmoothWarmingUp
SmoothBursty是恒定速率令牌桶实现
SmoothWarmingUp是逐渐加速知道稳定的令牌桶实现

源码解析

SmoothRateLimiter

SmoothRateLimiter是令牌桶抽象，其有四个关键的属性：

/** 当前存储的许可证。 */
double storedPermits;

/** 存储许可证的最大数量。 */
double maxPermits;

/**
 * 两个单位请求之间的间隔，以我们的稳定速率。
 * 例如，每秒 5 个许可的稳定速率具有 200 毫秒的稳定间隔。
 */
double stableIntervalMicros;

/**
 * 授予下一个请求（无论其大小如何）的时间。
 * 在批准请求后，这将在将来进一步推送。大请求比小请求更进一步。
 */
private long nextFreeTicketMicros = 0L; // could be either in the past or future

他们的作用可以看注释。此外还有两个内部类，实现此抽象：

SmoothBursty（平滑突发）：以恒定的速率生成令牌。
SmoothWarmingUp（顺利热身）：令牌生成的速度逐渐提升，直到达到一个稳定的值。

SmoothBursty恒速

首先来看下恒定速率生成令牌的实现。其使用方法是：

// 创建一个限流器，设置每秒放置的令牌数为1个
RateLimiter rateLimiter = RateLimiter.create(1);

RateLimiter#create：

public static RateLimiter create(double permitsPerSecond) {
  /*
   * The default RateLimiter configuration can save the unused permits of up to one second. This
   * is to avoid unnecessary stalls in situations like this: A RateLimiter of 1qps, and 4 threads,
   * all calling acquire() at these moments:
   *
   * T0 at 0 seconds
   * T1 at 1.05 seconds
   * T2 at 2 seconds
   * T3 at 3 seconds
   *
   * Due to the slight delay of T1, T2 would have to sleep till 2.05 seconds, and T3 would also
   * have to sleep till 3.05 seconds.
   */
  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实例，默认的maxBurstSeconds是1，其中SleepingStopwatch是guava实现的一个时钟类。
代码第三行调用了RateLimiter#setRate：

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

abstract void doSetRate(double permitsPerSecond, long nowMicros);

方法签名为：更新此 RateLimiter的稳定速率， permitsPerSecond 即构造 RateLimiter的工厂方法中提供的参数。当前受限制的线程不会因此调用而被唤醒，因此它们不会遵守新的速率，只有后续请求才会。
但请注意，由于每个请求都会偿还（如有必要，通过等待）前一个请求的成本，这意味着调用setRate后的下一个请求将不受新速率的影响，它将支付前一个请求的成本，这是根据以前的速率计算的。

此方法调用了抽象方法doSetRate，这里的实现是SmoothRateLimiter提供的，来看SmoothRateLimiter#doSetRate源码：

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

此方法：

调用 resync(nowMicros) 对 storedPermits 与 nextFreeTicketMicros 进行了调整——如果当前时间晚于 nextFreeTicketMicros，则计算这段时间内产生的令牌数，累加到 storedPermits 上，并更新下次可获取令牌时间 nextFreeTicketMicros 为当前时间。
计算stableIntervalMicros的值，此值代表产生令牌的时间间隔，1/permitsPerSecond(每秒几个令牌)
调用doSetRate(double, double)

其将输入的permitsPerSecond转换为速度，传给SmoothRateLimiter的doSetRate(permitsPerSecond, stableIntervalMicros)方法，doSetRate(permitsPerSecond, stableIntervalMicros)方法是抽象方法

abstract void doSetRate(double permitsPerSecond, double stableIntervalMicros);

由SmoothBursty的实现为：

@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 {
        storedPermits =
        (oldMaxPermits == 0.0)
        ? 0.0 // initial state
        : storedPermits * maxPermits / oldMaxPermits;
    }
}

此方法计算maxPermits的值maxBurstSeconds * permitsPerSecond。maxBurstSeconds是new SmoothBursty(SleepingStopwatch stopwatch, double maxBurstSeconds)构造方法传进来的，这里离默认是1。

获取令牌

acquire(int)

acquire(int)方法在获取不到令牌时阻塞等待，直到获取到令牌。
获取令牌方法为RateLimiter#acquire(int):

// 从中 RateLimiter获取给定数量的令牌，阻塞直到请求可以被批准。告诉睡眠时间（如果有）
@CanIgnoreReturnValue
public double acquire(int permits) {
  long microsToWait = reserve(permits);
  stopwatch.sleepMicrosUninterruptibly(microsToWait);
  return 1.0 * microsToWait / SECONDS.toMicros(1L);
}

调用RateLimiter#reserve 获取需要等待的时间：

// RateLimiter 保留给定数量的令牌以供将来使用，并返回使用这些令牌需要等待的微秒数。
final long reserve(int permits) {
  checkPermits(permits);
  synchronized (mutex()) {
    return reserveAndGetWaitLength(permits, stopwatch.readMicros());
  }
}

调用RateLimiter#reserveAndGetWaitLength：

//	保留令牌并返回使用者需要等待的时间。
final long reserveAndGetWaitLength(int permits, long nowMicros) {
    long momentAvailable = reserveEarliestAvailable(permits, nowMicros);
    return max(momentAvailable - nowMicros, 0);
}

调用RateLimiter#reserveEarliestAvailable：

abstract long reserveEarliestAvailable(int permits, long nowMicros);

是抽象方法，具体实现为SmoothRateLimiter#reserveEarliestAvailable：

// 更新下次可取令牌时间点与存储的令牌数，返回本次可取令牌的时间点
@Override
final long reserveEarliestAvailable(int requiredPermits, long nowMicros) {
    resync(nowMicros); // 如果nextFreeTicket小于当前时间，更新当前存储的令牌数，和下次可使用令牌的时间为now
    // nextFreeTicketMicros表示下一个可以分配令牌的时间点，这个值返回后，
    // 上一层的函数会调用stopwatch.sleepMicrosUninterruptibly(microsToWait);
    // 即阻塞到这个分配的时间点
    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;
}
// 如果nextFreeTicket小于当前时间，更新当前存储的令牌数，和下次可使用令牌的时间为now
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;
    }
}

上面是获取令牌的关键方法：

tryAcquire(int,long,TimeUnit)

指定时间内尝试获取令牌，获取到或获取超时返回。

public boolean tryAcquire(int permits, long timeout, TimeUnit unit) {
    // 将timeout换算成微秒
    long timeoutMicros = max(unit.toMicros(timeout), 0);
    checkPermits(permits);
    long microsToWait;
    synchronized (mutex()) {
        long nowMicros = stopwatch.readMicros();
        // 判断是否可以在timeoutMicros时间范围内获取令牌
        if (!canAcquire(nowMicros, timeoutMicros)) {
            return false;
        } else {
            // 获取令牌，并返回需要等待的毫秒数
            microsToWait = reserveAndGetWaitLength(permits, nowMicros);
        }
    }
    // 等待microsToWait时间
    stopwatch.sleepMicrosUninterruptibly(microsToWait);
    return true;
}

private boolean canAcquire(long nowMicros, long timeoutMicros) {
    return queryEarliestAvailable(nowMicros) - timeoutMicros <= nowMicros;
}

// 返回可用令牌的最早时间
abstract long queryEarliestAvailable(long nowMicros);

// SmoothRateLimiter#queryEarliestAvailable
final long queryEarliestAvailable(long nowMicros) {
	// 授予下一个请求（无论其大小如何）的时间
	return nextFreeTicketMicros;
}

该方法执行下面三步：

判断能否在指定超时时间内获取到令牌，通过 nextFreeTicketMicros <= nowMicros + timeoutMicros 是否为true来判断，即当前时间+超时时间在可取令牌时间之后，则可取（预消费的特性），否则不可获取。
如果不可获取，立即返回false。
如果可获取，则调用 reserveAndGetWaitLength(permits, nowMicros) 来更新下次可取令牌时间点与当前存储的令牌数，返回等待时间（逻辑与acquire(int)相同），并阻塞等待相应的时间，返回true。

存量桶系数

令牌桶算法中，多余的令牌会放到桶里。这个桶的容量是有上限的，决定这个容量的就是存量桶系数，默认为 1.0，即默认存量桶的容量是 1.0 倍的限流值。推荐设置 0.6~1.5 之间。
存量桶系数的影响有两方面：

突发流量第一个周期放过的请求数。如存量桶系数等于 0.6，第一个周期最多放过 1.6 倍限流值的请求数。
影响误杀率。存量桶系数越大，越能容忍流量不均衡问题。误杀率：服务限流是对单机进行限流，线上场景经常会用单机限流模拟集群限流。由于机器之间的秒级流量不够均衡，所以很容易出现误限。例如两台服务器，总限流值 20，每台限流 10，某一秒两台服务器的流量分别是 5、15，这时其中一台就限流了 5 个请求。减小误杀率的两个办法：
- 拉长限流周期。
- 使用令牌桶算法，并且调出较好的存量桶系数。

小结

RateLimiter 令牌桶的实现并不是起一个线程不断往桶里放令牌，而是以一种延迟计算的方式（参考resync函数），在每次获取令牌之前计算该段时间内可以产生多少令牌，将产生的令牌加入令牌桶中并更新数据来实现，比起一个线程来不断往桶里放令牌高效得多。（想想如果需要针对每个用户限制某个接口的访问，则针对每个用户都得创建一个RateLimiter，并起一个线程来控制令牌存放的话，如果在线用户数有几十上百万，起线程来控制是一件多么恐怖的事情）