Instrument request/response sizes

service/src/main/java/org/whispersystems/textsecuregcm/metrics/MetricsHttpChannelListener.java

@@ -59,6 +59,10 @@ public class MetricsHttpChannelListener implements HttpChannel.Listener, Contain
   public static final String REQUEST_COUNTER_NAME = MetricsRequestEventListener.REQUEST_COUNTER_NAME;
   public static final String REQUESTS_BY_VERSION_COUNTER_NAME = MetricsRequestEventListener.REQUESTS_BY_VERSION_COUNTER_NAME;
   @VisibleForTesting
+  static final String RESPONSE_BYTES_COUNTER_NAME = MetricsRequestEventListener.RESPONSE_BYTES_COUNTER_NAME;
+  @VisibleForTesting
+  static final String REQUEST_BYTES_COUNTER_NAME = MetricsRequestEventListener.REQUEST_BYTES_COUNTER_NAME;
+  @VisibleForTesting
   static final String URI_INFO_PROPERTY_NAME = MetricsHttpChannelListener.class.getName() + ".uriInfo";
 
   @VisibleForTesting
@@ -141,6 +145,9 @@ public class MetricsHttpChannelListener implements HttpChannel.Listener, Contain
 
     meterRegistry.counter(REQUEST_COUNTER_NAME, tags).increment();
 
+    meterRegistry.counter(RESPONSE_BYTES_COUNTER_NAME, tags).increment(request.getResponse().getContentCount());
+    meterRegistry.counter(REQUEST_BYTES_COUNTER_NAME, tags).increment(request.getContentRead());
+
     UserAgentTagUtil.getClientVersionTag(requestInfo.userAgent(), clientReleaseManager).ifPresent(
         clientVersionTag -> meterRegistry.counter(REQUESTS_BY_VERSION_COUNTER_NAME,
             Tags.of(clientVersionTag, UserAgentTagUtil.getPlatformTag(requestInfo.userAgent()))).increment());

service/src/main/java/org/whispersystems/textsecuregcm/metrics/MetricsRequestEventListener.java

@@ -15,8 +15,11 @@ import io.micrometer.core.instrument.Tags;
 import org.glassfish.jersey.server.ContainerResponse;
 import org.glassfish.jersey.server.monitoring.RequestEvent;
 import org.glassfish.jersey.server.monitoring.RequestEventListener;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.storage.ClientReleaseManager;
 import org.whispersystems.textsecuregcm.util.logging.UriInfoUtil;
+import org.whispersystems.websocket.WebSocketResourceProvider;
 
 import javax.annotation.Nullable;
 import java.util.ArrayList;
