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

Move RateLimiter interface to return Mono<Response>

Browse Source
This commit is contained in:
Spencer Gibb
2017-08-09 13:05:34 -06:00
parent 539721fa7a
commit 4d22218dbd
5 changed files with 30 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -17,24 +17,23 @@
package org.springframework.cloud.gateway.filter.factory;
import java.util.Arrays;
import java.util.List;
import org.springframework.beans.BeansException;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter.Response;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.http.HttpStatus;
import org.springframework.tuple.Tuple;
import org.springframework.web.server.WebFilter;
import java.util.Arrays;
import java.util.List;
/**
* 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 RequestRateLimiterWebFilterFactory implements WebFilterFactory, ApplicationContextAware {
public class RequestRateLimiterWebFilterFactory
implements WebFilterFactory, ApplicationContextAware {
public static final String REPLENISH_RATE_KEY = "replenishRate";
public static final String BURST_CAPACITY_KEY = "burstCapacity";
@@ -52,10 +51,10 @@ public class RequestRateLimiterWebFilterFactory implements WebFilterFactory, App
this.context = context;
}
@Override
public List<String> argNames() {
return Arrays.asList(REPLENISH_RATE_KEY, BURST_CAPACITY_KEY, KEY_RESOLVER_NAME_KEY);
return Arrays.asList(REPLENISH_RATE_KEY, BURST_CAPACITY_KEY,
KEY_RESOLVER_NAME_KEY);
}
@SuppressWarnings("unchecked")
@@ -70,18 +69,16 @@ public class RequestRateLimiterWebFilterFactory implements WebFilterFactory, App
String beanName = args.getString(KEY_RESOLVER_NAME_KEY);
KeyResolver keyResolver = this.context.getBean(beanName, KeyResolver.class);
return (exchange, chain) ->
keyResolver.resolve(exchange).flatMap(key -> {
Response response = rateLimiter.isAllowed(key, replenishRate, capacity);
//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();
});
return (exchange, chain) -> keyResolver.resolve(exchange)
.flatMap(key -> rateLimiter.isAllowed(key, replenishRate, capacity)
.flatMap(response -> {
if (response.isAllowed()) {
return chain.filter(exchange);
}
exchange.getResponse()
.setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
return exchange.getResponse().setComplete();
}));
}
}

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

@@ -1,10 +1,12 @@
package org.springframework.cloud.gateway.filter.ratelimit;
import reactor.core.publisher.Mono;
/**
* @author Spencer Gibb
*/
public interface RateLimiter {
Response isAllowed(String id, long replenishRate, long burstCapacity);
Mono<Response> isAllowed(String id, long replenishRate, long burstCapacity);
class Response {
private final boolean allowed;

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

@@ -5,7 +5,6 @@ import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
@@ -41,8 +40,7 @@ public class RedisRateLimiter implements RateLimiter {
*/
@Override
// TODO: signature? params (tuple?).
// TODO: change to Mono<?>
public Response isAllowed(String id, long replenishRate, long burstCapacity) {
public Mono<Response> isAllowed(String id, long replenishRate, long burstCapacity) {
try {
// Make a unique key per user.
@@ -110,7 +108,7 @@ public class RedisRateLimiter implements RateLimiter {
}
return response;
}).block();
});
}
catch (Exception e) {
@@ -121,6 +119,6 @@ public class RedisRateLimiter implements RateLimiter {
*/
log.error("Error determining if user allowed from redis", e);
}
return new Response(true, -1);
return Mono.just(new Response(true, -1));
}
}

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

@@ -67,7 +67,7 @@ public class RequestRateLimiterWebFilterFactoryTests extends BaseWebClientTests
KEY_RESOLVER_NAME_KEY, keyResolverName);
when(rateLimiter.isAllowed(key, replenishRate, burstCapacity))
.thenReturn(new Response(allowed, 1));
.thenReturn(Mono.just(new Response(allowed, 1)));
MockServerHttpRequest request = MockServerHttpRequest.get("/").build();

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

@@ -38,13 +38,13 @@ public class RedisRateLimiterTests extends BaseWebClientTests {
// Bursts work
for (int i = 0; i < burstCapacity; i++) {
Response response = rateLimiter.isAllowed(id, replenishRate, burstCapacity);
Response response = rateLimiter.isAllowed(id, replenishRate, burstCapacity).block();
assertThat(response.isAllowed()).as("Burst # %s is allowed", i).isTrue();
}
Response response = rateLimiter.isAllowed(id, replenishRate, burstCapacity);
Response response = rateLimiter.isAllowed(id, replenishRate, burstCapacity).block();
if (response.isAllowed()) { //TODO: sometimes there is an off by one error
response = rateLimiter.isAllowed(id, replenishRate, burstCapacity);
response = rateLimiter.isAllowed(id, replenishRate, burstCapacity).block();
}
assertThat(response.isAllowed()).as("Burst # %s is not allowed", burstCapacity).isFalse();
@@ -52,11 +52,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);
response = rateLimiter.isAllowed(id, replenishRate, burstCapacity).block();
assertThat(response.isAllowed()).as("steady state # %s is allowed", i).isTrue();
}
response = rateLimiter.isAllowed(id, replenishRate, burstCapacity);
response = rateLimiter.isAllowed(id, replenishRate, burstCapacity).block();
assertThat(response.isAllowed()).as("steady state # %s is allowed", replenishRate).isFalse();
}