Scope disconnection request listeners to a single connection

service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java

@@ -626,7 +626,8 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         () -> dynamicConfigurationManager.getConfiguration().getSvrbStatusCodesToIgnoreForAccountDeletion());
     SecureStorageClient secureStorageClient = new SecureStorageClient(storageCredentialsGenerator,
         storageServiceExecutor, storageServiceRetryExecutor, config.getSecureStorageServiceConfiguration());
-    DisconnectionRequestManager disconnectionRequestManager = new DisconnectionRequestManager(pubsubClient, disconnectionRequestListenerExecutor);
+    final GrpcClientConnectionManager grpcClientConnectionManager = new GrpcClientConnectionManager();
+    DisconnectionRequestManager disconnectionRequestManager = new DisconnectionRequestManager(pubsubClient, grpcClientConnectionManager, disconnectionRequestListenerExecutor);
     ProfilesManager profilesManager = new ProfilesManager(profiles, cacheCluster, asyncCdnS3Client, config.getCdnConfiguration().bucket());
     MessagesCache messagesCache = new MessagesCache(messagesCluster, messageDeliveryScheduler,
         messageDeletionAsyncExecutor, clock);
@@ -675,8 +676,6 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     MessageDeliveryLoopMonitor messageDeliveryLoopMonitor =
         config.logMessageDeliveryLoops() ? new RedisMessageDeliveryLoopMonitor(rateLimitersCluster) : new NoopMessageDeliveryLoopMonitor();
 
-    disconnectionRequestManager.addListener(webSocketConnectionEventManager);
-
     final RegistrationLockVerificationManager registrationLockVerificationManager = new RegistrationLockVerificationManager(
         accountsManager, disconnectionRequestManager, svr2CredentialsGenerator, registrationRecoveryPasswordsManager,
         pushNotificationManager, rateLimiters);
@@ -816,10 +815,6 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         config.getAppleDeviceCheck().teamId(),
         config.getAppleDeviceCheck().bundleId());
 
-    final GrpcClientConnectionManager grpcClientConnectionManager = new GrpcClientConnectionManager();
-
-    disconnectionRequestManager.addListener(grpcClientConnectionManager);
-
     final ManagedDefaultEventLoopGroup localEventLoopGroup = new ManagedDefaultEventLoopGroup();
 
     final RemoteDeprecationFilter remoteDeprecationFilter = new RemoteDeprecationFilter(dynamicConfigurationManager);
@@ -1001,7 +996,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         config.idlePrimaryDeviceReminderConfiguration().minIdleDuration(), Clock.systemUTC()));
     webSocketEnvironment.setConnectListener(
         new AuthenticatedConnectListener(accountsManager, receiptSender, messagesManager, messageMetrics, pushNotificationManager,
-            pushNotificationScheduler, webSocketConnectionEventManager, websocketScheduledExecutor,
+            pushNotificationScheduler, webSocketConnectionEventManager, disconnectionRequestManager, websocketScheduledExecutor,
             messageDeliveryScheduler, clientReleaseManager, messageDeliveryLoopMonitor, experimentEnrollmentManager));
     webSocketEnvironment.jersey().register(new RateLimitByIpFilter(rateLimiters));
     webSocketEnvironment.jersey().register(new RequestStatisticsFilter(TrafficSource.WEBSOCKET));

service/src/main/java/org/whispersystems/textsecuregcm/auth/DisconnectionRequestListener.java

@@ -5,20 +5,15 @@
 
 package org.whispersystems.textsecuregcm.auth;
 
-import java.util.Collection;
-import java.util.UUID;
-
 /**
- * A disconnection request listener receives and handles requests to close authenticated client network connections.
+ * A disconnection request listener receives and handles a request to close an authenticated network connection for a
+ * specific client.
  */
 public interface DisconnectionRequestListener {
 
   /**
-   * Handles a request to close authenticated network connections for one or more authenticated devices. Requests are
+   * Handles a request to close an authenticated network connection for a specific authenticated device. Requests are
    * dispatched on dedicated threads, and implementations may safely block.
-   *
-   * @param accountIdentifier the account identifier for which to close authenticated connections
-   * @param deviceIds the device IDs within the identified account for which to close authenticated connections
    */
-  void handleDisconnectionRequest(UUID accountIdentifier, Collection<Byte> deviceIds);
+  void handleDisconnectionRequest();
 }

