基于 Redis 的分布式令牌桶、漏桶、滑动窗口实现
- 漏桶
- 实现
- 功能性测试
- 滑动窗口
- 实现
- 功能性测试
- 令牌桶
- 单机
- 分布式
- 功能性测试
- 参考
本文的限流工具都只有功能性测试（见正文），未进行过高并发和大流量下的性能测试，生产环境下的性能未知，仅供参考。完整源码详见 github。
漏桶
漏桶是最简单的限流工具，设计思路为：如果时间间隔达到规定的时间间隔，则允许通过，否则返回失败。
实现
LeakyLimiter 类中有四个属性，最核心的是 intervalNanos，表示时间间隔。如下所示：
private final RedisService redisService;
// 漏桶唯一标识
private final String name;
// 分布式互斥锁
private final RLock lock;
// 每两滴水之间的时间间隔
private final long intervalNanos;
redisService 用于操作缓存，lock 表示分布式锁。
上锁和解锁的方法如下所示：
/**
* 尝试获取锁
* @return 获取成功返回 true
*/
private boolean lock() {
try {
// 等待 100 秒，获得锁 100 秒后自动解锁
return this.lock.tryLock(100, 100, TimeUnit.SECONDS);
} catch (InterruptedException e) {
e.printStackTrace();
}
return false;
}
/**
* 释放锁
*/
private void unlock() {
this.lock.unlock();
}
在 acquire 方法中检查是否已达到时间间隔，如下所示：
/**
* 尝试通过漏桶
*
* @return 获取成功返回 true，失败返回 false
*/
@Override
public boolean acquire() {
while (true) {
if (lock()) {
try {
return tryAcquire();
} finally {
unlock();
}
}
}
}
private boolean tryAcquire() {
long recent = getRecent();
long now = System.nanoTime();
if (now - recent >= this.intervalNanos) {
resync(now);
return true;
} else {
log.info("Acquire LeakyLimiter[" + this.name + "] failed.");
return false;
}
}
getRecent 的作用是获取当前时间，resync 是同步缓存中的最新时间戳。
private long getRecent() {
Long recent = redisService.get(LeakyBucketKey.leakyBucket, this.name, Long.class);
if (recent == null) {
recent = System.nanoTime();
resync(recent);
return recent - intervalNanos;
}
return recent;
}
private void resync(long now) {
redisService.setwe(LeakyBucketKey.leakyBucket, this.name, now, LeakyBucketKey.leakyBucket.expireSeconds());
}
功能性测试
测试漏桶功能，多线程同时请求通过漏桶，只有一个线程能通过。代码如下所示：
@Test
public void getLeakyLimiter() {
LeakyLimiterFactory factory = new LeakyLimiterFactory();
LeakyLimiterConfig config = new LeakyLimiterConfig("testLeakyLimiter", 1, redissonService.getRLock("testLeakylock"), redisService);
final LeakyLimiter leakyLimiter = factory.getLeakyLimiter(config);
final int N = 3;
Runnable task = new Runnable() {
@Override
public void run() {
if (leakyLimiter.acquire()) {
System.out.println(Thread.currentThread().getName() + " passed.");
} else {
System.out.println(Thread.currentThread().getName() + " failed.");
}
}
};
Executor executor = Executors.newFixedThreadPool(N);
for (int i = 0; i < N; i++) {
executor.execute(task);
}
try {
Thread.sleep(2 * 1000);
} catch (Exception e) {
e.printStackTrace();
}
Executor executor2 = Executors.newFixedThreadPool(N);
for (int i = 0; i < N; i++) {
executor2.execute(task);
}
}
滑动窗口
计数器限流是统计一段时间间隔内的请求数，如果达到了阈值，则拒绝后面的请求。滑动窗口在此基础上将时间间隔进行细分，让请求更平滑地执行。
如果计数器的时间间隔为 1s，限制请求数为 1000，考虑如下情况：在前一秒的最后 100ms 通过请求数 1000，下一秒的前 100ms 通过请求数 1000，实际上在 200ms 的时间内通过了 2000 个请求，远远超过了限流器的原始设计。
滑动窗口把 1s 的时间段分成更小的部分，例如 10 份，当时间到达后一秒的前 100ms 时，滑动窗口的范围是前一秒的后 900ms 和后一面的前 100ms，这时候窗口范围内已经达到了限制请求数，不会允许此时的 1000 个请求通过。无论何时，窗口范围内都只允许最大 1000 个请求。
实现
滑动窗口使用链表实现，链表的每一个节点是 Node 类的实例。Node 表示一小段时间间隔，类中有三个属性，分别代表“起始时间”、“终止时间”、“时间段内计数”。如果滑动窗口已经完全经过该时间段，可以把该段删除。
Node 节点如下所示：
public class Node {
private long startTime;
private long endTime;
private long count;
// getter and setter
// ...
}
Window 是在缓存中传递的载体，包括以下属性：
// 唯一标识
private String name;
// 滑动窗口
private LinkedList<Node> slots;
// 时间间隔
private long intervalNanos;
// 窗口大小
private long windowSize;
// 流量限制
private long limit;
除了 getter 和 setter 之外，tryAcquire 方法也在此类中实现。tryAcquire 主要有三个步骤，首先删除已经无效的节点，然后统计在滑动窗口范围内的计数，如果已经达到计数限制，则返回请求失败，否则进入最后一步，更新链表状态。如果当前时间戳所在节点已经存在，把对应节点的计数加一即可，如果不存在，先创造一个包含当前时间戳的节点，然后再把计数加一。
// 尝试获取
public boolean tryAcquire(long tokens) {
long now = System.nanoTime();
// 删除已经过时的节点
long earliestWindowStartTime = now - intervalNanos * windowSize;
while (!slots.isEmpty() && slots.getFirst().getEndTime() < earliestWindowStartTime) {
slots.removeFirst();
}
long count = 0;
// 当前所有窗口的计数
for (Node node : slots) {
count += node.getCount();
}
// 如果达到计数限制，返回 false，表示获取失败
if (count + tokens > limit) {
return false;
}
// 允许获取，更新计数
// 如果当前时间点已经有了节点，在其所属节点（最后一个节点）上累加，否则先创建一个再累加。
Node lastNode = slots.isEmpty() ? null : slots.getLast();
long lastEndTime = (lastNode == null) ? now : lastNode.getEndTime();
if (now >= lastEndTime) {
long startTime = now - (now - lastEndTime) % intervalNanos;
long endTime = startTime + intervalNanos;
slots.add(new Node(startTime, endTime, tokens));
} else {
lastNode.addCount(tokens);
}
return true;
}
操作类 SlidingWindowLimiter 进行了进一步的封装，putDefaultWindow 在缓存中没有对应窗口的时候放入默认窗口，getWindow 获取缓存中的窗口，setWindow 用于更新窗口状态，acquire 用于获取请求许可。
/**
* 放入新的（默认）窗口
* 必须在 lock 内调用
* @return 返回 Window 实例
*/
public Window putDefaultWindow() {
if (!redisService.exists(WindowKey.window, this.name)) {
Window window = new Window(name, new LinkedList<>(), intervalNanos, windowSize, limit);
// 存入缓存，设置有效时间
redisService.setwe(WindowKey.window, this.name, window, WindowKey.window.expireSeconds());
}
return redisService.get(WindowKey.window, this.name, Window.class);
}
/**
* 从缓存获取窗口
* @return 从缓存获取到的窗口
*/
private Window getWindow() {
return redisService.get(WindowKey.window, this.name, Window.class);
}
/**
* 更新
* @param window 新的窗口
*/
private void setWindow(Window window) {
redisService.setwe(WindowKey.window, this.name, window, WindowKey.window.expireSeconds());
}
public boolean acquire(long tokens) {
while (true) {
if (lock()) {
Window window = getWindow();
if (window == null) {
window = putDefaultWindow();
}
boolean success = window.tryAcquire(tokens);
try {
setWindow(window);
return success;
} finally {
unlock();
}
}
}
}
功能性测试
代码如下所示：
@RunWith(SpringRunner.class)
@SpringBootTest(classes = MainApplication.class)
public class SlidingWindowLimiterFactoryTest {
@Autowired
RedisService redisService;
@Autowired
RedissonService redissonService;
private static Map<Boolean, String> map = new ConcurrentHashMap<>();
static {
map.putIfAbsent(true, "passed");
map.putIfAbsent(false, "failed");
}
@Test
public void getSlidingWindowLimiter() {
SlidingWindowLimiterFactory factory = new SlidingWindowLimiterFactory();
SlidingWindowLimiterConfig config = new SlidingWindowLimiterConfig("testSlidingWindowLimiter", 2, 10, redissonService.getRLock("testWindowLock"), redisService);
SlidingWindowLimiter limiter = factory.getSlidingWindowLimiter(config);
System.out.println("Main thread " + map.get(limiter.acquire()) + " at first time."); // passed
try {
System.out.println("After sleep 500 millis--------");
Thread.sleep(500);
} catch (Exception e) {
e.printStackTrace();
}
System.out.println("Main thread " + map.get(limiter.acquire()) + " at second time."); // passed
System.out.println("Main thread " + map.get(limiter.acquire()) + " at third time."); // failed
try {
System.out.println("After sleep 600 millis--------");
Thread.sleep(600);
} catch (Exception e) {
e.printStackTrace();
}
System.out.println("Main thread " + map.get(limiter.acquire()) + " at forth time."); // passed
try {
System.out.println("After sleep 500 millis--------");
Thread.sleep(500);
} catch (Exception e) {
e.printStackTrace();
}
System.out.println("Main thread " + map.get(limiter.acquire()) + " at fifth time."); // passed
System.out.println("Main thread " + map.get(limiter.acquire()) + " at sixth time."); // failed
System.out.println("Main thread " + map.get(limiter.acquire()) + " at seventh time."); // failed
try {
System.out.println("After sleep 1100 millis--------");
Thread.sleep(1100);
} catch (Exception e) {
e.printStackTrace();
}
System.out.println("Main thread " + map.get(limiter.acquire()) + " at eighth time."); // passed
System.out.println("Main thread " + map.get(limiter.acquire()) + " at ninth time."); //passed
}
}
令牌桶
单机
Google 开源工具包 Guava 提供了限流工具类 RateLimiter，该类是令牌桶算法的一个具体实现。
RateLimiter有两种限流方式，分别对应两个类。一种是令牌生成速度恒定的方式，对应 SmoothBursty 类，一种是令牌初始速度缓慢，慢慢提升最后维持在一个稳定值的方式，对应 SmoothWarmingUp 类。之后的源码分析以 SmoothBursty 为例。
在 RateLimiter 的 create 函数中 “实例化” 了RateLimiter 类（这个说法不正确，因为 RateLimiter 是抽象类，不能实例化），实际上是直接实例化 SmoothBursty 类，而 SmoothBursty 继承自 SmoothRateLimiter 类，SmoothRateLimiter 类继承自 RateLimiter 类。
RateLimiter rateLimiter = new SmoothBursty(stopwatch, 1.0D);
rateLimiter.setRate(permitsPerSecond);
return rateLimiter;
此处的 1.0D 表示 maxBurstSeconds 设置为 1.0，即只能创建桶大小为 permitsPerSecond*1 的 SmoothBursty 对象。
SmoothBursty 中主要有以下几个字段：
// 桶中当前存放的令牌个数。
double storedPermits;
// 桶中最多存放多少个令牌。
double maxPermits;
// 加入令牌的平均时间。即每两个令牌之间的时间间隔。
// Micros 表示微秒，Millis 表示毫秒。
double stableIntervalMicros;
// 下一次请求可以获取令牌的起始时间。
private long nextFreeTicketMicros;
// 桶中最多存放多少秒的令牌数。
final double maxBurstSeconds;
在成员函数中最重要的是 resync 函数，用于同步，其作用主要是更新桶中令牌数和下次可获取令牌时间。
private void resync(long nowMicros) {
// 如果当前时间比下一次可获取令牌的时间还要晚，说明上次请求的令牌已经完全结清了，本次请求可以不用等待
// 且在上一次请求完成到当前时间内，桶中还会匀速放入令牌，进入 if 块，计算出令牌数
if (nowMicros > this.nextFreeTicketMicros) {
// 计算上一次请求完成到当前时间内，已经补充的令牌数
// 然后将补充的令牌数加上原来的令牌数，即为当前桶中的令牌数
this.storedPermits = Math.min(this.maxPermits, this.storedPermits + (double)(nowMicros - this.nextFreeTicketMicros) / this.stableIntervalMicros);
// 时间更新为当前时间
this.nextFreeTicketMicros = nowMicros;
}
}
需要获取令牌时调用 acquire 函数即可。
acquire 函数主要依赖以下几个函数实现
public double acquire(int permits) {
// 计算本次请求需要休眠多久才能拿到令牌，时间单位是微秒
long microsToWait = this.reserve(permits);
// 开始休眠
this.stopwatch.sleepMicrosUninterruptibly(microsToWait);
// 返回需要等待的时间
return 1.0D * (double)microsToWait / (double)TimeUnit.SECONDS.toMicros(1L);
}
final long reserve(int permits) {
checkPermits(permits);
// 上锁
synchronized(this.mutex()) {
return this.reserveAndGetWaitLength(permits, this.stopwatch.readMicros());
}
}
final long reserveAndGetWaitLength(int permits, long nowMicros) {
long momentAvailable = this.reserveEarliestAvailable(permits, nowMicros);
return Math.max(momentAvailable - nowMicros, 0L);
}
final long reserveEarliestAvailable(int requiredPermits, long nowMicros) {
// 在处理本次请求之前，先调用上面提到的 resync 方法，同步桶中的令牌数和下次可获取令牌时间，将桶内数据同步到最新状态
this.resync(nowMicros);
// 如果上次请求还没补齐，returnValue 为下次可获取时间，否则为当前时间
long returnValue = this.nextFreeTicketMicros;
double storedPermitsToSpend = Math.min((double)requiredPermits, this.storedPermits);
// 缺少的令牌数
double freshPermits = (double)requiredPermits - storedPermitsToSpend;
// storedPermitsToWaitTime 函数返回值恒为 0
// waitMicros 表示需要等待多长时间
long waitMicros = this.storedPermitsToWaitTime(this.storedPermits, storedPermitsToSpend) + (long)(freshPermits * this.stableIntervalMicros);
// 更新下次可请求时间
this.nextFreeTicketMicros += waitMicros;
// 减少桶中的现存令牌数
this.storedPermits -= storedPermitsToSpend;
return returnValue;
}
可以看到 acquire 函数主要对 nextFreeTicketMicros 和 storedPermits 这两个属性修改，一次请求可以预先支付令牌，那么下一次请求则需要相应的等待。
tryAcquire 相关方法在 acquire 基础上加上了超时限制，思路与 acquire 类似。
分布式
在实现分布式 RateLimiter 之前，需要确认 Redis 环境已配置完成。令牌的添加存储获取等均需要用到 Redis，除此之外，还会用到 Redis 实现的分布式互斥锁。
基于 redis 的分布式互斥锁用于跨 JVM 实现互斥以达到控制资源访问的目的，在 Redisson 框架已经实现，用户只需要将 Redisson 工具添加到项目依赖即可。
将单机令牌桶改造成分布式令牌桶的要点包括以下两点：
- 令牌桶属性（对象）保存在 redis 中，对令牌桶属性的读取和更改都在 redis 中进行。
- 使用分布式互斥锁保障数据读写的安全
PermitBucket
定义 PermitBucket 类作为存储令牌的桶，之后将会作为此令牌桶保存在缓存中的载体。
类中包含以下属性：
/**
* 唯一标识
*/
private String name;
/**
* 最大存储令牌数
*/
private long maxPermits;
/**
* 当前存储令牌数
*/
private long storedPermits;
/**
* 每两次添加令牌之间的时间间隔（逐个添加令牌），单位为纳秒
*/
private long intervalNanos;
/**
* 上次更新的时间
*/
private long lastUpdateTime;
除了 getter 和 setter 之外，还有一个 reSync 函数，它的功能是同步令牌桶的状态，也就是根据当前时间和上一次时间戳的间隔，更新令牌桶中当前令牌数。如下所示：
/**
* 更新当前持有的令牌数
* 若当前时间晚于 lastUpdateTime，则计算该段时间内可以生成多少令牌，将生成的令牌加入令牌桶中并更新数据
*
* @param now 当前时间
*/
public void reSync(long now, long storedPermitsToSpend) {
if (now > lastUpdateTime) {
long newStoredPermits = Math.min(maxPermits, storedPermits + (now - lastUpdateTime) / intervalNanos - storedPermitsToSpend);
// now 距离 lastUpdateTime 很短时，防止 lastUpdateTime 变了而 storedPermits 没变
if (newStoredPermits != storedPermits) {
storedPermits = newStoredPermits;
lastUpdateTime = now;
}
}
}
RateLimiter
PermitBucket 作为实体对象用于保存令牌桶状态，而获取令牌等操作在服务类 PermitLimiter 中实现。
此类中的 RedisService 属性用于缓存操作，RLock 属性用于分布式锁，其他属性如下所示：
/**
* 唯一标识
*/
private String name;
/**
* 最大存储令牌数
*/
private long maxPermits;
/**
* 当前存储令牌数
*/
private long storedPermits;
/**
* 每两次添加令牌之间的时间间隔（逐个添加令牌），单位为纳秒
*/
private long intervalNanos;
/**
* 上次更新的时间
*/
private long lastUpdateTime;
类中所有的方法都只有简单的计算，耗时很短，且必须串行执行，使用分布式锁保证在整个分布式系统中方法串行化执行：
/**
* 尝试获取锁
* @return 获取成功返回 true
*/
private boolean lock() {
try {
// 等待 100 秒，获得锁 100 秒后自动解锁
return lock.tryLock(100, 100, TimeUnit.SECONDS);
} catch (InterruptedException e) {
e.printStackTrace();
}
return false;
}
/**
* 释放锁
*/
private void unlock() {
lock.unlock();
}
acquire 方法是此类的核心方法，用于获取令牌，成功返回 true，失败返回 false。
/**
* 尝试获取 permits 个令牌
*
* @return 获取成功返回 true，失败返回 false
*/
public boolean acquire(long permits) {
checkPermits(permits);
while (true) {
if (lock()) {
long wait;
try {
wait = canAcquire(permits);
if (wait <= 0L) {
return doAcquire(permits);
}
else {
return false;
}
} finally {
unlock();
}
}
}
}
acquire 相当于令牌桶的入口，从 acquire 可以看出，所有的操作都使用分布式锁保护。
canAcquire 用于当前是否有令牌可以获取，如果有，调用 doAcquire 执行后续操作，如果没有直接返回 false，获取失败。
在 canAcquire 中先从缓存中获取令牌桶实体，调用令牌桶的 reSync 方法更新其状态，更新之后把它保存到缓存里。函数返回值为需要等待的时间。如果可以马上获取，返回 0。
/**
* 当前是否可以获取到令牌，如果获取不到，至少需要等多久
* @param permits 请求的令牌数
* @return 等待时间，单位是纳秒。为 0 表示可以马上获取
*/
private long canAcquire(long permits){
PermitBucket bucket = getBucket();
long now = System.nanoTime();
bucket.reSync(now, 0L);
setBucket(bucket);
if (permits <= bucket.getStoredPermits()) {
return 0L;
}
else {
return (permits - bucket.getStoredPermits()) * bucket.getIntervalNanos();
}
}
doAcquire 函数在令牌桶中减去响应的令牌数，并再次更新令牌桶状态。
/**
* 确认可以获取，就获取 permits 个令牌，更新缓存
* @param permits 请求 token 个令牌
* @return 需要等待的时间
*/
private boolean doAcquire(long permits) {
PermitBucket bucket = getBucket();
if (permits > bucket.getStoredPermits())
return false;
// 当前时间
long now = System.nanoTime();
if (now > bucket.getLastUpdateTime()) {
// 可以消耗的令牌数/需要消耗的令牌数
long storedPermitsToSpend = Math.min(permits, bucket.getStoredPermits());
// 更新一下
bucket.reSync(now, storedPermitsToSpend);
// 缓存中更新桶的状态
setBucket(bucket);
return true;
}
return false;
}
acquireTillSuccess 不断尝试获取令牌直到成功，线程可能会多次进入休眠状态，相当于阻塞了整个线程，所以不推荐使用此方法。
/**
* 获取成功或超时才返回
* @param permits 获取的令牌数
* @param timeout 超时时间，单位为秒
*/
public boolean acquireTillSuccess(long permits, long timeout) {
checkPermits(permits);
long start = System.nanoTime();
long timeoutNanos = TimeUnit.SECONDS.toNanos(timeout);
while (true) {
long wait = 0L;
if (lock()) {
try {
wait = canAcquire(permits);
if (wait <= 0L && doAcquire(permits)) {
return true;
}
} finally {
unlock();
}
}
try {
Thread.sleep(TimeUnit.NANOSECONDS.toMillis(wait));
} catch (Exception e) {
log.info(e.toString());
}
if (System.nanoTime() - start > timeoutNanos)
return false;
}
}
除了上面提到的 acquire 系列方法外，此类还提供了手动添加令牌功能，用于支持瞬时流量：
/**
* 添加指定数量令牌
* @param permits 要添加的令牌数
*/
public void addPermits(long permits) {
checkPermits(permits);
while (true) {
if (lock()) {
try {
PermitBucket bucket = getBucket();
long now = System.nanoTime();
bucket.reSync(now, 0L);
long newPermits = calculateAddPermits(bucket, permits);
bucket.setStoredPermits(newPermits);
setBucket(bucket);
return;
} finally {
unlock();
}
}
}
}
/**
* 计算添加之后桶里的令牌数
* @param bucket 桶
* @param addPermits 添加的令牌数
* @return
*/
private long calculateAddPermits(PermitBucket bucket, long addPermits) {
long newPermits = bucket.getStoredPermits() + addPermits;
if (newPermits > bucket.getMaxPermits()) {
newPermits = bucket.getMaxPermits();
}
return newPermits;
}
功能性测试
代码如下所示：
@Test
public void getPermitLimiter() {
PermitLimiterFactory factory = new PermitLimiterFactory();
PermitLimiterConfig config = new PermitLimiterConfig("testPermitLimiter", 1, 1000, redissonService.getRLock("testPermitLock"), redisService);
PermitLimiter permitLimiter = factory.getPermitLimiter(config);
if (permitLimiter.acquire()) {
System.out.println("Main thread passed at first time.");
} else {
System.out.println("Main thread failed at first time.");
}
if (permitLimiter.acquire()) {
System.out.println("Main thread passed at second time.");
} else {
System.out.println("Main thread failed at second time.");
}
System.out.println("Before first added: " + permitLimiter.getBucket().getStoredPermits());
permitLimiter.addPermits(100);
System.out.println("After added 100 permits: " + permitLimiter.getBucket().getStoredPermits());
if (permitLimiter.acquire()) {
System.out.println("Main thread passed at third time.");
} else {
System.out.println("Main thread failed at third time.");
}
if (permitLimiter.acquire()) {
System.out.println("Main thread passed at forth time.");
} else {
System.out.println("Main thread failed at forth time.");
}
System.out.println("Before second added: " + permitLimiter.getBucket().getStoredPermits());
permitLimiter.addPermits(500);
System.out.println("After added 500 permits: " + permitLimiter.getBucket().getStoredPermits());
}
每一秒钟产生 1 个令牌，最大令牌数限制为 1000。添加令牌之前只有第一次能成功获取令牌，添加之后，每一次都能成功获取，且剩余令牌数符合预期。
参考
- 基于Redis的分布式令牌桶限流器
- Guava RateLimiter源码解析