springboot api限流

项目上有需要对三方调用系统限制调用频率的需求，项目不大没有用微服务，也就没用网关那些，查阅资料，使用谷歌 guava 实现，令牌桶模式。

1、引入包

<dependency>
            <groupId>com.google.guava</groupId>
            <artifactId>guava</artifactId>
            <version>30.1-jre</version>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-aop</artifactId>
        </dependency>

2、自定义注解声明

@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.METHOD})
@Documented
public @interface CurrentLimit {
    /**
     * 资源的key,唯一
     * 作用：不同的接口，不同的流量控制
     */
    String key() default "";

    /**
     * 最多的访问限制次数
     */
    double permitsPerSecond();

    /**
     * 获取令牌最大等待时间
     */
    long timeout();

    /**
     * 获取令牌最大等待时间,单位(例:分钟/秒/毫秒) 默认:毫秒
     */
    TimeUnit timeunit() default TimeUnit.MILLISECONDS;

    /**
     * 得不到令牌的提示语
     */
    String msg() default "操作频繁，请稍后再试！";

}

3、注解功能实现

@Slf4j
@Aspect
@Component
public class CurrentLimitAOP {

    /**
     * 不同的接口，不同的流量控制
     * map的key为 Limiter.key
     */
    private final Map<String, RateLimiter> limitMap = Maps.newConcurrentMap();

    @Around("@annotation(com.leenleda.data.synchronism.dataflow.aop.CurrentLimit)")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        //拿limit的注解
        CurrentLimit currentLimit = method.getAnnotation(CurrentLimit.class);
        if (currentLimit != null) {
            //key作用：不同的接口，不同的流量控制
            String key = currentLimit.key();
            RateLimiter rateLimiter = null;
            //验证缓存是否有命中key
            if (!limitMap.containsKey(key)) {
                // 创建令牌桶
                rateLimiter = RateLimiter.create(currentLimit.permitsPerSecond());
                limitMap.put(key, rateLimiter);
            }
            rateLimiter = limitMap.get(key);
            // 拿令牌
            boolean acquire = rateLimiter.tryAcquire(currentLimit.timeout(), currentLimit.timeunit());
            // 拿不到命令，直接返回异常提示
            if (!acquire) {
                this.responseFail(currentLimit.msg());
                return null;
            }
        }
        return joinPoint.proceed();
    }

    private void responseFail(String msg)  {
        HttpServletResponse response=((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getResponse();
        response.setStatus(HttpStatus.SERVICE_UNAVAILABLE.value());
        response.setContentType(MediaType.APPLICATION_JSON_VALUE);
        response.setCharacterEncoding("UTF-8");
        try {
            response.getWriter().write(msg);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}

4、使用，限制500ms只能有1次请求。

    @PostMapping("/take-data")
    @CurrentLimit(key = "take-data", permitsPerSecond = 1, timeout = 500)
    public void takeData(HttpServletRequest request) {
    
    }

其中返回值部分可以根据需求来返回统一格式。使用过程中，一个方法一个key，看功能实现就能看到。