Fix flaky MessageMetricsTest

Make the MeterRegistry in MessageMetrics configurable

service/src/test/java/org/whispersystems/textsecuregcm/controllers/MessageControllerTest.java

@@ -119,6 +119,7 @@ import org.whispersystems.textsecuregcm.limits.CardinalityEstimator;
 import org.whispersystems.textsecuregcm.limits.RateLimiter;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.mappers.RateLimitExceededExceptionMapper;
+import org.whispersystems.textsecuregcm.metrics.MessageMetrics;
 import org.whispersystems.textsecuregcm.providers.MultiRecipientMessageProvider;
 import org.whispersystems.textsecuregcm.push.MessageSender;
 import org.whispersystems.textsecuregcm.push.NotPushRegisteredException;
@@ -212,7 +213,7 @@ class MessageControllerTest {
           new MessageController(rateLimiters, cardinalityEstimator, messageSender, receiptSender, accountsManager,
               messagesManager, pushNotificationManager, reportMessageManager, multiRecipientMessageExecutor,
               messageDeliveryScheduler, ReportSpamTokenProvider.noop(), mock(ClientReleaseManager.class), dynamicConfigurationManager,
-              serverSecretParams, SpamChecker.noop(), clock))
+              serverSecretParams, SpamChecker.noop(), new MessageMetrics(), clock))
       .build();
 
   @BeforeEach

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

@@ -13,11 +13,9 @@ import static org.mockito.Mockito.when;
 
 import io.micrometer.core.instrument.Counter;
 import io.micrometer.core.instrument.Meter;
-import io.micrometer.core.instrument.Metrics;
 import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
 import java.util.Optional;
 import java.util.UUID;
