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，看功能實現就能看到。