使用 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