Bucket4j 限流

最近需要对用户进行限流，经过一番查找，发现 Bucket4j 官网推荐的第三方文章有一篇详细的分布式限流讲解 如何在 Java 中通过 Bucket4j 提供速率限制 - DZone Java，分布式暂时还用不上，因此参考着写了一个单机的。

依赖

<dependency>
    <groupId>com.bucket4j</groupId>
    <artifactId>bucket4j_jdk8-caffeine</artifactId>
    <version>8.0.1</version>
</dependency>
<dependency>
    <groupId>com.github.ben-manes.caffeine</groupId>
    <artifactId>caffeine</artifactId>
    <version>3.1.1</version>
</dependency>

实现

限流要满足在 m 个单位时间段仅允许 n 次访问，因为限流注解设计如下（即 time 个 unit 内限制最多 limit 次访问）：

import java.lang.annotation.*;

@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
@Documented
public @interface RateLimiter {
    TimeUnit unit();

    long time();

    long limit();
}

接着就是使用 Spring 拦截器进行限流

public class RateLimiterHandlerInterceptor implements HandlerInterceptor {
    private final ProxyManager<String> proxyManager;

    public RateLimiterHandlerInterceptor() {
        Caffeine<String, RemoteBucketState> builder = Caffeine.newBuilder()
                .removalListener((key, graph, cause) -> {});
        proxyManager = new CaffeineProxyManager<>(builder, Duration.ofMinutes(10));
    }

    @Override
    public boolean preHandle(@NonNull HttpServletRequest request, @NonNull HttpServletResponse response, @NonNull Object handler) {
        if (handler instanceof HandlerMethod) {
            HandlerMethod handlerMethod = (HandlerMethod) handler;
            RateLimiter rateLimiter = handlerMethod.getMethodAnnotation(RateLimiter.class);;
            if (rateLimiter != null) {
                String key = "";
                Bucket bucket = proxyManager.builder().build(key, () -> bucketConfigurationByRateLimiter(rateLimiter));
                if (!bucket.tryConsume(1)) {
                    response.setStatus(429);
                    return false;
                }
            }
        }
        return true;
    }
    
    private BucketConfiguration bucketConfigurationByRateLimiter(RateLimiter rateLimiter) {
        TimeUnit timeUnit = rateLimiter.unit();
        long time = rateLimiter.time();
        long limit = rateLimiter.limit();
        Bandwidth bandwidth = Bandwidth.simple(limit, TimeUnit.SECONDS.equals(timeUnit) ? Duration.ofSeconds(time) : Duration.ofMinutes(time));
        return BucketConfiguration.builder()
        	.addLimit(bandwidth)
       		.build();
    }
}

上面的 key （也可以是一个对象，注意重写 equals 和 hashcode）就是限流的粒度，对用户限流就更改为用户 id，对接口限流就更改为接口全限定名，另外对复杂接口也可以用策略模式来获取 key。

接着注册拦截器后在需要限流的接口加上 @RateLimiter 就可以了。

总结

上面只介绍了一个最简单的限流方法，开头的分布式限流文章中作者也介绍了对一个接口多个粒度的限流方法，强烈推荐！

如果限流的场景不复杂的话也可以试试现成的 Starter MarcGiffing/bucket4j-spring-boot-starter: Spring Boot Starter for Bucket4j (github.com)