From 1ba00a66eb8aff99dc54a2a8af8eb66d0da64b27 Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Wed, 9 Mar 2022 16:54:43 -0500
Subject: [PATCH] Pass the reason for displacement to presence displacement
 listeners

---
 .../push/ClientPresenceManager.java             | 11 ++++++-----
 .../push/DisplacedPresenceListener.java         |  2 +-
 .../websocket/WebSocketConnection.java          |  2 +-
 .../push/ClientPresenceManagerTest.java         | 17 ++++++++---------
 4 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/push/ClientPresenceManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/push/ClientPresenceManager.java
index 317cb889d..88bea9049 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/push/ClientPresenceManager.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/push/ClientPresenceManager.java
@@ -158,7 +158,7 @@ public class ClientPresenceManager extends RedisClusterPubSubAdapter<String, Str
     try (final Timer.Context ignored = setPresenceTimer.time()) {
       final String presenceKey = getPresenceKey(accountUuid, deviceId);
 
-      displacePresence(presenceKey);
+      displacePresence(presenceKey, true);
 
       displacementListenersByPresenceKey.put(presenceKey, displacementListener);
 
@@ -182,7 +182,7 @@ public class ClientPresenceManager extends RedisClusterPubSubAdapter<String, Str
     final String presenceKey = getPresenceKey(accountUuid, deviceId);
 
     if (isLocallyPresent(accountUuid, deviceId)) {
-      displacePresence(presenceKey);
+      displacePresence(presenceKey, false);
     }
 
     // If connected locally, we still need to clean up the presence key.
@@ -190,11 +190,11 @@ public class ClientPresenceManager extends RedisClusterPubSubAdapter<String, Str
     presenceCluster.useCluster(connection -> connection.sync().del(presenceKey));
   }
 
