Bucket4j: Spring Boot Manual Integration

<dependencyManagement>
    <dependencies>
        <dependency>
            <groupId>com.bucket4j</groupId>
            <artifactId>bucket4j_jdk17-bom</artifactId>
            <version>8.21.0</version>
            <type>pom</type>
            <scope>import</scope>
        </dependency>
        <!-- import this before spring-cloud-dependencies, which manages Resilience4j 2.3.0 -->
        <dependency>
            <groupId>io.github.resilience4j</groupId>
            <artifactId>resilience4j-bom</artifactId>
            <version>2.4.0</version>
            <type>pom</type>
            <scope>import</scope>
        </dependency>
    </dependencies>
</dependencyManagement>

<dependencies>
    <dependency><groupId>com.bucket4j</groupId><artifactId>bucket4j_jdk17-core</artifactId></dependency>
    <dependency><groupId>com.bucket4j</groupId><artifactId>bucket4j_jdk17-lettuce</artifactId></dependency>
    <dependency><groupId>com.bucket4j</groupId><artifactId>bucket4j_jdk17-postgresql</artifactId></dependency>
    <dependency><groupId>com.bucket4j</groupId><artifactId>bucket4j_jdk17-caffeine</artifactId></dependency>

    <dependency><groupId>org.springframework.boot</groupId><artifactId>spring-boot-starter-webmvc</artifactId></dependency>
    <dependency><groupId>org.springframework.boot</groupId><artifactId>spring-boot-starter-security</artifactId></dependency>
    <dependency><groupId>org.springframework.boot</groupId><artifactId>spring-boot-starter-data-redis</artifactId></dependency>
    <dependency><groupId>org.springframework.boot</groupId><artifactId>spring-boot-starter-jdbc</artifactId></dependency>
    <dependency><groupId>org.springframework.boot</groupId><artifactId>spring-boot-starter-flyway</artifactId></dependency>
    <dependency><groupId>org.flywaydb</groupId><artifactId>flyway-database-postgresql</artifactId></dependency>
    <dependency><groupId>org.springframework.boot</groupId><artifactId>spring-boot-starter-validation</artifactId></dependency>
    <dependency><groupId>org.springframework.boot</groupId><artifactId>spring-boot-starter-aspectj</artifactId></dependency>
    <dependency><groupId>io.github.resilience4j</groupId><artifactId>resilience4j-spring-boot4</artifactId><version>2.4.0</version></dependency>
</dependencies>

bookshelf:
  rate-limit:
    default-plan: free
    plans:
      free:    { capacity: 20,   period: 1m }
      pro:     { capacity: 200,  period: 1m }
      partner: { capacity: 2000, period: 1m }
      catalog-import: { capacity: 2, period: 1m }
      isbn:    { capacity: 1000, period: 1d, refill: intervally }
    api-keys:
      key-free-demo: free
      key-pro-demo: pro
    users:
      alice: pro
    costs:
      "[/books/search]": 5
    on-backend-failure: closed

@Validated
@ConfigurationProperties("bookshelf.rate-limit")
public record RateLimitProperties(
        @DefaultValue("free") @NotBlank String defaultPlan,
        @NotEmpty Map<String, @Valid PlanLimit> plans,
        @DefaultValue Map<String, String> apiKeys,
        @DefaultValue Map<String, String> users,
        @DefaultValue Map<String, Long> costs,
        @DefaultValue("closed") FailureMode onBackendFailure,
        Integer filterOrder) {

    public enum FailureMode { OPEN, CLOSED }
    public enum Refill { GREEDY, INTERVALLY }

    public record PlanLimit(@Positive long capacity, @DefaultValue("1m") Duration period,
                            @DefaultValue("greedy") Refill refill) {

        public BucketConfiguration toConfiguration() {
            return BucketConfiguration.builder()
                    .addLimit(limit -> {
                        var capacityStage = limit.capacity(capacity);
                        return refill == Refill.GREEDY
                                ? capacityStage.refillGreedy(capacity, period)
                                : capacityStage.refillIntervally(capacity, period);
                    })
                    .build();
        }

        /** What ProxyManager.builder().build(key, supplier) takes: called only when the bucket must be created. */
        public Supplier<BucketConfiguration> supplier() {
            BucketConfiguration configuration = toConfiguration();
            return () -> configuration;
        }
    }

    public long costOf(String path) { return costs.getOrDefault(path, 1L); }
    public String planOfUser(String username) { return users.getOrDefault(username, defaultPlan); }
    public String planOfApiKey(String apiKey) { return apiKeys.get(apiKey); }
}

