Add rate-limit module: Bucket4j 8.21 token buckets on Redis in a Spring Boot 4.1 filter, per-user/per-key/global limits, 429 with Retry-After, two instances sharing one bucket

Co-Authored-By: Claude Sonnet 5.5 <[email protected]>
Claude-Session: https://claude.ai/code/session_017ic4b7rnC7jppGGdHGESvy

rate-limit/src/main/java/com/ankurm/ratelimit/Caller.java

@@ -0,0 +1,11 @@
+package com.ankurm.ratelimit;
+
+/**
+ * Who is calling, reduced to the three things the limiter needs.
+ *
+ * @param scope     user, api-key or anonymous
+ * @param bucketKey the Redis key; never contains a raw API key
+ * @param limit     the limit that applies to this caller
+ */
+public record Caller(String scope, String bucketKey, RateLimitProperties.Limit limit) {
+}

rate-limit/src/main/java/com/ankurm/ratelimit/Decision.java

@@ -0,0 +1,15 @@
+package com.ankurm.ratelimit;
+
+/**
+ * The outcome of one rate-limit check, already in the units the HTTP headers use.
+ *
+ * @param allowed           whether the request may proceed
+ * @param scope             which bucket decided: user, api-key, anonymous or global
+ * @param limit             capacity of that bucket
+ * @param remaining         tokens left after this request
+ * @param retryAfterSeconds when rejected: whole seconds until enough tokens have refilled, never 0
+ * @param resetSeconds      whole seconds until that bucket would be full again
+ */
+public record Decision(boolean allowed, String scope, long limit, long remaining,
+                       long retryAfterSeconds, long resetSeconds) {
+}

rate-limit/src/main/java/com/ankurm/ratelimit/DemoController.java

@@ -0,0 +1,34 @@
+package com.ankurm.ratelimit;
+
+import java.util.Map;
+
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.RestController;
+
+@RestController
+class DemoController {
+
+    private final String instance;
+
+    DemoController(@Value("${instance.name:solo}") String instance) {
+        this.instance = instance;
+    }
+
+    @GetMapping("/api/hello")
+    Map<String, String> hello() {
+        return Map.of("message", "hello", "instance", instance);
+    }
+
+    /** Costs 5 tokens, see {@code ratelimit.costs} in application.yml. */
+    @GetMapping("/api/report")
+    Map<String, String> report() {
+        return Map.of("report", "expensive", "instance", instance);
+    }
+
+    /** Not under /api/, so the filter skips it. */
+    @GetMapping("/health")
+    Map<String, String> health() {
+        return Map.of("status", "UP");
+    }
+}

rate-limit/src/main/java/com/ankurm/ratelimit/RateLimitApplication.java

@@ -0,0 +1,17 @@
+package com.ankurm.ratelimit;
+
+import org.springframework.boot.SpringApplication;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+import org.springframework.boot.context.properties.EnableConfigurationProperties;
+
+/**
+ * Rate limiting with Bucket4j 8.21 and Redis. See README.md for the endpoint and test tables.
+ */
+@SpringBootApplication
+@EnableConfigurationProperties(RateLimitProperties.class)
+public class RateLimitApplication {
+
+    public static void main(String[] args) {
+        SpringApplication.run(RateLimitApplication.class, args);
+    }
+}

rate-limit/src/main/java/com/ankurm/ratelimit/RateLimitFilter.java

