SpringBoo+令牌桶+AOP实现限流

一、 pom.xml

        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-aop</artifactId>
        </dependency>
        <dependency>
            <groupId>com.google.guava</groupId>
            <artifactId>guava</artifactId>
            <version>30.1-jre</version>
        </dependency>

二、aop

@Aspect
@Component
public class LimitAop {
	/**
	 * 不同的接口，不同的流量控制
	 * map的key为 Limiter.key
	 */
	private final Map<String, RateLimiter> limitMap = Maps.newConcurrentMap();

	@Around("@annotation(com.example.springboot_rate_limiter.utils.LimitCustom)")
	public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
		MethodSignature signature = (MethodSignature) joinPoint.getSignature();
		Method method = signature.getMethod();
		//拿limit的注解
		LimitCustom limit = method.getAnnotation(LimitCustom.class);
		if (limit != null) {
			//key作用：不同的接口，不同的流量控制
			String key = limit.key();
			RateLimiter rateLimiter;
			//验证缓存是否有命中key
			if (!limitMap.containsKey(key)) {
				// 创建令牌桶
				rateLimiter = RateLimiter.create(limit.permitsPerSecond());
				limitMap.put(key, rateLimiter);
				System.out.println("新建了令牌桶= " + key + "，容量=" + limit.permitsPerSecond() + "");
			}
			rateLimiter = limitMap.get(key);
			// 拿令牌
			boolean acquire = rateLimiter.tryAcquire(limit.timeout(), limit.timeunit());
			// 拿不到命令，直接返回异常提示
			if (!acquire) {
				System.out.println("令牌桶={" + key + "}，获取令牌失败!");
				this.responseFail(limit.msg());
				return null;
			}
		}
		return joinPoint.proceed();
	}

}

三、自定义元注解

@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.METHOD})
@Documented
public @interface LimitCustom {

	/**
	 * 资源的key,唯一
	 * 作用：不同的接口，不同的流量控制
	 */
	String key() default "";

	/**
	 * 最多的访问限制次数
	 */
	double permitsPerSecond () ;

	/**
	 * 获取令牌最大等待时间
	 */
	long timeout();

	/**
	 * 获取令牌最大等待时间,单位(例:分钟/秒/毫秒) 默认:毫秒
	 */
	TimeUnit timeunit() default TimeUnit.MILLISECONDS;

	/**
	 * 得不到令牌的提示语
	 */
	String msg() default "系统繁忙,请稍后再试.";

四、测试类

	@GetMapping("/test1")
	@LimitCustom(key = "testLimiter1",permitsPerSecond = 1,timeout = 500, msg = "当前排队人数较多，请稍后再试！")
	public String testLimiter1() {
		return "令牌桶testLimiter1获取令牌成功!";
	}

	@GetMapping("/test2")
	@LimitCustom(key = "testLimiter2",permitsPerSecond = 1,timeout = 500, msg = "系统繁忙，请稍后再试！")
	public String testLimiter2() {
		return "令牌桶testLimiter2获取令牌成功!";
	}