guava RateLimiter限流实战

最近在工作的开发中，遇到一个需要对api接口限流的功能，以防止发生系统被恶意请求攻击，导致应用性能下降，甚至整个服务崩溃的情况。

一、限流算法

常用的限流算法有：漏桶算法和令牌桶算法；
漏桶算法的大致思想是将请求放入一个漏桶中，漏桶以一定的速度来处理请求，当请求过大时漏桶溢出，如下图，不管外部请求速度有多快，都会以一个恒定的速度来处理。

在有的应用场景下，不仅需要限定请求速度，还要求允许某种程度的突发传输，显然漏桶算法并不适用。令牌桶算法的原理是系统会以一个恒定的速度往桶里放入令牌，而如果请求需要被处理，则需要先从桶里获取一个令牌，当桶里没有令牌可取时，则拒绝服务。 令牌桶的另外一个好处是可以方便的改变速度。 一旦需要提高速率，则按需提高放入桶中的令牌的速率，如下图；

二、guava RateLimiter

guava包的RateLimiter类是一种令牌算法的工程实现，经常用于限制对一些物理资源或者逻辑资源的访问速率。它有两种令牌生成方式：一种是稳定的分配令牌，另一种是有一个预热期，在预热器内，每秒分配的令牌数会稳定地增长直到达到稳定的速率。RateLimiter和Java中的信号量(java.util.concurrent.Semaphore)类似，与Semaphore 相比，Semaphore 限制了并发访问的数量而不是使用速率。
RateLimiter主要api方法：

double acquire(); // 阻塞直到获取一个许可，返回被限制的睡眠等待时间，单位秒

double acquire(int permits); // 阻塞直到获取permits个许可，返回被限制的睡眠等待时间，单位秒

boolean tryAcquire();  // 尝试获取一个许可

boolean tryAcquire(int permits);  // 尝试获取permits个许可

boolean tryAcquire(long timeout, TimeUnit unit);  // 尝试获取一个许可，最多等待timeout时间

boolean tryAcquire(int permits, long timeout, TimeUnit unit);  // 尝试获取permits个许可，最多等待timeout时间

demo：
下面的demo是基于spring aop的方式来实现接口限流。仅用于参考使用，具体使用方案可根据具体业务自由发挥。
demo下载地址
封装一个令牌管理类：

public class Shaping {

    private static final ConcurrentMap<String, RateLimiter> resourceLimiterMap = Maps.newConcurrentMap();

    /**
     * 初始化令牌桶
     * @param resource
     * @param qps
     */
    public static void updateResourceQps(String resource, double qps) {
        RateLimiter limiter = resourceLimiterMap.get(resource);
        if (limiter == null) {
            limiter = RateLimiter.create(qps);
            RateLimiter putByOtherThread = resourceLimiterMap.putIfAbsent(resource, limiter);
            if (putByOtherThread != null) {
                limiter = putByOtherThread;
            }
        }
        limiter.setRate(qps);
    }

    /**
     * 尝试获得令牌
     * @param resource
     */
    public static void tryAcquire(String resource){
        RateLimiter limiter = resourceLimiterMap.get(resource);
        if (limiter == null) {
            return;
        }
        if (!limiter.tryAcquire()) {
            throw new RuntimeException(resource+" 接口访问太频繁");
        }
    }

    public static void removeResource(String resource) {
        resourceLimiterMap.remove(resource);
    }

}

系统启动后，使用RequestMappingHandlerMapping遍历接口方法，为每一个接口设置RateLimiter

    @Autowired
    private RequestMappingHandlerMapping handlerMapping;

    /**
     * 为每一个接口初始化一个令牌桶
     */
    @PostConstruct
    private void init() {
        Map<RequestMappingInfo, HandlerMethod> map = handlerMapping.getHandlerMethods();
        for (Map.Entry<RequestMappingInfo, HandlerMethod> entry : map.entrySet()) {
            String mapping = entry.getKey().getPatternsCondition().toString();
            String[] methodPattern = mapping.replaceAll("\\[|\\]", "").split("\\|\\|");
            if (ArrayUtils.isNotEmpty(methodPattern)) {
                for (String m : methodPattern) {
                    Shaping.updateResourceQps(m, 10d);
                }
            }
        }
    }

aop拦截controller调用

@Around("execution(public * cn.com.bucket.controller.TestController.*(..))")
    public Object round(ProceedingJoinPoint jp) {
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
        try {
            //获得获取令牌
            Shaping.tryAcquire(request.getRequestURI());
           return jp.proceed();
        } catch (Throwable throwable) {
            throwable.printStackTrace();
            return throwable.getMessage();
        }
    }

编写一个controller接口

    @RequestMapping("api/test1")
    @ResponseBody
    public Object test1() {
        return "test1........";
    }

至此所有准备工作已经完成。让我们来用一个例子测试：

public static void main(String[] args) {
        CyclicBarrier barrier = new CyclicBarrier(15); //栅栏
        for (int i=0;i<15;i++){
            new Thread(() -> {
                try {
                    barrier.await();//等待15个线程同时开启
                    CloseableHttpClient httpclient = HttpClients.createDefault();
                    HttpPost httppost = new HttpPost("http://localhost:8080/api/test1");
                    CloseableHttpResponse response = httpclient.execute(httppost);
                    System.out.println(EntityUtils.toString(response.getEntity()));
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }).start();

        }
    }

可以看到运行结果中，部分请求获取令牌失败：

test1........
/api/test1 接口访问太频繁
test1........
/api/test1 接口访问太频繁
test1........
test1........
/api/test1 接口访问太频繁
test1........
/api/test1 接口访问太频繁
test1........
test1........
test1........
test1........
test1........
test1........