public class RateLimitFilter extends OncePerRequestFilter {

    private final RateLimitService limits;
    private final ClientIdentifier clients;
    private final RateLimitProperties properties;
    private final HandlerExceptionResolver exceptionResolver;

    // constructor omitted

    @Override
    protected boolean shouldNotFilter(HttpServletRequest request) {
        return request.getRequestURI().startsWith("/actuator/health");     // probes are never limited
    }

    @Override
    protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, FilterChain chain)
            throws ServletException, IOException {
        Client client = clients.identify(request);
        try {
            Decision decision = limits.check(client.key(), client.plan(), properties.costOf(request.getRequestURI()));
            if (!decision.allowed()) {
                throw new RateLimitExceededException(decision);
            }
            if (decision.remaining() >= 0) {
                response.setHeader("X-Rate-Limit-Remaining", String.valueOf(decision.remaining()));
            }
        } catch (RateLimitExceededException | RateLimitUnavailableException e) {
            // a filter runs before DispatcherServlet: hand the exception to the same @ExceptionHandler methods
            exceptionResolver.resolveException(request, response, null, e);
            return;
        }
        chain.doFilter(request, response);
    }
}

@Bean
@ConditionalOnProperty(prefix = "bookshelf.rate-limit", name = "mechanism", havingValue = "filter", matchIfMissing = true)
FilterRegistrationBean<RateLimitFilter> rateLimitFilter(RateLimitService limits, ClientIdentifier clients,
        RateLimitProperties properties, SecurityFilterProperties security,
        @Qualifier("handlerExceptionResolver") HandlerExceptionResolver resolver) {
    var registration = new FilterRegistrationBean<>(new RateLimitFilter(limits, clients, properties, resolver));
    // Spring Security's chain is registered at spring.security.filter.order (default -100): +1 runs after it
    int order = properties.filterOrder() != null ? properties.filterOrder() : security.getOrder() + 1;
    registration.setOrder(order);
    registration.addUrlPatterns("/*");
    return registration;
}

@Retention(RetentionPolicy.RUNTIME) @Target(ElementType.METHOD)
public @interface RateLimitCost { long value(); }

public class RateLimitInterceptor implements HandlerInterceptor {
    // fields and constructor as in the filter

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
        if (!(handler instanceof HandlerMethod method)) {
            return true;                                   // static resources, error pages
        }
        RateLimitCost cost = method.getMethodAnnotation(RateLimitCost.class);
        Client client = clients.identify(request);
        Decision decision = limits.check(client.key(), client.plan(), cost == null ? 1 : cost.value());
        if (!decision.allowed()) {
            throw new RateLimitExceededException(decision);       // MVC maps it through @ExceptionHandler
        }
        if (decision.remaining() >= 0) {
            response.setHeader("X-Rate-Limit-Remaining", String.valueOf(decision.remaining()));
        }
        return true;
    }
}

@Bean
@ConditionalOnProperty(prefix = "bookshelf.rate-limit", name = "mechanism", havingValue = "interceptor")
WebMvcConfigurer rateLimitInterceptor(RateLimitService limits, ClientIdentifier clients) {
    return new WebMvcConfigurer() {
        @Override
        public void addInterceptors(InterceptorRegistry registry) {
            registry.addInterceptor(new RateLimitInterceptor(limits, clients)).addPathPatterns("/books/**", "/account/**");
        }
    };
}