@@ -0,0 +1,121 @@
+package com.ankurm.ratelimit;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
+import java.util.HexFormat;
+
+import jakarta.servlet.FilterChain;
+import jakarta.servlet.ServletException;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpServletResponse;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.springframework.stereotype.Component;
+import org.springframework.web.filter.OncePerRequestFilter;
+
+/**
+ * The servlet filter that applies the limits. It runs before Spring MVC, so a rejected request
+ * never reaches a controller, a message converter or a database.
+ */
+@Component
+public class RateLimitFilter extends OncePerRequestFilter {
+
+    private static final Logger log = LoggerFactory.getLogger(RateLimitFilter.class);
+
+    private final RateLimiter limiter;
+    private final RateLimitProperties props;
+
+    RateLimitFilter(RateLimiter limiter, RateLimitProperties props) {
+        this.limiter = limiter;
+        this.props = props;
+    }
+
+    @Override
+    protected boolean shouldNotFilter(HttpServletRequest request) {
+        return !request.getRequestURI().startsWith("/api/");
+    }
+
+    @Override
+    protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, FilterChain chain)
+            throws ServletException, IOException {
+        Caller caller = identify(request);
+        if (caller == null) {
+            problem(response, 401, "Unauthorized", "Unknown API key", null);
+            return;
+        }
+        int cost = props.costs().getOrDefault(request.getRequestURI(), 1);
+
+        Decision decision;
+        try {
+            decision = limiter.check(caller, cost);
+        } catch (RuntimeException ex) {
+            if (props.failOpen()) {
+                log.warn("Rate limiter unavailable, letting the request through: {}", ex.toString());
+                response.setHeader("X-RateLimit-Status", "unavailable");
+                chain.doFilter(request, response);
+            } else {
+                log.warn("Rate limiter unavailable, rejecting the request: {}", ex.toString());
+                problem(response, 503, "Service Unavailable", "Rate limiter unavailable", null);
+            }
+            return;
+        }
+
+        response.setHeader("X-RateLimit-Limit", Long.toString(decision.limit()));
+        response.setHeader("X-RateLimit-Remaining", Long.toString(decision.remaining()));
+        response.setHeader("X-RateLimit-Reset", Long.toString(decision.resetSeconds()));
+        response.setHeader("X-RateLimit-Scope", decision.scope());
+        if (decision.allowed()) {
+            chain.doFilter(request, response);
+        } else {
+            response.setHeader("Retry-After", Long.toString(decision.retryAfterSeconds()));
+            problem(response, 429, "Too Many Requests",
+                    "Rate limit exceeded for scope '" + decision.scope() + "'. Retry in "
+                            + decision.retryAfterSeconds() + " second(s).", decision.retryAfterSeconds());
+        }
+    }
+
+    /**
+     * API key first, then user, then address. {@code X-User-Id} stands in for the authenticated
+     * principal so the demo needs no login; a real application takes the name from
+     * {@code request.getUserPrincipal()} and ignores any such header.
+     */
+    private Caller identify(HttpServletRequest request) {
+        String apiKey = request.getHeader("X-API-Key");
+        if (apiKey != null) {
+            String plan = props.apiKeys().get(apiKey);
+            if (plan == null) {
+                return null;   // unknown key: 401, and no bucket is created for it
+            }
+            return new Caller("api-key", "rl:key:" + sha256Prefix(apiKey), props.plans().get(plan));
+        }
+        String user = request.getUserPrincipal() != null
+                ? request.getUserPrincipal().getName() : request.getHeader("X-User-Id");
+        if (user != null && !user.isBlank()) {
+            return new Caller("user", "rl:user:" + user, props.user());
+        }
+        return new Caller("anonymous", "rl:ip:" + request.getRemoteAddr(), props.anonymous());
+    }
+
+    private static String sha256Prefix(String value) {
+        try {
+            byte[] digest = MessageDigest.getInstance("SHA-256").digest(value.getBytes(StandardCharsets.UTF_8));
+            return HexFormat.of().formatHex(digest, 0, 8);
+        } catch (NoSuchAlgorithmException e) {
+            throw new IllegalStateException(e);
+        }
+    }
+
+    private static void problem(HttpServletResponse response, int status, String title, String detail,
+                                Long retryAfter) throws IOException {
+        response.setStatus(status);
+        response.setContentType("application/problem+json");
+        String body = "{\"type\":\"about:blank\",\"title\":\"" + title + "\",\"status\":" + status
+                + ",\"detail\":\"" + detail + "\"" + (retryAfter != null ? ",\"retryAfterSeconds\":" + retryAfter : "")
+                + "}";
+        // Bytes, not a Writer: a Writer would append ;charset=ISO-8859-1 to the media type.
+        response.getOutputStream().write(body.getBytes(StandardCharsets.UTF_8));
+    }
+}

