From 4d22218dbdad426b029d2012a5b90d6e50119e3b Mon Sep 17 00:00:00 2001 From: Spencer Gibb Date: Wed, 9 Aug 2017 13:05:34 -0600 Subject: [PATCH] Move RateLimiter interface to return Mono --- .../RequestRateLimiterWebFilterFactory.java | 39 +++++++++---------- .../gateway/filter/ratelimit/RateLimiter.java | 4 +- .../filter/ratelimit/RedisRateLimiter.java | 8 ++-- ...questRateLimiterWebFilterFactoryTests.java | 2 +- .../ratelimit/RedisRateLimiterTests.java | 10 ++--- 5 files changed, 30 insertions(+), 33 deletions(-) diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterWebFilterFactory.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterWebFilterFactory.java index 9f54f089..c35656ae 100644 --- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterWebFilterFactory.java +++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterWebFilterFactory.java @@ -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 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(); + })); } } diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RateLimiter.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RateLimiter.java index 4499cbcf..53c5904a 100644 --- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RateLimiter.java +++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RateLimiter.java @@ -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 isAllowed(String id, long replenishRate, long burstCapacity); class Response { private final boolean allowed; diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiter.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiter.java index 3c8ef013..3101e66b 100644 --- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiter.java +++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiter.java @@ -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 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)); } } diff --git a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterWebFilterFactoryTests.java b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterWebFilterFactoryTests.java index 53b7abdf..2af4b1ff 100644 --- a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterWebFilterFactoryTests.java +++ b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterWebFilterFactoryTests.java @@ -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(); diff --git a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterTests.java b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterTests.java index 9d5f5046..508de657 100644 --- a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterTests.java +++ b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterTests.java @@ -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(); }