Guava RateLimiter 使用

RateLimiter令牌桶算法。令牌桶算法的思想是以固定的速率生成令牌，在请求之前都需要从令牌桶里获取足令牌。当令牌数量不足的时候，请求将被阻塞等待或者返回。RateLimiter常用于限制访问资源的速率。

配置文件
ratelimit.properties
# 限流模块的配置

# 是否开启限流
ratelimit.doRateLimit=false

# 配置超时时间（配置将等待获取，不配置将直接返回）,单位毫秒
ratelimit.waitTimeout=20

# 服务限流保护，服务每秒允许的TPS（需评估单个服务所允许的最大TPS）
ratelimit.permitsPerSecond=200

RateLimitConfig

/**
 * 限流配置信息
 * 
 */
@Data
@Component
@ConfigurationProperties(prefix = "ratelimit")
@PropertySource(value = "classpath:ratelimit.properties")
public class RateLimitConfig implements Serializable {
    
    private static final long serialVersionUID = 1L;
    private boolean           doRateLimit      = false;
    private long              waitTimeout;
    private long              permitsPerSecond;

}

/**
 * 限流服务接口
 * 
 */
public interface IRateLimitService {

    /**
     * 尝试获取许可证，获取1个，立即返回非阻塞
     * 
     * @return
     */
    boolean tryAcquire();

    /**
     * 尝试获取多个许可证，立即返回非阻塞
     * 
     * @param permits
     * @return
     */
    boolean tryAcquire(int permits);

    /**
     * 阻塞获取许可证，获取1个，若超过timeout未获取到许可证，则返回false
     * 
     * @param timeout
     * @return
     */
    boolean acquire(long timeout);

    /**
     * 阻塞获取多个许可证，若超过timeout未获取到许可证，则返回false
     * 
     * @param permits
     * @param timeout
     * @return
     */
    boolean acquire(int permits, long timeout);

}

/**
 * Guava RateLimiter的限流实现
 * 
 */
@Service
public class RateLimitServiceImpl implements IRateLimitService {

    private RateLimitConfig config;

    private RateLimiter rateLimiter;

    @Autowired
    public RateLimitServiceImpl(RateLimitConfig config) {
        this.config = config;
        this.rateLimiter = RateLimiter.create(this.config.getPermitsPerSecond());
    }

    @Override
    public boolean tryAcquire() {
        return this.tryAcquire(1);
    }

    @Override
    public boolean tryAcquire(int permits) {
        return rateLimiter.tryAcquire(permits);
    }

    @Override
    public boolean acquire(long timeout) {
        return this.acquire(1, timeout);
    }

    @Override
    public boolean acquire(int permits, long timeout) {
        long start = System.currentTimeMillis();
        for (;;) {
            boolean tryAcquire = rateLimiter.tryAcquire(permits);
            if (tryAcquire) {
                return true;
            }
            long end = System.currentTimeMillis();
            if ((end - start) >= timeout) {
                return false;
            }
        }
    }

}

切面拦截
/**
 * 限流切面
 * 
 */
@Aspect
@Order(-1)
@Component
public class RateLimitAspect {
    private static final Logger LOG = LoggerFactory.getLogger(RateLimitAspect.class);

    private RateLimitConfig config;

    private IRateLimitService rateLimitService;

    @Autowired
    public RateLimitAspect(RateLimitConfig config, IRateLimitService rateLimitService) {
        this.config = config;
        this.rateLimitService = rateLimitService;
    }

    @Pointcut("execution(public * xxx.xxx.*Controller.*(..))")
    public void executionMethod() {}

    @Around(value = "executionMethod()")
    public Object doRateLimit(ProceedingJoinPoint pjp) throws Throwable {
        if (LOG.isDebugEnabled()) {
            LOG.debug("进入限流处理切面!");
        }
        Object result = null;
        // 判断是否限流
        try {
            if (config.isDoRateLimit()) {
                // 开启限流
                boolean acquireResult = false;
                // 1.查看是否配置超时时间
                if (config.getWaitTimeout() == 0L) {
                    // 2.获取令牌
                    acquireResult = rateLimitService.tryAcquire();
                } else {
                    // 2.获取令牌，超时时间内获取令牌
                    acquireResult = rateLimitService.acquire(config.getWaitTimeout());
                }
                if (acquireResult) {
                    // 3.成功获取令牌，放行
                    result = pjp.proceed();
                } else {
                    // 3.失败获取令牌，返回错误码 429 => to many requests
                    result = ResponseEntity.status(HttpStatus.TOO_MANY_REQUESTS).build();
                }
            } else {
                // 无开启限流，直接放行
                result = pjp.proceed();
            }
        } catch (Throwable e) {
            throw e;
        }

        if (LOG.isDebugEnabled()) {
            LOG.debug("限流处理切面结束!");
        }
        return result;
    }

}