Java 接口限流

内容：

1、限流原理 -- 令牌桶算法

令牌桶算法的原理是系统会以一个恒定的速度（每秒生成一个令牌）往桶里放入令牌。当有访问者（针对于 IP）要访问接口时，则需要先从桶里获取一个令牌，当桶里没有令牌可取时，则拒绝服务。当桶满时，新添加的令牌被丢弃或拒绝。

2、知识点

Springboot
Guava -- RateLimiter
Interceptor（拦截器）

3、具体实现

1）先写一个限流 Service -- LoadingCacheService，代码如下：

@Service
public class LoadingCacheService {

    private final Logger logger = LoggerFactory.getLogger(LoadingCacheService.class);

    private LoadingCache<String, RateLimiter> ipRequestCaches = CacheBuilder.newBuilder()
            // 设置缓存上限
            .maximumSize(10000)
            // 设置一分钟对象没有被读/写访问则对象从内存中删除
            .expireAfterAccess(1, TimeUnit.MINUTES)
            // CacheLoader 类实现自动加载
            .build(new CacheLoader<String, RateLimiter>() {
                @Override
                public RateLimiter load(String s) {
                    // 新的 IP 初始化 (限流每秒生成 2 个令牌)
                    return RateLimiter.create(2);
                }
            });

    public boolean hasToken(HttpServletRequest request) {
        try {
            String ip = this.getIPAddress(request);
            String url = request.getRequestURL().toString();
            String key = "req_limit_".concat(url).concat(ip);
            // 有则返回，没有就添加后获取
            RateLimiter limiter = ipRequestCaches.get(key);

            return limiter.tryAcquire();
        } catch (Exception e) {
            logger.error("获取令牌异常：", e);
        }
        return false;
    }

    /**
     * 获取当前网络 ip
     *
     * @param request HttpServletRequest
     * @return 真实的 ip 地址
     */
    private String getIPAddress(HttpServletRequest request) {
        String ipAddress = request.getHeader("HTTP_X_FORWARDED_FOR");
        if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
            ipAddress = request.getHeader("x-forwarded-for");
        }
        if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
            ipAddress = request.getHeader("Proxy-Client-IP");
        }
        if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
            ipAddress = request.getHeader("WL-Proxy-Client-IP");
        }
        if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
            ipAddress = request.getRemoteAddr();
        }
        // 对于通过多个代理的情况，第一个 IP 为客户端真实 IP，多个 IP 按照','分割
        // "***.***.***.***".length() = 15
        if (ipAddress != null && ipAddress.length() > 15) {
            if (ipAddress.indexOf(",") > 0) {
                ipAddress = ipAddress.substring(0, ipAddress.indexOf(","));
            }
        }
        return ipAddress;
    }

}

以上代码要注意的点：

Guava 用到了缓存，感兴趣的同学，可以自己深入学习一下；
RateLimiter.create(2) 意思就是每秒生成两个令牌，如果改为 3 ，就是每秒生成 3 个；
仅仅靠 request.getRemoteAddr() 有可能获取不到用户的真实 IP ，需要用 getIPAddress() 方法。

2）写一个拦截器组件 -- RequestInterceptor

@Component
public class RequestInterceptor implements HandlerInterceptor {

    @Resource
    private LoadingCacheService loadingCacheService;

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object o) throws IOException {
        if (loadingCacheService.hasToken(request)) {
            return true;
        }
        outputError(response);
        return false;
    }

    @Override
    public void postHandle(HttpServletRequest request, HttpServletResponse response, Object o, ModelAndView modelAndView) {

    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object o, Exception e) {

    }

    private void outputError(HttpServletResponse response) throws IOException {
        response.setContentType("application/json;charset=UTF-8");
        response.setStatus(429);
        ServletOutputStream outputStream = null;
        try {
            outputStream = response.getOutputStream();
            outputStream.write(JSONObject.toJSONString(new ErrorRes("请求太频繁！")).getBytes("UTF-8"));
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (outputStream != null) {
                outputStream.flush();
                outputStream.close();
            }
        }
    }

}

3）编写拦截器配置类 -- InterceptorConfig，并注册刚才编写的拦截器 -- RequestInterceptor

@Configuration
public class InterceptorConfig extends WebMvcConfigurerAdapter {

    @Resource
    private RequestInterceptor requestInterceptor;

    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(requestInterceptor).addPathPatterns("/api/v1/**");
        // 注册
        super.addInterceptors(registry);
    }

}

以上代码要注意的点：

addPathPatterns 里面的内容，就是要拦截的接口；
要拦截多个地方，可以用逗号隔开，比如：addPathPatterns("/api/v1/**", "/api/v2/**") 。

4、结语

如果我的博客你看到了这里，我想说明一下，我一般会在开头就先写实现的具体代码，而在最后进行总结。

之前在网上搜集过一些资料，还有用到自定义注解的，可以参考：https://blog.csdn.net/u013476435/article/details/82180663。而我没有用的原因是：那些注解都用到了 aop，大部分在超流了以后，会通过抛异常的形式来处理，但我想要的时候通过返回给用户一个“请求太频繁”的提示，来达到目的。

当然，还没有考虑到就是恶意攻击。那就得再另起一篇来说明了，比如：增加 IP 黑名单等等。但就我们公司目前的业务，暂时是通过手动配置 IP 黑名单来处理的，还没有在程序中限制。关于这块，以后如果有用到的话，我会进行补充。在这写出来，也是给以后的自己提个醒！

目录：

内容：

猜你喜欢