✨ 使用 Guava 限流实现

xkcoding · xkcoding · commit ae4dd40327ab · 2019-09-12T15:28:58.000+08:00
diff --git a/spring-boot-demo-ratelimit-guava/pom.xml b/spring-boot-demo-ratelimit-guava/pom.xml
@@ -28,11 +28,32 @@
       <artifactId>spring-boot-starter-web</artifactId>
     </dependency>
 
+    <dependency>
+      <groupId>org.springframework.boot</groupId>
+      <artifactId>spring-boot-starter-aop</artifactId>
+    </dependency>
+
+    <dependency>
+      <groupId>cn.hutool</groupId>
+      <artifactId>hutool-all</artifactId>
+    </dependency>
+
+    <dependency>
+      <groupId>com.google.guava</groupId>
+      <artifactId>guava</artifactId>
+    </dependency>
+
     <dependency>
       <groupId>org.springframework.boot</groupId>
       <artifactId>spring-boot-starter-test</artifactId>
       <scope>test</scope>
     </dependency>
+
+    <dependency>
+      <groupId>org.projectlombok</groupId>
+      <artifactId>lombok</artifactId>
+      <optional>true</optional>
+    </dependency>
   </dependencies>
 
   <build>
diff --git a/spring-boot-demo-ratelimit-guava/src/main/java/com/xkcoding/ratelimit/guava/annotation/RateLimiter.java b/spring-boot-demo-ratelimit-guava/src/main/java/com/xkcoding/ratelimit/guava/annotation/RateLimiter.java
@@ -0,0 +1,43 @@
+package com.xkcoding.ratelimit.guava.annotation;
+
+import org.springframework.core.annotation.AliasFor;
+import org.springframework.core.annotation.AnnotationUtils;
+
+import java.lang.annotation.*;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * <p>
+ * 限流注解，添加了 {@link AliasFor} 必须通过 {@link AnnotationUtils} 获取，才会生效
+ *
+ * @author yangkai.shen
+ * @date Created in 2019/9/12 14:14
+ * @see AnnotationUtils
+ * </p>
+ */
+@Target(ElementType.METHOD)
+@Retention(RetentionPolicy.RUNTIME)
+@Documented
+public @interface RateLimiter {
+    int NOT_LIMITED = 0;
+
+    /**
+     * qps
+     */
+    @AliasFor("qps") double value() default NOT_LIMITED;
+
+    /**
+     * qps
+     */
+    @AliasFor("value") double qps() default NOT_LIMITED;
+
+    /**
+     * 超时时长
+     */
+    int timeout() default 0;
+
+    /**
+     * 超时时间单位
+     */
+    TimeUnit timeUnit() default TimeUnit.MILLISECONDS;
+}
diff --git a/spring-boot-demo-ratelimit-guava/src/main/java/com/xkcoding/ratelimit/guava/aspect/RateLimiterAspect.java b/spring-boot-demo-ratelimit-guava/src/main/java/com/xkcoding/ratelimit/guava/aspect/RateLimiterAspect.java
@@ -0,0 +1,52 @@
+package com.xkcoding.ratelimit.guava.aspect;
+
+import com.xkcoding.ratelimit.guava.annotation.RateLimiter;
+import lombok.extern.slf4j.Slf4j;
+import org.aspectj.lang.ProceedingJoinPoint;
+import org.aspectj.lang.annotation.Around;
+import org.aspectj.lang.annotation.Aspect;
+import org.aspectj.lang.annotation.Pointcut;
+import org.aspectj.lang.reflect.MethodSignature;
+import org.springframework.core.annotation.AnnotationUtils;
+import org.springframework.stereotype.Component;
+
+import java.lang.reflect.Method;
+
+/**
+ * <p>
+ * 限流切面
+ * </p>
+ *
+ * @author yangkai.shen
+ * @date Created in 2019/9/12 14:27
+ */
+@Slf4j
+@Aspect
+@Component
+public class RateLimiterAspect {
+    private static final com.google.common.util.concurrent.RateLimiter RATE_LIMITER = com.google.common.util.concurrent.RateLimiter.create(Double.MAX_VALUE);
+
+    @Pointcut("@annotation(com.xkcoding.ratelimit.guava.annotation.RateLimiter)")
+    public void rateLimit() {
+
+    }
+
+    @Around("rateLimit()")
+    public Object pointcut(ProceedingJoinPoint point) throws Throwable {
+        MethodSignature signature = (MethodSignature) point.getSignature();
+        Method method = signature.getMethod();
+        // 通过 AnnotationUtils.findAnnotation 获取 RateLimiter 注解
+        RateLimiter rateLimiter = AnnotationUtils.findAnnotation(method, RateLimiter.class);
+        if (rateLimiter != null && rateLimiter.qps() > RateLimiter.NOT_LIMITED) {
+            double qps = rateLimiter.qps();
+            log.debug("【{}】的QPS设置为: {}", method.getName(), qps);
+            // 重新设置 QPS
+            RATE_LIMITER.setRate(qps);
+            // 尝试获取令牌
+            if (!RATE_LIMITER.tryAcquire(rateLimiter.timeout(), rateLimiter.timeUnit())) {
+                throw new RuntimeException("手速太快了，慢点儿吧~");
+            }
+        }
+        return point.proceed();
+    }
+}
diff --git a/spring-boot-demo-ratelimit-guava/src/main/java/com/xkcoding/ratelimit/guava/controller/TestController.java b/spring-boot-demo-ratelimit-guava/src/main/java/com/xkcoding/ratelimit/guava/controller/TestController.java
@@ -0,0 +1,33 @@
+package com.xkcoding.ratelimit.guava.controller;
+
+import cn.hutool.core.lang.Dict;
+import com.xkcoding.ratelimit.guava.annotation.RateLimiter;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.RestController;
+
+/**
+ * <p>
+ * 测试
+ * </p>
+ *
+ * @author yangkai.shen
+ * @date Created in 2019/9/12 14:22
+ */
+@Slf4j
+@RestController
+public class TestController {
+
+    @RateLimiter(value = 1.0, timeout = 300)
+    @GetMapping("/test1")
+    public Dict test1() {
+        log.info("【test1】被执行了。。。。。");
+        return Dict.create().set("msg", "hello,world!").set("description", "别想一直看到我，不信你快速刷新看看~");
+    }
+
+    @GetMapping("/test2")
+    public Dict test2() {
+        log.info("【test2】被执行了。。。。。");
+        return Dict.create().set("msg", "hello,world!").set("description", "我一直都在，卟离卟弃");
+    }
+}
diff --git a/spring-boot-demo-ratelimit-guava/src/main/java/com/xkcoding/ratelimit/guava/handler/GlobalExceptionHandler.java b/spring-boot-demo-ratelimit-guava/src/main/java/com/xkcoding/ratelimit/guava/handler/GlobalExceptionHandler.java
@@ -0,0 +1,22 @@
+package com.xkcoding.ratelimit.guava.handler;
+
+import cn.hutool.core.lang.Dict;
+import org.springframework.web.bind.annotation.ExceptionHandler;
+import org.springframework.web.bind.annotation.RestControllerAdvice;
+
+/**
+ * <p>
+ * 全局异常拦截
+ * </p>
+ *
+ * @author yangkai.shen
+ * @date Created in 2019/9/12 15:00
+ */
+@RestControllerAdvice
+public class GlobalExceptionHandler {
+
+    @ExceptionHandler(RuntimeException.class)
+    public Dict handler(RuntimeException ex) {
+        return Dict.create().set("msg", ex.getMessage());
+    }
+}
diff --git a/spring-boot-demo-ratelimit-guava/src/main/resources/application.yml b/spring-boot-demo-ratelimit-guava/src/main/resources/application.yml
@@ -2,3 +2,6 @@ server:
   port: 8080
   servlet:
     context-path: /demo
+logging:
+  level:
+    com.xkcoding: debug
