在应用中有很多场景需要限制流量防止突然的流量峰值影响服务器、数据库的性能。在谷歌提供的Guava包中就有这么一个工具用于对流量进行限制。

令牌桶算法

令牌桶算法是往一个固定大小的令牌桶中定时发放令牌,请求只有拿到令牌之后才能继续访问,当没有拿到令牌则阻塞或直接返回。

使用

guava中的RateLimiter实现了令牌桶算法,使用也很简单

    public void test() throws InterruptedException {
		//创建limiter
        RateLimiter limiter=RateLimiter.create(2);
		//使用limiter
        System.out.println(limiter.acquire());        
        System.out.println(limiter.acquire());
        System.out.println(limiter.acquire());
        Thread.sleep(1000);
        System.out.println(limiter.acquire());
        System.out.println(limiter.acquire());
    }

limiter有两个方法acquire和tryAcquire,tryAcquire是直接返回是否获得了令牌,返回类型为bool,而acquire是阻塞的形式获得令牌,返回的结果是等待的时间(单位秒)。所以上面的例子,创建了一个每秒两个令牌的limiter,第一次acquire会立即获得,之后的两次会需要等待半秒钟,在sleep1秒之后,令牌桶会满,此时就可以连续两次无需等待获得令牌

0.0
0.498939
0.472683
0.0
0.0

实现

guava并没有和原理一样使用另一个线程对令牌桶中加令牌,而是记录了令牌的状态,在每次调用acquire的时候,刷新令牌数,并计算是否可以获得令牌以及需要等待的时间,再强制sleep来实现流量的限制

  final long reserveEarliestAvailable(int requiredPermits, long nowMicros) {
	//根据当前时间刷新令牌状态
    resync(nowMicros);
	//返回下一个令牌的时间
    long returnValue = nextFreeTicketMicros;
	//当前已有的permit
    double storedPermitsToSpend = min(requiredPermits, this.storedPermits);
	//需要等待的permit
    double freshPermits = requiredPermits - storedPermitsToSpend;
	//算出需要等待的时间
    long waitMicros =
        storedPermitsToWaitTime(this.storedPermits, storedPermitsToSpend)
            + (long) (freshPermits * stableIntervalMicros);

    //保存令牌状态
	this.nextFreeTicketMicros = LongMath.saturatedAdd(nextFreeTicketMicros, waitMicros);
    this.storedPermits -= storedPermitsToSpend;
	//返回需要下一个令牌的时间
    return returnValue;
  }

在外层则会根据返回值进行sleep。
对于tryAcquire则是判断是否可以获得令牌,如果有令牌,则会获取该令牌,此时如果一次需要获得多张令牌,但无法获得全部令牌的情况下会将计算出下一个令牌的时间,但是当前请求会被放行

  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();
	  //如果不能在timeout时间内获得一张令牌则返回false
      if (!canAcquire(nowMicros, timeoutMicros)) {
        return false;
      } else {
	  //不然则获得全部令牌
        microsToWait = reserveAndGetWaitLength(permits, nowMicros);
      }
    }
	//此处等待时间应该是0
    stopwatch.sleepMicrosUninterruptibly(microsToWait);
    return true;
  }

由此看出,如果需要几秒发放一个请求,可以通过acquire、tryAcquire可以通过获取多个permit(令牌)来实现。