From 92d36b725fe2ad85e5cf150c9c197bf0eae96080 Mon Sep 17 00:00:00 2001 From: Jon Chambers Date: Mon, 27 Jun 2022 17:45:06 -0400 Subject: [PATCH] Allow presence keys to expire if not periodically renewed --- .../push/ClientPresenceManager.java | 22 +++++-- .../AuthenticatedConnectListener.java | 19 ++++++ .../src/main/resources/lua/renew_presence.lua | 7 +++ .../push/ClientPresenceManagerTest.java | 58 +++++++++++++++++++ 4 files changed, 101 insertions(+), 5 deletions(-) create mode 100644 service/src/main/resources/lua/renew_presence.lua 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 47359650d..317cb889d 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/push/ClientPresenceManager.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/push/ClientPresenceManager.java @@ -55,6 +55,7 @@ public class ClientPresenceManager extends RedisClusterPubSubAdapter pubSubConnection; private final ClusterLuaScript clearPresenceScript; + private final ClusterLuaScript renewPresenceScript; private final ExecutorService keyspaceNotificationExecutorService; private final ScheduledExecutorService scheduledExecutorService; @@ -71,6 +72,7 @@ public class ClientPresenceManager extends RedisClusterPubSubAdapter connection.sync().upstream().commands().unsubscribe(getManagerPresenceChannel(managerId))); } - public void setPresent(final UUID accountUuid, final long deviceId, - final DisplacedPresenceListener displacementListener) { + public void setPresent(final UUID accountUuid, final long deviceId, final DisplacedPresenceListener displacementListener) { + try (final Timer.Context ignored = setPresenceTimer.time()) { final String presenceKey = getPresenceKey(accountUuid, deviceId); @@ -164,13 +166,18 @@ public class ClientPresenceManager extends RedisClusterPubSubAdapter commands = connection.sync(); commands.sadd(connectedClientSetKey, presenceKey); - commands.set(presenceKey, managerId); + commands.setex(presenceKey, PRESENCE_EXPIRATION_SECONDS, managerId); }); subscribeForRemotePresenceChanges(presenceKey); } } + public void renewPresence(final UUID accountUuid, final long deviceId) { + renewPresenceScript.execute(List.of(getPresenceKey(accountUuid, deviceId)), + List.of(managerId, String.valueOf(PRESENCE_EXPIRATION_SECONDS))); + } + public void disconnectPresence(final UUID accountUuid, final long deviceId) { final String presenceKey = getPresenceKey(accountUuid, deviceId); @@ -295,6 +302,11 @@ public class ClientPresenceManager extends RedisClusterPubSubAdapter apnFallbackManager.cancel(auth.getAccount(), device)); + final AtomicReference> renewPresenceFutureReference = new AtomicReference<>(); + context.addListener(new WebSocketSessionContext.WebSocketEventListener() { @Override public void onWebSocketClose(WebSocketSessionContext context, int statusCode, String reason) { openWebsocketCounter.dec(); timer.stop(); + final ScheduledFuture renewPresenceFuture = renewPresenceFutureReference.get(); + + if (renewPresenceFuture != null) { + renewPresenceFuture.cancel(false); + } + connection.stop(); RedisOperation.unchecked( @@ -94,6 +107,12 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener { connection.start(); clientPresenceManager.setPresent(auth.getAccount().getUuid(), device.getId(), connection); messagesManager.addMessageAvailabilityListener(auth.getAccount().getUuid(), device.getId(), connection); + + renewPresenceFutureReference.set(scheduledExecutorService.scheduleAtFixedRate(() -> RedisOperation.unchecked(() -> + clientPresenceManager.renewPresence(auth.getAccount().getUuid(), device.getId())), + RENEW_PRESENCE_INTERVAL_MINUTES, + RENEW_PRESENCE_INTERVAL_MINUTES, + TimeUnit.MINUTES)); } catch (final Exception e) { log.warn("Failed to initialize websocket", e); context.getClient().close(1011, "Unexpected error initializing connection"); diff --git a/service/src/main/resources/lua/renew_presence.lua b/service/src/main/resources/lua/renew_presence.lua new file mode 100644 index 000000000..71b47c869 --- /dev/null +++ b/service/src/main/resources/lua/renew_presence.lua @@ -0,0 +1,7 @@ +local presenceKey = KEYS[1] +local presenceUuid = ARGV[1] +local expireSeconds = ARGV[2] + +if redis.call("GET", presenceKey) == presenceUuid then + redis.call("EXPIRE", presenceKey, expireSeconds) +end 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 5eb0b5dd6..632541f6c 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/push/ClientPresenceManagerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/push/ClientPresenceManagerTest.java @@ -196,6 +196,64 @@ class ClientPresenceManagerTest { .sismember(ClientPresenceManager.MANAGER_SET_KEY, missingPeerId))); } + @Test + void testInitialPresenceExpiration() { + final UUID accountUuid = UUID.randomUUID(); + final long deviceId = 1; + + clientPresenceManager.setPresent(accountUuid, deviceId, NO_OP); + + { + final int ttl = REDIS_CLUSTER_EXTENSION.getRedisCluster().withCluster(connection -> + connection.sync().ttl(ClientPresenceManager.getPresenceKey(accountUuid, deviceId)).intValue()); + + assertTrue(ttl > 0); + } + } + + @Test + void testRenewPresence() { + final UUID accountUuid = UUID.randomUUID(); + final long deviceId = 1; + + final String presenceKey = ClientPresenceManager.getPresenceKey(accountUuid, deviceId); + + REDIS_CLUSTER_EXTENSION.getRedisCluster().useCluster(connection -> + connection.sync().set(presenceKey, clientPresenceManager.getManagerId())); + + { + final int ttl = REDIS_CLUSTER_EXTENSION.getRedisCluster().withCluster(connection -> + connection.sync().ttl(presenceKey).intValue()); + + assertEquals(-1, ttl); + } + + clientPresenceManager.renewPresence(accountUuid, deviceId); + + { + final int ttl = REDIS_CLUSTER_EXTENSION.getRedisCluster().withCluster(connection -> + connection.sync().ttl(presenceKey).intValue()); + + assertTrue(ttl > 0); + } + } + + @Test + void testExpiredPresence() { + final UUID accountUuid = UUID.randomUUID(); + final long deviceId = 1; + + clientPresenceManager.setPresent(accountUuid, deviceId, NO_OP); + + assertTrue(clientPresenceManager.isPresent(accountUuid, deviceId)); + + // Hackily set this key to expire immediately + REDIS_CLUSTER_EXTENSION.getRedisCluster().useCluster(connection -> + connection.sync().expire(ClientPresenceManager.getPresenceKey(accountUuid, deviceId), 0)); + + assertFalse(clientPresenceManager.isPresent(accountUuid, deviceId)); + } + private void addClientPresence(final String managerId) { final String clientPresenceKey = ClientPresenceManager.getPresenceKey(UUID.randomUUID(), 7);