From e74ad2b555ee1ccd51d1c788c1da6f896f3a2380 Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Wed, 19 Aug 2020 14:29:38 -0400
Subject: [PATCH] Make RedisClusterMessagesCache a Managed class.

---
 .../textsecuregcm/WhisperServerService.java           |  1 +
 .../storage/RedisClusterMessagesCache.java            | 11 ++++++++++-
 .../storage/RedisClusterMessagesCacheTest.java        |  6 +++++-
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index 90722f948..3fb675c95 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -349,6 +349,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     environment.lifecycle().manage(messagePersister);
     environment.lifecycle().manage(accountDatabaseCrawler);
     environment.lifecycle().manage(remoteConfigsManager);
+    environment.lifecycle().manage(clusterMessagesCache);
     environment.lifecycle().manage(clusterMessagePersister);
     environment.lifecycle().manage(clientPresenceManager);
 
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/RedisClusterMessagesCache.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/RedisClusterMessagesCache.java
index 2e6987d0f..45eefde4f 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/RedisClusterMessagesCache.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/RedisClusterMessagesCache.java
@@ -2,6 +2,7 @@ package org.whispersystems.textsecuregcm.storage;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.protobuf.InvalidProtocolBufferException;
+import io.dropwizard.lifecycle.Managed;
 import io.lettuce.core.ScriptOutputType;
 import io.lettuce.core.cluster.SlotHash;
 import io.lettuce.core.cluster.event.ClusterTopologyChangedEvent;
@@ -36,7 +37,7 @@ import java.util.concurrent.ExecutorService;
 
 import static com.codahale.metrics.MetricRegistry.name;
 
-public class RedisClusterMessagesCache extends RedisClusterPubSubAdapter<String, String> implements UserMessagesCache {
+public class RedisClusterMessagesCache extends RedisClusterPubSubAdapter<String, String> implements UserMessagesCache, Managed {
 
     private final FaultTolerantRedisCluster                     redisCluster;
     private final FaultTolerantPubSubConnection<String, String> pubSubConnection;
@@ -90,7 +91,10 @@ public class RedisClusterMessagesCache extends RedisClusterPubSubAdapter<String,
         this.getItemsScript           = ClusterLuaScript.fromResource(redisCluster, "lua/get_items.lua",             ScriptOutputType.MULTI);
         this.removeQueueScript        = ClusterLuaScript.fromResource(redisCluster, "lua/remove_queue.lua",          ScriptOutputType.STATUS);
         this.getQueuesToPersistScript = ClusterLuaScript.fromResource(redisCluster, "lua/get_queues_to_persist.lua", ScriptOutputType.MULTI);
+    }
 
+    @Override
+    public void start() {
         pubSubConnection.usePubSubConnection(connection -> {
             connection.addListener(this);
             connection.getResources().eventBus().get()
@@ -102,6 +106,11 @@ public class RedisClusterMessagesCache extends RedisClusterPubSubAdapter<String,
         });
     }
 
+    @Override
+    public void stop() {
+        pubSubConnection.usePubSubConnection(connection -> connection.sync().masters().commands().unsubscribe());
+    }
+
     private void resubscribeAll() {
         final Set<String> queueNames;
 
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/RedisClusterMessagesCacheTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/RedisClusterMessagesCacheTest.java
index 4202fe521..0a7d1e4c3 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/RedisClusterMessagesCacheTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/RedisClusterMessagesCacheTest.java
@@ -35,14 +35,18 @@ public class RedisClusterMessagesCacheTest extends AbstractMessagesCacheTest {
 
         notificationExecutorService = Executors.newSingleThreadExecutor();
         messagesCache               = new RedisClusterMessagesCache(getRedisCluster(), notificationExecutorService);
+
+        messagesCache.start();
     }
 
     @Override
     public void tearDown() throws Exception {
-        super.tearDown();
+        messagesCache.stop();
 
         notificationExecutorService.shutdown();
         notificationExecutorService.awaitTermination(1, TimeUnit.SECONDS);
+
+        super.tearDown();
     }
 
     @Override