service/src/main/java/org/whispersystems/textsecuregcm/auth/DisconnectionRequestManager.java

@@ -5,21 +5,27 @@
 
 package org.whispersystems.textsecuregcm.auth;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.protobuf.InvalidProtocolBufferException;
 import io.dropwizard.lifecycle.Managed;
 import io.lettuce.core.pubsub.RedisPubSubAdapter;
 import io.micrometer.core.instrument.Counter;
 import io.micrometer.core.instrument.Metrics;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.UUID;
 import java.util.concurrent.CompletionStage;
-import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executor;
 import javax.annotation.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.auth.grpc.AuthenticatedDevice;
+import org.whispersystems.textsecuregcm.grpc.net.GrpcClientConnectionManager;
 import org.whispersystems.textsecuregcm.identity.IdentityType;
 import org.whispersystems.textsecuregcm.metrics.MetricsUtil;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantPubSubConnection;
@@ -37,11 +43,11 @@ import org.whispersystems.textsecuregcm.util.UUIDUtil;
 public class DisconnectionRequestManager extends RedisPubSubAdapter<byte[], byte[]> implements Managed {
 
   private final FaultTolerantRedisClient pubSubClient;
+  private final GrpcClientConnectionManager grpcClientConnectionManager;
   private final Executor listenerEventExecutor;
 
-  // We expect just a couple listeners to get added at startup time and not at all at steady-state. There are several
-  // reasonable ways to model this, but a copy-on-write list gives us good flexibility with minimal performance cost.
-  private final List<DisconnectionRequestListener> listeners = new CopyOnWriteArrayList<>();
+  private final Map<AccountIdentifierAndDeviceId, List<DisconnectionRequestListener>> listeners =
+      new ConcurrentHashMap<>();
 
   @Nullable
   private FaultTolerantPubSubConnection<byte[], byte[]> pubSubConnection;
@@ -56,10 +62,14 @@ public class DisconnectionRequestManager extends RedisPubSubAdapter<byte[], byte
 
   private static final Logger logger = LoggerFactory.getLogger(DisconnectionRequestManager.class);
 
+  private record AccountIdentifierAndDeviceId(UUID accountIdentifier, byte deviceId) {}
+
   public DisconnectionRequestManager(final FaultTolerantRedisClient pubSubClient,
+      final GrpcClientConnectionManager grpcClientConnectionManager,
       final Executor listenerEventExecutor) {
 
     this.pubSubClient = pubSubClient;
+    this.grpcClientConnectionManager = grpcClientConnectionManager;
     this.listenerEventExecutor = listenerEventExecutor;
   }
 
@@ -85,13 +95,41 @@ public class DisconnectionRequestManager extends RedisPubSubAdapter<byte[], byte
   }
 
   /**
-   * Adds a listener for disconnection requests. Listeners will receive all broadcast disconnection requests regardless
-   * of whether the device in connection is connected to this server.
+   * Adds a listener for disconnection requests for a specific authenticated device.
    *
+   * @param accountIdentifier TODO
+   * @param deviceId TODO
    * @param listener the listener to register
    */
