From cb6a231066d847ffa7cdf70dfbf106a5870120dc Mon Sep 17 00:00:00 2001 From: Spencer Gibb Date: Wed, 9 Aug 2017 12:10:08 -0600 Subject: [PATCH] Initial move off of lua --- .../config/GatewayAutoConfiguration.java | 28 ++---- .../gateway/filter/ratelimit/RateLimiter.java | 2 +- .../filter/ratelimit/RedisRateLimiter.java | 94 ++++++++++++++----- .../ratelimit/RedisRateLimiterTests.java | 2 +- 4 files changed, 81 insertions(+), 45 deletions(-) diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/config/GatewayAutoConfiguration.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/config/GatewayAutoConfiguration.java index cca972b9..f3417a8b 100644 --- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/config/GatewayAutoConfiguration.java +++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/config/GatewayAutoConfiguration.java @@ -19,7 +19,6 @@ package org.springframework.cloud.gateway.config; import java.util.List; -import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.boot.actuate.endpoint.Endpoint; import org.springframework.boot.autoconfigure.AutoConfigureAfter; import org.springframework.boot.autoconfigure.AutoConfigureBefore; @@ -78,18 +77,12 @@ import org.springframework.cloud.gateway.route.RouteLocator; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; import org.springframework.context.annotation.Primary; -import org.springframework.core.io.ClassPathResource; -import org.springframework.data.redis.core.RedisTemplate; -import org.springframework.data.redis.core.StringRedisTemplate; -import org.springframework.data.redis.core.script.DefaultRedisScript; -import org.springframework.data.redis.core.script.RedisScript; -import org.springframework.scripting.support.ResourceScriptSource; +import org.springframework.data.redis.core.ReactiveRedisTemplate; import com.netflix.hystrix.HystrixObservableCommand; import reactor.core.publisher.Flux; import reactor.ipc.netty.http.client.HttpClient; -import reactor.ipc.netty.resources.PoolResources; import rx.RxReactiveStreams; /** @@ -109,8 +102,8 @@ public class GatewayAutoConfiguration { @ConditionalOnMissingBean public HttpClient httpClient() { return HttpClient.create(opts -> { - opts.poolResources(PoolResources.elastic("proxy")); - // opts.disablePool(); //TODO: why do I need this again? + // opts.poolResources(PoolResources.elastic("proxy")); + opts.disablePool(); //TODO: why do I need this again? }); } @@ -321,20 +314,11 @@ public class GatewayAutoConfiguration { } - @ConditionalOnClass(RedisTemplate.class) + @ConditionalOnClass(ReactiveRedisTemplate.class) protected static class GatewayRedisConfiguration { @Bean - public RedisScript redistRequestRateLimiterScript() { - DefaultRedisScript redisScript = new DefaultRedisScript<>(); - redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("META-INF/scripts/request_rate_limiter.lua"))); - redisScript.setResultType(List.class); - return redisScript; - } - - @Bean - public RedisRateLimiter redisRateLimiter(StringRedisTemplate redisTemplate, - @Qualifier("redistRequestRateLimiterScript") RedisScript redisScript) { - return new RedisRateLimiter(redisTemplate, redisScript); + public RedisRateLimiter redisRateLimiter(ReactiveRedisTemplate redisTemplate) { + return new RedisRateLimiter(redisTemplate); } } 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 cd85b1d8..4499cbcf 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 @@ -4,7 +4,7 @@ package org.springframework.cloud.gateway.filter.ratelimit; * @author Spencer Gibb */ public interface RateLimiter { - Response isAllowed(String id, int replenishRate, int burstCapacity); + Response 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 100ff958..1751f3a1 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 @@ -1,13 +1,19 @@ package org.springframework.cloud.gateway.filter.ratelimit; +import java.time.Duration; +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; -import org.springframework.data.redis.core.StringRedisTemplate; -import org.springframework.data.redis.core.script.RedisScript; +import org.springframework.data.redis.core.ReactiveRedisTemplate; -import java.time.Instant; -import java.util.Arrays; -import java.util.List; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.util.function.Tuple2; /** * See https://stripe.com/blog/rate-limiters and @@ -18,12 +24,10 @@ import java.util.List; public class RedisRateLimiter implements RateLimiter { private Log log = LogFactory.getLog(getClass()); - private final StringRedisTemplate redisTemplate; - private final RedisScript script; + private final ReactiveRedisTemplate redisTemplate; - public RedisRateLimiter(StringRedisTemplate redisTemplate, RedisScript script) { + public RedisRateLimiter(ReactiveRedisTemplate redisTemplate) { this.redisTemplate = redisTemplate; - this.script = script; } /** @@ -37,29 +41,77 @@ public class RedisRateLimiter implements RateLimiter { @Override //TODO: signature? params (tuple?). //TODO: change to Mono - public Response isAllowed(String id, int replenishRate, int burstCapacity) { + public Response isAllowed(String id, long replenishRate, long burstCapacity) { try { // Make a unique key per user. - String prefix = "request_rate_limiter." + id; + String key = "request_rate_limiter." + id; // You need two Redis keys for Token Bucket. - List keys = Arrays.asList(prefix + ".tokens", prefix + ".timestamp"); + // String tokensKey = key + ".tokens"; + // String timestampKey = key + ".timestamp"; // The arguments to the LUA script. time() returns unixtime in seconds. - Object[] args = new String[]{ replenishRate+"", burstCapacity +"", Instant.now().getEpochSecond()+"", "1"}; - // allowed, tokens_left = redis.eval(SCRIPT, keys, args) - List results = this.redisTemplate.execute(this.script, keys, args); + long now = Instant.now().getEpochSecond(); + int requested = 1; - boolean allowed = new Long(1L).equals(results.get(0)); - Long tokensLeft = (Long) results.get(1); + double fillTime = (double)burstCapacity / (double)replenishRate; + int ttl = (int)Math.floor(fillTime * 2); - Response response = new Response(allowed, tokensLeft); + Mono booleanMono = this.redisTemplate.hasKey(key); + Boolean hasKey = booleanMono.block(); - if (log.isDebugEnabled()) { - log.debug("response: "+response); + Mono> valuesMono; + if (hasKey) { + valuesMono = this.redisTemplate.opsForHash().multiGet(key, Arrays.asList("tokens", "timestamp")); + } else { + valuesMono = Mono.just(new ArrayList<>()); } - return response; + Mono responseMono = valuesMono.map(objects -> { + Long lastTokens = null; + + if (objects.size() >= 1) { + lastTokens= (Long) objects.get(0); + } + if (lastTokens == null) { + lastTokens = burstCapacity; + } + + Long lastRefreshed = null; + if (objects.size() >= 2) { + lastRefreshed = (Long) objects.get(1); + } + if (lastRefreshed == null) { + lastRefreshed = 0L; + } + + long delta = Math.max(0, (now - lastRefreshed)); + long filledTokens = Math.min(burstCapacity, lastTokens + (delta * replenishRate)); + boolean allowed = filledTokens >= requested; + long newTokens = filledTokens; + if (allowed) { + newTokens = filledTokens - requested; + } + + HashMap updated = new HashMap<>(); + updated.put("tokens", newTokens); + updated.put("timestamp", now); + Mono putAllMono = this.redisTemplate.opsForHash().putAll(key, updated); + Mono expireMono = this.redisTemplate.expire(key, Duration.ofSeconds(ttl)); + + Flux> zip = Flux.zip(putAllMono, expireMono); + Tuple2 objects1 = zip.blockLast(); + + Response response = new Response(allowed, newTokens); + + if (log.isDebugEnabled()) { + log.debug("response: " + response); + } + + return response; + }); + + return responseMono.block(); } catch (Exception e) { /* We don't want a hard dependency on Redis to allow traffic. 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 22d758f4..9d5f5046 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 @@ -33,7 +33,7 @@ public class RedisRateLimiterTests extends BaseWebClientTests { public void requestRateLimiterWebFilterFactoryWorks() throws Exception { String id = UUID.randomUUID().toString(); - int replenishRate = 10; + int replenishRate = 1;//10; int burstCapacity = 2 * replenishRate; // Bursts work