自定义注解方式对接口进行限流实现

一、引入依赖

<dependency>
    <groupId>com.google.guava</groupId>
    <artifactId>guava</artifactId>
    <version>18.0</version>
</dependency>

如果已经引入可忽略上面依赖

<dependency>
    <groupId>io.springfox</groupId>
    <artifactId>springfox-staticdocs</artifactId>
    <version>2.6.1</version>
    <scope>test</scope>
</dependency>

二、书写限流自定义注解

package com.xxx.xxx.xxx;

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;

/**
 * 限流注解
 * 
 * 限流值 默认1秒一次请求；<br/>
 * 例如： <br/>
 * rateLimiterNumber=0.5 这个是配置2秒通过1个请求；<br/>
 * rateLimiterNumber=2 这个是配置1秒通过2个请求；<br/>
 * 
 * 
 * @author zxl
 * @date 2021/03/09 10:38
 */
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface RateLimiterAnnotation {
    /**
     * 限流值 默认1秒一次请求；<br/>
     * 例如： <br/>
     * setLimiterValue=0.5 这个是配置2秒通过1个请求；<br/>
     * setLimiterValue=2 这个是配置1秒通过2个请求；<br/>
     * 
     * @return 限流值
     */
    double setLimiterValue() default 1.0;

    /**
     * 获取超时时间 默认:500
     */
    long timeOut() default 500;

    /**
     * 超时时间单位 默认:毫秒
     */
    TimeUnit timeUnit() default TimeUnit.MILLISECONDS;

    /**
     * 无法获取返回提示信息
     */
    String msg() default "亲，服务器快被挤爆了，请稍后再试！";
}

三、书写限流拦截器

package com.xxx.xxx.xxx;

/**
 * @author zxl
 * @date 2021/03/09 10:31
 */

import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.PrintWriter;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 限流拦截器
 * 
 * @author zxl
 * @date 2021/03/09 10:31
 **/
@Slf4j
@Component
public class RateLimiterInterceptor implements HandlerInterceptor {

    /**
     * 不同的方法存放不同的令牌桶
     */
    private final Map<String, RateLimiter> rateLimiterMap = new ConcurrentHashMap<>();

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
        PrintWriter out = null;
        try {
            if (handler instanceof HandlerMethod) {
                HandlerMethod handlerMethod = (HandlerMethod)handler;
                RateLimiterAnnotation rateLimit = handlerMethod.getMethodAnnotation(RateLimiterAnnotation.class);
                // 判断是否有注解
                if (rateLimit != null) {
                    // 获取请求url
                    String url = request.getRequestURI();
                    RateLimiter rateLimiter;
                    // 判断map集合中是否有创建好的令牌桶
                    if (!rateLimiterMap.containsKey(url)) {
                        // 创建令牌桶,以n r/s往桶中放入令牌
                        rateLimiter = RateLimiter.create(rateLimit.setLimiterValue());
                        rateLimiterMap.put(url, rateLimiter);
                    }
                    rateLimiter = rateLimiterMap.get(url);
                    // 获取令牌
                    boolean acquire = rateLimiter.tryAcquire(rateLimit.timeOut(), rateLimit.timeUnit());
                    if (acquire) {
                        return true;
                    } else {
                        log.warn("请求被限流,url:{}", request.getServletPath());
                        response.setCharacterEncoding("UTF-8");
                        response.setContentType("application/json; charset=utf-8");

                        out = response.getWriter();
                        out.println("{\"status\":" + 500 + " ,\"message\" :\"" + rateLimit.msg() + "\"}");
                        return false;

                    }
                }
            }
            return true;
        } catch (Exception e) {
            e.printStackTrace();
            log.error("请求url:{}" + request.getServletPath(), e);
            return false;
        } finally {
            if (out != null) {
                out.flush();
                out.close();
            }
        }
    }
}

四、配置限流拦截器

package com.xxx.xxx.configuration;

import com.xxx.xxx.xxx.RateLimiterInterceptor;
import org.springframework.boot.web.servlet.MultipartConfigFactory;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.util.unit.DataSize;
import org.springframework.web.servlet.config.annotation.InterceptorRegistry;
import org.springframework.web.servlet.config.annotation.ResourceHandlerRegistry;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;

import javax.annotation.Resource;
import javax.servlet.MultipartConfigElement;

/***
 * Description:限流拦截器配置
 * 
 * @author zxl
 */
@Configuration
public class RateLimiterConfig implements WebMvcConfigurer {

    /**
     * 限流配置
     */
    @Resource
    private RateLimiterInterceptor rateLimiterInterceptor;

    /**
     * 限流配置
     *
     * @param registry
     */
    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(rateLimiterInterceptor).addPathPatterns("/**");
    }

}

五、接口使用

@RateLimiterAnnotation(setLimiterValue = 0.1, timeOut = 200, timeUnit = TimeUnit.MILLISECONDS, msg = "服务器繁忙,请稍后再试")
    @PostMapping("/RateLimiter/test")
    public String test() {
        return "";
    }

记录一下方便查找