From 584fd06b88499f01d97e6c5d51a9e1c3c8985b09 Mon Sep 17 00:00:00 2001
From: Chris Eager <chris@signal.org>
Date: Thu, 17 Oct 2024 16:39:06 -0500
Subject: [PATCH] Add metric for closed connection age to KeepAliveController

---
 .../controllers/KeepAliveController.java      | 22 +++++++++++++------
 .../websocket/WebSocketClient.java            |  7 +++---
 2 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeepAliveController.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeepAliveController.java
index d36f44917..1d91c3a13 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeepAliveController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeepAliveController.java
@@ -10,7 +10,10 @@ import static com.codahale.metrics.MetricRegistry.name;
 import io.dropwizard.auth.Auth;
 import io.micrometer.core.instrument.Metrics;
 import io.micrometer.core.instrument.Tags;
+import io.micrometer.core.instrument.Timer;
 import io.swagger.v3.oas.annotations.tags.Tag;
+import java.time.Duration;
+import java.time.Instant;
 import java.util.Optional;
 import javax.ws.rs.GET;
 import javax.ws.rs.Path;
@@ -33,8 +36,9 @@ public class KeepAliveController {
 
   private final ClientPresenceManager clientPresenceManager;
 
-  private static final String NO_LOCAL_SUBSCRIPTION_COUNTER_NAME = name(KeepAliveController.class,
-      "noLocalSubscription");
+  private static final String CLOSED_CONNECTION_AGE_DISTRIBUTION_NAME = name(KeepAliveController.class,
+      "closedConnectionAge");
+
 
   public KeepAliveController(final ClientPresenceManager clientPresenceManager) {
     this.clientPresenceManager = clientPresenceManager;
@@ -46,16 +50,20 @@ public class KeepAliveController {
 
     maybeAuth.ifPresent(auth -> {
       if (!clientPresenceManager.isLocallyPresent(auth.getAccount().getUuid(), auth.getAuthenticatedDevice().getId())) {
+
+        final Duration age = Duration.between(context.getClient().getCreated(), Instant.now());
+
         logger.debug("***** No local subscription found for {}::{}; age = {}ms, User-Agent = {}",
-            auth.getAccount().getUuid(), auth.getAuthenticatedDevice().getId(),
-            System.currentTimeMillis() - context.getClient().getCreatedTimestamp(),
+            auth.getAccount().getUuid(), auth.getAuthenticatedDevice().getId(), age.toMillis(),
             context.getClient().getUserAgent());
 
         context.getClient().close(1000, "OK");
 
-        Metrics.counter(NO_LOCAL_SUBSCRIPTION_COUNTER_NAME,
-                Tags.of(UserAgentTagUtil.getPlatformTag(context.getClient().getUserAgent())))
-            .increment();
+        Timer.builder(CLOSED_CONNECTION_AGE_DISTRIBUTION_NAME)
+            .tags(Tags.of(UserAgentTagUtil.getPlatformTag(context.getClient().getUserAgent())))
+            .publishPercentileHistogram(true)
+            .register(Metrics.globalRegistry)
+            .record(age);
       }
     });
 
diff --git a/websocket-resources/src/main/java/org/whispersystems/websocket/WebSocketClient.java b/websocket-resources/src/main/java/org/whispersystems/websocket/WebSocketClient.java
index ba89c53ee..ee281dba1 100644
--- a/websocket-resources/src/main/java/org/whispersystems/websocket/WebSocketClient.java
+++ b/websocket-resources/src/main/java/org/whispersystems/websocket/WebSocketClient.java
@@ -7,6 +7,7 @@ package org.whispersystems.websocket;
 import com.google.common.net.HttpHeaders;
 import java.nio.ByteBuffer;
 import java.security.SecureRandom;
+import java.time.Instant;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -31,7 +32,7 @@ public class WebSocketClient {
   private final RemoteEndpoint remoteEndpoint;
   private final WebSocketMessageFactory messageFactory;
   private final Map<Long, CompletableFuture<WebSocketResponseMessage>> pendingRequestMapper;
-  private final long created;
+  private final Instant created;
 
   public WebSocketClient(Session session, RemoteEndpoint remoteEndpoint, WebSocketMessageFactory messageFactory,
                          Map<Long, CompletableFuture<WebSocketResponseMessage>> pendingRequestMapper) {
@@ -39,7 +40,7 @@ public class WebSocketClient {
     this.remoteEndpoint = remoteEndpoint;
     this.messageFactory = messageFactory;
     this.pendingRequestMapper = pendingRequestMapper;
-    this.created = System.currentTimeMillis();
+    this.created = Instant.now();
   }
 
   public CompletableFuture<WebSocketResponseMessage> sendRequest(String verb, String path,
@@ -78,7 +79,7 @@ public class WebSocketClient {
     return session.getUpgradeRequest().getHeader(HttpHeaders.USER_AGENT);
   }
 
-  public long getCreatedTimestamp() {
+  public Instant getCreated() {
     return this.created;
   }