-import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.whispersystems.textsecuregcm.entities.MessageProtos;
@@ -33,6 +31,7 @@ class MessageMetricsTest {
   private final UUID aci = UUID.fromString("11111111-1111-1111-1111-111111111111");
   private final UUID pni = UUID.fromString("22222222-2222-2222-2222-222222222222");
   private final UUID otherUuid = UUID.fromString("99999999-9999-9999-9999-999999999999");
+  private MessageMetrics messageMetrics;
   private SimpleMeterRegistry simpleMeterRegistry;
 
   @BeforeEach
@@ -42,35 +41,28 @@ class MessageMetricsTest {
     when(account.isIdentifiedBy(any())).thenReturn(false);
     when(account.isIdentifiedBy(new AciServiceIdentifier(aci))).thenReturn(true);
     when(account.isIdentifiedBy(new PniServiceIdentifier(pni))).thenReturn(true);
-    Metrics.globalRegistry.clear();
     simpleMeterRegistry = new SimpleMeterRegistry();
-    Metrics.globalRegistry.add(simpleMeterRegistry);
-  }
-
-  @AfterEach
-  void teardown() {
-    Metrics.globalRegistry.remove(simpleMeterRegistry);
-    Metrics.globalRegistry.clear();
+    messageMetrics = new MessageMetrics(simpleMeterRegistry);
   }
 
   @Test
   void measureAccountOutgoingMessageUuidMismatches() {
 
     final OutgoingMessageEntity outgoingMessageToAci = createOutgoingMessageEntity(new AciServiceIdentifier(aci));
-    MessageMetrics.measureAccountOutgoingMessageUuidMismatches(account, outgoingMessageToAci);
+    messageMetrics.measureAccountOutgoingMessageUuidMismatches(account, outgoingMessageToAci);
 
     Optional<Counter> counter = findCounter(simpleMeterRegistry);
 
     assertTrue(counter.isEmpty());
 
     final OutgoingMessageEntity outgoingMessageToPni = createOutgoingMessageEntity(new PniServiceIdentifier(pni));
-    MessageMetrics.measureAccountOutgoingMessageUuidMismatches(account, outgoingMessageToPni);
+    messageMetrics.measureAccountOutgoingMessageUuidMismatches(account, outgoingMessageToPni);
     counter = findCounter(simpleMeterRegistry);
 
     assertTrue(counter.isEmpty());
 
     final OutgoingMessageEntity outgoingMessageToOtherUuid = createOutgoingMessageEntity(new AciServiceIdentifier(otherUuid));
-    MessageMetrics.measureAccountOutgoingMessageUuidMismatches(account, outgoingMessageToOtherUuid);
+    messageMetrics.measureAccountOutgoingMessageUuidMismatches(account, outgoingMessageToOtherUuid);
     counter = findCounter(simpleMeterRegistry);
 
     assertEquals(1.0, counter.map(Counter::count).orElse(0.0));
@@ -83,26 +75,26 @@ class MessageMetricsTest {
   @Test
   void measureAccountEnvelopeUuidMismatches() {
     final MessageProtos.Envelope envelopeToAci = createEnvelope(new AciServiceIdentifier(aci));
-    MessageMetrics.measureAccountEnvelopeUuidMismatches(account, envelopeToAci);
+    messageMetrics.measureAccountEnvelopeUuidMismatches(account, envelopeToAci);
 
     Optional<Counter> counter = findCounter(simpleMeterRegistry);
 
     assertTrue(counter.isEmpty());
 
     final MessageProtos.Envelope envelopeToPni = createEnvelope(new PniServiceIdentifier(pni));
-    MessageMetrics.measureAccountEnvelopeUuidMismatches(account, envelopeToPni);
+    messageMetrics.measureAccountEnvelopeUuidMismatches(account, envelopeToPni);
     counter = findCounter(simpleMeterRegistry);
 
     assertTrue(counter.isEmpty());
 
     final MessageProtos.Envelope envelopeToOtherUuid = createEnvelope(new AciServiceIdentifier(otherUuid));
-    MessageMetrics.measureAccountEnvelopeUuidMismatches(account, envelopeToOtherUuid);
+    messageMetrics.measureAccountEnvelopeUuidMismatches(account, envelopeToOtherUuid);
     counter = findCounter(simpleMeterRegistry);
 
     assertEquals(1.0, counter.map(Counter::count).orElse(0.0));
 
     final MessageProtos.Envelope envelopeToNull = createEnvelope(null);
-    MessageMetrics.measureAccountEnvelopeUuidMismatches(account, envelopeToNull);
+    messageMetrics.measureAccountEnvelopeUuidMismatches(account, envelopeToNull);
     counter = findCounter(simpleMeterRegistry);
 
     assertEquals(1.0, counter.map(Counter::count).orElse(0.0));

service/src/test/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnectionIntegrationTest.java

@@ -46,6 +46,7 @@ import org.mockito.stubbing.Answer;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
 import org.whispersystems.textsecuregcm.entities.MessageProtos;
 import org.whispersystems.textsecuregcm.entities.MessageProtos.Envelope;
+import org.whispersystems.textsecuregcm.metrics.MessageMetrics;
 import org.whispersystems.textsecuregcm.push.ReceiptSender;
 import org.whispersystems.textsecuregcm.redis.RedisClusterExtension;
 import org.whispersystems.textsecuregcm.storage.Account;
@@ -124,6 +125,7 @@ class WebSocketConnectionIntegrationTest {
     final WebSocketConnection webSocketConnection = new WebSocketConnection(
         mock(ReceiptSender.class),
         new MessagesManager(messagesDynamoDb, messagesCache, reportMessageManager, sharedExecutorService),
+        new MessageMetrics(),
         new AuthenticatedAccount(account, device),
         device,
         webSocketClient,
@@ -209,6 +211,7 @@ class WebSocketConnectionIntegrationTest {
     final WebSocketConnection webSocketConnection = new WebSocketConnection(
         mock(ReceiptSender.class),
         new MessagesManager(messagesDynamoDb, messagesCache, reportMessageManager, sharedExecutorService),
+        new MessageMetrics(),
         new AuthenticatedAccount(account, device),
         device,
         webSocketClient,
@@ -275,6 +278,7 @@ class WebSocketConnectionIntegrationTest {
     final WebSocketConnection webSocketConnection = new WebSocketConnection(
         mock(ReceiptSender.class),
         new MessagesManager(messagesDynamoDb, messagesCache, reportMessageManager, sharedExecutorService),
+        new MessageMetrics(),
         new AuthenticatedAccount(account, device),
         device,
         webSocketClient,

service/src/test/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnectionTest.java

@@ -49,6 +49,7 @@ import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
 import org.eclipse.jetty.websocket.api.UpgradeRequest;
+import org.jetbrains.annotations.NotNull;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -56,6 +57,7 @@ import org.mockito.stubbing.Answer;
 import org.whispersystems.textsecuregcm.auth.AccountAuthenticator;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
 import org.whispersystems.textsecuregcm.identity.AciServiceIdentifier;
+import org.whispersystems.textsecuregcm.metrics.MessageMetrics;
 import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
 import org.whispersystems.textsecuregcm.push.PushNotificationManager;
 import org.whispersystems.textsecuregcm.push.ReceiptSender;
@@ -121,7 +123,7 @@ class WebSocketConnectionTest {
     WebSocketAccountAuthenticator webSocketAuthenticator =
         new WebSocketAccountAuthenticator(accountAuthenticator, mock(PrincipalSupplier.class));
     AuthenticatedConnectListener connectListener = new AuthenticatedConnectListener(receiptSender, messagesManager,
-        mock(PushNotificationManager.class), mock(ClientPresenceManager.class),
+        new MessageMetrics(), mock(PushNotificationManager.class), mock(ClientPresenceManager.class),
         retrySchedulingExecutor, messageDeliveryScheduler, clientReleaseManager);
     WebSocketSessionContext sessionContext = mock(WebSocketSessionContext.class);
 
@@ -201,8 +203,7 @@ class WebSocketConnectionTest {
           return future;
         });
 
-    WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager,
-        auth, device, client, retrySchedulingExecutor, Schedulers.immediate(), clientReleaseManager);
+    WebSocketConnection connection = webSocketConnection(client);
 
     connection.start();
     verify(client, times(3)).sendRequest(eq("PUT"), eq("/api/v1/message"), any(List.class),
@@ -229,8 +230,7 @@ class WebSocketConnectionTest {
   @Test
   public void testOnlineSend() {
     final WebSocketClient client = mock(WebSocketClient.class);
-    final WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor, Schedulers.immediate(), clientReleaseManager);
+    final WebSocketConnection connection = webSocketConnection(client);
 
     final UUID accountUuid = UUID.randomUUID();
 
@@ -349,8 +349,7 @@ class WebSocketConnectionTest {
           return future;
         });
 
-    WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager,
-        auth, device, client, retrySchedulingExecutor, Schedulers.immediate(), clientReleaseManager);
+    WebSocketConnection connection = webSocketConnection(client);
 
     connection.start();
 
@@ -373,8 +372,7 @@ class WebSocketConnectionTest {
   @Test
   void testProcessStoredMessageConcurrency() {
     final WebSocketClient client = mock(WebSocketClient.class);
-    final WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor, Schedulers.immediate(), clientReleaseManager);
+    final WebSocketConnection connection = webSocketConnection(client);
 
     when(account.getNumber()).thenReturn("+18005551234");
     when(account.getUuid()).thenReturn(UUID.randomUUID());
@@ -438,8 +436,7 @@ class WebSocketConnectionTest {
   @Test
   void testProcessStoredMessagesMultiplePages() {
     final WebSocketClient client = mock(WebSocketClient.class);
-    final WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor, Schedulers.immediate(), clientReleaseManager);
+    final WebSocketConnection connection = webSocketConnection(client);
 
     when(account.getNumber()).thenReturn("+18005551234");
     final UUID accountUuid = UUID.randomUUID();
@@ -487,8 +484,7 @@ class WebSocketConnectionTest {
   @Test
   void testProcessStoredMessagesMultiplePagesBackpressure() {
     final WebSocketClient client = mock(WebSocketClient.class);
-    final WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor, Schedulers.immediate(), clientReleaseManager);
+    final WebSocketConnection connection = webSocketConnection(client);
 
     when(account.getNumber()).thenReturn("+18005551234");
     final UUID accountUuid = UUID.randomUUID();
@@ -571,8 +567,7 @@ class WebSocketConnectionTest {
   @Test
   void testProcessStoredMessagesContainsSenderUuid() {
     final WebSocketClient client = mock(WebSocketClient.class);
-    final WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor, Schedulers.immediate(), clientReleaseManager);
+    final WebSocketConnection connection = webSocketConnection(client);
 
     when(account.getNumber()).thenReturn("+18005551234");
     final UUID accountUuid = UUID.randomUUID();
@@ -630,11 +625,15 @@ class WebSocketConnectionTest {
     verify(client).sendRequest(eq("PUT"), eq("/api/v1/queue/empty"), any(List.class), eq(Optional.empty()));
   }
 
+  private @NotNull WebSocketConnection webSocketConnection(final WebSocketClient client) {
+    return new WebSocketConnection(receiptSender, messagesManager, new MessageMetrics(), auth, device, client,
+        retrySchedulingExecutor, Schedulers.immediate(), clientReleaseManager);
+  }
+
   @Test
   void testProcessStoredMessagesSingleEmptyCall() {
     final WebSocketClient client = mock(WebSocketClient.class);
-    final WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor, Schedulers.immediate(), clientReleaseManager);
+    final WebSocketConnection connection = webSocketConnection(client);
 
     final UUID accountUuid = UUID.randomUUID();
 
@@ -662,8 +661,7 @@ class WebSocketConnectionTest {
   @Test
   public void testRequeryOnStateMismatch() {
     final WebSocketClient client = mock(WebSocketClient.class);
-    final WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor, Schedulers.immediate(), clientReleaseManager);
+    final WebSocketConnection connection = webSocketConnection(client);
     final UUID accountUuid = UUID.randomUUID();
 
     when(account.getNumber()).thenReturn("+18005551234");
@@ -718,8 +716,7 @@ class WebSocketConnectionTest {
   @Test
   void testProcessCachedMessagesOnly() {
     final WebSocketClient client = mock(WebSocketClient.class);
-    final WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor, Schedulers.immediate(), clientReleaseManager);
+    final WebSocketConnection connection = webSocketConnection(client);
 
     final UUID accountUuid = UUID.randomUUID();
 
@@ -750,8 +747,7 @@ class WebSocketConnectionTest {
   @Test
   void testProcessDatabaseMessagesAfterPersist() {
     final WebSocketClient client = mock(WebSocketClient.class);
-    final WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor, Schedulers.immediate(), clientReleaseManager);
+    final WebSocketConnection connection = webSocketConnection(client);
 
     final UUID accountUuid = UUID.randomUUID();
 
@@ -797,8 +793,7 @@ class WebSocketConnectionTest {
     final WebSocketClient client = mock(WebSocketClient.class);
     when(client.isOpen()).thenReturn(true);
 
-    WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor, Schedulers.immediate(), clientReleaseManager);
+    WebSocketConnection connection = webSocketConnection(client);
     connection.start();
 
     verify(retrySchedulingExecutor, times(WebSocketConnection.MAX_CONSECUTIVE_RETRIES)).schedule(any(Runnable.class),
@@ -821,8 +816,7 @@ class WebSocketConnectionTest {
     final WebSocketClient client = mock(WebSocketClient.class);
     when(client.isOpen()).thenReturn(false);
 
-    WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor, Schedulers.immediate(), clientReleaseManager);
+    WebSocketConnection connection = webSocketConnection(client);
     connection.start();
 
     verify(retrySchedulingExecutor, never()).schedule(any(Runnable.class), anyLong(), any());
@@ -855,8 +849,7 @@ class WebSocketConnectionTest {
     when(messagesManager.delete(any(), anyByte(), any(), any())).thenReturn(
         CompletableFuture.completedFuture(Optional.empty()));
 
-    WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor, messageDeliveryScheduler, clientReleaseManager);
+    WebSocketConnection connection = webSocketConnection(client);
 
     connection.start();
 
@@ -912,8 +905,7 @@ class WebSocketConnectionTest {
     when(messagesManager.delete(any(), anyByte(), any(), any())).thenReturn(
         CompletableFuture.completedFuture(Optional.empty()));
 
-    WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor, Schedulers.immediate(), clientReleaseManager);
+    WebSocketConnection connection = webSocketConnection(client);
 
     connection.start();