From 38293efe759f879d611f80ebaf1ced9a2563dec7 Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Thu, 6 Aug 2020 12:30:12 -0400
Subject: [PATCH] Keep a running count of the number of open websockets.

---
 .../textsecuregcm/websocket/AuthenticatedConnectListener.java | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/websocket/AuthenticatedConnectListener.java b/service/src/main/java/org/whispersystems/textsecuregcm/websocket/AuthenticatedConnectListener.java
index 083b5b2d4..ecd6143ed 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/websocket/AuthenticatedConnectListener.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/websocket/AuthenticatedConnectListener.java
@@ -1,5 +1,6 @@
 package org.whispersystems.textsecuregcm.websocket;
 
+import com.codahale.metrics.Counter;
 import com.codahale.metrics.MetricRegistry;
 import com.codahale.metrics.SharedMetricRegistries;
 import com.codahale.metrics.Timer;
@@ -29,6 +30,7 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
   private static final MetricRegistry metricRegistry               = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
   private static final Timer          durationTimer                = metricRegistry.timer(name(WebSocketConnection.class, "connected_duration"                 ));
   private static final Timer          unauthenticatedDurationTimer = metricRegistry.timer(name(WebSocketConnection.class, "unauthenticated_connection_duration"));
+  private static final Counter        openWebsocketCounter         = metricRegistry.counter(name(WebSocketConnection.class, "open_websockets"));
 
   private final PushSender         pushSender;
   private final ReceiptSender      receiptSender;
@@ -64,6 +66,7 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
                                                                   .setContent(ByteString.copyFrom(connectionId.getBytes()))
                                                                   .build();
 
+      openWebsocketCounter.inc();
       RedisOperation.unchecked(() -> apnFallbackManager.cancel(account, device));
       pubSubManager.publish(address, connectMessage);
       pubSubManager.subscribe(address, connection);
@@ -71,6 +74,7 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
       context.addListener(new WebSocketSessionContext.WebSocketEventListener() {
         @Override
         public void onWebSocketClose(WebSocketSessionContext context, int statusCode, String reason) {
+          openWebsocketCounter.dec();
           pubSubManager.unsubscribe(address, connection);
           timer.stop();
         }