From cfe659ca8f61acddcceef7ee5ba0c149f7537a03 Mon Sep 17 00:00:00 2001
From: lilong <lilong@axzo.cn>
Date: Thu, 21 Mar 2024 17:10:06 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E5=A2=9E=E5=8A=A0=E8=8E=B7=E5=8F=96?=
 =?UTF-8?q?=E4=BB=A4=E7=89=8C=E7=AD=89=E5=BE=85=E8=B6=85=E6=97=B6=E6=96=B9?=
 =?UTF-8?q?=E6=B3=95?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../cn/axzo/pokonyan/client/RateLimiter.java  |  2 ++
 .../client/impl/RedisRateLimiterImpl.java     | 19 +++++++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/src/main/java/cn/axzo/pokonyan/client/RateLimiter.java b/src/main/java/cn/axzo/pokonyan/client/RateLimiter.java
index f93af78..36643cd 100644
--- a/src/main/java/cn/axzo/pokonyan/client/RateLimiter.java
+++ b/src/main/java/cn/axzo/pokonyan/client/RateLimiter.java
@@ -18,6 +18,8 @@ public interface RateLimiter {
      */
     boolean tryAcquire();
 
+    boolean tryAcquire(long timeoutMillis);
+
     /**
      * 获取窗口类型
      *
diff --git a/src/main/java/cn/axzo/pokonyan/client/impl/RedisRateLimiterImpl.java b/src/main/java/cn/axzo/pokonyan/client/impl/RedisRateLimiterImpl.java
index 01d537f..2ea1490 100644
--- a/src/main/java/cn/axzo/pokonyan/client/impl/RedisRateLimiterImpl.java
+++ b/src/main/java/cn/axzo/pokonyan/client/impl/RedisRateLimiterImpl.java
@@ -12,6 +12,7 @@ import org.redisson.api.RateType;
 import org.redisson.api.RedissonClient;
 
 import java.util.Objects;
+import java.util.concurrent.TimeUnit;
 
 @Slf4j
 public class RedisRateLimiterImpl implements RateLimiter {
@@ -26,6 +27,12 @@ public class RedisRateLimiterImpl implements RateLimiter {
 
     private RRateLimiter rateLimiter;
 
+    /**
+     * 默认超时时间（毫秒）
+     */
+    private static final long DEFAULT_TIME_OUT_MILLIS = 5 * 1000;
+
+
     @Builder
     RedisRateLimiterImpl(RedissonClient redissonClient,
                          WindowType windowType,
@@ -56,6 +63,11 @@ public class RedisRateLimiterImpl implements RateLimiter {
         return rateLimiterWorker.tryAcquire();
     }
 
+    @Override
+    public boolean tryAcquire(long timeoutMillis) {
+        return rateLimiterWorker.tryAcquire(timeoutMillis);
+    }
+
     @Override
     public WindowType getWindowType() {
         return windowType;
@@ -89,6 +101,11 @@ public class RedisRateLimiterImpl implements RateLimiter {
         public boolean tryAcquire() {
             return rateLimiter.tryAcquire(1);
         }
+
+        @Override
+        public boolean tryAcquire(long timeoutMillis) {
+            return rateLimiter.tryAcquire(1, timeoutMillis, TimeUnit.MILLISECONDS);
+        }
     }
 
     interface RateLimiterWorker {
@@ -98,5 +115,7 @@ public class RedisRateLimiterImpl implements RateLimiter {
          * @return 如果获取成功则返回true, 失败则为false
          */
         boolean tryAcquire();
+
+        boolean tryAcquire(long timeoutMillis);
     }
 }