From 0bb3b90dec7cb838b00d8f915938e0cf027e5302 Mon Sep 17 00:00:00 2001 From: Manish M Demblani Date: Sat, 20 Sep 2025 00:04:45 +0400 Subject: [PATCH] enhance RateLimitException to return extra information --- .../starter/context/RateLimitException.java | 22 +++++++++++++++++++ .../config/aspect/RateLimitAspect.java | 20 ++++++++++++----- .../method/method/MethodRateLimitTest.java | 15 ++++++++----- 3 files changed, 46 insertions(+), 11 deletions(-) diff --git a/bucket4j-spring-boot-starter-context/src/main/java/com/giffing/bucket4j/spring/boot/starter/context/RateLimitException.java b/bucket4j-spring-boot-starter-context/src/main/java/com/giffing/bucket4j/spring/boot/starter/context/RateLimitException.java index 96f6e74a..0d1f87bc 100644 --- a/bucket4j-spring-boot-starter-context/src/main/java/com/giffing/bucket4j/spring/boot/starter/context/RateLimitException.java +++ b/bucket4j-spring-boot-starter-context/src/main/java/com/giffing/bucket4j/spring/boot/starter/context/RateLimitException.java @@ -5,4 +5,26 @@ * {@link RateLimiting} annotation. */ public class RateLimitException extends RuntimeException { + private final long retryAfterNanoSeconds; + private final long remainingTokens; + private final String configurationName; + + public RateLimitException(Long retryAfterNanoSeconds, Long remainingTokens, String configurationName) { + super("Rate limit exceeded for configuration: " + configurationName + ". Retry after: " + (retryAfterNanoSeconds != null ? retryAfterNanoSeconds : 0) + "ns. Remaining tokens: " + (remainingTokens != null ? remainingTokens : 0)); + this.retryAfterNanoSeconds = retryAfterNanoSeconds != null ? retryAfterNanoSeconds : 0; + this.remainingTokens = remainingTokens != null ? remainingTokens : 0; + this.configurationName = configurationName; + } + + public long getRetryAfterNanoSeconds() { + return retryAfterNanoSeconds; + } + + public long getRemainingTokens() { + return remainingTokens; + } + + public String getConfigurationName() { + return configurationName; + } } diff --git a/bucket4j-spring-boot-starter/src/main/java/com/giffing/bucket4j/spring/boot/starter/config/aspect/RateLimitAspect.java b/bucket4j-spring-boot-starter/src/main/java/com/giffing/bucket4j/spring/boot/starter/config/aspect/RateLimitAspect.java index c9032423..ac733b4f 100644 --- a/bucket4j-spring-boot-starter/src/main/java/com/giffing/bucket4j/spring/boot/starter/config/aspect/RateLimitAspect.java +++ b/bucket4j-spring-boot-starter/src/main/java/com/giffing/bucket4j/spring/boot/starter/config/aspect/RateLimitAspect.java @@ -8,6 +8,7 @@ import com.giffing.bucket4j.spring.boot.starter.context.properties.RateLimit; import com.giffing.bucket4j.spring.boot.starter.service.RateLimitService; import com.giffing.bucket4j.spring.boot.starter.utils.RateLimitAopUtils; +import com.giffing.bucket4j.spring.boot.starter.context.RateLimitException; import jakarta.annotation.PostConstruct; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; @@ -125,7 +126,11 @@ public Object processMethodsAnnotatedWithRateLimitAnnotation(ProceedingJoinPoint } else if (fallbackMethod != null) { return fallbackMethod.invoke(joinPoint.getTarget(), joinPoint.getArgs()); } else { - throw new RateLimitException(); + throw new RateLimitException( + consumedResult.retryAfterNanoSeconds, + consumedResult.remainingLimit, + rateLimitAnnotation.name() + ); } return methodResult; @@ -144,13 +149,14 @@ private static void performPostRateLimit(RateLimitService.RateLimitConfigresult< private static RateLimitConsumedResult performRateLimit(RateLimitService.RateLimitConfigresult rateLimitConfigResult, Method method, Map params, RateLimit annotationRateLimit) { boolean allConsumed = true; Long remainingLimit = null; + Long retryAfterNanoSeconds = null; for (RateLimitCheck rl : rateLimitConfigResult.getRateLimitChecks()) { var wrapper = rl.rateLimit(new ExpressionParams<>(method).addParams(params), annotationRateLimit); if (wrapper != null && wrapper.getRateLimitResult() != null) { var rateLimitResult = wrapper.getRateLimitResult(); - if (rateLimitResult.isConsumed()) { - remainingLimit = RateLimitService.getRemainingLimit(remainingLimit, rateLimitResult); - } else { + remainingLimit = RateLimitService.getRemainingLimit(remainingLimit, rateLimitResult); + retryAfterNanoSeconds = rateLimitResult.getNanosToWaitForRefill(); + if (!rateLimitResult.isConsumed()) { allConsumed = false; break; } @@ -158,11 +164,13 @@ private static RateLimitConsumedResult performRateLimit(RateLimitService.RateLim } if (allConsumed) { log.debug("rate-limit-remaining;limit:{}", remainingLimit); + } else { + log.debug("rate-limit-consumed;limit:{};rate-limit-retry-after;limit:{}", remainingLimit, retryAfterNanoSeconds); } - return new RateLimitConsumedResult(allConsumed, remainingLimit); + return new RateLimitConsumedResult(allConsumed, remainingLimit, retryAfterNanoSeconds); } - private record RateLimitConsumedResult(boolean allConsumed, Long remainingLimit) { + private record RateLimitConsumedResult(boolean allConsumed, Long remainingLimit, Long retryAfterNanoSeconds) { } /* diff --git a/examples/general-tests/src/main/java/com/giffing/bucket4j/spring/boot/starter/general/tests/method/method/MethodRateLimitTest.java b/examples/general-tests/src/main/java/com/giffing/bucket4j/spring/boot/starter/general/tests/method/method/MethodRateLimitTest.java index cdb68d1a..8bfc15f6 100644 --- a/examples/general-tests/src/main/java/com/giffing/bucket4j/spring/boot/starter/general/tests/method/method/MethodRateLimitTest.java +++ b/examples/general-tests/src/main/java/com/giffing/bucket4j/spring/boot/starter/general/tests/method/method/MethodRateLimitTest.java @@ -80,14 +80,19 @@ public void assert_no_rate_limit_with_skip_condition_matches() { @Test public void assert_rate_limit_with_cache_key() { - for(int i = 0; i < 5; i++) { - // rate limit by parameter value + for (int i = 0; i < 5; i++) { testService.withCacheKey("key1"); testService.withCacheKey("key2"); - // all tokens consumed } - assertThrows(RateLimitException.class, () -> testService.withCacheKey("key1")); - assertThrows(RateLimitException.class, () -> testService.withCacheKey("key2")); + RateLimitException ex1 = assertThrows(RateLimitException.class, () -> testService.withCacheKey("key1")); + assertTrue(ex1.getRetryAfterNanoSeconds() >= 0); + assertTrue(ex1.getRemainingTokens() >= 0); + assertNotNull(ex1.getConfigurationName()); + + RateLimitException ex2 = assertThrows(RateLimitException.class, () -> testService.withCacheKey("key2")); + assertTrue(ex2.getRetryAfterNanoSeconds() >= 0); + assertTrue(ex2.getRemainingTokens() >= 0); + assertNotNull(ex2.getConfigurationName()); } @Test