-  public void addListener(final DisconnectionRequestListener listener) {
-    listeners.add(listener);
+  public void addListener(final UUID accountIdentifier, final byte deviceId, final DisconnectionRequestListener listener) {
+    listeners.compute(new AccountIdentifierAndDeviceId(accountIdentifier, deviceId), (_, existingListeners) -> {
+      final List<DisconnectionRequestListener> listeners =
+          existingListeners == null ? new ArrayList<>() : existingListeners;
+
+      listeners.add(listener);
+
+      return listeners;
+    });
+  }
+
+  /**
+   * Removes a listener for disconnection requests for a specific authenticated device.
+   *
+   * @param accountIdentifier TODO
+   * @param deviceId TODO
+   * @param listener the listener to remove
+   */
+  public void removeListener(final UUID accountIdentifier, final byte deviceId, final DisconnectionRequestListener listener) {
+    listeners.computeIfPresent(new AccountIdentifierAndDeviceId(accountIdentifier, deviceId), (_, existingListeners) -> {
+      existingListeners.remove(listener);
+
+      return existingListeners.isEmpty() ? null : existingListeners;
+    });
+  }
+
+  @VisibleForTesting
+  List<DisconnectionRequestListener> getListeners(final UUID accountIdentifier, final byte deviceId) {
+    return listeners.getOrDefault(new AccountIdentifierAndDeviceId(accountIdentifier, deviceId), Collections.emptyList());
   }
 
   /**
@@ -154,12 +192,17 @@ public class DisconnectionRequestManager extends RedisPubSubAdapter<byte[], byte
       return;
     }
 
-    for (final DisconnectionRequestListener listener : listeners) {
-      try {
-        listenerEventExecutor.execute(() -> listener.handleDisconnectionRequest(accountIdentifier, deviceIds));
-      } catch (final Exception e) {
-        logger.warn("Listener failed to handle disconnection request", e);
-      }
-    }
+    deviceIds.forEach(deviceId -> {
+      grpcClientConnectionManager.closeConnection(new AuthenticatedDevice(accountIdentifier, deviceId));
+
+      listeners.getOrDefault(new AccountIdentifierAndDeviceId(accountIdentifier, deviceId), Collections.emptyList())
+          .forEach(listener -> listenerEventExecutor.execute(() -> {
+            try {
+              listener.handleDisconnectionRequest();
+            } catch (final Exception e) {
+              logger.warn("Listener failed to handle disconnection request", e);
+            }
+          }));
+    });
   }
 }

service/src/main/java/org/whispersystems/textsecuregcm/grpc/net/GrpcClientConnectionManager.java

@@ -10,19 +10,16 @@ import io.netty.channel.local.LocalChannel;
 import io.netty.util.AttributeKey;
 import java.net.InetAddress;
 import java.util.ArrayList;
-import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Optional;
-import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import javax.annotation.Nullable;
 import org.apache.commons.lang3.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-import org.whispersystems.textsecuregcm.auth.DisconnectionRequestListener;
 import org.whispersystems.textsecuregcm.auth.grpc.AuthenticatedDevice;
 import org.whispersystems.textsecuregcm.grpc.ChannelNotFoundException;
 import org.whispersystems.textsecuregcm.grpc.RequestAttributes;
@@ -45,7 +42,7 @@ import org.whispersystems.textsecuregcm.util.ClosableEpoch;
  * Methods for requesting connection closure accept an {@link AuthenticatedDevice} to identify the connection and may
  * be called from any application code.
  */
-public class GrpcClientConnectionManager implements DisconnectionRequestListener {
+public class GrpcClientConnectionManager {
 
   private final Map<LocalAddress, Channel> remoteChannelsByLocalAddress = new ConcurrentHashMap<>();
   private final Map<AuthenticatedDevice, List<Channel>> remoteChannelsByAuthenticatedDevice = new ConcurrentHashMap<>();
@@ -62,7 +59,7 @@ public class GrpcClientConnectionManager implements DisconnectionRequestListener
   static final AttributeKey<ClosableEpoch> EPOCH_ATTRIBUTE_KEY =
       AttributeKey.valueOf(GrpcClientConnectionManager.class, "epoch");
 
-  private static OutboundCloseErrorMessage SERVER_CLOSED =
+  private static final OutboundCloseErrorMessage SERVER_CLOSED =
       new OutboundCloseErrorMessage(OutboundCloseErrorMessage.Code.SERVER_CLOSED, "server closed");
 
   private static final Logger log = LoggerFactory.getLogger(GrpcClientConnectionManager.class);
@@ -268,11 +265,4 @@ public class GrpcClientConnectionManager implements DisconnectionRequestListener
           }));
     });
   }
