Introduce EnumMapUtil

signalapp · Oct 10, 2024 · 240a406 · 240a406
1 parent 38d25f9
commit 240a406
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 34 deletions.
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/DeviceController.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/DeviceController.java
@@ -16,15 +16,12 @@
 import io.swagger.v3.oas.annotations.media.Schema;
 import io.swagger.v3.oas.annotations.responses.ApiResponse;
 import java.time.Duration;
-import java.util.Arrays;
 import java.util.EnumMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.function.Function;
-import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import javax.validation.Valid;
 import javax.validation.constraints.Max;
@@ -72,6 +69,7 @@
 import org.whispersystems.textsecuregcm.storage.Device.DeviceCapabilities;
 import org.whispersystems.textsecuregcm.storage.DeviceSpec;
 import org.whispersystems.textsecuregcm.storage.LinkDeviceTokenAlreadyUsedException;
+import org.whispersystems.textsecuregcm.util.EnumMapUtil;
 import org.whispersystems.textsecuregcm.util.ExceptionUtils;
 import org.whispersystems.textsecuregcm.util.LinkDeviceToken;
 import org.whispersystems.textsecuregcm.util.ua.ClientPlatform;
@@ -116,15 +114,8 @@ public DeviceController(final AccountsManager accounts,
     this.rateLimiters = rateLimiters;
     this.maxDeviceConfiguration = maxDeviceConfiguration;
 
-    linkedDeviceListenersByPlatform = Arrays.stream(ClientPlatform.values())
-        .collect(Collectors.toMap(
-            Function.identity(),
-            clientPlatform -> buildGauge(clientPlatform.name().toLowerCase()),
-            (a, b) -> {
-              throw new AssertionError("Duplicate client platform enumeration key");
-            },
-            () -> new EnumMap<>(ClientPlatform.class)
-        ));
+    linkedDeviceListenersByPlatform =
+        EnumMapUtil.toEnumMap(ClientPlatform.class, clientPlatform -> buildGauge(clientPlatform.name().toLowerCase()));
 
     linkedDeviceListenersForUnrecognizedPlatforms = buildGauge("unknown");
   }

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/metrics/OpenWebSocketCounter.java b/service/src/main/java/org/whispersystems/textsecuregcm/metrics/OpenWebSocketCounter.java
@@ -4,16 +4,13 @@
 import io.micrometer.core.instrument.Tag;
 import io.micrometer.core.instrument.Tags;
 import io.micrometer.core.instrument.Timer;
+import org.whispersystems.textsecuregcm.util.EnumMapUtil;
 import org.whispersystems.textsecuregcm.util.ua.ClientPlatform;
 import org.whispersystems.textsecuregcm.util.ua.UnrecognizedUserAgentException;
 import org.whispersystems.textsecuregcm.util.ua.UserAgentUtil;
 import org.whispersystems.websocket.session.WebSocketSessionContext;
-import java.util.Arrays;
-import java.util.EnumMap;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.function.Function;
-import java.util.stream.Collectors;
 
 public class OpenWebSocketCounter {
 
@@ -28,27 +25,13 @@ public OpenWebSocketCounter(final String openWebSocketGaugeName, final String du
   }
 
   public OpenWebSocketCounter(final String openWebSocketGaugeName, final String durationTimerName, final Tags tags) {
-    openWebsocketsByClientPlatform = Arrays.stream(ClientPlatform.values())
-            .collect(Collectors.toMap(
-                Function.identity(),
-                clientPlatform -> buildGauge(openWebSocketGaugeName, clientPlatform.name().toLowerCase(), tags),
-                (a, b) -> {
-                  throw new AssertionError("Duplicate client platform enumeration key");
-                },
-                () -> new EnumMap<>(ClientPlatform.class)
-            ));
+    openWebsocketsByClientPlatform = EnumMapUtil.toEnumMap(ClientPlatform.class,
+        clientPlatform -> buildGauge(openWebSocketGaugeName, clientPlatform.name().toLowerCase(), tags));
 
     openWebsocketsFromUnknownPlatforms = buildGauge(openWebSocketGaugeName, "unknown", tags);
 
-    durationTimersByClientPlatform = Arrays.stream(ClientPlatform.values())
-        .collect(Collectors.toMap(
-            clientPlatform -> clientPlatform,
-            clientPlatform -> buildTimer(durationTimerName, clientPlatform.name().toLowerCase(), tags),
-            (a, b) -> {
-              throw new AssertionError("Duplicate client platform enumeration key");
-            },
-            () -> new EnumMap<>(ClientPlatform.class)
-        ));
+    durationTimersByClientPlatform = EnumMapUtil.toEnumMap(ClientPlatform.class,
+        clientPlatform -> buildTimer(durationTimerName, clientPlatform.name().toLowerCase(), tags));
 
     durationTimerForUnknownPlatforms = buildTimer(durationTimerName, "unknown", tags);
   }

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/util/EnumMapUtil.java b/service/src/main/java/org/whispersystems/textsecuregcm/util/EnumMapUtil.java
@@ -0,0 +1,24 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.util;
+
+import java.util.Arrays;
+import java.util.EnumMap;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+public class EnumMapUtil {
+
+  private EnumMapUtil() {}
+
+  public static <E extends Enum<E>, V> EnumMap<E, V> toEnumMap(final Class<E> enumClass, final Function<E, V> valueMapper) {
+    return Arrays.stream(enumClass.getEnumConstants())
+        .collect(Collectors.toMap(Function.identity(), valueMapper, (a, b) -> {
+              throw new AssertionError("Duplicate enumeration key");
+            },
+            () -> new EnumMap<>(enumClass)));
+  }
+}