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........