Merge pull request #502 from KamToHung/fix-501

Redis限流支持cluster模式
dromara · Nov 11, 2024 · bb30df3 · bb30df3
2 parents 8ae462b + 212ff6a
commit bb30df3
Show file tree

Hide file tree

Showing 9 changed files with 165 additions and 44 deletions.
diff --git a/...a/org/dromara/dynamictp/extension/limiter/redis/ratelimiter/AbstractRedisRateLimiter.java b/...a/org/dromara/dynamictp/extension/limiter/redis/ratelimiter/AbstractRedisRateLimiter.java
@@ -23,7 +23,6 @@
 import org.springframework.data.redis.core.script.RedisScript;
 import org.springframework.scripting.support.ResourceScriptSource;
 
-import java.time.Instant;
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
@@ -61,22 +60,12 @@ public RedisScript<List<Long>> getScript() {
         return script;
     }
 
-    @Override
-    public List<String> getKeys(final String key) {
-        String cacheKey = PREFIX + ":" + key;
-        return Collections.singletonList(cacheKey);
-    }
-
-    @Override
-    public List<Long> isAllowed(String key, long windowSize, int limit) {
+    public List<Object> isAllowed(String key, long windowSize, int limit) {
         RedisScript<?> script = this.getScript();
         List<String> keys = this.getKeys(key);
-
+        String[] values = this.getArgs(key, windowSize, limit);
         return Collections.unmodifiableList((List) Objects.requireNonNull(stringRedisTemplate.execute(script, keys,
-                doubleToString(windowSize), doubleToString(limit), doubleToString(Instant.now().getEpochSecond()))));
+                values)));
     }
 
-    private String doubleToString(final double param) {
-        return String.valueOf(param);
-    }
 }
diff --git a/...g/dromara/dynamictp/extension/limiter/redis/ratelimiter/NotifyRedisRateLimiterFilter.java b/...g/dromara/dynamictp/extension/limiter/redis/ratelimiter/NotifyRedisRateLimiterFilter.java
@@ -18,8 +18,6 @@
 package org.dromara.dynamictp.extension.limiter.redis.ratelimiter;
 
 import lombok.extern.slf4j.Slf4j;
-import lombok.val;
-import org.apache.commons.collections4.CollectionUtils;
 import org.dromara.dynamictp.common.pattern.filter.Invoker;
 import org.dromara.dynamictp.core.notifier.chain.filter.NotifyFilter;
 import org.dromara.dynamictp.core.notifier.context.BaseNotifyCtx;
