注册中心 Eureka 源码解析 —— 基于令牌桶算法的 RateLimiter

public class RateLimiter {

    /**
     * 速率单位转换成毫秒
     */
    private final long rateToMsConversion;

    public RateLimiter(TimeUnit averageRateUnit) {
        switch (averageRateUnit) {
            case SECONDS: // 秒级
                rateToMsConversion = 1000;
                break;
            case MINUTES: // 分钟级
                rateToMsConversion = 60 * 1000;
                break;
            default:
                throw new IllegalArgumentException("TimeUnit of " + averageRateUnit + " is not supported");
        }
    }
}

// RateLimiter.java
/**
* 获取令牌( Token )
*
* @param burstSize 令牌桶上限
* @param averageRate 令牌再装平均速率
* @return 是否获取成功
*/
public boolean acquire(int burstSize, long averageRate) {
   return acquire(burstSize, averageRate, System.currentTimeMillis());
}

public boolean acquire(int burstSize, long averageRate, long currentTimeMillis) {
   if (burstSize <= 0 || averageRate <= 0) { // Instead of throwing exception, we just let all the traffic go
       return true;
   }

   // 填充 令牌
   refillToken(burstSize, averageRate, currentTimeMillis);
   // 消费 令牌
   return consumeToken(burstSize);
}

  1: /**
  2:  * 速率单位转换成毫秒
  3:  */
  4: private final long rateToMsConversion;
  5: 
  6: /**
  7:  * 消耗令牌数
  8:  */
  9: private final AtomicInteger consumedTokens = new AtomicInteger();
 10: /**
 11:  * 最后填充令牌的时间
 12:  */
 13: private final AtomicLong lastRefillTime = new AtomicLong(0);
 14: 
 15: private void refillToken(int burstSize, long averageRate, long currentTimeMillis) {
 16:     // 获得 最后填充令牌的时间
 17:     long refillTime = lastRefillTime.get();
 18:     // 获得 过去多少毫秒
 19:     long timeDelta = currentTimeMillis - refillTime;
 20: 
 21:     // 计算 可填充最大令牌数量
 22:     long newTokens = timeDelta * averageRate / rateToMsConversion;
 23:     if (newTokens > 0) {
 24:         // 计算 新的填充令牌的时间
 25:         long newRefillTime = refillTime == 0
 26:                 ? currentTimeMillis
 27:                 : refillTime + newTokens * rateToMsConversion / averageRate;
 28:         // CAS 保证有且仅有一个线程进入填充
 29:         if (lastRefillTime.compareAndSet(refillTime, newRefillTime)) {
 30:             while (true) { // 死循环，直到成功
 31:                 // 计算 填充令牌后的已消耗令牌数量
 32:                 int currentLevel = consumedTokens.get();
 33:                 int adjustedLevel = Math.min(currentLevel, burstSize); // In case burstSize decreased
 34:                 int newLevel = (int) Math.max(0, adjustedLevel - newTokens);
 35:                 // CAS 避免和正在消费令牌的线程冲突
 36:                 if (consumedTokens.compareAndSet(currentLevel, newLevel)) {
 37:                     return;
 38:                 }
 39:             }
 40:         }
 41:     }
 42: }

  1: private boolean consumeToken(int burstSize) {
  2:     while (true) { // 死循环，直到没有令牌，或者获取令牌成功
  3:         // 没有令牌
  4:         int currentLevel = consumedTokens.get();
  5:         if (currentLevel >= burstSize) {
  6:             return false;
  7:         }
  8:         // CAS 避免和正在消费令牌或者填充令牌的线程冲突
  9:         if (consumedTokens.compareAndSet(currentLevel, currentLevel + 1)) {
 10:             return true;
 11:         }
 12:     }
 13: }

  1: @Override
  2: public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
  3:     // 获得 Target
  4:     Target target = getTarget(request);
  5: 
  6:     // Other Target ，不做限流
  7:     if (target == Target.Other) {
  8:         chain.doFilter(request, response);
  9:         return;
 10:     }
 11: 
 12:     HttpServletRequest httpRequest = (HttpServletRequest) request;
 13:     // 判断是否被限流
 14:     if (isRateLimited(httpRequest, target)) {
 15:         // TODO[0012]：监控相关，跳过
 16:         incrementStats(target);
 17:         // 如果开启限流，返回 503 状态码
 18:         if (serverConfig.isRateLimiterEnabled()) {
 19:             ((HttpServletResponse) response).setStatus(HttpServletResponse.SC_SERVICE_UNAVAILABLE);
 20:             return;
 21:         }
 22:     }
 23:     chain.doFilter(request, response);
 24: }

class InstanceInfoReplicator implements Runnable {

    /**
     * RateLimiter
     */
    private final RateLimiter rateLimiter;
    /**
     * 令牌桶上限，默认：2
     */
    private final int burstSize;
    /**
     * 令牌再装平均速率，默认：60 * 2 / 30 = 4
     */
    private final int allowedRatePerMinute;

    InstanceInfoReplicator(DiscoveryClient discoveryClient, InstanceInfo instanceInfo, int replicationIntervalSeconds, int burstSize) {
        // ... 省略其他代码

        this.rateLimiter = new RateLimiter(TimeUnit.MINUTES);
        this.replicationIntervalSeconds = replicationIntervalSeconds;
        this.burstSize = burstSize;

        this.allowedRatePerMinute = 60 * this.burstSize / this.replicationIntervalSeconds;
        logger.info("InstanceInfoReplicator onDemand update allowed rate per min is {}", allowedRatePerMinute);
    }

    public boolean onDemandUpdate() {
        if (rateLimiter.acquire(burstSize, allowedRatePerMinute)) { // 限流
            scheduler.submit(new Runnable() {
                @Override
                public void run() {
                    logger.debug("Executing on-demand update of local InstanceInfo");
                    // 取消任务
                    Future latestPeriodic = scheduledPeriodicRef.get();
                    if (latestPeriodic != null && !latestPeriodic.isDone()) {
                        logger.debug("Canceling the latest scheduled update, it will be rescheduled at the end of on demand update");
                        latestPeriodic.cancel(false);
                    }
                    // 再次调用
                    InstanceInfoReplicator.this.run();
                }
            });
            return true;
        } else {
            logger.warn("Ignoring onDemand update due to rate limiter");
            return false;
        }
    }

}