Fix flaky MessageMetricsTest

Make the MeterRegistry in MessageMetrics configurable

service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java

@@ -161,6 +161,7 @@ import org.whispersystems.textsecuregcm.mappers.RateLimitExceededExceptionMapper
 import org.whispersystems.textsecuregcm.mappers.RegistrationServiceSenderExceptionMapper;
 import org.whispersystems.textsecuregcm.mappers.ServerRejectedExceptionMapper;
 import org.whispersystems.textsecuregcm.mappers.SubscriptionProcessorExceptionMapper;
+import org.whispersystems.textsecuregcm.metrics.MessageMetrics;
 import org.whispersystems.textsecuregcm.metrics.MetricsApplicationEventListener;
 import org.whispersystems.textsecuregcm.metrics.MetricsHttpChannelListener;
 import org.whispersystems.textsecuregcm.metrics.MetricsUtil;
@@ -858,6 +859,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     final MetricsHttpChannelListener metricsHttpChannelListener = new MetricsHttpChannelListener(clientReleaseManager,
         Set.of(websocketServletPath, provisioningWebsocketServletPath, "/health-check"));
     metricsHttpChannelListener.configure(environment);
+    final MessageMetrics messageMetrics = new MessageMetrics();
 
     environment.jersey().register(new BufferingInterceptor());
     environment.jersey().register(new VirtualExecutorServiceProvider("managed-async-virtual-thread-"));
@@ -874,7 +876,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     webSocketEnvironment.jersey().register(new VirtualExecutorServiceProvider("managed-async-websocket-virtual-thread-"));
     webSocketEnvironment.setAuthenticator(new WebSocketAccountAuthenticator(accountAuthenticator, new AccountPrincipalSupplier(accountsManager)));
     webSocketEnvironment.setConnectListener(
-        new AuthenticatedConnectListener(receiptSender, messagesManager, pushNotificationManager,
+        new AuthenticatedConnectListener(receiptSender, messagesManager, messageMetrics, pushNotificationManager,
             clientPresenceManager, websocketScheduledExecutor, messageDeliveryScheduler, clientReleaseManager));
     webSocketEnvironment.jersey()
         .register(new WebsocketRefreshApplicationEventListener(accountsManager, clientPresenceManager));
@@ -939,6 +941,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
       log.info("Registered spam filter: {}", filter.getClass().getName());
     });
 