-  private void displacePresence(final String presenceKey) {
+  private void displacePresence(final String presenceKey, final boolean connectedElsewhere) {
     final DisplacedPresenceListener displacementListener = displacementListenersByPresenceKey.get(presenceKey);
 
     if (displacementListener != null) {
-      displacementListener.handleDisplacement();
+      displacementListener.handleDisplacement(connectedElsewhere);
     }
 
     clearPresence(presenceKey);
@@ -287,12 +287,13 @@ public class ClientPresenceManager extends RedisClusterPubSubAdapter<String, Str
       if ("set".equals(message) || "del".equals(message)) {
         // for "set", another process has overwritten this presence key, which means the client has connected to another host.
         // for "del", another process has indicated the client should be disconnected
+        final boolean connectedElsewhere = "set".equals(message);
 
         // At this point, we're on a Lettuce IO thread and need to dispatch to a separate thread before making
         // synchronous Lettuce calls to avoid deadlocking.
         keyspaceNotificationExecutorService.execute(() -> {
           try {
-            displacePresence(channel.substring("__keyspace@0__:".length()));
+            displacePresence(channel.substring("__keyspace@0__:".length()), connectedElsewhere);
             remoteDisplacementMeter.mark();
           } catch (final Exception e) {
             log.warn("Error displacing presence", e);
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/push/DisplacedPresenceListener.java b/service/src/main/java/org/whispersystems/textsecuregcm/push/DisplacedPresenceListener.java
index 879d4895f..88258be70 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/push/DisplacedPresenceListener.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/push/DisplacedPresenceListener.java
@@ -12,5 +12,5 @@ package org.whispersystems.textsecuregcm.push;
 @FunctionalInterface
 public interface DisplacedPresenceListener {
 
-    void handleDisplacement();
+    void handleDisplacement(boolean connectedElsewhere);
 }
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java b/service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java
index 52c688934..b14930bb7 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java
@@ -382,7 +382,7 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
   }
 
   @Override
-  public void handleDisplacement() {
+  public void handleDisplacement(final boolean connectedElsewhere) {
     Metrics.counter(DISPLACEMENT_COUNTER_NAME, List.of(UserAgentTagUtil.getPlatformTag(client.getUserAgent()))).increment();
 
     try {
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/push/ClientPresenceManagerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/push/ClientPresenceManagerTest.java
index 632541f6c..2c628f510 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/push/ClientPresenceManagerTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/push/ClientPresenceManagerTest.java
@@ -39,7 +39,7 @@ class ClientPresenceManagerTest {
   private ScheduledExecutorService presenceRenewalExecutorService;
   private ClientPresenceManager clientPresenceManager;
 
-  private static final DisplacedPresenceListener NO_OP = () -> {
+  private static final DisplacedPresenceListener NO_OP = connectedElsewhere -> {
   };
 
   @BeforeEach
@@ -94,7 +94,7 @@ class ClientPresenceManagerTest {
     final long deviceId = 1;
 
     final AtomicInteger displacementCounter = new AtomicInteger(0);
-    final DisplacedPresenceListener displacementListener = displacementCounter::incrementAndGet;
+    final DisplacedPresenceListener displacementListener = connectedElsewhere -> displacementCounter.incrementAndGet();
 
     clientPresenceManager.setPresent(accountUuid, deviceId, displacementListener);
 
@@ -114,7 +114,7 @@ class ClientPresenceManagerTest {
 
     clientPresenceManager.start();
 
-    clientPresenceManager.setPresent(accountUuid, deviceId, () -> displaced.complete(null));
+    clientPresenceManager.setPresent(accountUuid, deviceId, connectedElsewhere -> displaced.complete(null));
 
     REDIS_CLUSTER_EXTENSION.getRedisCluster().useCluster(
         connection -> connection.sync().set(ClientPresenceManager.getPresenceKey(accountUuid, deviceId),
@@ -132,7 +132,7 @@ class ClientPresenceManagerTest {
 
     clientPresenceManager.start();
 
-    clientPresenceManager.setPresent(accountUuid, deviceId, () -> displaced.complete(null));
+    clientPresenceManager.setPresent(accountUuid, deviceId, connectedElsewhere -> displaced.complete(null));
 
     clientPresenceManager.getPubSubConnection()
         .usePubSubConnection(connection -> connection.getResources().eventBus()
@@ -336,11 +336,10 @@ class ClientPresenceManagerTest {
       final long deviceId = 1L;
 
       final CompletableFuture<?> displaced = new CompletableFuture<>();
-      final DisplacedPresenceListener listener1 = () -> displaced.complete(null);
+      final DisplacedPresenceListener listener1 = connectedElsewhere -> displaced.complete(null);
       server1.setPresent(uuid1, deviceId, listener1);
 
-      server2.setPresent(uuid1, deviceId, () -> {
-      });
+      server2.setPresent(uuid1, deviceId, connectedElsewhere -> {});
 
       assertTimeoutPreemptively(Duration.ofSeconds(10), displaced::join);
     }
@@ -351,7 +350,7 @@ class ClientPresenceManagerTest {
       final long deviceId = 1L;
 
       final CompletableFuture<?> displaced = new CompletableFuture<>();
-      final DisplacedPresenceListener listener1 = () -> displaced.complete(null);
+      final DisplacedPresenceListener listener1 = connectedElsewhere -> displaced.complete(null);
       server1.setPresent(uuid1, deviceId, listener1);
 
       server1.disconnectPresence(uuid1, deviceId);
@@ -365,7 +364,7 @@ class ClientPresenceManagerTest {
       final long deviceId = 1L;
 
       final CompletableFuture<?> displaced = new CompletableFuture<>();
-      final DisplacedPresenceListener listener1 = () -> displaced.complete(null);
+      final DisplacedPresenceListener listener1 = connectedElsewhere -> displaced.complete(null);
       server1.setPresent(uuid1, deviceId, listener1);
 
       server2.disconnectPresence(uuid1, deviceId);