微服务实战 Spring Cloud Alibaba（二十一）：使用注解的方式实现单个接口的令牌桶限流

上一篇文章《微服务实战spring cloud alibaba（二十）利用guava框架实现单个接口的令牌桶限流》我们介绍了使用guava框架实现单个接口的令牌桶限流。但是大家想象一下，真实的生产业务，有直接在controller里面编写令牌筒的判断逻辑的吗？

所以说这肯定是行不通的，那真实环境是怎么做的呢？其实这里肯定是使用注解的啦，spring里面的aop利用起来非常的方便。下面我们来实战一下：

一、创建一个maven项目，引入如下的依赖

<dependency>
			<groupId>org.springframework.boot</groupId>
			<artifactId>spring-boot-starter</artifactId>
		</dependency>
		<dependency>
			<groupId>org.springframework.boot</groupId>
			<artifactId>spring-boot-starter-web</artifactId>
		</dependency>
		<dependency>
			<groupId>com.google.guava</groupId>
			<artifactId>guava</artifactId>
			<version>30.1-jre</version>
		</dependency>
		<!-- log start -->
		<dependency>
			<groupId>org.slf4j</groupId>
			<artifactId>slf4j-api</artifactId>
			<version>1.6.6</version>
		</dependency>
		<dependency>
			<groupId>org.slf4j</groupId>
			<artifactId>slf4j-log4j12</artifactId>
			<version>1.6.6</version>
		</dependency>
		<dependency>
			<groupId>org.projectlombok</groupId>
			<artifactId>lombok</artifactId>
			<version>${lombok.version}</version>
			<scope>provided</scope>
		</dependency>
		<!-- OKHttp3依赖 -->
		<dependency>
			<groupId>com.squareup.okhttp3</groupId>
			<artifactId>okhttp</artifactId>
			<version>3.8.1</version>
		</dependency>
		<dependency>
			<groupId>org.springframework.boot</groupId>
			<artifactId>spring-boot-starter-aop</artifactId>
		</dependency>
		<!-- https://mvnrepository.com/artifact/com.alibaba/fastjson -->
		<dependency>
			<groupId>com.alibaba</groupId>
			<artifactId>fastjson</artifactId>
			<version>1.2.83</version>
		</dependency>

二、定义一个注解RateLimitAnnotation

使用注解实现限流，那么首先我们需要自定义一个注解，完整代码示例如下：

package com.demo.api.aop;

import java.lang.annotation.Documented;
import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;
import java.util.concurrent.TimeUnit;

/**
 * 自定义限流的注解
 * @author Administrator
 *
 */
@Retention(RetentionPolicy.RUNTIME)
@Target({ ElementType.METHOD })
@Documented
public @interface RateLimitAnnotation {

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

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

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

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

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

三、实现自定义注解RateLimitAspect

自定义好了注解之后，我们需要实现这个注解里面的逻辑，完整示例代码如下

package com.demo.api.aop;

import java.io.IOException;
import java.lang.reflect.Method;
import java.util.Map;

import javax.servlet.http.HttpServletResponse;

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 org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import com.alibaba.fastjson.JSON;
import com.demo.api.model.BaseResponse;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;

import lombok.extern.slf4j.Slf4j;

/**
 * 实现自定义限流注解的代码逻辑
 * 
 * @author Administrator
 *
 */
@Slf4j
@Aspect
@Component
public class RateLimitAspect {

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

	@Around("@annotation(com.demo.api.aop.RateLimitAnnotation)")
	public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
		MethodSignature signature = (MethodSignature) joinPoint.getSignature();
		Method method = signature.getMethod();
		// 拿limit的注解
		RateLimitAnnotation limit = method.getAnnotation(RateLimitAnnotation.class);
		if (limit != null) {
			// key作用：不同的接口，不同的流量控制
			String key = limit.key();
			RateLimiter rateLimiter = null;
			// 验证缓存是否有命中key
			if (!limitMap.containsKey(key)) {
				// 创建令牌桶
				rateLimiter = RateLimiter.create(limit.permitsPerSecond());
				limitMap.put(key, rateLimiter);
				log.info("当前key：{} 的令牌桶找不到，所以在这里新创建一个令牌筒：{}  每秒钟的容量是：{}",
						new Object[] { key, key, limit.permitsPerSecond() });
			}
			rateLimiter = limitMap.get(key);
			// 拿令牌
			boolean acquire = rateLimiter.tryAcquire(limit.timeout(), limit.timeunit());
			// 拿不到命令，直接返回异常提示
			if (!acquire) {
				log.debug("当前令牌桶：{}，获取令牌失败，准备返回错误信息", key);
				this.responseFail(limit.msg());
				return null;
			}
		}
		return joinPoint.proceed();
	}

	/**
	 * 直接向前端抛出异常
	 * 
	 * @param msg 提示信息
	 */
	private void responseFail(String msg) {
		try {
			HttpServletResponse response = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes())
					.getResponse();
			// 使用setCharacterEncoding方法设置输出内容使用UTF-8进行编码
			response.setCharacterEncoding("UTF-8");
			
			BaseResponse rs = BaseResponse.fail(msg);
			String jsonString = JSON.toJSONString(rs);
			
			log.info("注解信息是：{}",jsonString);
			response.getWriter().write(jsonString);
		} catch (Exception e) {
			log.error(e.getMessage(), e);
		}
	}
	
}

四、编写一个接口，添加刚才自定义的注解

@GetMapping("/testLimiter1")
	@RateLimitAnnotation(key = "testLimiter1", permitsPerSecond = 2, timeout = 500, timeunit = TimeUnit.MILLISECONDS, msg = "当前排队人数较多，请稍后再试！")
	public BaseResponse testLimiter1() {
		return BaseResponse.success("请求成功，获取令牌成功，时间：" + LocalDateTime.now().format(dtf));
	}

这里我们定义的是每秒钟放2个令牌。