Signal/Signal-Server
1
0
Fork 0
mirror of https://github.com/signalapp/Signal-Server synced 2026-04-22 02:48:01 +01:00
Code Issues Packages Projects Releases Wiki Activity

migrate token bucket redis record format from json to hash: phase 1

Browse Source
This commit is contained in:
Sergey Skrobotov
2023-03-15 14:32:14 -07:00
parent ebf8aa7b15
commit 483e444174
6 changed files with 277 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
service/src/test/java/org/whispersystems/textsecuregcm/limits/RateLimitersLuaScriptTest.java
View File

@@ -19,6 +19,7 @@ import java.util.Optional;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
import org.whispersystems.textsecuregcm.redis.RedisClusterExtension;
import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
@@ -60,7 +61,51 @@ public class RateLimitersLuaScriptTest {
final RateLimiter rateLimiter = limiters.forDescriptor(descriptor);
rateLimiter.validate("test", 25);
rateLimiter.validate("test", 25);
assertThrows(Exception.class, () -> rateLimiter.validate("test", 25));
assertThrows(RateLimitExceededException.class, () -> rateLimiter.validate("test", 25));
}
@Test
public void testFormatMigration() throws Exception {
final RateLimiters.For descriptor = RateLimiters.For.REGISTRATION;
final FaultTolerantRedisCluster redisCluster = REDIS_CLUSTER_EXTENSION.getRedisCluster();
final RateLimiters limiters = new RateLimiters(
Map.of(descriptor.id(), new RateLimiterConfig(60, 60)),
dynamicConfig,
RateLimiters.defaultScript(redisCluster),
redisCluster,
Clock.systemUTC());
final RateLimiter rateLimiter = limiters.forDescriptor(descriptor);
// embedding an existing value in the old format
redisCluster.useCluster(c -> c.sync().set(
StaticRateLimiter.bucketName(descriptor.id(), "test"),
serializeToOldBucketValueFormat(new TokenBucket(60, 60, 30, System.currentTimeMillis() + 10000))
));
assertThrows(RateLimitExceededException.class, () -> rateLimiter.validate("test", 40));
// embedding an existing value in the old format
redisCluster.useCluster(c -> c.sync().set(
StaticRateLimiter.bucketName(descriptor.id(), "test1"),
serializeToOldBucketValueFormat(new TokenBucket(60, 60, 30, System.currentTimeMillis() + 10000))
));
rateLimiter.validate("test1", 20);
assertThrows(RateLimitExceededException.class, () -> rateLimiter.validate("test1", 20));
// embedding an existing value in the new format
redisCluster.useCluster(c -> c.sync().hset(
StaticRateLimiter.bucketName(descriptor.id(), "test2"),
Map.of("s", "30", "t", String.valueOf(System.currentTimeMillis() + 10000))
));
assertThrows(RateLimitExceededException.class, () -> rateLimiter.validate("test2", 40));
// embedding an existing value in the new format
redisCluster.useCluster(c -> c.sync().hset(
StaticRateLimiter.bucketName(descriptor.id(), "test3"),
Map.of("s", "30", "t", String.valueOf(System.currentTimeMillis() + 10000))
));
rateLimiter.validate("test3", 20);
assertThrows(RateLimitExceededException.class, () -> rateLimiter.validate("test3", 20));
}
@Test
@@ -117,6 +162,19 @@ public class RateLimitersLuaScriptTest {
assertEquals(750L, decodeBucket(key).orElseThrow().spaceRemaining);
}
private String serializeToOldBucketValueFormat(final TokenBucket bucket) {
try {
return SystemMapper.jsonMapper().writeValueAsString(Map.of(
"bucketSize", bucket.bucketSize,
"leakRatePerMillis", bucket.leakRatePerMillis,
"spaceRemaining", bucket.spaceRemaining,
"lastUpdateTimeMillis", bucket.lastUpdateTimeMillis
));
} catch (JsonProcessingException e) {
throw new RuntimeException(e);
}
}
private Optional<TokenBucket> decodeBucket(final String key) {
try {
final String json = redisCommandsHandler.get(key);