SpringBoot基于RateLimiter实现请求限流

自定义注解

@Documented
@Inherited
@Target({ElementType.METHOD,ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
public @interface RequestLimit {

    int maxQPS() default 1;
}

切面处理

@Slf4j
@Aspect
@Component
public class RateLimitAspect {

    private final Map<String, RateLimiter> rateLimiterMap = Maps.newConcurrentMap();

    private static final String lock = "lock";

    @Before("@annotation(com.go.relax.core.config.RequestLimit) && " +
            "(@annotation(org.springframework.web.bind.annotation.GetMapping))||" +
            "@annotation(org.springframework.web.bind.annotation.PostMapping))||" +
            "@annotation(org.springframework.web.bind.annotation.PutMapping))||" +
            "@annotation(org.springframework.web.bind.annotation.DeleteMapping))||" +
            "@annotation(org.springframework.web.bind.annotation.RequestMapping))")
    public void checkRateLimit(JoinPoint joinPoint) throws Throwable {
        String className = joinPoint.getTarget().getClass().getCanonicalName();
        String methodName = joinPoint.getSignature().getName();

        RequestLimit requestLimit = ((MethodSignature) joinPoint.getSignature()).getMethod().getAnnotation(RequestLimit.class);
        if (requestLimit == null) {
            return;
        }

        int maxQPS = requestLimit.maxQPS();
        String rateLimitKey = className + "." + methodName;
        RateLimiter rateLimiter = rateLimiterMap.get(rateLimitKey);
        if (rateLimiter == null) {
            synchronized (lock) {
                rateLimiter = rateLimiterMap.get(rateLimitKey);
                if (rateLimiter == null) {
                    rateLimiter = RateLimiter.create(maxQPS);
                    rateLimiterMap.put(rateLimitKey, rateLimiter);
                }
            }
        }
        ServletRequestAttributes servletRequestAttributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        Preconditions.checkArgument(servletRequestAttributes != null, "未获取到请求信息");
        HttpServletRequest httpServletRequest = servletRequestAttributes.getRequest();
        String requestMethod = httpServletRequest.getMethod();
        String requestURI = httpServletRequest.getRequestURI();

        boolean tryAcquire = rateLimiter.tryAcquire(100, TimeUnit.MILLISECONDS);
        if (!tryAcquire) {
            log.warn("请求被限流 {}:{} maxQPS:{}", requestMethod, requestURI, maxQPS);
            throw new RateLimitException();
        }
        log.info("请求未被限流 {}:{} maxQPS:{}", requestMethod, requestURI, maxQPS);
    }
}

接口限流配置

@Slf4j
@RestController
@RequestMapping("/adConfig")
public class AdConfigController {

    @Resource
    private AdConfigManager adConfigManager;

    @RequestLimit
    @GetMapping("/list")
    public Object queryList() {
        return Response.successResult(adConfigManager.queryList());
    }

}