springboot+RateLimiter+AOP自定义注解限流

RateLimiter简介

RateLimiter是Guava库中的一个限流器，它提供如下功能：
(1)基于PPS进行限流
(2)基于PPS限流的同时提供热启动

springboot集成RateLimiter

pom.xml引入

<!--引入guava依赖-->
<dependency>
	<groupId>com.google.guava</groupId>
	<artifactId>guava</artifactId>
	<version>31.1-jre</version>
</dependency>

RateLimiter常用api

(1)RateLimiter.create(double permitsPerSecond); ===> 指定速率，每秒产生permitsPerSecond个令牌
(2)rateLimiter.acquire(int permits); ===> 返回获取permits个令牌所需的时间
(3)rateLimiter.tryAcquire(1); ===> 获取1个令牌，返回false获取失败，true获取成功
(4)rateLimiter.tryAcquire(1, 2, TimeUnit.SECONDS); ===> 获取1个令牌，最多等待两秒。返回false获取失败，true获取成功

package com.mry.springboottools.controller;

import com.google.common.util.concurrent.RateLimiter;

import java.util.concurrent.TimeUnit;

/**
 * RateLimiter测试类
 */
public class RateLimiterTest {

    public static void main(String[] args) {
        //指定速率，每秒产生一个令牌  
        RateLimiter rateLimiter = RateLimiter.create(1);
        System.out.println(rateLimiter.getRate());
        //修改为每秒产生2个令牌
        rateLimiter.setRate(2);
        System.out.println(rateLimiter.getRate());

        //while (true) {  
        //获取2个令牌所需要的时间  
        double acquire = rateLimiter.acquire(2);
        //输出结果， 第一次0秒，后面每次等待接近0.5秒的时间  
        //0.0  
        //0.995198  
        //0.9897  
        //0.999413  
        //0.998272  
        //0.99202  
        //0.993757  
        //0.996198  
        //0.99523  
        //0.99532  
        //0.992674  
        System.out.println(acquire);
        // }  
        //获取1个令牌  
        boolean result = rateLimiter.tryAcquire(1);
        System.out.println(result);
        //获取1个令牌,最多等待两秒  
        boolean result2 = rateLimiter.tryAcquire(1, 2, TimeUnit.SECONDS);
        System.out.println(result2);

    }
}

代码实现

自定义注解Limiter

package com.mry.springboottools.annotation;

import java.lang.annotation.*;
import java.util.concurrent.TimeUnit;

/**
 * 限流注解
 *
 * @author
 * @date
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface Limiter {

    /**
     * 不进行限流
     */
    int NOT_LIMITED = 0;

    /**
     * qps (每秒并发量)
     */
    double qps() default NOT_LIMITED;

    /**
     * 超时时长,默认不等待
     */
    int timeout() default 0;

    /**
     * 超时时间单位,默认毫秒
     */
    TimeUnit timeUnit() default TimeUnit.MILLISECONDS;

    /**
     * 返回错误信息
     */
    String msg() default "系统忙，请稍后再试";

}

限流切面

package com.mry.springboottools.aspect;

import com.google.common.util.concurrent.RateLimiter;
import com.mry.springboottools.annotation.Limiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;
import java.lang.reflect.Method;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * 限流切面
 * @author
 * @date
 */
@Slf4j
@Aspect
@Component
public class RateLimiterAspect {
    /**
     * key: 类全路径+方法名
     */
    private static final ConcurrentMap<String, RateLimiter> RATE_LIMITER_CACHE = new ConcurrentHashMap<>();


    @Around("@within(limiter) || @annotation(limiter)")
    public Object pointcut(ProceedingJoinPoint point, Limiter limiter) throws Throwable {
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        if (limiter != null && limiter.qps() > Limiter.NOT_LIMITED) {
            double qps = limiter.qps();
            //这个key可以根据具体需求配置,例如根据ip限制,或用户
            String key = method.getDeclaringClass().getName() + method.getName();
            if (RATE_LIMITER_CACHE.get(key) == null) {
                // 初始化 QPS
                RATE_LIMITER_CACHE.put(key, RateLimiter.create(qps));
            }
            // 尝试获取令牌
            if (RATE_LIMITER_CACHE.get(key) != null && !RATE_LIMITER_CACHE.get(key).tryAcquire(limiter.timeout(), limiter.timeUnit())) {
                log.error("触发限流操作{}", key);
                throw new RuntimeException(limiter.msg());
            }
        }
        return point.proceed();
    }

}

验证

package com.mry.springboottools.controller;

import com.google.common.util.concurrent.RateLimiter;
import com.mry.springboottools.annotation.Limiter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import java.util.concurrent.TimeUnit;

/**
 * 测试限流
 * @author
 * @date
 */
@RestController
@RequestMapping("limiter")
@Slf4j
public class RateLimiterController {

    @GetMapping
    @Limiter(qps = 1, msg = "您已被限流！")
    public String getUserName() {
        String userName = "mry";
        log.info("userName = {}", userName);
        return userName;
    }
}

接口请求：

日志输出：