cleaned up code

Author: sophia-chen-ttd

src/main/java/com/uid2/admin/salt/KeyIdGenerator.java

@@ -5,32 +5,42 @@
 import java.util.Arrays;
 import java.util.concurrent.atomic.AtomicInteger;
 
+/**
+ * Intended outcomes of KeyIdGenerator:
+ * - Key ids are always consecutive, starting from 0
+ * - When the last allocated key id reaches 16777215, the next key id will wrap around to 0
+ * - Continuing to increment from the highest key id will result in monotonic incrementation of key ids for all newly rotated buckets
+ *
+ * Assumptions:
+ * - The latest assigned key ids are from the latest updated buckets
+ * - Key ids from these buckets will always be consecutive (apart from wraparound) as they have not rotated again after last assignment
+ */
 public class KeyIdGenerator {
     private static final int MAX_KEY_ID = 16777215; // 3 bytes
     private final AtomicInteger lastActiveKeyId;
 
-    public KeyIdGenerator(SaltEntry[] salts) {
-        this.lastActiveKeyId = new AtomicInteger(getLastActiveKeyId(salts));
+    public KeyIdGenerator(SaltEntry[] buckets) {
+        this.lastActiveKeyId = new AtomicInteger(getLastActiveKeyId(buckets));
     }
 
-    private int getLastActiveKeyId(SaltEntry[] salts) {
-        long maxLastUpdated = Arrays.stream(salts).mapToLong(SaltEntry::lastUpdated).max().orElse(0);
-        int[] lastActiveKeyIds = Arrays.stream(salts)
+    private static int getLastActiveKeyId(SaltEntry[] buckets) {
+        long maxLastUpdated = Arrays.stream(buckets).mapToLong(SaltEntry::lastUpdated).max().orElse(0);
+        int[] lastActiveKeyIdsSorted = Arrays.stream(buckets)
                 .filter(s -> s.lastUpdated() == maxLastUpdated && s.currentKey() != null)
                 .mapToInt(s -> s.currentKey().id())
                 .sorted()
                 .toArray();
 
-        if (lastActiveKeyIds.length == 0) return MAX_KEY_ID; // so that next ID will start at 0
+        if (lastActiveKeyIdsSorted.length == 0) return MAX_KEY_ID; // so that next ID will start at 0
 
-        int highestId = lastActiveKeyIds[lastActiveKeyIds.length - 1];
+        int highestId = lastActiveKeyIdsSorted[lastActiveKeyIdsSorted.length - 1];
 
         if (highestId < MAX_KEY_ID) return highestId;
 
-        // Wrapped case - find the last consecutive ID
-        for (int i = 0; i < lastActiveKeyIds.length - 1; i++) {
-            if (lastActiveKeyIds[i + 1] - lastActiveKeyIds[i] > 1) {
-                return lastActiveKeyIds[i];
+        // Wrapped case - find the last consecutive ID from 0
+        for (int i = 0; i < lastActiveKeyIdsSorted.length - 1; i++) {
+            if (lastActiveKeyIdsSorted[i + 1] - lastActiveKeyIdsSorted[i] > 1) {
+                return lastActiveKeyIdsSorted[i];
             }
         }
 

src/main/java/com/uid2/admin/salt/SaltRotation.java

@@ -115,23 +115,23 @@ private SaltEntry[] rotateSalts(SaltEntry[] oldSalts, List<SaltEntry> saltsToRot
         return updatedSalts;
     }
 
-    private SaltEntry updateSalt(SaltEntry oldSalt, TargetDate targetDate, boolean shouldRotate, KeyIdGenerator keyIdGenerator) throws Exception {
-        var lastUpdated = shouldRotate ? targetDate.asEpochMs() : oldSalt.lastUpdated();
-        var refreshFrom = calculateRefreshFrom(oldSalt, targetDate);
-        var currentSalt = calculateCurrentSalt(oldSalt, shouldRotate);
-        var previousSalt = calculatePreviousSalt(oldSalt, shouldRotate, targetDate);
-        var currentKey = calculateCurrentKey(oldSalt, shouldRotate, keyIdGenerator);
-        var previousKey = calculatePreviousKey(oldSalt,shouldRotate, targetDate);
+    private SaltEntry updateSalt(SaltEntry oldBucket, TargetDate targetDate, boolean shouldRotate, KeyIdGenerator keyIdGenerator) throws Exception {
+        var lastUpdated = shouldRotate ? targetDate.asEpochMs() : oldBucket.lastUpdated();
+        var refreshFrom = calculateRefreshFrom(oldBucket, targetDate);
+        var currentSalt = calculateCurrentSalt(oldBucket, shouldRotate);
+        var previousSalt = calculatePreviousSalt(oldBucket, shouldRotate, targetDate);
+        var currentKeySalt = calculateCurrentKeySalt(oldBucket, shouldRotate, keyIdGenerator);
+        var previousKeySalt = calculatePreviousKeySalt(oldBucket,shouldRotate, targetDate);
 
         return new SaltEntry(
-                oldSalt.id(),
-                oldSalt.hashedId(),
+                oldBucket.id(),
+                oldBucket.hashedId(),
                 lastUpdated,
                 currentSalt,
                 refreshFrom,
                 previousSalt,
-                currentKey,
-                previousKey
+                currentKeySalt,
+                previousKeySalt
         );
     }
 
@@ -141,9 +141,15 @@ private long calculateRefreshFrom(SaltEntry salt, TargetDate targetDate) {
     }
 
     private String calculateCurrentSalt(SaltEntry salt, boolean shouldRotate) throws Exception {
-        return shouldRotate ?
-                ENABLE_V4_RAW_UID ? null : this.keyGenerator.generateRandomKeyString(32)
-            : salt.currentSalt();
+        if (shouldRotate) {
+            if (ENABLE_V4_RAW_UID) {
+                return null;
+            }
+            else {
+                return this.keyGenerator.generateRandomKeyString(32);
+            }
+        }
+        return salt.currentSalt();
     }
 
     private String calculatePreviousSalt(SaltEntry salt, boolean shouldRotate, TargetDate targetDate) {
@@ -156,7 +162,7 @@ private String calculatePreviousSalt(SaltEntry salt, boolean shouldRotate, Targe
         return null;
     }
 
-    private SaltEntry.KeyMaterial calculateCurrentKey(SaltEntry salt, boolean shouldRotate, KeyIdGenerator keyIdGenerator) throws Exception {
+    private SaltEntry.KeyMaterial calculateCurrentKeySalt(SaltEntry bucket, boolean shouldRotate, KeyIdGenerator keyIdGenerator) throws Exception {
         if (shouldRotate) {
             if (ENABLE_V4_RAW_UID) {
                 return new SaltEntry.KeyMaterial(
@@ -168,15 +174,15 @@ private SaltEntry.KeyMaterial calculateCurrentKey(SaltEntry salt, boolean should
                 return null;
             }
         }
-        return salt.currentKey();
+        return bucket.currentKey();
     }
 
-    private SaltEntry.KeyMaterial calculatePreviousKey(SaltEntry salt, boolean shouldRotate, TargetDate targetDate) {
+    private SaltEntry.KeyMaterial calculatePreviousKeySalt(SaltEntry bucket, boolean shouldRotate, TargetDate targetDate) {
         if (shouldRotate) {
-            return salt.currentKey();
+            return bucket.currentKey();
         }
-        if (targetDate.saltAgeInDays(salt) < 90) {
-            return salt.previousKey();
+        if (targetDate.saltAgeInDays(bucket) < 90) {
+            return bucket.previousKey();
         }
         return null;
     }