Support ReactiveRedisTemplate execution of RedisScript.
This commit is contained in:
@@ -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 {
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 }
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user