springMvc自定义注解实现拦截器限流

前言

本文适用于springMvc, 当然也包括springboot

代码

自定义拦截器注解:

@Retention(RUNTIME)
@Target(METHOD)
public @interface AccessLimit {
    // controller url 必填
    String pathInfoC();
    // 最大并发量
    int processQuantityC() default 100;
    // 开启限流功能
    boolean processFlagC() default true;
    // url 必填
    String pathInfo();
    // 最大并发量
    int processQuantity() default 100;
    // 开启限流功能
    boolean processFlag() default true;
}

拦截器注册及设置拦截的url:

@Configuration
public class WebConfig extends WebMvcConfigurerAdapter {

    @Autowired
    private AccessInterceptor interceptor;

    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(interceptor).addPathPatterns("/mediaData/*");
    }
}

拦截器处理, 这里限流用redis做的:

@Component
public class AccessInterceptor extends HandlerInterceptorAdapter {

    @Autowired
    private RedisUtil redisUtil;

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {

        if(handler instanceof HandlerMethod){
            HandlerMethod hm = (HandlerMethod) handler;

            String errorJsonMsg = "{\"success\":\"false\",\"resultCode\":\"-1\",\"resultMsg\":\"Too many requests in the same period of time\"}";

            AccessLimit accessLimit = hm.getMethodAnnotation(AccessLimit.class);
            if(accessLimit == null){
                return true;
            }

            String pathInfoC = accessLimit.pathInfoC();
            int processQuantityC = accessLimit.processQuantityC();
            boolean processFlagC = accessLimit.processFlagC();
            String pathInfo = accessLimit.pathInfo();
            int processQuantity = accessLimit.processQuantity();
            boolean processFlag = accessLimit.processFlag();

            String accessAllCount = "hw:accessAllCount:" + pathInfoC;
            if (processFlagC) {
                request.setAttribute("accessAllCount", accessAllCount);
                if (redisUtil.incr(accessAllCount) > processQuantityC) {
                    redisUtil.decr(accessAllCount);
                    PrintWriter pw = response.getWriter();
                    pw.write(errorJsonMsg);
                    pw.flush();
                    pw.close();
                    return false;
                }
            }

            if (processFlag) {
                String accessCount = "hw:accessCount:" + pathInfoC + pathInfo;
                request.setAttribute("accessCount", accessCount);
                if (redisUtil.incr(accessCount) > processQuantity) {
                    redisUtil.decr(accessCount);
                    redisUtil.decr(accessAllCount);
                    PrintWriter pw = response.getWriter();
                    pw.write(errorJsonMsg);
                    pw.flush();
                    pw.close();
                    return false;
                }
            }

            return true;
        }

        return true;
    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex)
            throws Exception {
        redisUtil.decr(String.valueOf(request.getAttribute("accessCount")));
        redisUtil.decr(String.valueOf(request.getAttribute("accessAllCount")));
    }
}

使用限流的controller:

@RestController
@RequestMapping("/mediaData")
public class MediaDataController extends BaseController {

    @Autowired
    private IMediaDataService mediaDataService;

    @PostMapping(value = "/doWork", produces = MediaType.APPLICATION_JSON_UTF8_VALUE)
    @AccessLimit(pathInfoC = "/mediaData", pathInfo = "/doWork")
    public String doWork(@RequestBody String inputJson) {
        String outputJson = "";
        try {
            outputJson = mediaDataService.doWork(inputJson);
        } catch (Exception e) {
            logger.error(e.getMessage());
            outputJson = mediaDataService.getErrorJson();
        }
        return outputJson;
    }
}