新增 rateExpression 参数 Spel 表达式支持

taptap · Mar 24, 2021 · c52149f · c52149f
1 parent 9c17977
commit c52149f
Show file tree

Hide file tree

Showing 6 changed files with 45 additions and 5 deletions.
diff --git a/src/main/java/com/taptap/ratelimiter/annotation/RateLimit.java b/src/main/java/com/taptap/ratelimiter/annotation/RateLimit.java
@@ -19,6 +19,12 @@
      */
     long rate();
 
+    /**
+     * 时间窗口流量数量表达式
+     * @return rateExpression
+     */
+    String rateExpression() default "";
+
     /**
      * 时间窗口，最小单位秒，如 2s，2h , 2d
      * @return rateInterval

diff --git a/src/main/java/com/taptap/ratelimiter/core/BizKeyProvider.java b/src/main/java/com/taptap/ratelimiter/core/BizKeyProvider.java
@@ -4,11 +4,16 @@
 import com.taptap.ratelimiter.annotation.RateLimitKey;
 import org.aspectj.lang.JoinPoint;
 import org.aspectj.lang.reflect.MethodSignature;
+import org.springframework.beans.factory.BeanFactory;
+import org.springframework.beans.factory.BeanFactoryAware;
+import org.springframework.beans.factory.config.ConfigurableBeanFactory;
+import org.springframework.context.expression.BeanFactoryResolver;
 import org.springframework.context.expression.MethodBasedEvaluationContext;
 import org.springframework.core.DefaultParameterNameDiscoverer;
 import org.springframework.core.ParameterNameDiscoverer;
 import org.springframework.expression.EvaluationContext;
 import org.springframework.expression.ExpressionParser;
+import org.springframework.expression.common.TemplateParserContext;
 import org.springframework.expression.spel.standard.SpelExpressionParser;
 import org.springframework.expression.spel.support.StandardEvaluationContext;
 import org.springframework.util.ObjectUtils;
@@ -23,12 +28,18 @@
  * @author kl (http://kailing.pub)
  * @since 2021/3/16
  */
-public class BizKeyProvider {
+public class BizKeyProvider implements BeanFactoryAware {
 
     private final ParameterNameDiscoverer nameDiscoverer = new DefaultParameterNameDiscoverer();
 
+    private static final TemplateParserContext PARSER_CONTEXT = new TemplateParserContext();
+
     private final ExpressionParser parser = new SpelExpressionParser();
 
+    private final StandardEvaluationContext evaluationContext = new StandardEvaluationContext();
+
+    private BeanFactory beanFactory;
+
     public String getKeyName(JoinPoint joinPoint, RateLimit rateLimit) {
         Method method = getMethod(joinPoint);
         List<String> definitionKeys = getSpelDefinitionKey(rateLimit.keys(), method, joinPoint.getArgs());
@@ -38,6 +49,17 @@ public String getKeyName(JoinPoint joinPoint, RateLimit rateLimit) {
         return StringUtils.collectionToDelimitedString(keyList,"","-","");
     }
 
+    public Long getRateValue(RateLimit rateLimit){
+        if (StringUtils.hasText(rateLimit.rateExpression())) {
+            String value = parser.parseExpression(resolve(rateLimit.rateExpression()), PARSER_CONTEXT)
+                    .getValue(String.class);
+            if (value != null) {
+                return Long.parseLong(value);
+            }
+        }
+        return rateLimit.rate();
+    }
+
     private Method getMethod(JoinPoint joinPoint) {
         MethodSignature signature = (MethodSignature) joinPoint.getSignature();
         Method method = signature.getMethod();
@@ -82,4 +104,15 @@ private List<String> getParameterKey(Parameter[] parameters, Object[] parameterV
         }
         return parameterKey;
     }
+
+    @SuppressWarnings("NullableProblems")
+    @Override
+    public void setBeanFactory(BeanFactory beanFactory){
+        this.beanFactory = beanFactory;
+        this.evaluationContext.setBeanResolver(new BeanFactoryResolver(beanFactory));
+    }
+
+    private String resolve(String value) {
+        return ((ConfigurableBeanFactory) this.beanFactory).resolveEmbeddedValue(value);
+    }
 }
diff --git a/src/main/java/com/taptap/ratelimiter/core/RateLimiterService.java b/src/main/java/com/taptap/ratelimiter/core/RateLimiterService.java
@@ -37,7 +37,7 @@ RateLimiterInfo getRateLimiterInfo(JoinPoint joinPoint, RateLimit rateLimit) {
                 logger.info("Gets the custom Key exception and degrades it to the default Key:{}", rateLimit, throwable);
             }
         }
-        long rate = rateLimit.rate();
+        long rate = bizKeyProvider.getRateValue(rateLimit);
         long rateInterval = DurationStyle.detectAndParse(rateLimit.rateInterval()).getSeconds();
         return new RateLimiterInfo(rateLimitKey, rate, rateInterval);
     }

diff --git a/src/main/resources/META-INF/ratelimiter-spring-boot-starter-rateLimit.lua b/src/main/resources/META-INF/ratelimiter-spring-boot-starter-rateLimit.lua
@@ -4,7 +4,6 @@
 -- Date: 2021/3/18
 -- Time: 11:17 上午
 -- To change this template use File | Settings | File Templates.
--- 限流器 lua 脚本，只是方便编写，实际读取的脚本不在这里
 local rateLimitKey = KEYS[1];
 local rate = tonumber(KEYS[2]);
 local rateInterval = tonumber(KEYS[3]);

diff --git a/src/test/java/com/taptap/ratelimiter/TestController.java b/src/test/java/com/taptap/ratelimiter/TestController.java
@@ -24,7 +24,7 @@ public String get(String name) {
     }
 
     @GetMapping("/get2")
-    @RateLimit(rate = 2, rateInterval = "10s")
+    @RateLimit(rate = 2, rateInterval = "10s",rateExpression = "${spring.ratelimiter.max}")
     public String get2() {
         return "get";
     }

diff --git a/src/test/resources/application.properties b/src/test/resources/application.properties
@@ -7,4 +7,6 @@ spring.ratelimiter.redis-database=1
 spring.ratelimiter.status-code=429
 spring.ratelimiter.response-body= {"code":429,"msg":"Too Many Requests"}
 
-logging.level.io.netty.resolver.dns = error
+logging.level.io.netty.resolver.dns = error
+
+spring.ratelimiter.max = 1