Use computeIfAbsent to record metrics

IABTechLab · Jun 11, 2024 · 149a399 · 149a399
1 parent 1e6cbd5
commit 149a399
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 24 deletions.
diff --git a/src/main/java/com/uid2/operator/service/EncryptedTokenEncoder.java b/src/main/java/com/uid2/operator/service/EncryptedTokenEncoder.java
@@ -14,13 +14,16 @@
 
 import java.time.Instant;
 import java.util.Base64;
+import java.util.HashMap;
+import java.util.Map;
 
 public class EncryptedTokenEncoder implements ITokenEncoder {
     private final KeyManager keyManager;
 
     public EncryptedTokenEncoder(KeyManager keyManager) {
         this.keyManager = keyManager;
     }
+    private final Map<String, Counter> _refreshTokenVersionCounters = new HashMap<>();
 
     public byte[] encode(AdvertisingToken t, Instant asOf) {
         final KeysetKey masterKey = this.keyManager.getMasterKey(asOf);
@@ -260,23 +263,24 @@ public AdvertisingToken decodeAdvertisingTokenV3orV4(Buffer b, byte[] bytes, Tok
         );
     }
 
+    private void recordRefreshTokenVersionCount(String siteId, TokenVersion tokenVersion) {
+        _refreshTokenVersionCounters.computeIfAbsent(siteId, id -> Counter
+                .builder(String.format("uid2_refresh_token_%s_served_count", tokenVersion.toString().toLowerCase()))
+                .description(String.format("Counter for the amount of refresh token %s served", tokenVersion.toString().toLowerCase()))
+                .tags("site_id", String.valueOf(siteId))
+                .register(Metrics.globalRegistry)).increment();
+
+    }
+
     public byte[] encode(RefreshToken t, Instant asOf) {
         final KeysetKey serviceKey = this.keyManager.getRefreshKey(asOf);
 
         switch (t.version) {
             case V2:
-                var v2Builder = Counter
-                        .builder("uid2_refresh_token_v2_served_count")
-                        .description("Counter for the amount of refresh token v2 served").tags(
-                                "timestamp", String.valueOf(asOf));
-                v2Builder.register(Metrics.globalRegistry).increment();
+                recordRefreshTokenVersionCount(String.valueOf(t.publisherIdentity.siteId), TokenVersion.V2);
                 return encodeV2(t, serviceKey);
             case V3:
-                var v3Builder = Counter
-                        .builder("uid2_refresh_token_v3_served_count")
-                        .description("Counter for the amount of refresh token v3 served").tags(
-                                "timestamp", String.valueOf(asOf));
-                v3Builder.register(Metrics.globalRegistry).increment();
+                recordRefreshTokenVersionCount(String.valueOf(t.publisherIdentity.siteId), TokenVersion.V3);
                 return encodeV3(t, serviceKey);
             default:
                 throw new ClientInputValidationException("RefreshToken version " + t.version + " not supported");

diff --git a/src/main/java/com/uid2/operator/vertx/UIDOperatorVerticle.java b/src/main/java/com/uid2/operator/vertx/UIDOperatorVerticle.java
@@ -114,6 +114,7 @@ public class UIDOperatorVerticle extends AbstractVerticle {
     public final static int MASTER_KEYSET_ID_FOR_SDKS = 9999999; //this is because SDKs have an issue where they assume keyset ids are always positive; that will be fixed.
     public final static long OPT_OUT_CHECK_CUTOFF_DATE = Instant.parse("2023-09-01T00:00:00.00Z").getEpochSecond();
     private final Handler<Boolean> saltRetrievalResponseHandler;
+    private final Map<String, Counter> _refreshTokenVersionCounters = new HashMap<>();
 
     private final int maxBidstreamLifetimeSeconds;
     private final int allowClockSkewSeconds;
@@ -1791,6 +1792,15 @@ public TokenVersion getRefreshTokenVersion(String s) {
         return null;
     }
 
+    private void recordRefreshTokenVersionCount(String siteId, TokenVersion tokenVersion) {
+        _refreshTokenVersionCounters.computeIfAbsent(siteId, id -> Counter
+                .builder(String.format("uid2_refresh_token_%s_received_count", tokenVersion.toString().toLowerCase()))
+                .description(String.format("Counter for the amount of refresh token %s received", tokenVersion.toString().toLowerCase()))
+                .tags("site_id", siteId)
+                .register(Metrics.globalRegistry)).increment();
+
+    }
+
     private RefreshResponse refreshIdentity(RoutingContext rc, String tokenStr) {
         final RefreshToken refreshToken;
         try {
@@ -1807,20 +1817,7 @@ private RefreshResponse refreshIdentity(RoutingContext rc, String tokenStr) {
         if (!AuthMiddleware.isAuthenticated(rc)) {
             rc.put(Const.RoutingContextData.SiteId, refreshToken.publisherIdentity.siteId);
         }
-        TokenVersion tokenVersion = this.getRefreshTokenVersion(tokenStr);
-        if (tokenVersion == TokenVersion.V2) {
-            var v2Builder = Counter
-                    .builder("uid2_refresh_token_v2_received_count")
-                    .description("Counter for the amount of refresh token v2 received").tags(
-                            "site_id", String.valueOf(rc.data().get(Const.RoutingContextData.SiteId)));
-            v2Builder.register(Metrics.globalRegistry).increment();
-        } else if (tokenVersion == TokenVersion.V3) {
-            var v3Builder = Counter
-                    .builder("uid2_refresh_token_v3_received_count")
-                    .description("Counter for the amount of refresh token v3 received").tags(
-                            "site_id", String.valueOf(rc.data().get(Const.RoutingContextData.SiteId)));
-            v3Builder.register(Metrics.globalRegistry).increment();
-        }
+        recordRefreshTokenVersionCount(String.valueOf(rc.data().get(Const.RoutingContextData.SiteId)), this.getRefreshTokenVersion(tokenStr));
 
         return this.idService.refreshIdentity(refreshToken);
     }

diff --git a/src/test/java/com/uid2/operator/TokenEncodingTest.java b/src/test/java/com/uid2/operator/TokenEncodingTest.java
@@ -94,6 +94,7 @@ public void testRefreshTokenEncoding(TokenVersion tokenVersion) {
 
         assertEquals(1, Metrics.globalRegistry
                 .get("uid2_refresh_token_v3_served_count")
+                .tag("site_id", "111")
                 .counter().count());
     }