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); } }