带你手撸令牌桶限流算法

带你手撸令牌桶限流算法

1.漏桶算法和令牌桶算法

漏桶算法和令牌桶算法都属于限流的基本算法，但是各自有各自的特点

1. 漏桶算法

   

   上图是从网上拷的一张算法示意图，其中，桶的体积表示能够处理请求的最大值，水龙头的水表示外部请求，漏下去的水表示处理的请求。

   也就是说，不管水龙头的水怎么往下流，开最大也好，关掉也好，都不会影响到桶往下滴水的速度，这也是漏桶算法最核心的一点，能够保证流量的平滑性（请求处理的速度基本一致），如果外部请求超过了桶的容积，那么水就流到桶外去了（这些请求也就都抛弃掉了）。因此，漏桶算法并不能满足对于突发流量高峰的场景，所以漏桶算法的具体实现不再阐述。

2. 令牌桶算法

   
   同样的，从网上拷的一张令牌桶的示意图。

   令牌桶相对于漏桶算法来说，首先多了一个缓存区，使其能够处理突发的流量高峰（就是削峰填谷的作用），并且令牌的生成是每时每刻都在生成的，而不像漏桶算法，如果一段时间一直没有请求，那么桶就一直是空的，这样就能够保证，在一段时间内，请求的处理的平均速度等于令牌生成的速度。

2.开始手撸令牌桶

关于令牌桶的算法，有很多很完善的实现，比如用的最多的Guava库中的RateLimiter就用到了令牌桶算法。但是，如果只是看源码，可能短时间内就忘得干干净净了，其实令牌桶的概念很简单，不如我们直接根据概念手撸一下，更能加深记忆。

通过上一步的令牌桶算法示意图，基本上可以先规划出所用到的数据结构。

首先，我们要规定一个速率，也就是要限制的qps，

然后，产生令牌也不能无限产生，要有一个可容纳令牌的最大值maxPermits，自然也要知道，目前有多少令牌可用currentPermits.

这样一来，数据结构就定义完成了，如下

// 令牌生成的速度
int limitQps;
// 最大令牌数
int maxPermits;
// 当前已有令牌数
int currentPermits;

对于，如何实现持续增加令牌数，有两种方法可以做，一个是起新的线程，持续增加，一个是惰性增加，在用到令牌的时候才去增加。如果起新的线程去增加，会涉及到对令牌资源增减的线程安全问题，所以先使用简单一点的方式实现，惰性增加。既然要惰性增加，那么就要知道本次请求距离上次请求经过了多长时间，因为知道经过了多长时间，才能知道要加多少个令牌，所以，要在数据结构里增加一个字段，用于记录上次请求的时间。那么，现在数据结构就变为了如下

// 令牌生成的速度
int limitQps;
// 最大令牌数
int maxPermits;
// 当前已有令牌数
int currentPermits;
// 上次请求时间
long lastRequestTimeStamp;

但是，之前有讲过，令牌桶算法可以应对突发流量的状况，也就是说就算所需令牌大于当前已有令牌，请求仍能够执行，那么，记录上次请求的时间显然是不准确的，因为，上个请求已经预支了以后请求的令牌了，所以，这里可以更改一下，变为记录下次请求能够执行的时间，这样，就算请求预支了令牌，以后的请求仍能够感知到。

// 令牌生成的速度
int limitQps;
// 最大令牌数
int maxPermits;
// 当前已有令牌数
int currentPermits;
// 下次请求可以执行的时间
long nextRequestExcuteTimeStamp;

在具有了以上数据结构后，就可以给出惰性增加令牌方法的实现

// 惰性增加令牌数
void increasePermits(long nowTimeStamp) {
  // 当前时间晚于下次请求可以执行的时间，也就意味着会有多余的令牌生成
 	if (nowTimeStamp > nextRequestExcuteTimeStamp) {
    	// 算一下晚了多少毫秒
      long interval = nowTimeStamp - nextRequestExcuteTimeStamp;
    	// 用时间间隔 * qps 得出这段间隔能生成多少令牌
    	int increasePermits = (limitQps * interval) / 1000;
      if ((increasePermits + currentPermits) <= maxPermits)
         currentPermits += increasePermits;
    		 // 因为令牌数已经刷新了，所以时间也要改一下
    		 nextRequestExcuteTimeStamp = nowTimeStamp;
      }
 	}
}

现在，增加令牌实现了，那我们就可以实现消耗令牌的方法了，也就是实际请求时的方法

/** 
 * @Description: 请求消耗令牌方法
 * @param permits	消耗令牌数 
 * @Author: lvqiushi 
 * @Date: 2020-03-17
 */
public void acquire(int permits) {
  long nowTimeStamp = System.currentTimeMillis();
  // 计算获取令牌要等多长时间
  long toWaitTime = getSpendTime(permits, nowTimeStamp);
  if (toWaitTime > 0) {
    Thread.sleep(toWaitTime);
  }
}

/** 
 * @Description: 要获取令牌所等待的时间
 * @param permits	消耗令牌数 
 * @Author: lvqiushi 
 * @Date: 2020-03-17
 */
public long getSpendTime(int permits, long nowTime) {
  long retrunTime = nextRequestExcuteTimeStamp;
   // 刷新当前令牌数
  increasePermits(nowTimeStamp);
  // 判断当前令牌够不够用，如果不够，需要多长时间
  long toWaitTime = 0L;
  if (permits > currentPermits) {
     // 计算额外获取的令牌要多次时间生成
     int toWaitPermits = permits - currentPermits;
     toWaitTime = toWaitPermits * (1 / limitQps);
     // 当前令牌数自然就消耗完了
     currentPermits = 0;
  } else {
     // 如果当前够用，就减掉令牌数
     currentPermits -= permits;
  }
  
  nextRequestExcuteTimeStamp += timeExpand;
  // 就算第一次有额外令牌消耗也不等待
  return max(retrunTime - nowTime, 0);
}

在getSpendTime方法中，涉及到了对临界资源的修改，所以要使用锁控制一下，也有两种方法，一种是整个方法加关键字synchronized锁一下，另一种是创建一个锁对象，用锁对象当锁synchronized一下。当然这两种方法都是针对单机版的，不过对于一般网关应用来说单机也够用了。

从以上代码，可以看出，令牌桶的概念其实真没有多少，以上代码再抽象优化一下，其实跟RateLimiter的SmoothBursty的限流实现就基本一致了。对于这些简单的算法，自己实现一下，可能就会记很久很久。