使用 bucket4j 进行限流

使用 bucket4j 进行限流

使用 bucket4j 进行限流。

bucket4j

介绍

TODO: 令牌桶、算法

用途

TODO

实现

基于注解

参考:Pixiv-Illustration-Collection-Backend

Maven 依赖:

<dependency>
    <groupId>com.github.vladimir-bukhtoyarov</groupId>
    <artifactId>bucket4j-core</artifactId>
    <version>4.10.0</version>
</dependency>
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-aop</artifactId>
</dependency>

注解

@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface RateLimit {
}

具体实现

@Aspect
@Order(1)
@Component
public class RateLimitProcessor implements HandlerInterceptor {
    private static final Logger logger = LoggerFactory.getLogger(RateLimitProcessor.class);

    /**
     * 最大容积为 2,每三秒生成一个令牌
     */
    private Bucket freeBucket = Bucket4j.builder().addLimit(Bandwidth.classic(2, Refill.intervally(1,
            Duration.ofSeconds(3)))).build();

    /**
     * 在使用 RateLimit 注解的方法上生效
     */
    @Pointcut(value = "@annotation(xyz.bolitao.asynclearn.aop.limit.RateLimit)")
    public void pointCut() {
    }

    @Around(value = "pointCut()")
    public Object handleLimit(ProceedingJoinPoint joinPoint) throws Throwable {
        Bucket requestBucket = this.freeBucket;
        ConsumptionProbe consumptionProbe = requestBucket.tryConsumeAndReturnRemaining(1);
        if (consumptionProbe.isConsumed()) {
            return joinPoint.proceed();
        }
        return ResponseEntity.status(HttpStatus.TOO_MANY_REQUESTS).body(new RateLimitException(HttpStatus.TOO_MANY_REQUESTS, "请求太频繁"));
    }
}

异常处理

基础异常类 BaseException

public class BaseException extends RuntimeException {
    private HttpStatus httpStatus;
    private String message;

    public BaseException() {
    }

    public BaseException(HttpStatus httpStatus, String message) {
        this.httpStatus = httpStatus;
        this.message = message;
    }
    // ... 省略 setter、getter、equals 等
}

限流异常类:

public class RateLimitException extends BaseException {
    private HttpStatus httpStatus;
    private String message;

    public RateLimitException(HttpStatus httpStatus, String message) {
        this.httpStatus = httpStatus;
        this.message = message;
    }
    // ... 省略部分
}

Controller

@RestController
public class LimitController {
    @RateLimit
    @GetMapping("/testLimit")
    public ResponseEntity<Object> testLimit() {
        return ResponseEntity.ok().body(new Date());
    }
}

测试

基于 Filter

参考:bucket4j/basic-usage.md at master · vladimir-bukhtoyarov/bucket4j

加载评论