rate-limit/src/main/java/com/ankurm/ratelimit/RateLimitProperties.java

@@ -0,0 +1,34 @@
+package com.ankurm.ratelimit;
+
+import java.time.Duration;
+import java.util.Map;
+
+import org.springframework.boot.context.properties.ConfigurationProperties;
+
+/**
+ * Every limit in the demo lives here, so a test can change one with a property override.
+ *
+ * @param failOpen      when Redis is unreachable: true lets the request through, false answers 503
+ * @param configVersion bumped by hand when a limit changes; see {@link RedisRateLimitConfig}
+ * @param user          per authenticated user
+ * @param anonymous     per client address, for callers with neither a user nor an API key
+ * @param global        one bucket shared by every caller of every instance
+ * @param plans         named limits an API key can be on
+ * @param apiKeys       API key to plan name (a real system would keep this in a database)
+ * @param costs         path to token cost; a path not listed here costs 1
+ */
+@ConfigurationProperties(prefix = "ratelimit")
+public record RateLimitProperties(
+        boolean failOpen,
+        long configVersion,
+        Limit user,
+        Limit anonymous,
+        Limit global,
+        Map<String, Limit> plans,
+        Map<String, String> apiKeys,
+        Map<String, Integer> costs) {
+
+    /** capacity tokens, refilled greedily at {@code capacity / period}. */
+    public record Limit(long capacity, Duration period) {
+    }
+}

rate-limit/src/main/java/com/ankurm/ratelimit/RateLimiter.java

@@ -0,0 +1,65 @@
+package com.ankurm.ratelimit;
+
+import io.github.bucket4j.BucketConfiguration;
+import io.github.bucket4j.ConsumptionProbe;
+import io.github.bucket4j.TokensInheritanceStrategy;
+import io.github.bucket4j.distributed.BucketProxy;
+import io.github.bucket4j.distributed.proxy.ProxyManager;
+
+import org.springframework.stereotype.Service;
+
+/**
+ * Decides whether a request may proceed. Two buckets are consulted: the caller's own
+ * (user, API key or address) and then the one global bucket.
+ *
+ * <p>The order matters. The caller's bucket goes first, so a caller who is already over their own
+ * limit never spends a global token. The price is the opposite case: when the global bucket
+ * rejects, the caller has already spent a token of their own.
+ */
+@Service
+public class RateLimiter {
+
+    static final String GLOBAL_KEY = "rl:global";
+
+    private final ProxyManager<String> buckets;
+    private final RateLimitProperties props;
+
+    RateLimiter(ProxyManager<String> buckets, RateLimitProperties props) {
+        this.buckets = buckets;
+        this.props = props;
+    }
+
+    public Decision check(Caller caller, int cost) {
+        Decision own = consume(caller.bucketKey(), caller.scope(), caller.limit(), cost);
+        if (!own.allowed()) {
+            return own;
+        }
+        Decision global = consume(GLOBAL_KEY, "global", props.global(), cost);
+        return global.allowed() ? own : global;
+    }
+
+    private Decision consume(String key, String scope, RateLimitProperties.Limit limit, int cost) {
+        ConsumptionProbe probe = bucket(key, limit).tryConsumeAndReturnRemaining(cost);
+        long retry = probe.isConsumed() ? 0 : Math.max(1, ceilSeconds(probe.getNanosToWaitForRefill()));
+        return new Decision(probe.isConsumed(), scope, limit.capacity(), probe.getRemainingTokens(),
+                retry, ceilSeconds(probe.getNanosToWaitForReset()));
+    }
+
+    /**
+     * Cheap to call on every request: nothing touches Redis until a token is consumed. The
+     * configuration supplier only runs when the key does not exist yet, which is why changing a
+     * limit in code does not change buckets that already exist. {@code
+     * withImplicitConfigurationReplacement} is what makes a new {@code configVersion} replace them.
+     */
+    BucketProxy bucket(String key, RateLimitProperties.Limit limit) {
+        return buckets.builder()
+                .withImplicitConfigurationReplacement(props.configVersion(), TokensInheritanceStrategy.PROPORTIONALLY)
+                .build(key, () -> BucketConfiguration.builder()
+                        .addLimit(l -> l.capacity(limit.capacity()).refillGreedy(limit.capacity(), limit.period()))
+                        .build());
+    }
+
+    private static long ceilSeconds(long nanos) {
+        return (nanos + 999_999_999L) / 1_000_000_000L;
+    }
+}

