Support ReactiveRedisTemplate execution of RedisScript.

This commit is contained in:
Spencer Gibb
2017-10-02 17:23:53 -04:00
parent f15c1e633c
commit b997943fb9
10 changed files with 125 additions and 64 deletions

View File

@@ -55,7 +55,6 @@ import org.springframework.cloud.gateway.filter.factory.SetStatusWebFilterFactor
import org.springframework.cloud.gateway.filter.factory.WebFilterFactory;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter;
import org.springframework.cloud.gateway.handler.FilteringWebHandler;
import org.springframework.cloud.gateway.handler.RoutePredicateHandlerMapping;
import org.springframework.cloud.gateway.handler.predicate.AfterRoutePredicateFactory;
@@ -81,12 +80,6 @@ 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.web.reactive.socket.client.ReactorNettyWebSocketClient;
import org.springframework.web.reactive.socket.client.WebSocketClient;
import org.springframework.web.reactive.socket.server.WebSocketService;
@@ -361,24 +354,6 @@ public class GatewayAutoConfiguration {
}
@ConditionalOnClass(RedisTemplate.class)
protected static class GatewayRedisConfiguration {
@Bean
public RedisScript<List> redisRequestRateLimiterScript() {
DefaultRedisScript<List> 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("redisRequestRateLimiterScript") RedisScript<List> redisScript) {
return new RedisRateLimiter(redisTemplate, redisScript);
}
}
@Configuration
@ConditionalOnClass(Health.class)
protected static class GatewayActuatorConfiguration {

View File

@@ -0,0 +1,63 @@
package org.springframework.cloud.gateway.config;
import java.util.List;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.autoconfigure.AutoConfigureAfter;
import org.springframework.boot.autoconfigure.AutoConfigureBefore;
import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.data.redis.RedisReactiveAutoConfiguration;
import org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.ResourceLoader;
import org.springframework.data.redis.connection.ReactiveRedisConnectionFactory;
import org.springframework.data.redis.core.ReactiveRedisTemplate;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.data.redis.serializer.RedisSerializationContext;
import org.springframework.data.redis.serializer.RedisSerializer;
import org.springframework.data.redis.serializer.StringRedisSerializer;
import org.springframework.scripting.support.ResourceScriptSource;
@Configuration
@AutoConfigureAfter(RedisReactiveAutoConfiguration.class)
@AutoConfigureBefore(GatewayAutoConfiguration.class)
@ConditionalOnBean(ReactiveRedisTemplate.class)
@ConditionalOnClass(RedisTemplate.class)
class GatewayRedisAutoConfiguration {
@Bean
@SuppressWarnings("unchecked")
public RedisScript redisRequestRateLimiterScript() {
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 ReactiveRedisTemplate<String, String> stringReactiveRedisTemplate(
ReactiveRedisConnectionFactory reactiveRedisConnectionFactory,
ResourceLoader resourceLoader) {
RedisSerializer<String> serializer = new StringRedisSerializer();
RedisSerializationContext<String , String> serializationContext = RedisSerializationContext
.<String, String>newSerializationContext()
.key(serializer)
.value(serializer)
.hashKey(serializer)
.hashValue(serializer)
.build();
return new ReactiveRedisTemplate<>(reactiveRedisConnectionFactory,
serializationContext);
}
@Bean
public RedisRateLimiter redisRateLimiter(ReactiveRedisTemplate<String, String> redisTemplate,
@Qualifier("redisRequestRateLimiterScript") RedisScript<String> redisScript) {
return new RedisRateLimiter(redisTemplate, redisScript);
}
}

View File

@@ -72,7 +72,7 @@ public class RequestRateLimiterWebFilterFactory implements WebFilterFactory, App
return (exchange, chain) ->
keyResolver.resolve(exchange).flatMap(key -> {
Response response = rateLimiter.isAllowed(key, replenishRate, capacity);
Response response = rateLimiter.isAllowed(key, replenishRate, capacity).block(); //FIXME: block()
//TODO: set some headers for rate, tokens left

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, int replenishRate, int burstCapacity);
Mono<Response> isAllowed(String id, int replenishRate, int burstCapacity);
class Response {
private final boolean allowed;

View File

@@ -1,13 +1,17 @@
package org.springframework.cloud.gateway.filter.ratelimit;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
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.ReactiveRedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import java.time.Instant;
import java.util.Arrays;
import java.util.List;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
/**
* See https://stripe.com/blog/rate-limiters and
@@ -18,26 +22,28 @@ import java.util.List;
public class RedisRateLimiter implements RateLimiter {
private Log log = LogFactory.getLog(getClass());
private final StringRedisTemplate redisTemplate;
private final RedisScript<List> script;
private final ReactiveRedisTemplate<String, String> redisTemplate;
private final RedisScript script;
public RedisRateLimiter(StringRedisTemplate redisTemplate, RedisScript<List> script) {
public RedisRateLimiter(ReactiveRedisTemplate<String, String> redisTemplate,
RedisScript script) {
this.redisTemplate = redisTemplate;
this.script = script;
}
/**
* 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.
* 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
* @return
*/
@Override
//TODO: signature? params (tuple?).
//TODO: change to Mono<?>
public Response isAllowed(String id, int replenishRate, int burstCapacity) {
// TODO: signature? params (tuple?).
@SuppressWarnings("unchecked")
public Mono<Response> isAllowed(String id, int replenishRate, int burstCapacity) {
try {
// Make a unique key per user.
@@ -47,26 +53,35 @@ 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.
Object[] args = new String[]{ replenishRate+"", burstCapacity +"", Instant.now().getEpochSecond()+"", "1"};
List<String> args = Arrays.asList(replenishRate + "", burstCapacity + "",
Instant.now().getEpochSecond() + "", "1");
// allowed, tokens_left = redis.eval(SCRIPT, keys, args)
List results = this.redisTemplate.execute(this.script, keys, args);
Flux<Long> flux = this.redisTemplate.execute(this.script, keys, args)
.log("redisratelimiter");
return flux.onErrorResume(throwable -> Flux.just(1L, -1L))
.reduce(new ArrayList<Long>(), (longs, l) -> {
longs.add(l);
return longs;
}).map(results -> {
boolean allowed = results.get(0) == 1L;
Long tokensLeft = results.get(1);
boolean allowed = new Long(1L).equals(results.get(0));
Long tokensLeft = (Long) results.get(1);
Response response = new Response(allowed, tokensLeft);
Response response = new Response(allowed, tokensLeft);
if (log.isDebugEnabled()) {
log.debug("response: "+response);
}
return response;
} catch (Exception e) {
/* We don't want a hard dependency on Redis to allow traffic.
Make sure to set an alert so you know if this is happening too much.
Stripe's observed failure rate is 0.01%. */
if (log.isDebugEnabled()) {
log.debug("response: " + response);
}
return response;
});
}
catch (Exception e) {
/*
* We don't want a hard dependency on Redis to allow traffic. Make sure to set
* an alert so you know if this is happening too much. Stripe's observed
* failure rate is 0.01%.
*/
log.error("Error determining if user allowed from redis", e);
}
return new Response(true, -1);
return Mono.just(new Response(true, -1));
}
}

View File

@@ -24,11 +24,13 @@ local delta = math.max(0, now-last_refreshed)
local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
local allowed = filled_tokens >= requested
local new_tokens = filled_tokens
local allowed_num = 0
if allowed then
new_tokens = filled_tokens - requested
allowed_num = 1
end
redis.call("setex", tokens_key, ttl, new_tokens)
redis.call("setex", timestamp_key, ttl, now)
return { allowed, new_tokens }
return { allowed_num, new_tokens }

View File

@@ -1,4 +1,5 @@
# Auto Configure
org.springframework.boot.autoconfigure.EnableAutoConfiguration=\
org.springframework.cloud.gateway.config.GatewayAutoConfiguration,\
org.springframework.cloud.gateway.config.GatewayLoadBalancerClientAutoConfiguration
org.springframework.cloud.gateway.config.GatewayLoadBalancerClientAutoConfiguration,\
org.springframework.cloud.gateway.config.GatewayRedisAutoConfiguration

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();

View File

@@ -13,6 +13,7 @@ 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 reactor.core.publisher.Mono;
import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.boot.test.context.SpringBootTest.WebEnvironment.RANDOM_PORT;
@@ -22,7 +23,7 @@ import static org.springframework.boot.test.context.SpringBootTest.WebEnvironmen
* @author Spencer Gibb
*/
@RunWith(SpringRunner.class)
@SpringBootTest(webEnvironment = RANDOM_PORT)
@SpringBootTest(webEnvironment = RANDOM_PORT, properties = "true")
@DirtiesContext
public class RedisRateLimiterTests extends BaseWebClientTests {
@@ -38,13 +39,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 +53,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();
}

View File

@@ -197,6 +197,7 @@ logging:
org.springframework.http.server.reactive: DEBUG
org.springframework.web.reactive: DEBUG
reactor.ipc.netty: DEBUG
redisratelimiter: DEBUG
management:
context-path: /admin
@@ -204,3 +205,4 @@ eureka:
client:
enabled: false
# port: 8081