-
-  @Override
-  public void handleDisconnectionRequest(final UUID accountIdentifier, final Collection<Byte> deviceIds) {
-    deviceIds.stream()
-        .map(deviceId -> new AuthenticatedDevice(accountIdentifier, deviceId))
-        .forEach(this::closeConnection);
-  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/push/WebSocketConnectionEventListener.java

@@ -22,11 +22,8 @@ public interface WebSocketConnectionEventListener {
   void handleMessagesPersisted();
 
   /**
-   * Indicates that the client's presence has been displaced and the listener should close the client's underlying
-   * network connection.
-   *
-   * @param connectedElsewhere if {@code true}, indicates that the client's presence has been displaced by another
-   *                           connection from the same client
+   * Indicates a newer instance of this client has started reading messages and the listener should close this client's
+   * underlying network connection.
    */
-  void handleConnectionDisplaced(boolean connectedElsewhere);
+  void handleConflictingMessageReader();
 }

service/src/main/java/org/whispersystems/textsecuregcm/push/WebSocketConnectionEventManager.java

@@ -17,11 +17,9 @@ import io.micrometer.core.instrument.Metrics;
 import io.micrometer.core.instrument.Tags;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
-import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
@@ -32,8 +30,6 @@ import java.util.function.Function;
 import javax.annotation.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-import org.whispersystems.textsecuregcm.auth.DisconnectionRequestListener;
-import org.whispersystems.textsecuregcm.entities.MessageProtos;
 import org.whispersystems.textsecuregcm.metrics.MetricsUtil;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantPubSubClusterConnection;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisClusterClient;
@@ -56,10 +52,9 @@ import org.whispersystems.textsecuregcm.util.Util;
  * servers, but cannot guarantee at-most-one behavior.
  *
  * @see WebSocketConnectionEventListener
- * @see org.whispersystems.textsecuregcm.storage.MessagesManager#insert(UUID, byte, MessageProtos.Envelope)
+ * @see org.whispersystems.textsecuregcm.storage.MessagesManager#insert(UUID, Map)
  */
-public class WebSocketConnectionEventManager extends RedisClusterPubSubAdapter<byte[], byte[]> implements Managed,
-    DisconnectionRequestListener {
+public class WebSocketConnectionEventManager extends RedisClusterPubSubAdapter<byte[], byte[]> implements Managed {
 
   private final AccountsManager accountsManager;
   private final PushNotificationManager pushNotificationManager;
@@ -145,7 +140,7 @@ public class WebSocketConnectionEventManager extends RedisClusterPubSubAdapter<b
   /**
    * Marks the given device as "present" for message delivery and registers a listener for new messages and conflicting
    * connections. If the given device already has a presence registered with this manager, that presence is displaced
-   * immediately and the listener's {@link WebSocketConnectionEventListener#handleConnectionDisplaced(boolean)} method is called.
+   * immediately and the listener's {@link WebSocketConnectionEventListener#handleConflictingMessageReader()} method is called.
    *
    * @param accountIdentifier the account identifier for the newly-connected device
    * @param deviceId the ID of the newly-connected device within the given account
@@ -189,7 +184,7 @@ public class WebSocketConnectionEventManager extends RedisClusterPubSubAdapter<b
         });
 
     if (displacedListener.get() != null) {
-      listenerEventExecutor.execute(() -> displacedListener.get().handleConnectionDisplaced(true));
+      listenerEventExecutor.execute(() -> displacedListener.get().handleConflictingMessageReader());
     }
 
     return subscribeFuture.get()
@@ -260,14 +255,6 @@ public class WebSocketConnectionEventManager extends RedisClusterPubSubAdapter<b
     return listenersByAccountAndDeviceIdentifier.containsKey(new AccountAndDeviceIdentifier(accountUuid, deviceId));
   }
 
-  @Override
-  public void handleDisconnectionRequest(final UUID accountIdentifier, final Collection<Byte> deviceIds) {
-    deviceIds.stream()
-        .map(deviceId -> listenersByAccountAndDeviceIdentifier.get(new AccountAndDeviceIdentifier(accountIdentifier, deviceId)))
-        .filter(Objects::nonNull)
-        .forEach(listener -> listener.handleConnectionDisplaced(false));
-  }
-
   @VisibleForTesting
   void resubscribe(final ClusterTopologyChangedEvent clusterTopologyChangedEvent) {
     final boolean[] changedSlots = RedisClusterUtil.getChangedSlots(clusterTopologyChangedEvent);
@@ -339,7 +326,7 @@ public class WebSocketConnectionEventManager extends RedisClusterPubSubAdapter<b
           // Only act on new connections to other event manager instances; we'll learn about displacements in THIS
           // instance when we update the listener map in `handleClientConnected`
           if (!this.serverId.equals(UUIDUtil.fromByteString(clientEvent.getClientConnected().getServerId()))) {
-            listenerEventExecutor.execute(() -> listener.handleConnectionDisplaced(true));
+            listenerEventExecutor.execute(listener::handleConflictingMessageReader);
           }
         }
 

service/src/main/java/org/whispersystems/textsecuregcm/websocket/AuthenticatedConnectListener.java

@@ -13,7 +13,9 @@ import java.util.concurrent.ScheduledExecutorService;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedDevice;
+import org.whispersystems.textsecuregcm.auth.DisconnectionRequestManager;
 import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
+import org.whispersystems.textsecuregcm.identity.IdentityType;
 import org.whispersystems.textsecuregcm.limits.MessageDeliveryLoopMonitor;
 import org.whispersystems.textsecuregcm.metrics.MessageMetrics;
 import org.whispersystems.textsecuregcm.metrics.OpenWebSocketCounter;
@@ -48,6 +50,7 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
   private final PushNotificationManager pushNotificationManager;
   private final PushNotificationScheduler pushNotificationScheduler;
   private final WebSocketConnectionEventManager webSocketConnectionEventManager;
+  private final DisconnectionRequestManager disconnectionRequestManager;
   private final ScheduledExecutorService scheduledExecutorService;
   private final Scheduler messageDeliveryScheduler;
   private final ClientReleaseManager clientReleaseManager;
@@ -58,17 +61,18 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
   private final OpenWebSocketCounter openUnauthenticatedWebSocketCounter;
 
   public AuthenticatedConnectListener(
-      AccountsManager accountsManager,
-      ReceiptSender receiptSender,
-      MessagesManager messagesManager,
-      MessageMetrics messageMetrics,
-      PushNotificationManager pushNotificationManager,
-      PushNotificationScheduler pushNotificationScheduler,
-      WebSocketConnectionEventManager webSocketConnectionEventManager,
-      ScheduledExecutorService scheduledExecutorService,
-      Scheduler messageDeliveryScheduler,
-      ClientReleaseManager clientReleaseManager,
-      MessageDeliveryLoopMonitor messageDeliveryLoopMonitor,
+      final AccountsManager accountsManager,
+      final ReceiptSender receiptSender,
+      final MessagesManager messagesManager,
+      final MessageMetrics messageMetrics,
+      final PushNotificationManager pushNotificationManager,
+      final PushNotificationScheduler pushNotificationScheduler,
+      final WebSocketConnectionEventManager webSocketConnectionEventManager,
+      final DisconnectionRequestManager disconnectionRequestManager,
+      final ScheduledExecutorService scheduledExecutorService,
+      final Scheduler messageDeliveryScheduler,
+      final ClientReleaseManager clientReleaseManager,
+      final MessageDeliveryLoopMonitor messageDeliveryLoopMonitor,
       final ExperimentEnrollmentManager experimentEnrollmentManager) {
 
     this.accountsManager = accountsManager;
@@ -78,6 +82,7 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
     this.pushNotificationManager = pushNotificationManager;
     this.pushNotificationScheduler = pushNotificationScheduler;
     this.webSocketConnectionEventManager = webSocketConnectionEventManager;
+    this.disconnectionRequestManager = disconnectionRequestManager;
     this.scheduledExecutorService = scheduledExecutorService;
     this.messageDeliveryScheduler = messageDeliveryScheduler;
     this.clientReleaseManager = clientReleaseManager;
@@ -104,7 +109,7 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
       final AuthenticatedDevice auth = context.getAuthenticated(AuthenticatedDevice.class);
 
       final Optional<Account> maybeAuthenticatedAccount = accountsManager.getByAccountIdentifier(auth.accountIdentifier());
-      final Optional<Device> maybeAuthenticatedDevice = maybeAuthenticatedAccount.flatMap(account -> account.getDevice(auth.deviceId()));;
+      final Optional<Device> maybeAuthenticatedDevice = maybeAuthenticatedAccount.flatMap(account -> account.getDevice(auth.deviceId()));
 
       if (maybeAuthenticatedAccount.isEmpty() || maybeAuthenticatedDevice.isEmpty()) {
         log.warn("{}:{} not found when opening authenticated WebSocket", auth.accountIdentifier(), auth.deviceId());
@@ -127,7 +132,15 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
           messageDeliveryLoopMonitor,
           experimentEnrollmentManager);
 
-      context.addWebsocketClosedListener((closingContext, statusCode, reason) -> {
+      disconnectionRequestManager.addListener(maybeAuthenticatedAccount.get().getIdentifier(IdentityType.ACI),
+          maybeAuthenticatedDevice.get().getId(),
+          connection);
+
+      context.addWebsocketClosedListener((_, _, _) -> {
+        disconnectionRequestManager.removeListener(maybeAuthenticatedAccount.get().getIdentifier(IdentityType.ACI),
+            maybeAuthenticatedDevice.get().getId(),
+            connection);
+
         // We begin the shutdown process by removing this client's "presence," which means it will again begin to
         // receive push notifications for inbound messages. We should do this first because, at this point, the
         // connection has already closed and attempts to actually deliver a message via the connection will not succeed.

service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java

@@ -37,6 +37,7 @@ import org.eclipse.jetty.util.StaticException;
 import org.reactivestreams.Publisher;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.auth.DisconnectionRequestListener;
 import org.whispersystems.textsecuregcm.controllers.MessageController;
 import org.whispersystems.textsecuregcm.entities.MessageProtos.Envelope;
 import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
@@ -65,7 +66,7 @@ import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.scheduler.Scheduler;
 
-public class WebSocketConnection implements WebSocketConnectionEventListener {
+public class WebSocketConnection implements WebSocketConnectionEventListener, DisconnectionRequestListener {
 
   private static final DistributionSummary messageTime = Metrics.summary(
       name(MessageController.class, "messageDeliveryDuration"));
@@ -506,24 +507,23 @@ public class WebSocketConnection implements WebSocketConnectionEventListener {
   }
 
   @Override
-  public void handleConnectionDisplaced(final boolean connectedElsewhere) {
+  public void handleConflictingMessageReader() {
+    closeConnection(4409, "Connected elsewhere");
+  }
+
+  @Override
+  public void handleDisconnectionRequest() {
+    closeConnection(4401, "Reauthentication required");
+  }
+
+  private void closeConnection(final int code, final String message) {
     final Tags tags = Tags.of(
         UserAgentTagUtil.getPlatformTag(client.getUserAgent()),
-        Tag.of("connectedElsewhere", String.valueOf(connectedElsewhere)));
+        // TODO We should probably just use the status code directly
+        Tag.of("connectedElsewhere", String.valueOf(code == 4409)));
 
     Metrics.counter(DISPLACEMENT_COUNTER_NAME, tags).increment();
 
-    final int code;
-    final String message;
-
-    if (connectedElsewhere) {
-      code = 4409;
-      message = "Connected elsewhere";
-    } else {
-      code = 4401;
-      message = "Reauthentication required";
-    }
-
     client.close(code, message);
   }
 

service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java

@@ -36,6 +36,7 @@ import org.whispersystems.textsecuregcm.controllers.SecureValueRecovery2Controll
 import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
 import org.whispersystems.textsecuregcm.controllers.SecureValueRecoveryBController;
 import org.whispersystems.textsecuregcm.experiment.PushNotificationExperimentSamples;
+import org.whispersystems.textsecuregcm.grpc.net.GrpcClientConnectionManager;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.metrics.MicrometerAwsSdkMetricPublisher;
 import org.whispersystems.textsecuregcm.push.APNSender;
@@ -254,7 +255,8 @@ record CommandDependencies(
         () -> dynamicConfigurationManager.getConfiguration().getSvrbStatusCodesToIgnoreForAccountDeletion());
     SecureStorageClient secureStorageClient = new SecureStorageClient(storageCredentialsGenerator,
         storageServiceExecutor, storageServiceRetryExecutor, configuration.getSecureStorageServiceConfiguration());
-    DisconnectionRequestManager disconnectionRequestManager = new DisconnectionRequestManager(pubsubClient, disconnectionRequestListenerExecutor);
+    GrpcClientConnectionManager grpcClientConnectionManager = new GrpcClientConnectionManager();
+    DisconnectionRequestManager disconnectionRequestManager = new DisconnectionRequestManager(pubsubClient, grpcClientConnectionManager, disconnectionRequestListenerExecutor);
     MessagesCache messagesCache = new MessagesCache(messagesCluster,
         messageDeliveryScheduler, messageDeletionExecutor, Clock.systemUTC());
     ProfilesManager profilesManager = new ProfilesManager(profiles, cacheCluster, asyncCdnS3Client,