rate-limit/src/main/java/com/ankurm/ratelimit/RedisRateLimitConfig.java

@@ -0,0 +1,58 @@
+package com.ankurm.ratelimit;
+
+import java.time.Duration;
+
+import io.github.bucket4j.TimeMeter;
+import io.github.bucket4j.distributed.ExpirationAfterWriteStrategy;
+import io.github.bucket4j.distributed.proxy.ProxyManager;
+import io.github.bucket4j.redis.lettuce.Bucket4jLettuce;
+import io.lettuce.core.RedisClient;
+import io.lettuce.core.api.StatefulRedisConnection;
+import io.lettuce.core.codec.ByteArrayCodec;
+import io.lettuce.core.codec.RedisCodec;
+import io.lettuce.core.codec.StringCodec;
+
+import org.springframework.beans.factory.ObjectProvider;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.data.redis.connection.lettuce.LettuceConnectionFactory;
+
+/**
+ * Wires Bucket4j to the Lettuce client that Spring Boot already configured, so
+ * {@code spring.data.redis.host/port/password/ssl} apply to the rate limiter too.
+ */
+@Configuration(proxyBeanMethods = false)
+class RedisRateLimitConfig {
+
+    /**
+     * A second connection on Boot's RedisClient, with String keys and byte[] values. Bucket4j
+     * stores each bucket as a binary blob, so the value codec must not touch the bytes.
+     */
+    @Bean(destroyMethod = "close")
+    StatefulRedisConnection<String, byte[]> rateLimitConnection(LettuceConnectionFactory factory) {
+        RedisClient client = (RedisClient) factory.getRequiredNativeClient();
+        StatefulRedisConnection<String, byte[]> connection =
+                client.connect(RedisCodec.of(StringCodec.UTF8, ByteArrayCodec.INSTANCE));
+        connection.setTimeout(Duration.ofMillis(500));
+        return connection;
+    }
+
+    /**
+     * The proxy manager is the only Redis-aware object. Everything else talks to
+     * {@link ProxyManager}, so the same code runs on Hazelcast, JDBC or Redisson by swapping this
+     * bean.
+     *
+     * <p>The clock is a bean so a test can move time without sleeping. Production gets the system
+     * clock.
+     */
+    @Bean
+    ProxyManager<String> rateLimitProxyManager(StatefulRedisConnection<String, byte[]> connection,
+                                               ObjectProvider<TimeMeter> clock) {
+        return Bucket4jLettuce.casBasedBuilder(connection)
+                // Delete a bucket's key once it has been full for 10 seconds: idle callers cost no memory.
+                .expirationAfterWrite(
+                        ExpirationAfterWriteStrategy.basedOnTimeForRefillingBucketUpToMax(Duration.ofSeconds(10)))
+                .clientClock(clock.getIfAvailable(() -> TimeMeter.SYSTEM_MILLISECONDS))
+                .build();
+    }
+}

rate-limit/src/main/resources/application.yml

@@ -0,0 +1,30 @@
+spring:
+  data:
+    redis:
+      host: localhost
+      port: 6379
+
+ratelimit:
+  fail-open: true
+  config-version: 1
+  user:
+    capacity: 10
+    period: 1m
+  anonymous:
+    capacity: 5
+    period: 1m
+  global:
+    capacity: 1000
+    period: 1m
+  plans:
+    free:
+      capacity: 5
+      period: 1m
+    pro:
+      capacity: 20
+      period: 1m
+  api-keys:
+    free-key-123: free
+    pro-key-456: pro
+  costs:
+    "[/api/report]": 5