com.ghp.request.limiter.impl.RateLimiter Maven / Gradle / Ivy

Go to download
package com.ghp.request.limiter.impl;

import com.ghp.request.annotation.RequestLimit;
import com.ghp.request.limiter.Limiter;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author ghp
 * @title 速率限流器
 * @description 基于谷歌 guava 工具包中提供的 RateLimiter 工具类实现，底层是采用令牌桶算法
 */
public class RateLimiter implements Limiter {

    /**
     * 记录所有被限流的接口
     * key是接口名，value是RateLimiter对象
     */
    private Map rateLimiterMap = new ConcurrentHashMap<>();

    /**
     * 采用令牌桶限流算法实现限流
     *
     * @param methodName
     * @param requestLimitAnnotation
     * return true-发生限流 false-未限流
     */
    @Override
    public boolean limit(String methodName, RequestLimit requestLimitAnnotation) {
        int qps = requestLimitAnnotation.qps();

        // 判断拦截的方法是否已存在，已存在则从 map 中取出已存在的，不存在则新建一个放入 map 中
        com.google.common.util.concurrent.RateLimiter rateLimiter = rateLimiterMap.computeIfAbsent(methodName, k ->
                com.google.common.util.concurrent.RateLimiter.create(qps));
        // 判断请求是否限流
        return !rateLimiter.tryAcquire();
    }
}