@@ -36,8 +34,6 @@
 @Slf4j
 public class NotifyRedisRateLimiterFilter implements NotifyFilter {
 
-    public static final int LUA_RES_REMAIN_INDEX = 2;
-
     @Resource
     private RedisRateLimiter<List<Long>> redisScriptRateLimiter;
 
@@ -49,27 +45,10 @@ public int getOrder() {
     @Override
     public void doFilter(BaseNotifyCtx context, Invoker<BaseNotifyCtx> nextFilter) {
         String notifyName = context.getExecutorWrapper().getThreadPoolName() + ":" + context.getNotifyItemEnum().getValue();
-        boolean checkResult = check(notifyName, context.getNotifyItem().getClusterLimit(),
-                context.getNotifyItem().getInterval());
+        boolean checkResult = redisScriptRateLimiter.check(notifyName, context.getNotifyItem().getInterval(), context.getNotifyItem().getClusterLimit());
         if (checkResult) {
             nextFilter.invoke(context);
         }
     }
 
-    private boolean check(String notifyName, int limit, long interval) {
-        try {
-            val res = redisScriptRateLimiter.isAllowed(notifyName, interval, limit);
-            if (CollectionUtils.isEmpty(res)) {
-                return true;
-            }
-            if (res.get(LUA_RES_REMAIN_INDEX) <= 0) {
-                log.debug("DynamicTp notify, trigger redis rate limit, limitKey:{}", res.get(0));
-                return false;
-            }
-            return true;
-        } catch (Exception e) {
-            log.error("DynamicTp notify, redis rate limit check failed, limitKey:{}", notifyName, e);
-            return true;
-        }
-    }
 }
diff --git a/...main/java/org/dromara/dynamictp/extension/limiter/redis/ratelimiter/RedisRateLimiter.java b/...main/java/org/dromara/dynamictp/extension/limiter/redis/ratelimiter/RedisRateLimiter.java
@@ -46,12 +46,23 @@ public interface RedisRateLimiter<T> {
     List<String> getKeys(String key);
 
     /**
-     * If allowed.
+     * Get args.
      *
      * @param key the key
      * @param windowSize the window size
      * @param limit the limit
+     * @return the args
+     */
+    String[] getArgs(String key, long windowSize, int limit);
+
+    /**
+     * check.
+     *
+     * @param name the key
+     * @param interval the interval
+     * @param limit the limit
      * @return the result
      */
-    T isAllowed(String key, long windowSize, int limit);
+    boolean check(String name, long interval, int limit);
+
 }
diff --git a/...a/org/dromara/dynamictp/extension/limiter/redis/ratelimiter/SlidingWindowRateLimiter.java b/...a/org/dromara/dynamictp/extension/limiter/redis/ratelimiter/SlidingWindowRateLimiter.java
@@ -17,29 +17,64 @@
 
 package org.dromara.dynamictp.extension.limiter.redis.ratelimiter;
 
+import lombok.extern.slf4j.Slf4j;
+import lombok.val;
+import org.apache.commons.collections4.CollectionUtils;
 import org.dromara.dynamictp.common.util.CommonUtil;
 import org.dromara.dynamictp.extension.limiter.redis.em.RateLimitEnum;
 import org.springframework.data.redis.core.StringRedisTemplate;
 
-import java.util.Arrays;
+import java.time.Instant;
+import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 
 /**
  * SlidingWindowRateLimiter related
  *
  * @author yanhom
  * @since 1.0.8
  **/
+@Slf4j
 public class SlidingWindowRateLimiter extends AbstractRedisRateLimiter {
 
+    public static final int LUA_RES_REMAIN_INDEX = 2;
+
     public SlidingWindowRateLimiter(StringRedisTemplate stringRedisTemplate) {
         super(RateLimitEnum.SLIDING_WINDOW.getScriptName(), stringRedisTemplate);
     }
 
     @Override
     public List<String> getKeys(final String key) {
         String cacheKey = CommonUtil.getInstance().getServiceName() + ":" + PREFIX + ":" + key;
+        return Collections.singletonList(cacheKey);
+    }
+
+    @Override
+    public String[] getArgs(String key, long windowSize, int limit) {
         String memberKey = CommonUtil.getInstance().getIp() + ":" + COUNTER.incrementAndGet();
-        return Arrays.asList(cacheKey, memberKey);
+        return new String[]{doubleToString(windowSize), doubleToString(limit), doubleToString(Instant.now().getEpochSecond()), memberKey};
     }
+
+    public boolean check(String name, long interval, int limit) {
+        try {
+            val res = isAllowed(name, interval, limit);
+            if (CollectionUtils.isEmpty(res)) {
+                return true;
+            }
+            if (Objects.isNull(res.get(LUA_RES_REMAIN_INDEX)) || (long) res.get(LUA_RES_REMAIN_INDEX) <= 0) {
+                log.debug("DynamicTp notify, trigger redis rate limit, limitKey:{}", res.get(0));
+                return false;
+            }
+            return true;
+        } catch (Exception e) {
+            log.error("DynamicTp notify, redis rate limit check failed, limitKey:{}", name, e);
+            return true;
+        }
+    }
+
+    private String doubleToString(final double param) {
+        return String.valueOf(param);
+    }
+
 }
diff --git a/extension/extension-limiter-redis/src/main/resources/scripts/sliding_window_rate_limiter.lua b/extension/extension-limiter-redis/src/main/resources/scripts/sliding_window_rate_limiter.lua
@@ -1,9 +1,9 @@
 local key = KEYS[1]
-local member = KEYS[2]
 
 local window_size = tonumber(ARGV[1])
 local limit = tonumber(ARGV[2])
 local timestamp = tonumber(ARGV[3])
+local member = ARGV[4]
 
 local accepted = 0
 local exists_key = redis.call('exists', key)

diff --git a/test/test-extension/pom.xml b/test/test-extension/pom.xml
@@ -3,13 +3,13 @@
     <modelVersion>4.0.0</modelVersion>
     <parent>
         <groupId>org.dromara.dynamictp</groupId>
-        <artifactId>dynamic-tp-all</artifactId>
+        <artifactId>dynamic-tp-test</artifactId>
         <version>${revision}</version>
-        <relativePath>../../pom.xml</relativePath>
+        <relativePath>../pom.xml</relativePath>
     </parent>
 
     <artifactId>dynamic-tp-test-extension</artifactId>
-    <packaging>jar</packaging>
+    <packaging>pom</packaging>
 
     <name>test-extension</name>
 

diff --git a/test/test-extension/test-extension-limiter-redis/pom.xml b/test/test-extension/test-extension-limiter-redis/pom.xml
@@ -0,0 +1,46 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+    <parent>
+        <groupId>org.dromara.dynamictp</groupId>
+        <artifactId>dynamic-tp-test-extension</artifactId>
+        <version>1.1.9.1</version>
+        <relativePath>../pom.xml</relativePath>
+    </parent>
+
+    <artifactId>test-extension-limiter-redis</artifactId>
+
+    <properties>
+        <maven.compiler.source>8</maven.compiler.source>
+        <maven.compiler.target>8</maven.compiler.target>
+        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
+    </properties>
+
+    <dependencies>
+        <dependency>
+            <groupId>org.dromara.dynamictp</groupId>
+            <artifactId>dynamic-tp-spring-boot-starter-common</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.springframework</groupId>
+            <artifactId>spring-aspects</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.dromara.dynamictp</groupId>
+            <artifactId>dynamic-tp-spring-boot-starter-extension-limiter-redis</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-pool2</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.springframework.boot</groupId>
+            <artifactId>spring-boot-starter-data-redis</artifactId>
+        </dependency>
+    </dependencies>
+
+</project>
diff --git a/...imiter-redis/src/test/java/org/dromara/dynamictp/test/extension/RedisRateLimiterTest.java b/...imiter-redis/src/test/java/org/dromara/dynamictp/test/extension/RedisRateLimiterTest.java
@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.dromara.dynamictp.test.extension;
+
+import lombok.val;
+import org.dromara.dynamictp.extension.limiter.redis.ratelimiter.RedisRateLimiter;
+import org.dromara.dynamictp.spring.EnableDynamicTp;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.test.context.junit.jupiter.SpringExtension;
+
+import javax.annotation.Resource;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+
+@SpringBootTest(classes = RedisRateLimiterTest.class)
+@ExtendWith(SpringExtension.class)
+@EnableAutoConfiguration
+@EnableDynamicTp
+class RedisRateLimiterTest {
+
+    @Resource
+    private RedisRateLimiter<List<Long>> redisScriptRateLimiter;
+
+    @Test
+    void testRedisRateLimiterCheck() throws InterruptedException {
+        for (int i = 0; i < 6; i++) {
+            TimeUnit.SECONDS.sleep(1);
+            val res = redisScriptRateLimiter.check("rate-limiter", 120, 5);
+            System.out.println(res);
+        }
+    }
+
+}
+
+
+
diff --git a/test/test-extension/test-extension-limiter-redis/src/test/resources/application.yml b/test/test-extension/test-extension-limiter-redis/src/test/resources/application.yml
@@ -0,0 +1,7 @@
+spring:
+  redis:
+    cluster:
+      nodes:
+        - 127.0.0.1:6379
+        - 127.0.0.1:6379
+        - 127.0.0.1:6379