+
     final List<Object> commonControllers = Lists.newArrayList(
         new AccountController(accountsManager, rateLimiters, turnTokenGenerator, registrationRecoveryPasswordsManager,
             usernameHashZkProofVerifier),
@@ -968,7 +971,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         new MessageController(rateLimiters, messageByteLimitCardinalityEstimator, messageSender, receiptSender,
             accountsManager, messagesManager, pushNotificationManager, reportMessageManager,
             multiRecipientMessageExecutor, messageDeliveryScheduler, reportSpamTokenProvider, clientReleaseManager,
-            dynamicConfigurationManager, zkSecretParams, spamChecker, Clock.systemUTC()),
+            dynamicConfigurationManager, zkSecretParams, spamChecker, messageMetrics, Clock.systemUTC()),
         new PaymentsController(currencyManager, paymentsCredentialsGenerator),
         new ProfileController(clock, rateLimiters, accountsManager, profilesManager, dynamicConfigurationManager,
             profileBadgeConverter, config.getBadges(), cdnS3Client, profileCdnPolicyGenerator, profileCdnPolicySigner,

service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java

@@ -167,6 +167,7 @@ public class MessageController {
   private final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager;
   private final ServerSecretParams serverSecretParams;
   private final SpamChecker spamChecker;
+  private final MessageMetrics messageMetrics;
   private final Clock clock;
 
   private static final int MAX_FETCH_ACCOUNT_CONCURRENCY = 8;
@@ -216,6 +217,7 @@ public class MessageController {
       final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager,
       final ServerSecretParams serverSecretParams,
       final SpamChecker spamChecker,
+      final MessageMetrics messageMetrics,
       final Clock clock) {
     this.rateLimiters = rateLimiters;
     this.messageByteLimitEstimator = messageByteLimitEstimator;
@@ -232,6 +234,7 @@ public class MessageController {
     this.dynamicConfigurationManager = dynamicConfigurationManager;
     this.serverSecretParams = serverSecretParams;
     this.spamChecker = spamChecker;
+    this.messageMetrics = messageMetrics;
     this.clock = clock;
   }
 
@@ -732,8 +735,8 @@ public class MessageController {
           final OutgoingMessageEntityList messages = new OutgoingMessageEntityList(envelopes
               .map(OutgoingMessageEntity::fromEnvelope)
               .peek(outgoingMessageEntity -> {
-                MessageMetrics.measureAccountOutgoingMessageUuidMismatches(auth.getAccount(), outgoingMessageEntity);
-                MessageMetrics.measureOutgoingMessageLatency(outgoingMessageEntity.serverTimestamp(), "rest", userAgent, clientReleaseManager);
+                messageMetrics.measureAccountOutgoingMessageUuidMismatches(auth.getAccount(), outgoingMessageEntity);
+                messageMetrics.measureOutgoingMessageLatency(outgoingMessageEntity.serverTimestamp(), "rest", userAgent, clientReleaseManager);
               })
               .collect(Collectors.toList()),
               messagesAndHasMore.second());

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

@@ -7,6 +7,8 @@ package org.whispersystems.textsecuregcm.metrics;
 
 import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
 
+import com.google.common.annotations.VisibleForTesting;
+import io.micrometer.core.instrument.MeterRegistry;
 import io.micrometer.core.instrument.Metrics;
 import io.micrometer.core.instrument.Tag;
 import io.micrometer.core.instrument.Timer;
@@ -30,13 +32,23 @@ public final class MessageMetrics {
       "mismatchedAccountEnvelopeUuid");
 
   public static final String DELIVERY_LATENCY_TIMER_NAME = name(MessageMetrics.class, "deliveryLatency");
+  private final MeterRegistry metricRegistry;
 
-  public static void measureAccountOutgoingMessageUuidMismatches(final Account account,
+  @VisibleForTesting
+  MessageMetrics(final MeterRegistry metricRegistry) {
+    this.metricRegistry = metricRegistry;
+  }
+
+  public MessageMetrics() {
+    this(Metrics.globalRegistry);
+  }
+
+  public void measureAccountOutgoingMessageUuidMismatches(final Account account,
       final OutgoingMessageEntity outgoingMessage) {
     measureAccountDestinationUuidMismatches(account, outgoingMessage.destinationUuid());
   }
 
-  public static void measureAccountEnvelopeUuidMismatches(final Account account,
+  public void measureAccountEnvelopeUuidMismatches(final Account account,
       final MessageProtos.Envelope envelope) {
     if (envelope.hasDestinationUuid()) {
       try {
@@ -47,16 +59,16 @@ public final class MessageMetrics {
     }
   }
 
-  private static void measureAccountDestinationUuidMismatches(final Account account, final ServiceIdentifier destinationIdentifier) {
+  private void measureAccountDestinationUuidMismatches(final Account account, final ServiceIdentifier destinationIdentifier) {
     if (!account.isIdentifiedBy(destinationIdentifier)) {
       // In all cases, this represents a mismatch between the account’s current PNI and its PNI when the message was
       // sent. This is an expected case, but if this metric changes significantly, it could indicate an issue to
       // investigate.
-      Metrics.counter(MISMATCHED_ACCOUNT_ENVELOPE_UUID_COUNTER_NAME).increment();
+      metricRegistry.counter(MISMATCHED_ACCOUNT_ENVELOPE_UUID_COUNTER_NAME).increment();
     }
   }
 
-  public static void measureOutgoingMessageLatency(final long serverTimestamp,
+  public void measureOutgoingMessageLatency(final long serverTimestamp,
       final String channel,
       final String userAgent,
       final ClientReleaseManager clientReleaseManager) {
@@ -70,7 +82,7 @@ public final class MessageMetrics {
     Timer.builder(DELIVERY_LATENCY_TIMER_NAME)
         .publishPercentileHistogram(true)
         .tags(tags)
-        .register(Metrics.globalRegistry)
+        .register(metricRegistry)
         .record(Duration.between(Instant.ofEpochMilli(serverTimestamp), Instant.now()));
   }
 }

service/src/main/java/org/whispersystems/textsecuregcm/websocket/AuthenticatedConnectListener.java

@@ -20,6 +20,7 @@ import java.util.concurrent.atomic.AtomicReference;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
+import org.whispersystems.textsecuregcm.metrics.MessageMetrics;
 import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
 import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
 import org.whispersystems.textsecuregcm.push.NotPushRegisteredException;
@@ -51,6 +52,7 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
 
   private final ReceiptSender receiptSender;
   private final MessagesManager messagesManager;
+  private final MessageMetrics messageMetrics;
   private final PushNotificationManager pushNotificationManager;
   private final ClientPresenceManager clientPresenceManager;
   private final ScheduledExecutorService scheduledExecutorService;
@@ -69,6 +71,7 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
 
   public AuthenticatedConnectListener(ReceiptSender receiptSender,
       MessagesManager messagesManager,
+      MessageMetrics messageMetrics,
       PushNotificationManager pushNotificationManager,
       ClientPresenceManager clientPresenceManager,
       ScheduledExecutorService scheduledExecutorService,
@@ -76,6 +79,7 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
       ClientReleaseManager clientReleaseManager) {
     this.receiptSender = receiptSender;
     this.messagesManager = messagesManager;
+    this.messageMetrics = messageMetrics;
     this.pushNotificationManager = pushNotificationManager;
     this.clientPresenceManager = clientPresenceManager;
     this.scheduledExecutorService = scheduledExecutorService;
@@ -138,7 +142,7 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
       final Device device = auth.getAuthenticatedDevice();
       final Timer.Sample sample = Timer.start();
       final WebSocketConnection connection = new WebSocketConnection(receiptSender,
-          messagesManager, auth, device,
+          messagesManager, messageMetrics, auth, device,
           context.getClient(),
           scheduledExecutorService,
           messageDeliveryScheduler,

service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java

@@ -109,6 +109,7 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
 
   private final ReceiptSender receiptSender;
   private final MessagesManager messagesManager;
+  private final MessageMetrics messageMetrics;
 
   private final AuthenticatedAccount auth;
   private final Device device;
@@ -141,6 +142,7 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
 
   public WebSocketConnection(ReceiptSender receiptSender,
       MessagesManager messagesManager,
+      MessageMetrics messageMetrics,
       AuthenticatedAccount auth,
       Device device,
       WebSocketClient client,
@@ -150,6 +152,7 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
 
     this(receiptSender,
         messagesManager,
+        messageMetrics,
         auth,
         device,
         client,
@@ -162,6 +165,7 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
   @VisibleForTesting
   WebSocketConnection(ReceiptSender receiptSender,
       MessagesManager messagesManager,
+      MessageMetrics messageMetrics,
       AuthenticatedAccount auth,
       Device device,
       WebSocketClient client,
@@ -172,6 +176,7 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
 
     this.receiptSender = receiptSender;
     this.messagesManager = messagesManager;
+    this.messageMetrics = messageMetrics;
     this.auth = auth;
     this.device = device;
     this.client = client;
@@ -208,7 +213,7 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
     sendMessageCounter.increment();
     sentMessageCounter.increment();
     bytesSentCounter.increment(body.map(bytes -> bytes.length).orElse(0));
-    MessageMetrics.measureAccountEnvelopeUuidMismatches(auth.getAccount(), message);
+    messageMetrics.measureAccountEnvelopeUuidMismatches(auth.getAccount(), message);
 
     // X-Signal-Key: false must be sent until Android stops assuming it missing means true
     return client.sendRequest("PUT", "/api/v1/message",
@@ -217,7 +222,7 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
           if (throwable != null) {
             sendFailuresCounter.increment();
           } else {
-            MessageMetrics.measureOutgoingMessageLatency(message.getServerTimestamp(), "websocket", client.getUserAgent(), clientReleaseManager);
+            messageMetrics.measureOutgoingMessageLatency(message.getServerTimestamp(), "websocket", client.getUserAgent(), clientReleaseManager);
           }
         }).thenCompose(response -> {
           final CompletableFuture<Void> result;