public Client identify(HttpServletRequest request) {
    Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
    if (authentication != null && authentication.isAuthenticated()
            && !(authentication instanceof AnonymousAuthenticationToken)) {
        return new Client("user:" + authentication.getName(), properties.planOfUser(authentication.getName()));
    }
    String apiKey = request.getHeader("X-API-Key");
    String plan = apiKey == null ? null : properties.planOfApiKey(apiKey);
    if (plan != null) {
        return new Client("key:" + fingerprint(apiKey), plan);      // never store the secret itself as a key
    }
    // unknown keys fall back to the address: random keys must not create unlimited buckets
    return new Client("ip:" + request.getRemoteAddr(), properties.defaultPlan());
}

static String fingerprint(String secret) {
    try {
        byte[] hash = MessageDigest.getInstance("SHA-256").digest(secret.getBytes(StandardCharsets.UTF_8));
        return HexFormat.of().formatHex(hash, 0, 8);
    } catch (NoSuchAlgorithmException e) {                          // every Java platform must provide SHA-256
        throw new IllegalStateException(e);
    }
}

public record Decision(boolean allowed, long remaining, long retryAfterSeconds, String plan) {

    static Decision of(ConsumptionProbe probe, String plan) {
        long seconds = Math.max(1, (long) Math.ceil(probe.getNanosToWaitForRefill() / (double) Duration.ofSeconds(1).toNanos()));
        return new Decision(probe.isConsumed(), probe.getRemainingTokens(), probe.isConsumed() ? 0 : seconds, plan);
    }
}

@RestControllerAdvice
class RateLimitExceptionHandler {

    @ExceptionHandler(RateLimitExceededException.class)
    ResponseEntity<ProblemDetail> tooManyRequests(RateLimitExceededException e) {
        Decision d = e.decision();
        ProblemDetail problem = ProblemDetail.forStatusAndDetail(HttpStatus.TOO_MANY_REQUESTS,
                "The %s plan allows fewer requests than you sent. Retry in %d seconds.".formatted(d.plan(), d.retryAfterSeconds()));
        problem.setType(URI.create("https://bookshelf.example/problems/rate-limit-exceeded"));
        problem.setTitle("Rate limit exceeded");
        problem.setProperty("plan", d.plan());
        problem.setProperty("retryAfterSeconds", d.retryAfterSeconds());
        return ResponseEntity.status(HttpStatus.TOO_MANY_REQUESTS)
                .header(HttpHeaders.RETRY_AFTER, String.valueOf(d.retryAfterSeconds()))
                .header("X-Rate-Limit-Remaining", "0")
                .header("X-Rate-Limit-Retry-After-Seconds", String.valueOf(d.retryAfterSeconds()))
                .body(problem);
    }

    @ExceptionHandler(RateLimitUnavailableException.class)
    ResponseEntity<ProblemDetail> storeUnavailable(RateLimitUnavailableException e) {
        ProblemDetail problem = ProblemDetail.forStatusAndDetail(HttpStatus.SERVICE_UNAVAILABLE,
                "The rate limiter cannot decide right now.");
        problem.setTitle("Rate limiter unavailable");
        return ResponseEntity.status(HttpStatus.SERVICE_UNAVAILABLE).header(HttpHeaders.RETRY_AFTER, "5").body(problem);
    }
}

HTTP/1.1 429
Retry-After: 3
X-Rate-Limit-Remaining: 0
X-Rate-Limit-Retry-After-Seconds: 3
Content-Type: application/problem+json

{"detail":"The free plan allows fewer requests than you sent. Retry in 3 seconds.","instance":"/books",
 "status":429,"title":"Rate limit exceeded","type":"https://bookshelf.example/problems/rate-limit-exceeded",
 "plan":"free","retryAfterSeconds":3}

