API简单限流实现

###背景

最近项目有个需求，需要对第三方接口调用加入调用次数限制。随设计自定义注解，使用拦截器拦截方法请求，将单位时间内的请求次数保存到redis。超出限制次数的请求直接拒绝或者异常处理。

1.自定义注解

@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.FIELD,ElementType.TYPE,ElementType.METHOD})
@Documented
public @interface AccessLimit {

    /**
     * 单位时间内允许访问的次数，默认60
     * @return
     */
    int maxCount() default 60;

    /**
     * 单位时间为1分钟,即默认限流为一分钟最大调用60次
     * @return
     */
    long time() default 1;
}

2.添加拦截器

@Configuration
@Slf4j
public class AccessLimitInterceptor implements HandlerInterceptor {

    @Autowired
    private RedisUtil redisUtil;

    @Autowired
    private MessageSource messageSource;

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        //请求输入方法
        if(handler instanceof HandlerMethod){
            HandlerMethod handlerMethod = (HandlerMethod) handler;
            Method method = handlerMethod.getMethod();
            if (!method.isAnnotationPresent(AccessLimit.class)) {
                return true;
            }
            AccessLimit accessLimit = method.getAnnotation(AccessLimit.class);
            if (accessLimit == null) {
                return true;
            }

            int maxCount = accessLimit.maxCount();
            long time = accessLimit.time();
            //存入redis中的key值
            String key = IpUtils.getIpAddr(request) + Constants.COLON + request.getRequestURI();
            try {
                //第一次访问将次数+1
                long visitTimes = redisUtil.increment(key,1);
                //第一次访问时设置过期时间
                if(1 == visitTimes){
                    redisUtil.setExpire(key,visitTimes,time, TimeUnit.MINUTES);
                }
                if(visitTimes > maxCount){
                    frequentRequest(response,messageSource);
                    return false;
                }
            }catch (Exception e){
                log.error("方法拦截中redis操作异常：" + e);
                return false;
            }
        }
        return true;
    }

    @Override
    public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler,
                           ModelAndView modelAndView) throws Exception {
        // TODO Auto-generated method stub
        HandlerInterceptor.super.postHandle(request, response, handler, modelAndView);
    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex)
            throws Exception {
        // TODO Auto-generated method stub
        HandlerInterceptor.super.afterCompletion(request, response, handler, ex);
    }

    /**
     * 接口请求频繁异常处理
     * @param response 请求响应
     * @param messageSource
     */
    private void frequentRequest(HttpServletResponse response,MessageSource messageSource){
        try {
            BaseResponse baseResponse = new BaseResponse();
            baseResponse.setState(ResultCode.FREQUENT_REQUEST_CODE.getValue());
            baseResponse.setStateDesc(InterUtils.interInfo(InterUtils.defLanguage, messageSource, "request.too.frequent"));
            response.setCharacterEncoding("utf-8");
            response.setContentType("application/json; charset=utf-8");
            PrintWriter writer = response.getWriter();
            JSONObject json = (JSONObject) JSONObject.toJSON(baseResponse);
            writer.write(json.toString());
        }catch (IOException e){
            log.error("IO异常:" + e.getMessage(), e);
        }
    }
}

3. 配置拦截器bean

@Configuration
public class WebMvcConfig implements WebMvcConfigurer {

    @Bean
    public AccessLimitInterceptor getAccessLimitInterceptor(){
        return new AccessLimitInterceptor();
    }

    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(getAccessLimitInterceptor()).addPathPatterns("/**");//.excludePathPatterns("/login");
    }
}

4.接口方法使用注解

@PostMapping("/api/sdk")
@AccessLimit(maxCount = 100,time = 60)
public String sdk(HttpServletRequest request){
	...
}

5.总结

从代码层面来考虑的话，此方式实现还是比较优雅的，对业务层也没有太多的耦合。且此种方式单体和分布式均适用，因为用户实际的访问次数都是存在redis容器里的，和应用的单体或分布式无关。