SpringCloud/spring-cloud-gateway
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Move RateLimiter to tuple args

Browse Source
This commit is contained in:
Spencer Gibb
2017-10-31 13:07:47 -04:00
parent 5bff2b22b3
commit 7659f89bc9
6 changed files with 60 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterGatewayFilterFactory.java
View File

@@ -17,9 +17,6 @@
package org.springframework.cloud.gateway.filter.factory;
import java.util.Arrays;
import java.util.List;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
@@ -27,69 +24,44 @@ import org.springframework.http.HttpStatus;
import org.springframework.tuple.Tuple;
/**
* User Request Rate Limiter filter.
* See https://stripe.com/blog/rate-limiters and
* User Request Rate Limiter filter. See https://stripe.com/blog/rate-limiters and
*/
public class RequestRateLimiterGatewayFilterFactory implements GatewayFilterFactory {
public static final String REPLENISH_RATE_KEY = "replenishRate";
public static final String BURST_CAPACITY_KEY = "burstCapacity";
public static final String KEY_RESOLVER_KEY = "keyResolver";
private final RateLimiter rateLimiter;
private final KeyResolver defaultKeyResolver;
public RequestRateLimiterGatewayFilterFactory(RateLimiter rateLimiter, KeyResolver defaultKeyResolver) {
public RequestRateLimiterGatewayFilterFactory(RateLimiter rateLimiter,
KeyResolver defaultKeyResolver) {
this.rateLimiter = rateLimiter;
this.defaultKeyResolver = defaultKeyResolver;
}
@Override
public List<String> argNames() {
return Arrays.asList(REPLENISH_RATE_KEY, BURST_CAPACITY_KEY, KEY_RESOLVER_KEY);
}
public GatewayFilter apply(int replenishRate, int burstCapacity) {
return apply(replenishRate, burstCapacity, this.defaultKeyResolver);
}
public GatewayFilter apply(int replenishRate, int burstCapacity, KeyResolver keyResolver) {
return (exchange, chain) ->
keyResolver.resolve(exchange).flatMap(key ->
//TODO: if key is empty?
rateLimiter.isAllowed(key, replenishRate, burstCapacity).flatMap(response -> {
//TODO: set some headers for rate, tokens left
if (response.isAllowed()) {
return chain.filter(exchange);
}
exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
return exchange.getResponse().setComplete();
}));
}
@SuppressWarnings("unchecked")
@Override
public GatewayFilter apply(Tuple args) {
// How many requests per second do you want a user to be allowed to do?
int replenishRate = args.getInt(REPLENISH_RATE_KEY);
// How much bursting do you want to allow?
int burstCapacity;
if (args.hasFieldName(BURST_CAPACITY_KEY)) {
burstCapacity = args.getInt(BURST_CAPACITY_KEY);
} else {
burstCapacity = 0;
}
KeyResolver keyResolver;
if (args.hasFieldName(KEY_RESOLVER_KEY)) {
keyResolver = args.getValue(KEY_RESOLVER_KEY, KeyResolver.class);
} else {
}
else {
keyResolver = defaultKeyResolver;
}
return apply(replenishRate, burstCapacity, keyResolver);
return (exchange, chain) -> keyResolver.resolve(exchange).flatMap(key ->
// TODO: if key is empty?
rateLimiter.isAllowed(key, args).flatMap(response -> {
// TODO: set some headers for rate, tokens left
if (response.isAllowed()) {
return chain.filter(exchange);
}
exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
return exchange.getResponse().setComplete();
}));
}
}

5
spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RateLimiter.java
View File

@@ -1,13 +1,14 @@
package org.springframework.cloud.gateway.filter.ratelimit;
import org.springframework.tuple.Tuple;
import reactor.core.publisher.Mono;
/**
* @author Spencer Gibb
*/
public interface RateLimiter {
//TODO: move ints to tuple
Mono<Response> isAllowed(String id, int replenishRate, int burstCapacity);
Mono<Response> isAllowed(String id, Tuple args);
class Response {
private final boolean allowed;

25
spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiter.java
View File

@@ -4,13 +4,13 @@ import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.logging.Level;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.data.redis.core.ReactiveRedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.tuple.Tuple;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
@@ -21,6 +21,9 @@ import reactor.core.publisher.Mono;
* @author Spencer Gibb
*/
public class RedisRateLimiter implements RateLimiter {
public static final String REPLENISH_RATE_KEY = "replenishRate";
public static final String BURST_CAPACITY_KEY = "burstCapacity";
private Log log = LogFactory.getLog(getClass());
private final ReactiveRedisTemplate<String, String> redisTemplate;
@@ -36,15 +39,23 @@ public class RedisRateLimiter implements RateLimiter {
* This uses a basic token bucket algorithm and relies on the fact that Redis scripts
* execute atomically. No other operations can run between fetching the count and
* writing the new count.
* @param replenishRate
* @param burstCapacity
* @param id
* @param args
* @return
*/
@Override
// TODO: signature? params (tuple?).
@SuppressWarnings("unchecked")
public Mono<Response> isAllowed(String id, int replenishRate, int burstCapacity) {
public Mono<Response> isAllowed(String id, Tuple args) {
// How many requests per second do you want a user to be allowed to do?
int replenishRate = args.getInt(REPLENISH_RATE_KEY);
// How much bursting do you want to allow?
int burstCapacity;
if (args.hasFieldName(BURST_CAPACITY_KEY)) {
burstCapacity = args.getInt(BURST_CAPACITY_KEY);
} else {
burstCapacity = 0;
}
try {
// Make a unique key per user.
@@ -54,10 +65,10 @@ public class RedisRateLimiter implements RateLimiter {
List<String> keys = Arrays.asList(prefix + ".tokens", prefix + ".timestamp");
// The arguments to the LUA script. time() returns unixtime in seconds.
List<String> args = Arrays.asList(replenishRate + "", burstCapacity + "",
List<String> scriptArgs = Arrays.asList(replenishRate + "", burstCapacity + "",
Instant.now().getEpochSecond() + "", "1");
// allowed, tokens_left = redis.eval(SCRIPT, keys, args)
Flux<List<Long>> flux = this.redisTemplate.execute(this.script, keys, args);
Flux<List<Long>> flux = this.redisTemplate.execute(this.script, keys, scriptArgs);
// .log("redisratelimiter", Level.FINER);
return flux.onErrorResume(throwable -> Flux.just(Arrays.asList(1L, -1L)))
.reduce(new ArrayList<Long>(), (longs, l) -> {

14
spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterGatewayFilterFactoryTests.java
View File

@@ -20,10 +20,14 @@ import org.springframework.mock.http.server.reactive.MockServerHttpRequest;
import org.springframework.mock.web.server.MockServerWebExchange;
import org.springframework.test.annotation.DirtiesContext;
import org.springframework.test.context.junit4.SpringRunner;
import org.springframework.tuple.Tuple;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.when;
import static org.springframework.boot.test.context.SpringBootTest.WebEnvironment.RANDOM_PORT;
import static org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter.BURST_CAPACITY_KEY;
import static org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter.REPLENISH_RATE_KEY;
import static org.springframework.tuple.TupleBuilder.tuple;
import reactor.core.publisher.Mono;
@@ -67,7 +71,8 @@ public class RequestRateLimiterGatewayFilterFactoryTests extends BaseWebClientTe
int replenishRate = 10;
int burstCapacity = 2 * replenishRate;
when(rateLimiter.isAllowed(key, replenishRate, burstCapacity))
Tuple args = tuple().of(REPLENISH_RATE_KEY, replenishRate, BURST_CAPACITY_KEY, burstCapacity);
when(rateLimiter.isAllowed(key, args))
.thenReturn(Mono.just(new Response(allowed, 1)));
@@ -77,10 +82,9 @@ public class RequestRateLimiterGatewayFilterFactoryTests extends BaseWebClientTe
when(this.filterChain.filter(exchange)).thenReturn(Mono.empty());
Mono<Void> response = filterFactory.apply(replenishRate, burstCapacity, keyResolver).filter(exchange, this.filterChain);
response.subscribe(aVoid -> {
assertThat(exchange.getResponse().getStatusCode()).isEqualTo(expectedStatus);
});
Mono<Void> response = filterFactory.apply(args).filter(exchange, this.filterChain);
response.subscribe(aVoid -> assertThat(exchange.getResponse().getStatusCode())
.isEqualTo(expectedStatus));
}

6
spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/PrincipalNameKeyResolverIntegrationTests.java
View File

@@ -33,8 +33,8 @@ import org.springframework.web.bind.annotation.RestController;
import static org.springframework.boot.test.context.SpringBootTest.WebEnvironment.DEFINED_PORT;
import static org.springframework.cloud.gateway.filter.factory.GatewayFilters.prefixPath;
import static org.springframework.cloud.gateway.filter.factory.RequestRateLimiterGatewayFilterFactory.BURST_CAPACITY_KEY;
import static org.springframework.cloud.gateway.filter.factory.RequestRateLimiterGatewayFilterFactory.REPLENISH_RATE_KEY;
import static org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter.BURST_CAPACITY_KEY;
import static org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter.REPLENISH_RATE_KEY;
import static org.springframework.cloud.gateway.handler.predicate.RoutePredicates.path;
import static org.springframework.tuple.TupleBuilder.tuple;
import static org.springframework.web.reactive.function.client.ExchangeFilterFunctions.basicAuthentication;
@@ -108,7 +108,7 @@ public class PrincipalNameKeyResolverIntegrationTests {
@Bean
RateLimiter rateLimiter() {
return (id, replenishRate, burstCapacity) -> Mono.just(new RateLimiter.Response(true, Long.MAX_VALUE));
return (id, args) -> Mono.just(new RateLimiter.Response(true, Long.MAX_VALUE));
}
@Bean

16
spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterTests.java
View File

@@ -13,9 +13,13 @@ import org.springframework.cloud.gateway.test.BaseWebClientTests;
import org.springframework.context.annotation.Import;
import org.springframework.test.annotation.DirtiesContext;
import org.springframework.test.context.junit4.SpringRunner;
import org.springframework.tuple.Tuple;
import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.boot.test.context.SpringBootTest.WebEnvironment.RANDOM_PORT;
import static org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter.BURST_CAPACITY_KEY;
import static org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter.REPLENISH_RATE_KEY;
import static org.springframework.tuple.TupleBuilder.tuple;
/**
* see https://gist.github.com/ptarjan/e38f45f2dfe601419ca3af937fff574d#file-1-check_request_rate_limiter-rb-L36-L62
@@ -36,15 +40,17 @@ public class RedisRateLimiterTests extends BaseWebClientTests {
int replenishRate = 10;
int burstCapacity = 2 * replenishRate;
Tuple args = tuple().of(REPLENISH_RATE_KEY, replenishRate, BURST_CAPACITY_KEY, burstCapacity);
// Bursts work
for (int i = 0; i < burstCapacity; i++) {
Response response = rateLimiter.isAllowed(id, replenishRate, burstCapacity).block();
Response response = rateLimiter.isAllowed(id, args).block();
assertThat(response.isAllowed()).as("Burst # %s is allowed", i).isTrue();
}
Response response = rateLimiter.isAllowed(id, replenishRate, burstCapacity).block();
Response response = rateLimiter.isAllowed(id, args).block();
if (response.isAllowed()) { //TODO: sometimes there is an off by one error
response = rateLimiter.isAllowed(id, replenishRate, burstCapacity).block();
response = rateLimiter.isAllowed(id, args).block();
}
assertThat(response.isAllowed()).as("Burst # %s is not allowed", burstCapacity).isFalse();
@@ -52,11 +58,11 @@ public class RedisRateLimiterTests extends BaseWebClientTests {
// # After the burst is done, check the steady state
for (int i = 0; i < replenishRate; i++) {
response = rateLimiter.isAllowed(id, replenishRate, burstCapacity).block();
response = rateLimiter.isAllowed(id, args).block();
assertThat(response.isAllowed()).as("steady state # %s is allowed", i).isTrue();
}
response = rateLimiter.isAllowed(id, replenishRate, burstCapacity).block();
response = rateLimiter.isAllowed(id, args).block();
assertThat(response.isAllowed()).as("steady state # %s is allowed", replenishRate).isFalse();
}