Use computeIfAbsent to record metrics

Author: cYKatherine

src/main/java/com/uid2/operator/service/EncryptedTokenEncoder.java

@@ -1,6 +1,7 @@
 package com.uid2.operator.service;
 
 import com.uid2.operator.model.*;
+import com.uid2.operator.util.Tuple;
 import com.uid2.operator.vertx.ClientInputValidationException;
 import com.uid2.shared.Const.Data;
 import com.uid2.shared.encryption.AesCbc;
@@ -11,16 +12,20 @@
 import io.vertx.core.buffer.Buffer;
 import io.micrometer.core.instrument.Counter;
 import io.micrometer.core.instrument.Metrics;
+import software.amazon.awssdk.services.kms.endpoints.internal.Value;
 
 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<Integer, Counter> _refreshTokenVersionCounters = new HashMap<>();
 
     public byte[] encode(AdvertisingToken t, Instant asOf) {
         final KeysetKey masterKey = this.keyManager.getMasterKey(asOf);
@@ -260,23 +265,24 @@ public AdvertisingToken decodeAdvertisingTokenV3orV4(Buffer b, byte[] bytes, Tok
         );
     }
 
+    private void recordRefreshTokenVersionCount(int 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(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(t.publisherIdentity.siteId, TokenVersion.V3);
                 return encodeV3(t, serviceKey);
             default:
                 throw new ClientInputValidationException("RefreshToken version " + t.version + " not supported");

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 {
@@ -1808,19 +1818,7 @@ private RefreshResponse refreshIdentity(RoutingContext rc, String tokenStr) {
             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)), tokenVersion);
 
         return this.idService.refreshIdentity(refreshToken);
     }

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());
     }