@@ -24,14 +27,19 @@ import java.util.List;
 import java.util.Optional;
 
 /**
- * Gathers and reports request-level metrics.
+ * Gathers and reports request-level metrics for WebSocket traffic only.
+ * For HTTP traffic, use {@link MetricsHttpChannelListener}.
  */
 public class MetricsRequestEventListener implements RequestEventListener {
 
+  private static final Logger logger = LoggerFactory.getLogger(MetricsRequestEventListener.class);
+
   private final ClientReleaseManager clientReleaseManager;
 
   public static final String REQUEST_COUNTER_NAME = MetricRegistry.name(MetricsRequestEventListener.class, "request");
   public static final String REQUESTS_BY_VERSION_COUNTER_NAME = MetricRegistry.name(MetricsRequestEventListener.class, "requestByVersion");
+  public static final String RESPONSE_BYTES_COUNTER_NAME = MetricRegistry.name(MetricsRequestEventListener.class, "responseBytes");
+  public static final String REQUEST_BYTES_COUNTER_NAME = MetricRegistry.name(MetricsRequestEventListener.class, "requestBytes");
 
   @VisibleForTesting
   static final String PATH_TAG = "path";
@@ -50,6 +58,10 @@ public class MetricsRequestEventListener implements RequestEventListener {
 
   public MetricsRequestEventListener(final TrafficSource trafficSource, final ClientReleaseManager clientReleaseManager) {
     this(trafficSource, Metrics.globalRegistry, clientReleaseManager);
+
+    if (trafficSource == TrafficSource.HTTP) {
+      logger.warn("Use {} for HTTP traffic", MetricsHttpChannelListener.class.getName());
+    }
   }
 
   @VisibleForTesting
@@ -85,6 +97,18 @@ public class MetricsRequestEventListener implements RequestEventListener {
 
         meterRegistry.counter(REQUEST_COUNTER_NAME, tags).increment();
 
+        Optional.ofNullable(event.getContainerRequest().getProperty(WebSocketResourceProvider.REQUEST_LENGTH_PROPERTY))
+            .filter(Integer.class::isInstance)
+            .map(Integer.class::cast)
+            .filter(bytes -> bytes >= 0)
+            .ifPresent(bytes -> meterRegistry.counter(REQUEST_BYTES_COUNTER_NAME, tags).increment(bytes));
+
+        Optional.ofNullable(event.getContainerRequest().getProperty(WebSocketResourceProvider.RESPONSE_LENGTH_PROPERTY))
+            .filter(Integer.class::isInstance)
+            .map(Integer.class::cast)
+            .filter(bytes -> bytes >= 0)
+            .ifPresent(bytes -> meterRegistry.counter(RESPONSE_BYTES_COUNTER_NAME, tags).increment(bytes));
+
         UserAgentTagUtil.getClientVersionTag(userAgent, clientReleaseManager)
             .ifPresent(clientVersionTag -> meterRegistry.counter(REQUESTS_BY_VERSION_COUNTER_NAME,
                     Tags.of(clientVersionTag, UserAgentTagUtil.getPlatformTag(userAgent)))

service/src/test/java/org/whispersystems/textsecuregcm/metrics/MetricsHttpChannelListenerIntegrationTest.java

@@ -46,6 +46,7 @@ import java.security.Principal;
 import java.time.Duration;
 import java.util.EnumSet;
 import java.util.HashSet;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
@@ -83,7 +84,9 @@ class MetricsHttpChannelListenerIntegrationTest {
 
   private static final TrafficSource TRAFFIC_SOURCE = TrafficSource.HTTP;
   private static final MeterRegistry METER_REGISTRY = mock(MeterRegistry.class);
-  private static final Counter COUNTER = mock(Counter.class);
+  private static final Counter REQUEST_COUNTER = mock(Counter.class);
+  private static final Counter RESPONSE_BYTES_COUNTER = mock(Counter.class);
+  private static final Counter REQUEST_BYTES_COUNTER = mock(Counter.class);
   private static final AtomicReference<CountDownLatch> COUNT_DOWN_LATCH_FUTURE_REFERENCE = new AtomicReference<>();
 
   private static final DropwizardAppExtension<Configuration> EXTENSION = new DropwizardAppExtension<>(
@@ -92,7 +95,9 @@ class MetricsHttpChannelListenerIntegrationTest {
   @AfterEach
   void teardown() {
     reset(METER_REGISTRY);
-    reset(COUNTER);
+    reset(REQUEST_COUNTER);
+    reset(RESPONSE_BYTES_COUNTER);
+    reset(REQUEST_BYTES_COUNTER);
   }
 
   @ParameterizedTest
@@ -105,11 +110,13 @@ class MetricsHttpChannelListenerIntegrationTest {
     COUNT_DOWN_LATCH_FUTURE_REFERENCE.set(countDownLatch);
 
     final ArgumentCaptor<Iterable<Tag>> tagCaptor = ArgumentCaptor.forClass(Iterable.class);
+    final Map<String, Counter> counterMap = Map.of(
+        MetricsHttpChannelListener.REQUEST_COUNTER_NAME, REQUEST_COUNTER,
+        MetricsHttpChannelListener.RESPONSE_BYTES_COUNTER_NAME, RESPONSE_BYTES_COUNTER,
+        MetricsHttpChannelListener.REQUEST_BYTES_COUNTER_NAME, REQUEST_BYTES_COUNTER
+    );
     when(METER_REGISTRY.counter(anyString(), any(Iterable.class)))
-        .thenAnswer(a -> MetricsHttpChannelListener.REQUEST_COUNTER_NAME.equals(a.getArgument(0, String.class))
+        .thenAnswer(a -> counterMap.getOrDefault(a.getArgument(0, String.class), mock(Counter.class)));
-            ? COUNTER
-            : mock(Counter.class))
-        .thenReturn(COUNTER);
 
     Client client = EXTENSION.client();
 
@@ -141,7 +148,7 @@ class MetricsHttpChannelListenerIntegrationTest {
     assertTrue(countDownLatch.await(1000, TimeUnit.MILLISECONDS));
 
     verify(METER_REGISTRY).counter(eq(MetricsHttpChannelListener.REQUEST_COUNTER_NAME), tagCaptor.capture());
-    verify(COUNTER).increment();
+    verify(REQUEST_COUNTER).increment();
 
     final Iterable<Tag> tagIterable = tagCaptor.getValue();
     final Set<Tag> tags = new HashSet<>();
@@ -186,11 +193,13 @@ class MetricsHttpChannelListenerIntegrationTest {
       COUNT_DOWN_LATCH_FUTURE_REFERENCE.set(countDownLatch);
 
       final ArgumentCaptor<Iterable<Tag>> tagCaptor = ArgumentCaptor.forClass(Iterable.class);
+      final Map<String, Counter> counterMap = Map.of(
+          MetricsHttpChannelListener.REQUEST_COUNTER_NAME, REQUEST_COUNTER,
+          MetricsHttpChannelListener.RESPONSE_BYTES_COUNTER_NAME, RESPONSE_BYTES_COUNTER,
+          MetricsHttpChannelListener.REQUEST_BYTES_COUNTER_NAME, REQUEST_BYTES_COUNTER
+      );
       when(METER_REGISTRY.counter(anyString(), any(Iterable.class)))
-          .thenAnswer(a -> MetricsHttpChannelListener.REQUEST_COUNTER_NAME.equals(a.getArgument(0, String.class))
+          .thenAnswer(a -> counterMap.getOrDefault(a.getArgument(0, String.class), mock(Counter.class)));
-              ? COUNTER
-              : mock(Counter.class))
-          .thenReturn(COUNTER);
 
       client.connect(new WebSocketListener() {
                        @Override
@@ -203,7 +212,7 @@ class MetricsHttpChannelListenerIntegrationTest {
       assertTrue(countDownLatch.await(1000, TimeUnit.MILLISECONDS));
 
       verify(METER_REGISTRY).counter(eq(MetricsHttpChannelListener.REQUEST_COUNTER_NAME), tagCaptor.capture());
-      verify(COUNTER).increment();
+      verify(REQUEST_COUNTER).increment();
 
       final Iterable<Tag> tagIterable = tagCaptor.getValue();
       final Set<Tag> tags = new HashSet<>();

service/src/test/java/org/whispersystems/textsecuregcm/metrics/MetricsHttpChannelListenerTest.java

@@ -39,6 +39,8 @@ class MetricsHttpChannelListenerTest {
   private MeterRegistry meterRegistry;
   private Counter requestCounter;
   private Counter requestsByVersionCounter;
+  private Counter responseBytesCounter;
+  private Counter requestBytesCounter;
   private ClientReleaseManager clientReleaseManager;
   private MetricsHttpChannelListener listener;
 
@@ -47,6 +49,8 @@ class MetricsHttpChannelListenerTest {
     meterRegistry = mock(MeterRegistry.class);
     requestCounter = mock(Counter.class);
     requestsByVersionCounter = mock(Counter.class);
+    responseBytesCounter = mock(Counter.class);
+    requestBytesCounter = mock(Counter.class);
 
     when(meterRegistry.counter(eq(MetricsHttpChannelListener.REQUEST_COUNTER_NAME), any(Iterable.class)))
         .thenReturn(requestCounter);
@@ -54,6 +58,12 @@ class MetricsHttpChannelListenerTest {
     when(meterRegistry.counter(eq(MetricsHttpChannelListener.REQUESTS_BY_VERSION_COUNTER_NAME), any(Iterable.class)))
         .thenReturn(requestsByVersionCounter);
 
+    when(meterRegistry.counter(eq(MetricsHttpChannelListener.RESPONSE_BYTES_COUNTER_NAME), any(Iterable.class)))
+        .thenReturn(responseBytesCounter);
+
+    when(meterRegistry.counter(eq(MetricsHttpChannelListener.REQUEST_BYTES_COUNTER_NAME), any(Iterable.class)))
+        .thenReturn(requestBytesCounter);
+
     clientReleaseManager = mock(ClientReleaseManager.class);
 
     listener = new MetricsHttpChannelListener(meterRegistry, clientReleaseManager, Collections.emptySet());
@@ -76,7 +86,9 @@ class MetricsHttpChannelListenerTest {
 
     final Response response = mock(Response.class);
     when(response.getStatus()).thenReturn(statusCode);
+    when(response.getContentCount()).thenReturn(1024L);
     when(request.getResponse()).thenReturn(response);
+    when(request.getContentRead()).thenReturn(512L);
     final ExtendedUriInfo extendedUriInfo = mock(ExtendedUriInfo.class);
     when(request.getAttribute(MetricsHttpChannelListener.URI_INFO_PROPERTY_NAME)).thenReturn(extendedUriInfo);
     when(extendedUriInfo.getMatchedTemplates()).thenReturn(List.of(new UriTemplate(path)));
@@ -87,6 +99,9 @@ class MetricsHttpChannelListenerTest {
 
     verify(requestCounter).increment();
 
+    verify(responseBytesCounter).increment(1024L);
+    verify(requestBytesCounter).increment(512L);
+
     verify(meterRegistry).counter(eq(MetricsHttpChannelListener.REQUEST_COUNTER_NAME), tagCaptor.capture());
 
     final Set<Tag> tags = new HashSet<>();
@@ -123,7 +138,9 @@ class MetricsHttpChannelListenerTest {
 
     final Response response = mock(Response.class);
     when(response.getStatus()).thenReturn(statusCode);
+    when(response.getContentCount()).thenReturn(1024L);
     when(request.getResponse()).thenReturn(response);
+    when(request.getContentRead()).thenReturn(512L);
     final ExtendedUriInfo extendedUriInfo = mock(ExtendedUriInfo.class);
     when(request.getAttribute(MetricsHttpChannelListener.URI_INFO_PROPERTY_NAME)).thenReturn(extendedUriInfo);
     when(extendedUriInfo.getMatchedTemplates()).thenReturn(List.of(new UriTemplate(path)));

service/src/test/java/org/whispersystems/textsecuregcm/metrics/MetricsRequestEventListenerTest.java

@@ -62,6 +62,8 @@ class MetricsRequestEventListenerTest {
 
   private MeterRegistry meterRegistry;
   private Counter counter;
+  private Counter responseBytesCounter;
+  private Counter requestBytesCounter;
   private MetricsRequestEventListener listener;
 
   private static final TrafficSource TRAFFIC_SOURCE = TrafficSource.HTTP;
@@ -70,6 +72,8 @@ class MetricsRequestEventListenerTest {
   void setup() {
     meterRegistry = mock(MeterRegistry.class);
     counter = mock(Counter.class);
+    responseBytesCounter = mock(Counter.class);
+    requestBytesCounter = mock(Counter.class);
 
     final ClientReleaseManager clientReleaseManager = mock(ClientReleaseManager.class);
     when(clientReleaseManager.isVersionActive(any(), any())).thenReturn(false);
@@ -91,6 +95,8 @@ class MetricsRequestEventListenerTest {
     when(request.getMethod()).thenReturn(method);
     when(request.getRequestHeader(HttpHeaders.USER_AGENT)).thenReturn(
         Collections.singletonList("Signal-Android/7.6.2 Android/34 libsignal/0.46.0"));
+    when(request.getProperty(WebSocketResourceProvider.REQUEST_LENGTH_PROPERTY)).thenReturn(512);
+    when(request.getProperty(WebSocketResourceProvider.RESPONSE_LENGTH_PROPERTY)).thenReturn(1024);
 
     final ContainerResponse response = mock(ContainerResponse.class);
     when(response.getStatus()).thenReturn(statusCode);
@@ -104,10 +110,17 @@ class MetricsRequestEventListenerTest {
     final ArgumentCaptor<Iterable<Tag>> tagCaptor = ArgumentCaptor.forClass(Iterable.class);
     when(meterRegistry.counter(eq(MetricsRequestEventListener.REQUEST_COUNTER_NAME), any(Iterable.class)))
         .thenReturn(counter);
+    when(meterRegistry.counter(eq(MetricsRequestEventListener.RESPONSE_BYTES_COUNTER_NAME), any(Iterable.class)))
+        .thenReturn(responseBytesCounter);
+    when(meterRegistry.counter(eq(MetricsRequestEventListener.REQUEST_BYTES_COUNTER_NAME), any(Iterable.class)))
+        .thenReturn(requestBytesCounter);
 
     listener.onEvent(event);
 
     verify(meterRegistry).counter(eq(MetricsRequestEventListener.REQUEST_COUNTER_NAME), tagCaptor.capture());
+    verify(counter).increment();
+    verify(responseBytesCounter).increment(1024L);
+    verify(requestBytesCounter).increment(512L);
 
     final Iterable<Tag> tagIterable = tagCaptor.getValue();
     final Set<Tag> tags = new HashSet<>();
@@ -155,6 +168,10 @@ class MetricsRequestEventListenerTest {
     final ArgumentCaptor<Iterable<Tag>> tagCaptor = ArgumentCaptor.forClass(Iterable.class);
     when(meterRegistry.counter(eq(MetricsRequestEventListener.REQUEST_COUNTER_NAME), any(Iterable.class)))
         .thenReturn(counter);
+    when(meterRegistry.counter(eq(MetricsRequestEventListener.RESPONSE_BYTES_COUNTER_NAME), any(Iterable.class)))
+        .thenReturn(responseBytesCounter);
+    when(meterRegistry.counter(eq(MetricsRequestEventListener.REQUEST_BYTES_COUNTER_NAME), any(Iterable.class)))
+        .thenReturn(requestBytesCounter);
 
     provider.onWebSocketConnect(session);
 
@@ -216,6 +233,10 @@ class MetricsRequestEventListenerTest {
     final ArgumentCaptor<Iterable<Tag>> tagCaptor = ArgumentCaptor.forClass(Iterable.class);
     when(meterRegistry.counter(eq(MetricsRequestEventListener.REQUEST_COUNTER_NAME), any(Iterable.class))).thenReturn(
         counter);
+    when(meterRegistry.counter(eq(MetricsRequestEventListener.RESPONSE_BYTES_COUNTER_NAME), any(Iterable.class)))
+        .thenReturn(responseBytesCounter);
+    when(meterRegistry.counter(eq(MetricsRequestEventListener.REQUEST_BYTES_COUNTER_NAME), any(Iterable.class)))
+        .thenReturn(requestBytesCounter);
 
     provider.onWebSocketConnect(session);
 

websocket-resources/src/main/java/org/whispersystems/websocket/WebSocketResourceProvider.java

@@ -173,18 +173,32 @@ public class WebSocketResourceProvider<T extends Principal> implements WebSocket
    * {@link org.whispersystems.websocket.ReusableAuth.MutableRef} for us to close when the request is finished
    */
   public static final String RESOLVED_PRINCIPAL_PROPERTY = WebSocketResourceProvider.class.getName() + ".resolvedPrincipal";
+
+  /**
+   * The property name where request byte count is stored for metrics collection
+   */
+  public static final String REQUEST_LENGTH_PROPERTY = WebSocketResourceProvider.class.getName() + ".requestBytes";
+
+  /**
+   * The property name where response byte count is stored for metrics collection
+   */
+  public static final String RESPONSE_LENGTH_PROPERTY = WebSocketResourceProvider.class.getName() + ".responseBytes";
+
   private void handleRequest(WebSocketRequestMessage requestMessage) {
     ContainerRequest containerRequest = new ContainerRequest(null, URI.create(requestMessage.getPath()),
         requestMessage.getVerb(), new WebSocketSecurityContext(new ContextPrincipal(context)),
         new MapPropertiesDelegate(new HashMap<>()), jerseyHandler.getConfiguration());
     containerRequest.headers(getCombinedHeaders(session.getUpgradeRequest().getHeaders(), requestMessage.getHeaders()));
 
+    final int requestBytes = requestMessage.getBody().map(body -> body.length).orElse(0);
+
     if (requestMessage.getBody().isPresent()) {
       containerRequest.setEntityStream(new ByteArrayInputStream(requestMessage.getBody().get()));
     }
 
     containerRequest.setProperty(remoteAddressPropertyName, remoteAddress);
     containerRequest.setProperty(REUSABLE_AUTH_PROPERTY, reusableAuth);
+    containerRequest.setProperty(REQUEST_LENGTH_PROPERTY, requestBytes);
 
     ByteArrayOutputStream responseBody = new ByteArrayOutputStream();
     CompletableFuture<ContainerResponse> responseFuture = (CompletableFuture<ContainerResponse>) jerseyHandler.apply(
@@ -203,6 +217,8 @@ public class WebSocketResourceProvider<T extends Principal> implements WebSocket
         })
         .thenAccept(response -> {
           try {
+            final int responseBytes = responseBody.size();
+            containerRequest.setProperty(RESPONSE_LENGTH_PROPERTY, responseBytes);
             sendResponse(requestMessage, response, responseBody);
           } catch (IOException e) {
             throw new RuntimeException(e);
@@ -213,6 +229,7 @@ public class WebSocketResourceProvider<T extends Principal> implements WebSocket
           logger.warn("Websocket Error: " + requestMessage.getVerb() + " " + requestMessage.getPath() + "\n"
               + requestMessage.getBody(), exception);
           try {
+            containerRequest.setProperty(RESPONSE_LENGTH_PROPERTY, 0);
             sendErrorResponse(requestMessage, Response.status(500).build());
           } catch (IOException e) {
             logger.warn("Failed to send error response", e);