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
This commit is contained in:
Claude
@@ -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) {
|
||||
}
|
||||
@@ -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) {
|
||||
}
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
@@ -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) {
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
@@ -0,0 +1,117 @@
|
||||
package com.ankurm.ratelimit;
|
||||
|
||||
import org.junit.jupiter.api.AfterAll;
|
||||
import org.junit.jupiter.api.BeforeAll;
|
||||
import org.junit.jupiter.api.MethodOrderer;
|
||||
import org.junit.jupiter.api.Order;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.TestMethodOrder;
|
||||
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.test.web.servlet.MockMvc;
|
||||
import org.springframework.mock.web.MockHttpServletResponse;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get;
|
||||
|
||||
/** Per-API-key limits by plan, token cost per endpoint, and what ends up in Redis. Writes output/03-api-keys.txt. */
|
||||
@TestMethodOrder(MethodOrderer.OrderAnnotation.class)
|
||||
class ApiKeyLimitTest extends RedisBackedTest {
|
||||
|
||||
static Transcript t;
|
||||
|
||||
@Autowired
|
||||
MockMvc mvc;
|
||||
|
||||
@BeforeAll
|
||||
static void open() {
|
||||
t = Transcript.open("03-api-keys.txt");
|
||||
}
|
||||
|
||||
@AfterAll
|
||||
static void close() {
|
||||
t.close();
|
||||
}
|
||||
|
||||
private MockHttpServletResponse call(String path, String key) throws Exception {
|
||||
var res = mvc.perform(get(path).header("X-API-Key", key)).andReturn().getResponse();
|
||||
t.line("GET %-11s X-API-Key: %-12s -> %d limit=%s remaining=%s retry-after=%s", path, key,
|
||||
res.getStatus(), res.getHeader("X-RateLimit-Limit"), res.getHeader("X-RateLimit-Remaining"),
|
||||
res.getHeader("Retry-After"));
|
||||
return res;
|
||||
}
|
||||
|
||||
@Test
|
||||
@Order(1)
|
||||
void planDecidesTheLimit() throws Exception {
|
||||
t.line("== free plan is 5 per minute, pro plan is 20 per minute");
|
||||
int free = 0;
|
||||
int pro = 0;
|
||||
for (int i = 0; i < 25; i++) {
|
||||
if (mvc.perform(get("/api/hello").header("X-API-Key", "free-key-123")).andReturn()
|
||||
.getResponse().getStatus() == 200) {
|
||||
free++;
|
||||
}
|
||||
if (mvc.perform(get("/api/hello").header("X-API-Key", "pro-key-456")).andReturn()
|
||||
.getResponse().getStatus() == 200) {
|
||||
pro++;
|
||||
}
|
||||
}
|
||||
t.line("25 requests each: free-key-123 allowed %d, pro-key-456 allowed %d", free, pro);
|
||||
assertThat(free).isEqualTo(5);
|
||||
assertThat(pro).isEqualTo(20);
|
||||
}
|
||||
|
||||
@Test
|
||||
@Order(2)
|
||||
void expensiveEndpointCostsFiveTokens() throws Exception {
|
||||
t.line("== /api/report costs 5 tokens, /api/hello costs 1");
|
||||
assertThat(call("/api/hello", "pro-key-456").getStatus()).isEqualTo(200);
|
||||
assertThat(call("/api/report", "pro-key-456").getStatus()).isEqualTo(200);
|
||||
assertThat(call("/api/report", "pro-key-456").getStatus()).isEqualTo(200);
|
||||
assertThat(call("/api/report", "pro-key-456").getStatus()).isEqualTo(200);
|
||||
var rejected = call("/api/report", "pro-key-456");
|
||||
assertThat(rejected.getStatus()).isEqualTo(429);
|
||||
t.line("-- 20 tokens = 1 hello + 3 reports (16), 4 left; a 4th report needs 5, so it waits for the refill");
|
||||
assertThat(rejected.getHeader("Retry-After")).isEqualTo("3");
|
||||
}
|
||||
|
||||
@Test
|
||||
@Order(3)
|
||||
void unknownKeyIs401AndCreatesNothing() throws Exception {
|
||||
t.line("== an unknown API key");
|
||||
var res = call("/api/hello", "made-up-key");
|
||||
assertThat(res.getStatus()).isEqualTo(401);
|
||||
t.line("body: %s", res.getContentAsString());
|
||||
t.line("keys in Redis after the 401: %s", REDIS.cli("DBSIZE"));
|
||||
assertThat(REDIS.cli("DBSIZE")).isEqualTo("0");
|
||||
}
|
||||
|
||||
@Test
|
||||
@Order(4)
|
||||
void whatRedisActuallyHolds() throws Exception {
|
||||
t.line("== what is in Redis after one request each from a key, a user and an anonymous caller");
|
||||
mvc.perform(get("/api/hello").header("X-API-Key", "free-key-123"));
|
||||
mvc.perform(get("/api/hello").header("X-User-Id", "alice"));
|
||||
mvc.perform(get("/api/hello"));
|
||||
mvc.perform(get("/api/hello"));
|
||||
t.line("$ redis-cli --scan | sort");
|
||||
String keys = REDIS.cli("--scan");
|
||||
java.util.List<String> sorted = keys.lines().sorted().toList();
|
||||
sorted.forEach(k -> t.line(k));
|
||||
assertThat(sorted).hasSize(4);
|
||||
assertThat(keys).doesNotContain("free-key-123");
|
||||
assertThat(sorted).anyMatch(k -> k.startsWith("rl:key:") && k.length() == "rl:key:".length() + 16);
|
||||
assertThat(sorted).contains("rl:user:alice", "rl:global", "rl:ip:127.0.0.1");
|
||||
|
||||
String keyKey = sorted.stream().filter(k -> k.startsWith("rl:key:")).findFirst().orElseThrow();
|
||||
t.line("$ redis-cli TYPE %s", keyKey);
|
||||
t.line(REDIS.cli("TYPE", keyKey));
|
||||
t.line("$ redis-cli STRLEN %s", keyKey);
|
||||
t.line(REDIS.cli("STRLEN", keyKey));
|
||||
long ttl = Long.parseLong(REDIS.cli("PTTL", "rl:user:alice"));
|
||||
t.line("$ redis-cli PTTL rl:user:alice (a bucket expires once it has been full for 10 s)");
|
||||
t.line("positive and at most 70000: %s", ttl > 0 && ttl <= 70_000);
|
||||
assertThat(ttl).isBetween(1L, 70_000L);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
package com.ankurm.ratelimit;
|
||||
|
||||
import java.net.URI;
|
||||
import java.net.http.HttpClient;
|
||||
import java.net.http.HttpRequest;
|
||||
import java.net.http.HttpResponse;
|
||||
|
||||
import io.github.bucket4j.distributed.proxy.ProxyManager;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import org.springframework.boot.builder.SpringApplicationBuilder;
|
||||
import org.springframework.context.ConfigurableApplicationContext;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* A bucket's configuration is stored in Redis with its tokens. Deploying a new limit does not change
|
||||
* buckets that already exist unless the configuration version is bumped. Writes output/06-config-change.txt.
|
||||
*/
|
||||
class ConfigChangeTest {
|
||||
|
||||
static final RedisServer REDIS = RedisServer.shared();
|
||||
static final HttpClient http = HttpClient.newHttpClient();
|
||||
|
||||
static ConfigurableApplicationContext boot(String... extraArgs) {
|
||||
String[] base = {"--server.port=0", "--spring.main.banner-mode=off", "--logging.level.root=WARN",
|
||||
"--spring.data.redis.port=" + REDIS.port()};
|
||||
String[] all = new String[base.length + extraArgs.length];
|
||||
System.arraycopy(base, 0, all, 0, base.length);
|
||||
System.arraycopy(extraArgs, 0, all, base.length, extraArgs.length);
|
||||
return new SpringApplicationBuilder(RateLimitApplication.class).run(all);
|
||||
}
|
||||
|
||||
static HttpResponse<String> alice(ConfigurableApplicationContext ctx) throws Exception {
|
||||
String port = ctx.getEnvironment().getProperty("local.server.port");
|
||||
return http.send(HttpRequest.newBuilder(URI.create("http://localhost:" + port + "/api/hello"))
|
||||
.header("X-User-Id", "alice").build(), HttpResponse.BodyHandlers.ofString());
|
||||
}
|
||||
|
||||
@SuppressWarnings("unchecked")
|
||||
static long storedCapacity(ConfigurableApplicationContext ctx) {
|
||||
ProxyManager<String> pm = ctx.getBean(ProxyManager.class);
|
||||
return pm.getProxyConfiguration("rl:user:alice").orElseThrow().getBandwidths()[0].getCapacity();
|
||||
}
|
||||
|
||||
@Test
|
||||
void newLimitNeedsANewConfigVersion() throws Exception {
|
||||
try (Transcript t = Transcript.open("06-config-change.txt")) {
|
||||
REDIS.flush();
|
||||
|
||||
t.line("== deploy 1: user limit 10, config-version 1. alice makes 3 requests.");
|
||||
try (ConfigurableApplicationContext v1 = boot("--ratelimit.user.capacity=10")) {
|
||||
for (int i = 0; i < 3; i++) {
|
||||
alice(v1);
|
||||
}
|
||||
t.line("stored capacity in Redis: %d", storedCapacity(v1));
|
||||
}
|
||||
|
||||
t.line("== deploy 2: user limit raised to 20 in application.yml, config-version still 1");
|
||||
try (ConfigurableApplicationContext v2 = boot("--ratelimit.user.capacity=20")) {
|
||||
var r = alice(v2);
|
||||
t.line("X-RateLimit-Limit header says %s", r.headers().firstValue("X-RateLimit-Limit").orElse("-"));
|
||||
t.line("stored capacity in Redis: %d", storedCapacity(v2));
|
||||
t.line("X-RateLimit-Remaining: %s", r.headers().firstValue("X-RateLimit-Remaining").orElse("-"));
|
||||
assertThat(r.headers().firstValue("X-RateLimit-Limit")).hasValue("20");
|
||||
assertThat(storedCapacity(v2)).isEqualTo(10);
|
||||
}
|
||||
|
||||
t.line("== deploy 3: limit 20 and config-version 2");
|
||||
try (ConfigurableApplicationContext v3 = boot("--ratelimit.user.capacity=20", "--ratelimit.config-version=2")) {
|
||||
var r = alice(v3);
|
||||
t.line("X-RateLimit-Limit header says %s", r.headers().firstValue("X-RateLimit-Limit").orElse("-"));
|
||||
t.line("stored capacity in Redis: %d", storedCapacity(v3));
|
||||
t.line("X-RateLimit-Remaining: %s", r.headers().firstValue("X-RateLimit-Remaining").orElse("-"));
|
||||
assertThat(storedCapacity(v3)).isEqualTo(20);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
package com.ankurm.ratelimit;
|
||||
|
||||
import java.time.Duration;
|
||||
|
||||
import io.github.bucket4j.distributed.BucketProxy;
|
||||
import org.junit.jupiter.api.AfterAll;
|
||||
import org.junit.jupiter.api.BeforeAll;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.mock.web.MockHttpServletResponse;
|
||||
import org.springframework.test.context.TestPropertySource;
|
||||
import org.springframework.test.web.servlet.MockMvc;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get;
|
||||
|
||||
/**
|
||||
* The global bucket, shrunk to 12 per minute so five users can hit it while each stays far under
|
||||
* their own 10 per minute. Writes output/04-global.txt.
|
||||
*/
|
||||
@TestPropertySource(properties = "ratelimit.global.capacity=12")
|
||||
class GlobalLimitTest extends RedisBackedTest {
|
||||
|
||||
@Autowired
|
||||
MockMvc mvc;
|
||||
|
||||
@Autowired
|
||||
RateLimiter limiter;
|
||||
|
||||
@Autowired
|
||||
RateLimitProperties props;
|
||||
|
||||
static Transcript t;
|
||||
|
||||
@BeforeAll
|
||||
static void open() {
|
||||
t = Transcript.open("04-global.txt");
|
||||
}
|
||||
|
||||
@AfterAll
|
||||
static void close() {
|
||||
t.close();
|
||||
}
|
||||
|
||||
private MockHttpServletResponse call(String user) throws Exception {
|
||||
var res = mvc.perform(get("/api/hello").header("X-User-Id", user)).andReturn().getResponse();
|
||||
t.line("GET /api/hello X-User-Id: %-3s -> %d limit=%s remaining=%s retry-after=%s scope=%s", user,
|
||||
res.getStatus(), res.getHeader("X-RateLimit-Limit"), res.getHeader("X-RateLimit-Remaining"),
|
||||
res.getHeader("Retry-After"), res.getHeader("X-RateLimit-Scope"));
|
||||
return res;
|
||||
}
|
||||
|
||||
@Test
|
||||
void globalBucketCapsEveryoneTogether() throws Exception {
|
||||
t.line("== global limit 12 per minute; four users send 3 requests each");
|
||||
for (String user : new String[] {"u1", "u2", "u3", "u4"}) {
|
||||
for (int i = 0; i < 3; i++) {
|
||||
assertThat(call(user).getStatus()).isEqualTo(200);
|
||||
}
|
||||
}
|
||||
t.line("== u5 has never called before, so their own bucket is full");
|
||||
var rejected = call("u5");
|
||||
assertThat(rejected.getStatus()).isEqualTo(429);
|
||||
assertThat(rejected.getHeader("X-RateLimit-Scope")).isEqualTo("global");
|
||||
assertThat(rejected.getHeader("X-RateLimit-Limit")).isEqualTo("12");
|
||||
assertThat(rejected.getHeader("Retry-After")).isEqualTo("5");
|
||||
|
||||
BucketProxy u5 = limiter.bucket("rl:user:u5", props.user());
|
||||
t.line("-- u5 was rejected by the global bucket, but their own bucket already paid: %d of 10 tokens left",
|
||||
u5.getAvailableTokens());
|
||||
assertThat(u5.getAvailableTokens()).isEqualTo(9);
|
||||
|
||||
clock.advance(Duration.ofSeconds(5));
|
||||
t.line("== five seconds later (Retry-After), u5 tries again");
|
||||
assertThat(call("u5").getStatus()).isEqualTo(200);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
package com.ankurm.ratelimit;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.concurrent.atomic.AtomicLong;
|
||||
|
||||
import io.github.bucket4j.TimeMeter;
|
||||
|
||||
/**
|
||||
* A clock the test moves by hand. Bucket4j's distributed buckets take their time from the client,
|
||||
* so handing the proxy manager this clock lets a test "wait" 6 seconds in zero milliseconds.
|
||||
*/
|
||||
public class ManualClock implements TimeMeter {
|
||||
|
||||
private final AtomicLong nanos = new AtomicLong(System.currentTimeMillis() * 1_000_000L);
|
||||
|
||||
public ManualClock advance(Duration d) {
|
||||
nanos.addAndGet(d.toNanos());
|
||||
return this;
|
||||
}
|
||||
|
||||
@Override
|
||||
public long currentTimeNanos() {
|
||||
return nanos.get();
|
||||
}
|
||||
|
||||
@Override
|
||||
public boolean isWallClockBased() {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
package com.ankurm.ratelimit;
|
||||
|
||||
import java.time.Duration;
|
||||
|
||||
import org.junit.jupiter.api.AfterAll;
|
||||
import org.junit.jupiter.api.BeforeAll;
|
||||
import org.junit.jupiter.api.MethodOrderer;
|
||||
import org.junit.jupiter.api.Order;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.TestMethodOrder;
|
||||
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.test.web.servlet.MockMvc;
|
||||
import org.springframework.test.web.servlet.MvcResult;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get;
|
||||
|
||||
/** Per-user limit: 10 requests per minute, greedy refill (one token every 6 s). Writes output/02-per-user.txt. */
|
||||
@TestMethodOrder(MethodOrderer.OrderAnnotation.class)
|
||||
class PerUserLimitTest extends RedisBackedTest {
|
||||
|
||||
static Transcript t;
|
||||
|
||||
@Autowired
|
||||
MockMvc mvc;
|
||||
|
||||
@BeforeAll
|
||||
static void open() {
|
||||
t = Transcript.open("02-per-user.txt");
|
||||
}
|
||||
|
||||
@AfterAll
|
||||
static void close() {
|
||||
t.close();
|
||||
}
|
||||
|
||||
private MvcResult call(String user) throws Exception {
|
||||
MvcResult r = mvc.perform(get("/api/hello").header("X-User-Id", user)).andReturn();
|
||||
var res = r.getResponse();
|
||||
t.line("GET /api/hello X-User-Id: %-5s -> %d limit=%s remaining=%s reset=%s retry-after=%s scope=%s",
|
||||
user, res.getStatus(), res.getHeader("X-RateLimit-Limit"), res.getHeader("X-RateLimit-Remaining"),
|
||||
res.getHeader("X-RateLimit-Reset"), res.getHeader("Retry-After"), res.getHeader("X-RateLimit-Scope"));
|
||||
return r;
|
||||
}
|
||||
|
||||
@Test
|
||||
@Order(1)
|
||||
void eleventhRequestInAMinuteGets429() throws Exception {
|
||||
t.line("== alice sends 12 requests with no pause");
|
||||
for (int i = 1; i <= 10; i++) {
|
||||
var res = call("alice").getResponse();
|
||||
assertThat(res.getStatus()).isEqualTo(200);
|
||||
assertThat(res.getHeader("X-RateLimit-Remaining")).isEqualTo(String.valueOf(10 - i));
|
||||
}
|
||||
MvcResult rejected = call("alice");
|
||||
assertThat(rejected.getResponse().getStatus()).isEqualTo(429);
|
||||
assertThat(rejected.getResponse().getHeader("Retry-After")).isEqualTo("6");
|
||||
assertThat(rejected.getResponse().getContentType()).isEqualTo("application/problem+json");
|
||||
call("alice");
|
||||
t.line("-- body of the 429:");
|
||||
t.line(rejected.getResponse().getContentAsString());
|
||||
}
|
||||
|
||||
@Test
|
||||
@Order(2)
|
||||
void aDifferentUserIsUnaffected() throws Exception {
|
||||
t.line("== alice is out of tokens; bob is a different bucket");
|
||||
for (int i = 0; i < 10; i++) {
|
||||
mvc.perform(get("/api/hello").header("X-User-Id", "alice"));
|
||||
}
|
||||
assertThat(call("alice").getResponse().getStatus()).isEqualTo(429);
|
||||
MvcResult bob = call("bob");
|
||||
assertThat(bob.getResponse().getStatus()).isEqualTo(200);
|
||||
assertThat(bob.getResponse().getHeader("X-RateLimit-Remaining")).isEqualTo("9");
|
||||
}
|
||||
|
||||
@Test
|
||||
@Order(3)
|
||||
void honouringRetryAfterWorks() throws Exception {
|
||||
t.line("== alice waits exactly Retry-After (6 s on the test clock) and tries again");
|
||||
for (int i = 0; i < 10; i++) {
|
||||
mvc.perform(get("/api/hello").header("X-User-Id", "alice"));
|
||||
}
|
||||
MvcResult rejected = call("alice");
|
||||
long wait = Long.parseLong(rejected.getResponse().getHeader("Retry-After"));
|
||||
clock.advance(Duration.ofSeconds(wait - 1));
|
||||
t.line("-- after %d s: still too early", wait - 1);
|
||||
assertThat(call("alice").getResponse().getStatus()).isEqualTo(429);
|
||||
clock.advance(Duration.ofSeconds(1));
|
||||
t.line("-- after %d s in total:", wait);
|
||||
assertThat(call("alice").getResponse().getStatus()).isEqualTo(200);
|
||||
assertThat(call("alice").getResponse().getStatus()).isEqualTo(429);
|
||||
}
|
||||
|
||||
@Test
|
||||
@Order(4)
|
||||
void healthIsNotLimited() throws Exception {
|
||||
t.line("== /health is outside /api/, so the filter skips it");
|
||||
for (int i = 0; i < 30; i++) {
|
||||
assertThat(mvc.perform(get("/health")).andReturn().getResponse().getStatus()).isEqualTo(200);
|
||||
}
|
||||
t.line("30 requests to /health -> all 200, X-RateLimit-Limit header present: %s",
|
||||
mvc.perform(get("/health")).andReturn().getResponse().containsHeader("X-RateLimit-Limit"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
package com.ankurm.ratelimit;
|
||||
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.boot.webmvc.test.autoconfigure.AutoConfigureMockMvc;
|
||||
import org.springframework.boot.test.context.SpringBootTest;
|
||||
import org.springframework.boot.test.context.TestConfiguration;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
import org.springframework.context.annotation.Import;
|
||||
import org.springframework.test.context.DynamicPropertyRegistry;
|
||||
import org.springframework.test.context.DynamicPropertySource;
|
||||
|
||||
/**
|
||||
* The full application, MockMvc in front of it, a real redis-server behind it and a clock the
|
||||
* test owns.
|
||||
*/
|
||||
@SpringBootTest
|
||||
@AutoConfigureMockMvc
|
||||
@Import(RedisBackedTest.ClockConfig.class)
|
||||
abstract class RedisBackedTest {
|
||||
|
||||
static final RedisServer REDIS = RedisServer.shared();
|
||||
|
||||
@TestConfiguration(proxyBeanMethods = false)
|
||||
static class ClockConfig {
|
||||
@Bean
|
||||
ManualClock manualClock() {
|
||||
return new ManualClock();
|
||||
}
|
||||
}
|
||||
|
||||
@DynamicPropertySource
|
||||
static void redis(DynamicPropertyRegistry registry) {
|
||||
registry.add("spring.data.redis.port", REDIS::port);
|
||||
}
|
||||
|
||||
@Autowired
|
||||
ManualClock clock;
|
||||
|
||||
@BeforeEach
|
||||
void flush() {
|
||||
REDIS.flush();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
package com.ankurm.ratelimit;
|
||||
|
||||
import java.net.URI;
|
||||
import java.net.http.HttpClient;
|
||||
import java.net.http.HttpRequest;
|
||||
import java.net.http.HttpResponse;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import org.springframework.boot.builder.SpringApplicationBuilder;
|
||||
import org.springframework.context.ConfigurableApplicationContext;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
|
||||
/**
|
||||
* What the limiter does when Redis dies under it, with ratelimit.fail-open true and false.
|
||||
* Each case starts its own redis-server and stops it. Writes output/07-redis-outage.txt.
|
||||
*/
|
||||
class RedisOutageTest {
|
||||
|
||||
static final HttpClient http = HttpClient.newHttpClient();
|
||||
|
||||
private HttpResponse<String> get(ConfigurableApplicationContext ctx) throws Exception {
|
||||
String port = ctx.getEnvironment().getProperty("local.server.port");
|
||||
return http.send(HttpRequest.newBuilder(URI.create("http://localhost:" + port + "/api/hello"))
|
||||
.header("X-User-Id", "dave").build(), HttpResponse.BodyHandlers.ofString());
|
||||
}
|
||||
|
||||
@Test
|
||||
void failOpenAndFailClosed() throws Exception {
|
||||
try (Transcript t = Transcript.open("07-redis-outage.txt")) {
|
||||
for (boolean failOpen : new boolean[] {true, false}) {
|
||||
RedisServer redis = RedisServer.start();
|
||||
try (ConfigurableApplicationContext ctx = new SpringApplicationBuilder(RateLimitApplication.class).run(
|
||||
"--server.port=0", "--spring.main.banner-mode=off", "--logging.level.root=OFF",
|
||||
"--spring.data.redis.port=" + redis.port(), "--ratelimit.fail-open=" + failOpen)) {
|
||||
t.line("== ratelimit.fail-open=%s", failOpen);
|
||||
var before = get(ctx);
|
||||
t.line("Redis up: %d X-RateLimit-Remaining=%s", before.statusCode(),
|
||||
before.headers().firstValue("X-RateLimit-Remaining").orElse("-"));
|
||||
assertThat(before.statusCode()).isEqualTo(200);
|
||||
|
||||
redis.close();
|
||||
Thread.sleep(300);
|
||||
|
||||
long start = System.nanoTime();
|
||||
var during = get(ctx);
|
||||
long millis = (System.nanoTime() - start) / 1_000_000;
|
||||
t.line("Redis down: %d X-RateLimit-Status=%s X-RateLimit-Remaining=%s", during.statusCode(),
|
||||
during.headers().firstValue("X-RateLimit-Status").orElse("-"),
|
||||
during.headers().firstValue("X-RateLimit-Remaining").orElse("-"));
|
||||
t.line("the request waited for the 500 ms command timeout: %s", millis >= 400 && millis < 5_000);
|
||||
assertThat(millis).isBetween(400L, 5_000L);
|
||||
assertThat(during.statusCode()).isEqualTo(failOpen ? 200 : 503);
|
||||
|
||||
if (failOpen) {
|
||||
assertThatThrownBy(() -> ctx.getBean(RateLimiter.class)
|
||||
.check(new Caller("user", "rl:user:dave", new RateLimitProperties.Limit(10,
|
||||
java.time.Duration.ofMinutes(1))), 1))
|
||||
.satisfies(e -> t.line("exception the filter caught: %s", e.getClass().getName()));
|
||||
}
|
||||
} finally {
|
||||
redis.close();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
package com.ankurm.ratelimit;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.net.InetSocketAddress;
|
||||
import java.net.ServerSocket;
|
||||
import java.net.Socket;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* A throwaway {@code redis-server} on a free port. Docker was not available where the module was
|
||||
* written, so the tests start the real Redis binary themselves. Nothing here is mocked.
|
||||
*/
|
||||
public final class RedisServer implements AutoCloseable {
|
||||
|
||||
private static RedisServer shared;
|
||||
|
||||
private final Process process;
|
||||
private final int port;
|
||||
|
||||
private RedisServer(Process process, int port) {
|
||||
this.process = process;
|
||||
this.port = port;
|
||||
}
|
||||
|
||||
/** One server per test JVM, stopped by a shutdown hook. */
|
||||
public static synchronized RedisServer shared() {
|
||||
if (shared == null) {
|
||||
shared = start();
|
||||
Runtime.getRuntime().addShutdownHook(new Thread(shared::close));
|
||||
}
|
||||
return shared;
|
||||
}
|
||||
|
||||
public static RedisServer start() {
|
||||
try {
|
||||
int port;
|
||||
try (ServerSocket s = new ServerSocket(0)) {
|
||||
port = s.getLocalPort();
|
||||
}
|
||||
Process p = new ProcessBuilder("redis-server", "--port", Integer.toString(port),
|
||||
"--save", "", "--appendonly", "no", "--bind", "127.0.0.1")
|
||||
.redirectOutput(ProcessBuilder.Redirect.DISCARD)
|
||||
.redirectError(ProcessBuilder.Redirect.DISCARD)
|
||||
.start();
|
||||
RedisServer server = new RedisServer(p, port);
|
||||
server.awaitReady();
|
||||
return server;
|
||||
} catch (IOException | InterruptedException e) {
|
||||
throw new IllegalStateException("could not start redis-server (is it on the PATH?)", e);
|
||||
}
|
||||
}
|
||||
|
||||
public int port() {
|
||||
return port;
|
||||
}
|
||||
|
||||
private void awaitReady() throws InterruptedException {
|
||||
for (int i = 0; i < 100; i++) {
|
||||
try (Socket s = new Socket()) {
|
||||
s.connect(new InetSocketAddress("127.0.0.1", port), 100);
|
||||
return;
|
||||
} catch (IOException ignored) {
|
||||
Thread.sleep(50);
|
||||
}
|
||||
}
|
||||
throw new IllegalStateException("redis-server did not start on port " + port);
|
||||
}
|
||||
|
||||
/** Runs redis-cli against this server and returns trimmed stdout. */
|
||||
public String cli(String... args) {
|
||||
try {
|
||||
List<String> cmd = new java.util.ArrayList<>(List.of("redis-cli", "-p", Integer.toString(port)));
|
||||
cmd.addAll(List.of(args));
|
||||
Process p = new ProcessBuilder(cmd).redirectErrorStream(true).start();
|
||||
String out = new String(p.getInputStream().readAllBytes(), StandardCharsets.UTF_8).trim();
|
||||
p.waitFor();
|
||||
return out;
|
||||
} catch (IOException | InterruptedException e) {
|
||||
throw new IllegalStateException(e);
|
||||
}
|
||||
}
|
||||
|
||||
public void flush() {
|
||||
cli("FLUSHALL");
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() {
|
||||
process.destroy();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
package com.ankurm.ratelimit;
|
||||
|
||||
import java.time.Duration;
|
||||
|
||||
import io.github.bucket4j.Bucket;
|
||||
import io.github.bucket4j.ConsumptionProbe;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* The token bucket with no Spring and no Redis: an in-memory Bucket4j bucket on a clock the test
|
||||
* moves. Writes output/01-token-bucket-basics.txt.
|
||||
*/
|
||||
class TokenBucketBasicsTest {
|
||||
|
||||
@Test
|
||||
void greedyVersusIntervallyRefill() {
|
||||
try (Transcript t = Transcript.open("01-token-bucket-basics.txt")) {
|
||||
ManualClock clock = new ManualClock();
|
||||
|
||||
// Capacity 5, refilled at 5 tokens per 10 seconds.
|
||||
Bucket greedy = Bucket.builder().withCustomTimePrecision(clock)
|
||||
.addLimit(l -> l.capacity(5).refillGreedy(5, Duration.ofSeconds(10))).build();
|
||||
ManualClock clock2 = new ManualClock();
|
||||
Bucket intervally = Bucket.builder().withCustomTimePrecision(clock2)
|
||||
.addLimit(l -> l.capacity(5).refillIntervally(5, Duration.ofSeconds(10))).build();
|
||||
|
||||
t.line("== bucket: capacity 5, refill 5 tokens per 10 s");
|
||||
t.line("-- burst: five calls in the same instant");
|
||||
for (int i = 1; i <= 6; i++) {
|
||||
ConsumptionProbe p = greedy.tryConsumeAndReturnRemaining(1);
|
||||
t.line("call %d: consumed=%s remaining=%d waitForRefill=%ds", i, p.isConsumed(),
|
||||
p.getRemainingTokens(), p.getNanosToWaitForRefill() / 1_000_000_000L);
|
||||
if (i <= 5) {
|
||||
assertThat(p.isConsumed()).isTrue();
|
||||
} else {
|
||||
assertThat(p.isConsumed()).isFalse();
|
||||
assertThat(p.getNanosToWaitForRefill()).isEqualTo(Duration.ofSeconds(2).toNanos());
|
||||
}
|
||||
}
|
||||
for (int i = 0; i < 5; i++) {
|
||||
intervally.tryConsume(1);
|
||||
}
|
||||
|
||||
t.line("-- greedy refill: tokens trickle back at 1 per 2 s");
|
||||
int elapsed = 0;
|
||||
for (int seconds : new int[] {1, 1, 2, 4}) {
|
||||
clock.advance(Duration.ofSeconds(seconds));
|
||||
elapsed += seconds;
|
||||
t.line("after %d s: available=%d", elapsed, greedy.getAvailableTokens());
|
||||
}
|
||||
assertThat(greedy.getAvailableTokens()).isEqualTo(4);
|
||||
|
||||
t.line("-- intervally refill: nothing, then everything");
|
||||
clock2.advance(Duration.ofMillis(9_999));
|
||||
t.line("after 9.999 s: available=%d", intervally.getAvailableTokens());
|
||||
assertThat(intervally.getAvailableTokens()).isZero();
|
||||
clock2.advance(Duration.ofMillis(1));
|
||||
t.line("after 10.000 s: available=%d", intervally.getAvailableTokens());
|
||||
assertThat(intervally.getAvailableTokens()).isEqualTo(5);
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void twoLimitsOnOneBucket() {
|
||||
try (Transcript t = Transcript.open("01b-burst-and-sustained.txt")) {
|
||||
ManualClock clock = new ManualClock();
|
||||
// A burst allowance of 3 per second AND a sustained ceiling of 10 per minute.
|
||||
Bucket bucket = Bucket.builder().withCustomTimePrecision(clock)
|
||||
.addLimit(l -> l.capacity(3).refillGreedy(3, Duration.ofSeconds(1)))
|
||||
.addLimit(l -> l.capacity(10).refillGreedy(10, Duration.ofMinutes(1)))
|
||||
.build();
|
||||
|
||||
t.line("== two limits: 3 per second AND 10 per minute (both must have a token)");
|
||||
int allowed = 0;
|
||||
for (int second = 0; second < 5; second++) {
|
||||
int thisSecond = 0;
|
||||
for (int call = 0; call < 5; call++) {
|
||||
if (bucket.tryConsume(1)) {
|
||||
thisSecond++;
|
||||
}
|
||||
}
|
||||
allowed += thisSecond;
|
||||
t.line("second %d: 5 attempts, %d allowed, running total %d", second, thisSecond, allowed);
|
||||
clock.advance(Duration.ofSeconds(1));
|
||||
}
|
||||
assertThat(allowed).isEqualTo(10);
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void oneBucketPerInstanceMultipliesTheLimit() {
|
||||
try (Transcript t = Transcript.open("01c-per-instance-buckets.txt")) {
|
||||
ManualClock clock = new ManualClock();
|
||||
// What an in-memory limiter does behind a load balancer: every JVM holds its own bucket.
|
||||
Bucket instanceA = Bucket.builder().withCustomTimePrecision(clock)
|
||||
.addLimit(l -> l.capacity(5).refillGreedy(5, Duration.ofMinutes(1))).build();
|
||||
Bucket instanceB = Bucket.builder().withCustomTimePrecision(clock)
|
||||
.addLimit(l -> l.capacity(5).refillGreedy(5, Duration.ofMinutes(1))).build();
|
||||
|
||||
t.line("== the same 5-per-minute limit held in memory by two instances");
|
||||
int allowed = 0;
|
||||
for (int i = 0; i < 12; i++) {
|
||||
Bucket target = i % 2 == 0 ? instanceA : instanceB;
|
||||
if (target.tryConsume(1)) {
|
||||
allowed++;
|
||||
}
|
||||
}
|
||||
t.line("12 requests alternating A, B: %d allowed (the limit is 5)", allowed);
|
||||
assertThat(allowed).isEqualTo(10);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
package com.ankurm.ratelimit;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* Collects the lines a test prints and writes them to {@code output/NN-name.txt}. The same test
|
||||
* asserts the same numbers, so a line in a transcript is something the build checks.
|
||||
*/
|
||||
public final class Transcript implements AutoCloseable {
|
||||
|
||||
private final Path file;
|
||||
private final List<String> lines = new ArrayList<>();
|
||||
|
||||
private Transcript(Path file) {
|
||||
this.file = file;
|
||||
}
|
||||
|
||||
public static Transcript open(String name) {
|
||||
return new Transcript(Path.of("output", name));
|
||||
}
|
||||
|
||||
public synchronized Transcript line(String format, Object... args) {
|
||||
String text = args.length == 0 ? format : String.format(format, args);
|
||||
lines.add(text);
|
||||
System.out.println(text);
|
||||
return this;
|
||||
}
|
||||
|
||||
@Override
|
||||
public synchronized void close() {
|
||||
try {
|
||||
Files.createDirectories(file.getParent());
|
||||
Files.write(file, lines);
|
||||
} catch (IOException e) {
|
||||
throw new IllegalStateException(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,167 @@
|
||||
package com.ankurm.ratelimit;
|
||||
|
||||
import java.net.URI;
|
||||
import java.net.http.HttpClient;
|
||||
import java.net.http.HttpRequest;
|
||||
import java.net.http.HttpResponse;
|
||||
import java.time.Duration;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.ExecutorService;
|
||||
import java.util.concurrent.Executors;
|
||||
import java.util.concurrent.Future;
|
||||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
|
||||
import org.junit.jupiter.api.AfterAll;
|
||||
import org.junit.jupiter.api.BeforeAll;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.MethodOrderer;
|
||||
import org.junit.jupiter.api.Order;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.TestMethodOrder;
|
||||
|
||||
import org.springframework.boot.builder.SpringApplicationBuilder;
|
||||
import org.springframework.context.ConfigurableApplicationContext;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* Two complete Spring Boot applications in one JVM, each with its own server port, its own
|
||||
* Lettuce connection and its own clock, both pointed at one redis-server.
|
||||
* Writes output/05-two-instances.txt.
|
||||
*/
|
||||
@TestMethodOrder(MethodOrderer.OrderAnnotation.class)
|
||||
class TwoInstancesTest {
|
||||
|
||||
static final RedisServer REDIS = RedisServer.shared();
|
||||
static Transcript t;
|
||||
static final HttpClient http = HttpClient.newHttpClient();
|
||||
|
||||
static final ManualClock clockA = new ManualClock();
|
||||
static final ManualClock clockB = new ManualClock();
|
||||
static ConfigurableApplicationContext a;
|
||||
static ConfigurableApplicationContext b;
|
||||
|
||||
@BeforeAll
|
||||
static void start() {
|
||||
t = Transcript.open("05-two-instances.txt");
|
||||
a = boot("A", clockA);
|
||||
b = boot("B", clockB);
|
||||
}
|
||||
|
||||
@AfterAll
|
||||
static void stop() {
|
||||
a.close();
|
||||
b.close();
|
||||
t.close();
|
||||
}
|
||||
|
||||
@BeforeEach
|
||||
void reset() {
|
||||
REDIS.flush();
|
||||
}
|
||||
|
||||
static ConfigurableApplicationContext boot(String name, ManualClock clock) {
|
||||
// Command-line arguments, not .properties(): the latter are defaults that application.yml overrides.
|
||||
return new SpringApplicationBuilder(RateLimitApplication.class)
|
||||
.initializers(ctx -> ((ConfigurableApplicationContext) ctx).getBeanFactory()
|
||||
.registerSingleton("rateLimitClock", clock))
|
||||
.run("--server.port=0", "--instance.name=" + name, "--spring.main.banner-mode=off",
|
||||
"--spring.data.redis.port=" + REDIS.port(), "--logging.level.root=WARN");
|
||||
}
|
||||
|
||||
static String url(ConfigurableApplicationContext ctx, String path) {
|
||||
return "http://localhost:" + ctx.getEnvironment().getProperty("local.server.port") + path;
|
||||
}
|
||||
|
||||
static HttpResponse<String> get(ConfigurableApplicationContext ctx, String path, String header, String value) {
|
||||
try {
|
||||
return http.send(HttpRequest.newBuilder(URI.create(url(ctx, path))).header(header, value).build(),
|
||||
HttpResponse.BodyHandlers.ofString());
|
||||
} catch (Exception e) {
|
||||
throw new IllegalStateException(e);
|
||||
}
|
||||
}
|
||||
|
||||
static String served(HttpResponse<String> r) {
|
||||
String body = r.body();
|
||||
int i = body.indexOf("\"instance\":\"");
|
||||
return i < 0 ? "-" : body.substring(i + 12, i + 13);
|
||||
}
|
||||
|
||||
@Test
|
||||
@Order(1)
|
||||
void oneBucketAcrossTwoInstances() {
|
||||
t.line("== free-key-123 allows 5 per minute; requests alternate between instance A and B");
|
||||
String[] order = {"A", "B", "A", "B", "A", "B", "A"};
|
||||
int allowed = 0;
|
||||
for (String which : order) {
|
||||
var r = get("A".equals(which) ? a : b, "/api/hello", "X-API-Key", "free-key-123");
|
||||
t.line("sent to %s -> %d remaining=%s (served by %s)", which, r.statusCode(),
|
||||
r.headers().firstValue("X-RateLimit-Remaining").orElse("-"), r.statusCode() == 200 ? served(r) : "-");
|
||||
if (r.statusCode() == 200) {
|
||||
allowed++;
|
||||
}
|
||||
}
|
||||
assertThat(allowed).isEqualTo(5);
|
||||
t.line("allowed in total: %d (not 10)", allowed);
|
||||
}
|
||||
|
||||
@Test
|
||||
@Order(2)
|
||||
void redisHoldsOneKey() {
|
||||
t.line("== one bucket in Redis, written by both instances");
|
||||
get(a, "/api/hello", "X-API-Key", "pro-key-456");
|
||||
get(b, "/api/hello", "X-API-Key", "pro-key-456");
|
||||
t.line("$ redis-cli --scan");
|
||||
List<String> keys = REDIS.cli("--scan").lines().sorted().toList();
|
||||
keys.forEach(k -> t.line(k));
|
||||
assertThat(keys).hasSize(2); // the key bucket and the global bucket
|
||||
}
|
||||
|
||||
@Test
|
||||
@Order(3)
|
||||
void twoHundredConcurrentRequestsAcrossBothInstances() throws Exception {
|
||||
t.line("== 200 concurrent requests for pro-key-456 (20 per minute), half to A and half to B");
|
||||
AtomicInteger ok = new AtomicInteger();
|
||||
AtomicInteger limited = new AtomicInteger();
|
||||
try (ExecutorService pool = Executors.newVirtualThreadPerTaskExecutor()) {
|
||||
List<Future<?>> futures = new ArrayList<>();
|
||||
for (int i = 0; i < 200; i++) {
|
||||
ConfigurableApplicationContext target = i % 2 == 0 ? a : b;
|
||||
futures.add(pool.submit(() -> {
|
||||
int status = get(target, "/api/hello", "X-API-Key", "pro-key-456").statusCode();
|
||||
(status == 200 ? ok : limited).incrementAndGet();
|
||||
}));
|
||||
}
|
||||
for (Future<?> f : futures) {
|
||||
f.get();
|
||||
}
|
||||
}
|
||||
t.line("200 sent: %d allowed, %d rejected with 429", ok.get(), limited.get());
|
||||
assertThat(ok.get()).isEqualTo(20);
|
||||
assertThat(limited.get()).isEqualTo(180);
|
||||
}
|
||||
|
||||
@Test
|
||||
@Order(4)
|
||||
void aClockThatRunsFastLeaksTokens() {
|
||||
t.line("== instance B's clock is 30 s ahead of A's (user carol: 10 per minute, one token per 6 s)");
|
||||
for (int i = 0; i < 10; i++) {
|
||||
get(a, "/api/hello", "X-User-Id", "carol");
|
||||
}
|
||||
var rejectedByA = get(a, "/api/hello", "X-User-Id", "carol");
|
||||
t.line("A, bucket empty -> %d", rejectedByA.statusCode());
|
||||
assertThat(rejectedByA.statusCode()).isEqualTo(429);
|
||||
|
||||
clockB.advance(Duration.ofSeconds(30));
|
||||
int allowedByB = 0;
|
||||
for (int i = 0; i < 8; i++) {
|
||||
if (get(b, "/api/hello", "X-User-Id", "carol").statusCode() == 200) {
|
||||
allowedByB++;
|
||||
}
|
||||
}
|
||||
t.line("B sees 30 s of refill that never happened on A's clock: %d of 8 allowed", allowedByB);
|
||||
assertThat(allowedByB).isEqualTo(5);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user