RateLimiter源码解读

旧城等待， 2022-12-29 09:47 83阅读 0赞

最近使用了Guava提供的限流器RateLimiter觉得十分好用，遍想一探究竟。

RateLimiter使用的限流算法为令牌桶，这里对令牌桶算法不做过多的解释，不懂的同学可以搜索一番。

#### 限流器创建 ####

public static RateLimiter create(double permitsPerSecond) { 
        return create(permitsPerSecond, SleepingStopwatch.createFromSystemTimer());
    }
    
    
    static RateLimiter create(double permitsPerSecond, SleepingStopwatch stopwatch) { 
        
        // 1.new SmoothBursty对象
        RateLimiter rateLimiter = new SmoothBursty(stopwatch, 1.0);
        // 2.设置 限流值
        rateLimiter.setRate(permitsPerSecond);
        return rateLimiter;
    }

SmoothBursty是RateLimiter的子类，它有一个参数名为maxBurstSeconds。这个参数的作用为，当限流器没有被使用时最多可以存放多少秒的令牌桶。默认的参数为1，即只可以保存1s产生的令牌。 下面来具体看一下限流值是如何被设置的。

public final void setRate(double permitsPerSecond) { 
        
        // 校验参数必须大于0
        checkArgument(
            permitsPerSecond > 0.0 && !Double.isNaN(permitsPerSecond), "rate must be positive");
        
        // 加锁设置限流值
        synchronized (mutex()) { 
          doSetRate(permitsPerSecond, stopwatch.readMicros());
        }
    }
    
    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为0）
        if (nowMicros > nextFreeTicketMicros) { 
            
          // 计算令牌数 (当前时间 - 下一次获取令牌的时间) / 冷启动时间
          double newPermits = (nowMicros - nextFreeTicketMicros) / coolDownIntervalMicros();
          
          // 计算可存储的令牌数 取最大的令牌数与存储的令牌数和新生成的令牌数之间较小的值
          storedPermits = min(maxPermits, storedPermits + newPermits);
            
          // 将下一次获取令牌数的时间设置为当前时间
          nextFreeTicketMicros = nowMicros;
        }
      }
    
    
    void doSetRate(double permitsPerSecond, double stableIntervalMicros) { 
          double oldMaxPermits = this.maxPermits;
        
          // 最大令牌数 = 最多可以储存的秒数 * 每秒产生的令牌数
          maxPermits = maxBurstSeconds * permitsPerSecond;
          if (oldMaxPermits == Double.POSITIVE_INFINITY) { 
            storedPermits = maxPermits;
          } else { 
            storedPermits =
                (oldMaxPermits == 0.0)
                    ? 0.0 // initial state
                    : storedPermits * maxPermits / oldMaxPermits;
          }
        }

在设置限流值的代码中，比较重要的方法是resync，该方法重新设置了下一次获取令牌数的时间nextFreeTicketMicros和可以获取的令牌数。在该方法内有一个coolDownIntervalMicros()方法，该方法是一个抽象方法具体的逻辑有子类实现SmoothBursty中的实现使用的是使用固定的时间间隔，即1s/限流值。参看和另一个类SmoothWarmingUp实现之差别可以得知SmoothWarmingUp提供的是带慢启动的限流器。

#### 限流器的使用 ####

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(nowMicros);
        long returnValue = nextFreeTicketMicros;
        double storedPermitsToSpend = min(requiredPermits, this.storedPermits);
        
        // 计算是否还需要新产生令牌
        double freshPermits = requiredPermits - storedPermitsToSpend;
        
        // 计算等待时间 如果 requiredPermits > storedPermits 则freshPermits 不为0 则需要等待
        // 如果 requiredPermits < storedPermits 则freshPermits 为0 则不需要等待
        long waitMicros =
            storedPermitsToWaitTime(this.storedPermits, storedPermitsToSpend)
                + (long) (freshPermits * stableIntervalMicros);
    
        // 计算下一次的获取令牌的时间
        this.nextFreeTicketMicros = LongMath.saturatedAdd(nextFreeTicketMicros, waitMicros);
        this.storedPermits -= storedPermitsToSpend;
        return returnValue;
    }

获取令牌的具体逻辑已经写到注释里面了，具体可以参看相关注释。这里说一下获取令牌的具体算法，短短几行代码就可以实现核心逻辑，真是代码的艺术。

1.  当当前时间 > 下一次获取令牌的时间时 则产生新的令牌 否则则不产生新的令牌
2.  需要的令牌数和存储的令牌数两者之间取小值和需要的令牌数做差，如果差值 > 0则说明当前存储的令牌数不足，需要等待时间产生新的令牌。如若 = 0 则说明存储的令牌数大于需要的令牌数，则不需要等待。