Spencer Gibb
8 years ago
7 changed files with 225 additions and 69 deletions
@ -0,0 +1,35 @@
@@ -0,0 +1,35 @@
|
||||
package org.springframework.cloud.gateway.filter.ratelimit; |
||||
|
||||
/** |
||||
* @author Spencer Gibb |
||||
*/ |
||||
public interface RateLimiter { |
||||
Response isAllowed(String id, int replenishRate, int capacity); |
||||
|
||||
class Response { |
||||
private final boolean allowed; |
||||
private final long tokensRemaining; |
||||
|
||||
public Response(boolean allowed, long tokensRemaining) { |
||||
this.allowed = allowed; |
||||
this.tokensRemaining = tokensRemaining; |
||||
} |
||||
|
||||
public boolean isAllowed() { |
||||
return allowed; |
||||
} |
||||
|
||||
public long getTokensRemaining() { |
||||
return tokensRemaining; |
||||
} |
||||
|
||||
@Override |
||||
public String toString() { |
||||
final StringBuffer sb = new StringBuffer("Response{"); |
||||
sb.append("allowed=").append(allowed); |
||||
sb.append(", tokensRemaining=").append(tokensRemaining); |
||||
sb.append('}'); |
||||
return sb.toString(); |
||||
} |
||||
} |
||||
} |
||||
@ -0,0 +1,71 @@
@@ -0,0 +1,71 @@
|
||||
package org.springframework.cloud.gateway.filter.ratelimit; |
||||
|
||||
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 java.time.Instant; |
||||
import java.util.Arrays; |
||||
import java.util.List; |
||||
|
||||
/** |
||||
* See https://stripe.com/blog/rate-limiters and
|
||||
* https://gist.github.com/ptarjan/e38f45f2dfe601419ca3af937fff574d#file-1-check_request_rate_limiter-rb-L11-L34
|
||||
* |
||||
* @author Spencer Gibb |
||||
*/ |
||||
public class RedisRateLimiter implements RateLimiter { |
||||
private Log log = LogFactory.getLog(getClass()); |
||||
|
||||
private final StringRedisTemplate redisTemplate; |
||||
private final RedisScript<List> script; |
||||
|
||||
public RedisRateLimiter(StringRedisTemplate redisTemplate, RedisScript<List> 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. |
||||
* @param replenishRate |
||||
* @param capacity |
||||
* @param id |
||||
* @return |
||||
*/ |
||||
@Override |
||||
//TODO: signature? params (tuple?). Return type, tokens left?
|
||||
public Response isAllowed(String id, int replenishRate, int capacity) { |
||||
|
||||
try { |
||||
// Make a unique key per user.
|
||||
String prefix = "request_rate_limiter." + id; |
||||
|
||||
// You need two Redis keys for Token Bucket.
|
||||
List<String> keys = Arrays.asList(prefix + ".tokens", prefix + ".timestamp"); |
||||
|
||||
// The arguments to the LUA script. time() returns unixtime in seconds.
|
||||
String[] args = new String[]{ replenishRate+"", capacity+"", Instant.now().getEpochSecond()+"", "1"}; |
||||
// allowed, tokens_left = redis.eval(SCRIPT, keys, args)
|
||||
List results = this.redisTemplate.execute(this.script, keys, args); |
||||
|
||||
boolean allowed = new Long(1L).equals(results.get(0)); |
||||
Long tokensLeft = (Long) results.get(1); |
||||
|
||||
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%. */ |
||||
log.error("Error determining if user allowed from redis", e); |
||||
} |
||||
return new Response(true, -1); |
||||
} |
||||
} |
||||
@ -0,0 +1,66 @@
@@ -0,0 +1,66 @@
|
||||
package org.springframework.cloud.gateway.filter.ratelimit; |
||||
|
||||
import java.util.UUID; |
||||
|
||||
import org.junit.Test; |
||||
import org.junit.runner.RunWith; |
||||
import org.springframework.beans.factory.annotation.Autowired; |
||||
import org.springframework.boot.SpringBootConfiguration; |
||||
import org.springframework.boot.autoconfigure.EnableAutoConfiguration; |
||||
import org.springframework.boot.test.context.SpringBootTest; |
||||
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter.Response; |
||||
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 static org.assertj.core.api.Assertions.assertThat; |
||||
import static org.springframework.boot.test.context.SpringBootTest.WebEnvironment.RANDOM_PORT; |
||||
|
||||
/** |
||||
* see https://gist.github.com/ptarjan/e38f45f2dfe601419ca3af937fff574d#file-1-check_request_rate_limiter-rb-L36-L62
|
||||
* @author Spencer Gibb |
||||
*/ |
||||
@RunWith(SpringRunner.class) |
||||
@SpringBootTest(webEnvironment = RANDOM_PORT) |
||||
@DirtiesContext |
||||
public class RedisRateLimiterTests extends BaseWebClientTests { |
||||
|
||||
@Autowired |
||||
private RedisRateLimiter rateLimiter; |
||||
|
||||
@Test |
||||
public void requestRateLimiterWebFilterFactoryWorks() throws Exception { |
||||
String id = UUID.randomUUID().toString(); |
||||
|
||||
int replenishRate = 10; |
||||
int capacity = 2 * replenishRate; |
||||
|
||||
// Bursts work
|
||||
for (int i = 0; i < capacity; i++) { |
||||
Response response = rateLimiter.isAllowed(id, replenishRate, capacity); |
||||
assertThat(response.isAllowed()).as("Burst # %s is allowed", i).isTrue(); |
||||
} |
||||
|
||||
Response response = rateLimiter.isAllowed(id, replenishRate, capacity); |
||||
assertThat(response.isAllowed()).as("Burst # %s is not allowed", capacity).isFalse(); |
||||
|
||||
Thread.sleep(1000); |
||||
|
||||
// # After the burst is done, check the steady state
|
||||
for (int i = 0; i < replenishRate; i++) { |
||||
response = rateLimiter.isAllowed(id, replenishRate, capacity); |
||||
assertThat(response.isAllowed()).as("steady state # %s is allowed", i).isTrue(); |
||||
} |
||||
|
||||
response = rateLimiter.isAllowed(id, replenishRate, capacity); |
||||
assertThat(response.isAllowed()).as("steady state # %s is allowed", replenishRate).isFalse(); |
||||
} |
||||
|
||||
@EnableAutoConfiguration |
||||
@SpringBootConfiguration |
||||
@Import(BaseWebClientTests.DefaultTestConfig.class) |
||||
public static class TestConfig { |
||||
|
||||
} |
||||
} |
||||
Loading…
Reference in new issue