Test rate limiters against a real cluster.

This commit is contained in:
Jon Chambers 2020-07-06 14:16:22 -04:00 committed by Jon Chambers
parent 6c1ba957bd
commit 933ce42d5a
2 changed files with 36 additions and 10 deletions

View File

@ -1,33 +1,55 @@
package org.whispersystems.textsecuregcm.limits; package org.whispersystems.textsecuregcm.limits;
import org.junit.After;
import org.junit.Before; import org.junit.Before;
import org.junit.Test; import org.junit.Test;
import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException; import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
import org.whispersystems.textsecuregcm.redis.AbstractRedisTest; import org.whispersystems.textsecuregcm.providers.RedisClientFactory;
import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster; import org.whispersystems.textsecuregcm.redis.AbstractRedisClusterTest;
import org.whispersystems.textsecuregcm.redis.ReplicatedJedisPool;
import redis.clients.jedis.Jedis; import redis.clients.jedis.Jedis;
import redis.embedded.RedisServer;
import java.io.IOException; import java.io.IOException;
import java.net.URISyntaxException;
import java.util.List;
import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeUnit;
import static org.junit.Assert.fail; import static org.junit.Assert.fail;
import static org.mockito.Mockito.mock;
public class RateLimiterTest extends AbstractRedisTest { public class RateLimiterTest extends AbstractRedisClusterTest {
private static final long NOW_MILLIS = System.currentTimeMillis(); private static final long NOW_MILLIS = System.currentTimeMillis();
private static final String KEY = "key"; private static final String KEY = "key";
private RedisServer redisServer;
private ReplicatedJedisPool replicatedJedisPool;
@FunctionalInterface @FunctionalInterface
private interface RateLimitedTask { private interface RateLimitedTask {
void run() throws RateLimitExceededException; void run() throws RateLimitExceededException;
} }
@Before @Before
public void clearCache() { public void clearCache() throws URISyntaxException, IOException {
try (final Jedis jedis = getReplicatedJedisPool().getWriteResource()) { redisServer = new RedisServer(AbstractRedisClusterTest.getNextRedisClusterPort());
redisServer.start();
final String redisUrl = "redis://127.0.0.1:" + redisServer.ports().get(0);
replicatedJedisPool = new RedisClientFactory("test-pool", redisUrl, List.of(redisUrl), new CircuitBreakerConfiguration()).getRedisClientPool();
try (final Jedis jedis = replicatedJedisPool.getWriteResource()) {
jedis.flushAll(); jedis.flushAll();
} }
getRedisCluster().useWriteCluster(connection -> connection.sync().flushall());
}
@After
public void stopServer() {
redisServer.stop();
} }
@Test @Test
@ -52,10 +74,12 @@ public class RateLimiterTest extends AbstractRedisTest {
final RateLimiter rateLimiter = buildRateLimiter(2, 8.333333333333334E-6); final RateLimiter rateLimiter = buildRateLimiter(2, 8.333333333333334E-6);
final String leakyBucketJson = "{\"bucketSize\":2,\"leakRatePerMillis\":8.333333333333334E-6,\"spaceRemaining\":0,\"lastUpdateTimeMillis\":" + (NOW_MILLIS - TimeUnit.MINUTES.toMillis(2)) + "}"; final String leakyBucketJson = "{\"bucketSize\":2,\"leakRatePerMillis\":8.333333333333334E-6,\"spaceRemaining\":0,\"lastUpdateTimeMillis\":" + (NOW_MILLIS - TimeUnit.MINUTES.toMillis(2)) + "}";
try (final Jedis jedis = getReplicatedJedisPool().getWriteResource()) { try (final Jedis jedis = replicatedJedisPool.getWriteResource()) {
jedis.set(rateLimiter.getBucketName(KEY), leakyBucketJson); jedis.set(rateLimiter.getBucketName(KEY), leakyBucketJson);
} }
getRedisCluster().useWriteCluster(connection -> connection.sync().set(rateLimiter.getBucketName(KEY), leakyBucketJson));
rateLimiter.validate(KEY, 1, NOW_MILLIS); rateLimiter.validate(KEY, 1, NOW_MILLIS);
assertRateLimitExceeded(() -> rateLimiter.validate(KEY, 1, NOW_MILLIS)); assertRateLimitExceeded(() -> rateLimiter.validate(KEY, 1, NOW_MILLIS));
} }
@ -65,10 +89,12 @@ public class RateLimiterTest extends AbstractRedisTest {
final RateLimiter rateLimiter = buildRateLimiter(2, 8.333333333333334E-6); final RateLimiter rateLimiter = buildRateLimiter(2, 8.333333333333334E-6);
final String leakyBucketJson = "{\"bucketSize\":2,\"leakRatePerMillis\":8.333333333333334E-6,\"spaceRemaining\":0,\"lastUpdateTimeMillis\":" + (NOW_MILLIS - TimeUnit.MINUTES.toMillis(1)) + "}"; final String leakyBucketJson = "{\"bucketSize\":2,\"leakRatePerMillis\":8.333333333333334E-6,\"spaceRemaining\":0,\"lastUpdateTimeMillis\":" + (NOW_MILLIS - TimeUnit.MINUTES.toMillis(1)) + "}";
try (final Jedis jedis = getReplicatedJedisPool().getWriteResource()) { try (final Jedis jedis = replicatedJedisPool.getWriteResource()) {
jedis.set(rateLimiter.getBucketName(KEY), leakyBucketJson); jedis.set(rateLimiter.getBucketName(KEY), leakyBucketJson);
} }
getRedisCluster().useWriteCluster(connection -> connection.sync().set(rateLimiter.getBucketName(KEY), leakyBucketJson));
assertRateLimitExceeded(() -> rateLimiter.validate(KEY, 1, NOW_MILLIS)); assertRateLimitExceeded(() -> rateLimiter.validate(KEY, 1, NOW_MILLIS));
} }
@ -83,6 +109,6 @@ public class RateLimiterTest extends AbstractRedisTest {
@SuppressWarnings("SameParameterValue") @SuppressWarnings("SameParameterValue")
private RateLimiter buildRateLimiter(final int bucketSize, final double leakRatePerMilli) throws IOException { private RateLimiter buildRateLimiter(final int bucketSize, final double leakRatePerMilli) throws IOException {
final double leakRatePerMinute = leakRatePerMilli * 60_000d; final double leakRatePerMinute = leakRatePerMilli * 60_000d;
return new RateLimiter(getReplicatedJedisPool(), mock(FaultTolerantRedisCluster.class), KEY, bucketSize, leakRatePerMinute); return new RateLimiter(replicatedJedisPool, getRedisCluster(), KEY, bucketSize, leakRatePerMinute);
} }
} }

View File

@ -143,7 +143,7 @@ public abstract class AbstractRedisClusterTest {
} }
} }
private static int getNextRedisClusterPort() throws IOException { public static int getNextRedisClusterPort() throws IOException {
final int MAX_ITERATIONS = 11_000; final int MAX_ITERATIONS = 11_000;
int port; int port;
for (int i = 0; i < MAX_ITERATIONS; i++) { for (int i = 0; i < MAX_ITERATIONS; i++) {