基于guava RateLimiter及spring aop 注解实现的访问速率限制

1、pom引入

        <dependency>
            <groupId>com.google.guava</groupId>
            <artifactId>guava</artifactId>
            <version>14.0.1</version>
        </dependency>

        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-aop</artifactId>
        </dependency>

2、注解类

@Documented
@Retention(RetentionPolicy.RUNTIME)
@Target({ ElementType.METHOD, ElementType.TYPE })
public @interface SpeedLimit {

    /**
     * 开启只对特定ip限制
     */
    boolean forIP() default false;

    /**
     * 请求次数
     */
    double num() default 1;

    /**
     * 时间（单位 s）
     */
    int time() default 1;

    /**
     * 阻塞等待
     */
    boolean block() default true;

}

3、主要服务类

@Slf4j
public class RequestCacheService {

    // 缓存最大数
    private static final int MAXIMIZE = 1000;
    //
    private static final int DEFAULT_RATE = 1;

    /**
     * 创建缓存器
     */
    private static final LoadingCache<String, RateLimiter> requestCaches = CacheBuilder.newBuilder()
            .maximumSize(MAXIMIZE)
            .expireAfterWrite(1, TimeUnit.MINUTES)
            .removalListener((notification) -> {
                log.info(notification.getKey() + " 令牌桶被移除了，原因： " + notification.getCause());
            })
            .build(new CacheLoader<String, RateLimiter>() {
                @Override
                public RateLimiter load(String key){
                    return RateLimiter.create(DEFAULT_RATE);
                }
            });

    /**
     * 访问速率限制
     *
     * @param key
     * @param num
     * @param time
     * @return
     * @throws ExecutionException
     */
    public static RateLimiter rateLimiter(String key, double num, int time) throws ExecutionException {
        double rate = Double.parseDouble(String.format("%.2f", num / time));
        if(0 == rate){
            rate = 0.01;
        }
        RateLimiter rateLimiter = requestCaches.get(key);
        if(rate != rateLimiter.getRate()){
            rateLimiter.setRate(rate);
        }
        return rateLimiter;
    }

    /**
     * 尝试获取令牌
     *
     * @param key
     * @param num
     * @param time
     * @param block
     * @return
     * @throws ExecutionException
     */
    public static boolean tryAcquire(String key, double num, int time, boolean block) throws ExecutionException {
        RateLimiter rateLimiter = rateLimiter(key, num, time);
        if(block){
            rateLimiter.acquire(1);
            return true;
        }
        return rateLimiter.tryAcquire();
    }
}

4、切面实现类

@Component
@Aspect
public class RateLimitAspect {

    @Resource
    private HttpServletRequest request;

    @Pointcut(" @annotation(com.example.bootdemo.aop.SpeedLimit) || " +
            "@within(com.example.bootdemo.aop.SpeedLimit)")
    public void pointcut() {


    }

    @Around("pointcut()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        String ipAddr = ServletUtil.getClientIP(request);
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        SpeedLimit annotation = signature.getMethod().getAnnotation(SpeedLimit.class);
        if(null == annotation){
            annotation = joinPoint.getTarget().getClass().getAnnotation(SpeedLimit.class);
        }
        boolean enableLimit = checkEnableLimit(ipAddr, annotation);
        try{
            if(enableLimit){
                return speedLimit(joinPoint, getKey(signature), annotation);
            }
            return joinPoint.proceed();
        }catch (Throwable e){
            throw e;
        }
    }

    /**
     * 速度限制
     *
     * @param joinPoint
     * @param key
     * @param annotation
     * @return
     * @throws Throwable
     */
    private Object speedLimit(ProceedingJoinPoint joinPoint, String key, SpeedLimit annotation) throws Throwable {
        double num = annotation.num();
        int time = annotation.time();
        boolean block = annotation.block();
        try{
            if (RequestCacheService.tryAcquire(key, num, time, block)) {
                // 获得令牌（不限制访问）
                return joinPoint.proceed();
            } else {
                // 未获得令牌（限制访问）
                return Result.FAIL();
            }
        }catch (Throwable e){
            throw e;
        }
    }

    /**
     * 检测是否开启限制
     *
     * @param ip
     * @param annotation
     */
    private boolean checkEnableLimit(String ip, SpeedLimit annotation) {
        if(annotation.forIP()){
            List<String> ips = getLimitIPs();
            if(null != ips && 0 != ips.size() && ips.contains(ip)){
                // 只对存在的ip进行限制
                return true;
            }
            return false;
        }
        // 默认限制所有
        return true;
    }

    /**
     * 获取key
     *
     * @param signature
     * @return
     */
    private String getKey(MethodSignature signature){
        return signature.getDeclaringTypeName() + "." + signature.getMethod().getName();
    }

    /**
     * 获取ip集合
     *
     * @return
     */
    private List<String> getLimitIPs() {
        return Lists.newArrayList("127.0.0.1");
    }


}

5、测试controller

@RestController
@RequestMapping("/limit")
public class LimitController {


    /**
     * 限制三秒内访问五次
     * 
     * @return
     */
    @RequestMapping("/test1")
    @SpeedLimit(num = 5, time = 3)
    public Result test1(){
        System.out.println("访问成功" + new SimpleDateFormat("HH:mm:ss").format(new Date()));
        return Result.OK();
    }
}