背景：SSM框架

需求：需要开发对外API接口，需要对API接口做节流处理，每秒只能处理1个请求

准备工作，pom导入aop包(如果添加后你的项目无法启动请检查你的spring包是否匹配)：

<!-- Spring AOP 日志管理需要导入的包 -->
<dependency>
      <groupId>org.springframework</groupId>
      <artifactId>spring-aop</artifactId>
      <version>5.0.11.RELEASE</version>
    </dependency>
    <dependency>
      <groupId>org.aspectj</groupId>
      <artifactId>aspectjrt</artifactId>
      <version>1.8.13</version>
    </dependency>
    <dependency>
      <groupId>org.aspectj</groupId>
      <artifactId>aspectjweaver</artifactId>
      <version>1.8.13</version>
    </dependency>

自定义一个限流用的注解，后面在需要限流的方法或接口上面只需添加该注解即可-需放在可被controller扫描的地方

import java.lang.annotation.*;
@Documented
@Target(value = ElementType.METHOD)
@Retention(value = RetentionPolicy.RUNTIME)
public @interface GuavaLimitRateAnnotation {

    // 每秒 1 个请求
    double limitCount() default 1d;

    // 限制类型
    String limitType();
}

创建限流 AOP 类-需放在可被controller扫描的地方

import com.google.common.util.concurrent.RateLimiter;
import com.xjt.utils.Filter;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.json.simple.JSONObject;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.lang.reflect.Method;
import java.util.Objects;
import java.util.logging.Logger;
@Aspect
@Component
public class GuavaLimitRateAspect {

    private static Logger logger = Logger.getLogger(String.valueOf(Filter.class));

    @Before("execution(@com.xjt.controller.GuavaLimitRateAnnotation * *(..))")
    public void limit(JoinPoint joinPoint) {
        // 1.获取当前方法
        Method currentMethod = getCurrentMethod(joinPoint);
        if (Objects.isNull(currentMethod)) {
            return;
        }
        // 2.从方法注解定义上获取限流的类型
        String limitType = currentMethod.getAnnotation(GuavaLimitRateAnnotation.class).limitType();
        double limitCount = currentMethod.getAnnotation(GuavaLimitRateAnnotation.class).limitCount();
        // 3.使用guava的令牌桶算法获取一个令牌，获取不到先等待
        RateLimiter rateLimiter = RateLimitHelper.getRateLimiter(limitType, limitCount);
        boolean b = rateLimiter.tryAcquire();
        if (b) {
            System.out.println("获取到令牌");
        } else {
            HttpServletResponse resp = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getResponse();
            JSONObject jsonObject = new JSONObject();
            jsonObject.put("success",false);
            jsonObject.put("msg","限流中");
            try {
                output(resp, jsonObject.toJSONString());
            } catch (Exception e) {
                logger.info("error,e:"+e);
            }
        }
    }

    public void output(HttpServletResponse response, String msg) throws IOException {
        response.setContentType("application/json;charset=UTF-8");
        ServletOutputStream outputStream = null;
        try {
            outputStream = response.getOutputStream();
            outputStream.write(msg.getBytes("UTF-8"));
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            assert outputStream != null;
            outputStream.flush();
            outputStream.close();
        }
    }
    private Method getCurrentMethod(JoinPoint joinPoint) {
        Method[] methods = joinPoint.getTarget().getClass().getMethods();
        Method target = null;
        for (Method method : methods) {
            if (method.getName().equals(joinPoint.getSignature().getName())) {
                target = method;
                break;
            }
        }
        return target;
    }
}

其中限流的核心 API 即为 RateLimiter 这个对象，涉及到的 RateLimitHelper 类如下：

import com.google.common.util.concurrent.RateLimiter;
import java.util.HashMap;
import java.util.Map;
public class RateLimitHelper {
    private RateLimitHelper(){}
    private static Map<String, RateLimiter> rateMap = new HashMap<>();
    public static RateLimiter getRateLimiter(String limitType, double limitCount ){
        RateLimiter rateLimiter = rateMap.get(limitType);
        if(rateLimiter == null){
            rateLimiter = RateLimiter.create(limitCount);
            rateMap.put(limitType,rateLimiter);
        }
        return rateLimiter;
    }
}

测试：

@Controller
@RequestMapping("/api")
public class LimitController {

    @RequestMapping("/api_get_test",method = RequestMethod.GET)
    @GuavaLimitRateAnnotation(limitType = "测试限流", limitCount = 1)
    @ResponseBody
    public void limitByGuava(HttpServletRequest request, HttpServletResponse response, UsersEntity usersEntity) {
        try {
            System.out.println("成功调用API/GET接口："+usersEntity.getApi_key());
            response.getWriter().write(net.sf.json.JSONObject.fromObject(usersEntity).toString());
        }catch (Exception e){
            e.printStackTrace();
        }
    }

}

image.png

完工！