@Bean
AsyncProxyManager<String> asyncProxyManager(RedisClient bucket4jRedisClient) {
    return Bucket4jLettuce.casBasedBuilder(bucket4jRedisClient)
            .expirationAfterWrite(ExpirationAfterWriteStrategy
                    .basedOnTimeForRefillingBucketUpToMax(Duration.ofMinutes(10)))
            .requestTimeout(Duration.ofMillis(250))
            .build()
            .withMapper((String key) -> key.getBytes(StandardCharsets.UTF_8))
            .asAsync();
}

@Component
@Order(Ordered.HIGHEST_PRECEDENCE + 10)
class RateLimitWebFilter implements WebFilter {

    private static final BucketConfiguration FREE = BucketConfiguration.builder()
            .addLimit(limit -> limit.capacity(5).refillGreedy(5, Duration.ofMinutes(1)))
            .build();

    private final AsyncProxyManager<String> proxyManager;
    private final RateLimitProperties properties;

    RateLimitWebFilter(AsyncProxyManager<String> proxyManager, RateLimitProperties properties) {
        this.proxyManager = proxyManager;
        this.properties = properties;
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
        String key = clientKey(exchange);
        return Mono.defer(() -> Mono.fromFuture(
                        proxyManager.builder()
                                .build(key, () -> CompletableFuture.completedFuture(FREE))
                                .tryConsumeAndReturnRemaining(1)))          // no thread parks on the network call
                .flatMap(probe -> probe.isConsumed() ? allow(exchange, chain, probe) : reject(exchange, probe));
    }

    private static Mono<Void> allow(ServerWebExchange exchange, WebFilterChain chain, ConsumptionProbe probe) {
        exchange.getResponse().getHeaders().set("X-Rate-Limit-Remaining", String.valueOf(probe.getRemainingTokens()));
        return chain.filter(exchange);
    }

    private static Mono<Void> reject(ServerWebExchange exchange, ConsumptionProbe probe) {
        long seconds = Math.max(1, (long) Math.ceil(probe.getNanosToWaitForRefill() / 1_000_000_000d));
        var response = exchange.getResponse();
        response.setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
        response.getHeaders().set(HttpHeaders.RETRY_AFTER, String.valueOf(seconds));
        response.getHeaders().setContentType(MediaType.APPLICATION_PROBLEM_JSON);
        String body = "{\"status\":429,\"title\":\"Rate limit exceeded\",\"retryAfterSeconds\":" + seconds + "}";
        return response.writeWith(Mono.just(response.bufferFactory().wrap(body.getBytes(StandardCharsets.UTF_8))));
    }
    private String clientKey(ServerWebExchange exchange) {
        String apiKey = exchange.getRequest().getHeaders().getFirst("X-API-Key");
        if (apiKey != null && properties.planOfApiKey(apiKey) != null) {
            return "key:" + ClientIdentifier.fingerprint(apiKey);   // issued keys only, never the secret itself
        }
        // unknown or missing keys fall back to the address: random keys must not create unlimited buckets
        InetSocketAddress address = exchange.getRequest().getRemoteAddress();
        return "ip:" + (address == null ? "unknown" : address.getAddress().getHostAddress());
    }
}

@Retention(RetentionPolicy.RUNTIME) @Target(ElementType.METHOD)
public @interface RateLimited {
    String plan();          // an entry of bookshelf.rate-limit.plans
    String key();           // SpEL over the arguments, e.g. "#libraryId"
    long tokens() default 1;
}

@Aspect
@Component
class RateLimitAspect {

    private final RateLimitService limits;
    private final ExpressionParser parser = new SpelExpressionParser();
    private final ParameterNameDiscoverer names = new DefaultParameterNameDiscoverer();

    @Around("@annotation(rateLimited)")
    Object limit(ProceedingJoinPoint call, RateLimited rateLimited) throws Throwable {
        var method = ((MethodSignature) call.getSignature()).getMethod();
        var context = new MethodBasedEvaluationContext(call.getTarget(), method, call.getArgs(), names);
        String subject = parser.parseExpression(rateLimited.key()).getValue(context, String.class);
        Decision decision = limits.check("method:" + rateLimited.plan() + ":" + subject, rateLimited.plan(), rateLimited.tokens());
        if (!decision.allowed()) {
            throw new RateLimitExceededException(decision);
        }
        return call.proceed();
    }
}

