Restore Redis retries for select operations

service/src/test/java/org/whispersystems/textsecuregcm/limits/DynamicRateLimiterTest.java

@@ -19,10 +19,13 @@ import java.time.Duration;
 import java.time.Instant;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import org.apache.commons.lang3.RandomStringUtils;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
@@ -36,6 +39,7 @@ import org.whispersystems.textsecuregcm.util.TestClock;
 class DynamicRateLimiterTest {
 
   private ClusterLuaScript validateRateLimitScript;
+  private ScheduledExecutorService retryExecutor;
 
   private static final TestClock CLOCK = TestClock.pinned(Instant.now());
 
@@ -44,10 +48,17 @@ class DynamicRateLimiterTest {
 
   @BeforeEach
   void setUp() throws IOException {
+    retryExecutor = Executors.newSingleThreadScheduledExecutor();
+
     validateRateLimitScript = ClusterLuaScript.fromResource(
         REDIS_CLUSTER_EXTENSION.getRedisCluster(), "lua/validate_rate_limit.lua", ScriptOutputType.INTEGER);
   }
 
+  @AfterEach
+  void tearDown() {
+    retryExecutor.shutdown();
+  }
+
   @ParameterizedTest
   @ValueSource(booleans = {true, false})
   void validate(final boolean failOpen) {
@@ -56,6 +67,7 @@ class DynamicRateLimiterTest {
         () -> new RateLimiterConfig(1, Duration.ofHours(1), failOpen),
         validateRateLimitScript,
         REDIS_CLUSTER_EXTENSION.getRedisCluster(),
+        retryExecutor,
         CLOCK);
 
     final String key = RandomStringUtils.insecure().nextAlphanumeric(16);
@@ -72,6 +84,7 @@ class DynamicRateLimiterTest {
         () -> new RateLimiterConfig(1, Duration.ofHours(1), failOpen),
         validateRateLimitScript,
         REDIS_CLUSTER_EXTENSION.getRedisCluster(),
+        retryExecutor,
         CLOCK);
 
     final String key = RandomStringUtils.insecure().nextAlphanumeric(16);
@@ -94,6 +107,7 @@ class DynamicRateLimiterTest {
         () -> new RateLimiterConfig(1, Duration.ofHours(1), failOpen),
         failingScript,
         REDIS_CLUSTER_EXTENSION.getRedisCluster(),
+        retryExecutor,
         CLOCK);
 
     final String key = RandomStringUtils.insecure().nextAlphanumeric(16);
@@ -116,6 +130,7 @@ class DynamicRateLimiterTest {
         () -> new RateLimiterConfig(1, Duration.ofHours(1), failOpen),
         failingScript,
         REDIS_CLUSTER_EXTENSION.getRedisCluster(),
+        retryExecutor,
         CLOCK);
 
     final String key = RandomStringUtils.insecure().nextAlphanumeric(16);
@@ -138,6 +153,7 @@ class DynamicRateLimiterTest {
         () -> new RateLimiterConfig(1, refillRate.get(), false),
         validateRateLimitScript,
         REDIS_CLUSTER_EXTENSION.getRedisCluster(),
+        retryExecutor,
         CLOCK);
 
     final String key = RandomStringUtils.insecure().nextAlphanumeric(16);
@@ -161,6 +177,7 @@ class DynamicRateLimiterTest {
         () -> new RateLimiterConfig(1, refillRate.get(), false),
         validateRateLimitScript,
         REDIS_CLUSTER_EXTENSION.getRedisCluster(),
+        retryExecutor,
         CLOCK);
 
     final String key = RandomStringUtils.insecure().nextAlphanumeric(16);
@@ -187,6 +204,7 @@ class DynamicRateLimiterTest {
         () -> new RateLimiterConfig(bucketSize.get(), Duration.ofMinutes(1), false),
         validateRateLimitScript,
         REDIS_CLUSTER_EXTENSION.getRedisCluster(),
+        retryExecutor,
         CLOCK);
 
     final String key = RandomStringUtils.insecure().nextAlphanumeric(16);
@@ -211,6 +229,7 @@ class DynamicRateLimiterTest {
         () -> new RateLimiterConfig(bucketSize.get(), Duration.ofMinutes(1), false),
         validateRateLimitScript,
         REDIS_CLUSTER_EXTENSION.getRedisCluster(),
+        retryExecutor,
         CLOCK);
 
     final String key = RandomStringUtils.insecure().nextAlphanumeric(16);
@@ -238,6 +257,7 @@ class DynamicRateLimiterTest {
         () -> new RateLimiterConfig(1, Duration.ofMinutes(1), failOpen.get()),
         failingScript,
         REDIS_CLUSTER_EXTENSION.getRedisCluster(),
+        retryExecutor,
         CLOCK);
 
     final String key = RandomStringUtils.insecure().nextAlphanumeric(16);
@@ -260,6 +280,7 @@ class DynamicRateLimiterTest {
         () -> new RateLimiterConfig(1, Duration.ofMinutes(1), failOpen.get()),
         failingScript,
         REDIS_CLUSTER_EXTENSION.getRedisCluster(),
+        retryExecutor,
         CLOCK);
 
     final String key = RandomStringUtils.insecure().nextAlphanumeric(16);

service/src/test/java/org/whispersystems/textsecuregcm/limits/RateLimitersLuaScriptTest.java

@@ -21,6 +21,7 @@ import java.time.Duration;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.ScheduledExecutorService;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -65,6 +66,7 @@ public class RateLimitersLuaScriptTest {
         dynamicConfig,
         RateLimiters.defaultScript(redisCluster),
         redisCluster,
+        mock(ScheduledExecutorService.class),
         Clock.systemUTC());
 
     final RateLimiter rateLimiter = limiters.forDescriptor(descriptor);
@@ -84,6 +86,7 @@ public class RateLimitersLuaScriptTest {
         dynamicConfig,
         RateLimiters.defaultScript(redisCluster),
         redisCluster,
+        mock(ScheduledExecutorService.class),
         Clock.systemUTC());
 
     final RateLimiter rateLimiter = limiters.forDescriptor(descriptor);
@@ -138,6 +141,7 @@ public class RateLimitersLuaScriptTest {
         dynamicConfig,
         RateLimiters.defaultScript(redisCluster),
         redisCluster,
+        mock(ScheduledExecutorService.class),
         Clock.systemUTC());
     when(redisCluster.withCluster(any())).thenThrow(new RedisException("fail"));
     final RateLimiter rateLimiter = limiters.forDescriptor(descriptor);

service/src/test/java/org/whispersystems/textsecuregcm/limits/RateLimitersTest.java

@@ -13,6 +13,7 @@ import static org.mockito.Mockito.when;
 import java.time.Duration;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.ScheduledExecutorService;
 import org.junit.jupiter.api.Test;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
 import org.whispersystems.textsecuregcm.redis.ClusterLuaScript;
@@ -36,7 +37,7 @@ public class RateLimitersTest {
   private final MutableClock clock = MockUtils.mutableClock(0);
 
   @Test
-  public void testValidateDuplicates() throws Exception {
+  public void testValidateDuplicates() {
     final TestDescriptor td1 = new TestDescriptor("id1");
     final TestDescriptor td2 = new TestDescriptor("id2");
     final TestDescriptor td3 = new TestDescriptor("id3");
@@ -47,6 +48,7 @@ public class RateLimitersTest {
         dynamicConfig,
         validateScript,
         redisCluster,
+        mock(ScheduledExecutorService.class),
         clock) {});
 
     new BaseRateLimiters<>(
@@ -54,12 +56,13 @@ public class RateLimitersTest {
         dynamicConfig,
         validateScript,
         redisCluster,
+        mock(ScheduledExecutorService.class),
         clock) {};
   }
 
   @Test
   void testUnchangingConfiguration() {
-    final RateLimiters rateLimiters = new RateLimiters(dynamicConfig, validateScript, redisCluster, clock);
+    final RateLimiters rateLimiters = new RateLimiters(dynamicConfig, validateScript, redisCluster, mock(ScheduledExecutorService.class), clock);
     final RateLimiter limiter = rateLimiters.getRateLimitResetLimiter();
     final RateLimiterConfig expected = RateLimiters.For.RATE_LIMIT_RESET.defaultConfig();
     assertEquals(expected, limiter.config());
@@ -78,7 +81,7 @@ public class RateLimitersTest {
 
     when(configuration.getLimits()).thenReturn(limitsConfigMap);
 
-    final RateLimiters rateLimiters = new RateLimiters(dynamicConfig, validateScript, redisCluster, clock);
+    final RateLimiters rateLimiters = new RateLimiters(dynamicConfig, validateScript, redisCluster, mock(ScheduledExecutorService.class), clock);
     final RateLimiter limiter = rateLimiters.getRateLimitResetLimiter();
 
     limitsConfigMap.put(RateLimiters.For.RATE_LIMIT_RESET.id(), initialRateLimiterConfig);
@@ -104,7 +107,7 @@ public class RateLimitersTest {
 
     when(configuration.getLimits()).thenReturn(mapForDynamic);
 
-    final RateLimiters rateLimiters = new RateLimiters(dynamicConfig, validateScript, redisCluster, clock);
+    final RateLimiters rateLimiters = new RateLimiters(dynamicConfig, validateScript, redisCluster, mock(ScheduledExecutorService.class), clock);
     final RateLimiter limiter = rateLimiters.forDescriptor(descriptor);
 
     // test only default is present