From 3e80669f4e10a14628c55e7eda4ca4b03fe4722b Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Thu, 3 Jul 2025 11:32:30 -0400
Subject: [PATCH] Reuse/extend lifetime of Redis singleton resources

---
 .../LocalFaultTolerantRedisClientFactory.java |  5 +--
 .../redis/RedisServerExtension.java           | 39 ++++++++++---------
 2 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/configuration/LocalFaultTolerantRedisClientFactory.java b/service/src/test/java/org/whispersystems/textsecuregcm/configuration/LocalFaultTolerantRedisClientFactory.java
index eff6297f9..118433190 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/configuration/LocalFaultTolerantRedisClientFactory.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/configuration/LocalFaultTolerantRedisClientFactory.java
@@ -33,9 +33,8 @@ public class LocalFaultTolerantRedisClientFactory implements FaultTolerantRedisC
     if (shutdownHookConfigured.compareAndSet(false, true)) {
       Runtime.getRuntime().addShutdownHook(new Thread(() -> {
         try {
-          REDIS_SERVER_EXTENSION.afterEach(null);
-          REDIS_SERVER_EXTENSION.afterAll(null);
-        } catch (Exception e) {
+          REDIS_SERVER_EXTENSION.close();
+        } catch (Throwable e) {
           throw new RuntimeException(e);
         }
       }));
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/redis/RedisServerExtension.java b/service/src/test/java/org/whispersystems/textsecuregcm/redis/RedisServerExtension.java
index ff6387951..8b9cde749 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/redis/RedisServerExtension.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/redis/RedisServerExtension.java
@@ -8,12 +8,10 @@ package org.whispersystems.textsecuregcm.redis;
 import static org.junit.jupiter.api.Assumptions.assumeFalse;
 
 import io.lettuce.core.RedisURI;
+import io.lettuce.core.resource.ClientResources;
 import java.io.IOException;
 import java.net.ServerSocket;
 import java.time.Duration;
-import io.lettuce.core.resource.ClientResources;
-import org.junit.jupiter.api.extension.AfterAllCallback;
-import org.junit.jupiter.api.extension.AfterEachCallback;
 import org.junit.jupiter.api.extension.BeforeAllCallback;
 import org.junit.jupiter.api.extension.BeforeEachCallback;
 import org.junit.jupiter.api.extension.ExtensionContext;
@@ -22,11 +20,12 @@ import org.whispersystems.textsecuregcm.configuration.RetryConfiguration;
 import redis.embedded.RedisServer;
 import redis.embedded.exceptions.EmbeddedRedisException;
 
-public class RedisServerExtension implements BeforeAllCallback, BeforeEachCallback, AfterAllCallback, AfterEachCallback {
+public class RedisServerExtension implements BeforeAllCallback, BeforeEachCallback, ExtensionContext.Store.CloseableResource {
 
   private static RedisServer redisServer;
+  private static ClientResources redisClientResources;
+
   private FaultTolerantRedisClient faultTolerantRedisClient;
-  private ClientResources redisClientResources;
 
   public static class RedisServerExtensionBuilder {
 
@@ -46,14 +45,18 @@ public class RedisServerExtension implements BeforeAllCallback, BeforeEachCallba
   public void beforeAll(final ExtensionContext context) throws Exception {
     assumeFalse(System.getProperty("os.name").equalsIgnoreCase("windows"));
 
-    redisServer = RedisServer.builder()
-        .setting("appendonly no")
-        .setting("save \"\"")
-        .setting("dir " + System.getProperty("java.io.tmpdir"))
-        .port(getAvailablePort())
-        .build();
+    if (redisServer == null) {
+      redisServer = RedisServer.builder()
+          .setting("appendonly no")
+          .setting("save \"\"")
+          .setting("dir " + System.getProperty("java.io.tmpdir"))
+          .port(getAvailablePort())
+          .build();
 
-    startWithRetries(3);
+      redisClientResources = ClientResources.builder().build();
+
+      startWithRetries(3);
+    }
   }
 
   public static RedisURI getRedisURI() {
@@ -62,7 +65,6 @@ public class RedisServerExtension implements BeforeAllCallback, BeforeEachCallba
 
   @Override
   public void beforeEach(final ExtensionContext context) {
-    redisClientResources = ClientResources.builder().build();
     final CircuitBreakerConfiguration circuitBreakerConfig = new CircuitBreakerConfiguration();
     circuitBreakerConfig.setWaitDurationInOpenState(Duration.ofMillis(500));
     faultTolerantRedisClient = new FaultTolerantRedisClient("test-redis-client",
@@ -76,15 +78,14 @@ public class RedisServerExtension implements BeforeAllCallback, BeforeEachCallba
   }
 
   @Override
-  public void afterEach(final ExtensionContext context) throws InterruptedException {
-    redisClientResources.shutdown().await();
-  }
-
-  @Override
-  public void afterAll(final ExtensionContext context) {
+  public void close() throws Throwable {
     if (redisServer != null) {
+      redisClientResources.shutdown().await();
       redisServer.stop();
     }
+
+    redisClientResources = null;
+    redisServer = null;
   }
 
   public FaultTolerantRedisClient getRedisClient() {