@Service
public class CatalogImportService {

    @RateLimited(plan = "catalog-import", key = "#libraryId")
    public String importCatalog(String libraryId) {
        return "imported " + libraryId;
    }
}

public class IsbnQuotaInterceptor implements ClientHttpRequestInterceptor {

    private final BucketProxy quota;       // built once, reused for every call
    private final Duration maxWait;

    public IsbnQuotaInterceptor(ProxyManager<String> proxyManager, RateLimitProperties.PlanLimit isbn, Duration maxWait) {
        this.quota = proxyManager.builder().build("quota:isbn", isbn.supplier());
        this.maxWait = maxWait;
    }

    @Override
    public ClientHttpResponse intercept(HttpRequest request, byte[] body, ClientHttpRequestExecution execution)
            throws IOException {
        if (maxWait.isZero()) {
            ConsumptionProbe probe = quota.tryConsumeAndReturnRemaining(1);
            if (!probe.isConsumed()) {
                long seconds = Math.max(1, (long) Math.ceil(probe.getNanosToWaitForRefill() / 1_000_000_000d));   // round up
                throw new QuotaExhaustedException(seconds);
            }
        } else {
            try {
                if (!quota.asBlocking().tryConsume(1, maxWait)) {      // parks the calling (virtual) thread, bounded
                    throw new QuotaExhaustedException(maxWait.toSeconds());
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                throw new IOException("Interrupted while waiting for the ISBN quota", e);
            }
        }
        return execution.execute(request, body);
    }
}

RestClient isbn = RestClient.builder().baseUrl("https://isbn.example")
        .requestInterceptor(new IsbnQuotaInterceptor(proxyManager, properties.plans().get("isbn"), Duration.ZERO))
        .build();

@Service
public class RateLimitService {

    private final ProxyManager<String> proxyManager;
    private final RateLimitProperties properties;
    private final CircuitBreaker circuitBreaker;

    RateLimitService(ProxyManager<String> proxyManager, RateLimitProperties properties, CircuitBreakerRegistry registry) {
        this.proxyManager = proxyManager;
        this.properties = properties;
        this.circuitBreaker = registry.circuitBreaker("bucket4j");
    }

    public Decision check(String key, String plan, long tokens) {
        var limit = properties.plans().get(plan);
        try {
            ConsumptionProbe probe = circuitBreaker.executeSupplier(() ->
                    proxyManager.builder().build(key, limit.supplier()).tryConsumeAndReturnRemaining(tokens));
            return Decision.of(probe, plan);
        } catch (RuntimeException storeFailure) {                 // includes CallNotPermittedException
            if (properties.onBackendFailure() == FailureMode.OPEN) {
                log.warn("Rate-limit store unavailable ({}), failing open", storeFailure.getClass().getSimpleName());
                return Decision.unchecked(plan);
            }
            throw new RateLimitUnavailableException(storeFailure);
        }
    }
}

resilience4j:
  circuitbreaker:
    instances:
      bucket4j:
        sliding-window-type: COUNT_BASED
        sliding-window-size: 4
        minimum-number-of-calls: 4
        failure-rate-threshold: 50
        wait-duration-in-open-state: 2s
        permitted-number-of-calls-in-half-open-state: 1

var config = BucketConfiguration.builder()
        .addLimit(limit -> limit.capacity(10).refillGreedy(10, Duration.ofSeconds(1))).build();
var bucket = proxyManager.builder().build("blocking:demo", () -> config);
AtomicInteger admitted = new AtomicInteger();
try (var executor = Executors.newVirtualThreadPerTaskExecutor()) {
    for (int i = 0; i < 50; i++) {
        executor.submit(() -> {
            try {
                if (bucket.asBlocking().tryConsume(1, Duration.ofSeconds(10))) {
                    admitted.incrementAndGet();
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        });
    }
}