added logs for key bucket count

Author: sophia-chen-ttd

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

@@ -62,6 +62,7 @@ public Result rotateSalts(
         logSaltAges("refreshable-salts", targetDate, refreshableSalts);
         logSaltAges("rotated-salts", targetDate, saltsToRotate);
         logSaltAges("total-salts", targetDate, Arrays.asList(postRotationSalts));
+        logKeyBucketCount(targetDate, preRotationSalts, postRotationSalts);
 
         var nextSnapshot = new SaltSnapshot(
                 nextEffective,
@@ -246,6 +247,20 @@ private void logSaltAges(String saltCountType, TargetDate targetDate, Collection
         }
     }
 
+    private void logKeyBucketCount(TargetDate targetDate, SaltEntry[] preRotationSalts, SaltEntry[] postRotationSalts) {
+        int newKeyBucketCounter = 0;
+        int totalKeyBucketCounter = 0;
+        for (int i = 0; i < preRotationSalts.length && i < postRotationSalts.length; i++) {
+            var oldSalt = preRotationSalts[i];
+            var updatedSalt = postRotationSalts[i];
+
+            if (updatedSalt.currentKey() != null) totalKeyBucketCounter++;
+            if (updatedSalt.currentKey() != null && oldSalt.currentSalt() != null) newKeyBucketCounter++;
+        }
+        LOGGER.info("salt_bucket_count_type={} target_date={} bucket_count={}", "new-key-buckets", targetDate, newKeyBucketCounter);
+        LOGGER.info("salt_bucket_count_type={} target_date={} bucket_count={}", "total-key-buckets", targetDate, totalKeyBucketCounter);
+    }
+
     @Getter
     public static final class Result {
         private final SaltSnapshot snapshot; // can be null if new snapshot is not needed

src/test/java/com/uid2/admin/salt/SaltRotationTest.java

@@ -389,7 +389,7 @@ void testLogFewSaltAgesOnRotation() throws Exception {
         var minAges = new Duration[]{Duration.ofDays(30), Duration.ofDays(60)};
         saltRotation.rotateSalts(lastSnapshot, minAges, 0.4, targetDate());
 
-        var actual = appender.list.stream().map(Object::toString).collect(Collectors.toSet());
+        var actual = appender.list.stream().map(Object::toString).filter(s -> s.contains("salt_count_type") || s.contains("Salt rotation complete")).collect(Collectors.toSet());
         assertThat(actual).isEqualTo(expected);
     }
 
@@ -429,7 +429,7 @@ void testLogManySaltAgesOnRotation() throws Exception {
         var minAges = new Duration[]{Duration.ofDays(30), Duration.ofDays(60)};
         saltRotation.rotateSalts(lastSnapshot, minAges, 0.2, targetDate());
 
-        var actual = appender.list.stream().map(Object::toString).collect(Collectors.toSet());
+        var actual = appender.list.stream().map(Object::toString).filter(s -> s.contains("salt_count_type") || s.contains("Salt rotation complete")).collect(Collectors.toSet());
         assertThat(actual).isEqualTo(expected);
     }
 
@@ -726,4 +726,35 @@ void testKeyRotationSaltToKeyRotation() throws Exception {
         assertThat(salts[0].currentSalt()).isNotNull();
         assertThat(salts[0].previousSalt()).isNull();
     }
+
+    @ParameterizedTest
+    @CsvSource({
+            "true, 1, 3",
+            "false, 0, 1"
+    })
+    void testKeyRotationLogKeyBuckets(boolean v4Enabled, int expectedNewKeyBuckets, int expectedTotalKeyBuckets) throws Exception {
+        saltRotation = new SaltRotation(keyGenerator, JsonObject.of(AdminConst.ENABLE_V4_RAW_UID, v4Enabled));
+
+        final Duration[] minAges = {
+                Duration.ofDays(30)
+        };
+
+        var willRefresh = targetDate();
+        var willNotRefresh = targetDate().plusDays(30);
+        var lastSnapshot = SaltSnapshotBuilder.start()
+                .entries(SaltBuilder.start().lastUpdated(targetDate().minusDays(60)).refreshFrom(willRefresh).currentSalt(),
+                        SaltBuilder.start().lastUpdated(targetDate().minusDays(60)).refreshFrom(willNotRefresh).currentSalt(),
+                        SaltBuilder.start().lastUpdated(targetDate().minusDays(60)).refreshFrom(willRefresh).currentKey(1),
+                        SaltBuilder.start().lastUpdated(targetDate().minusDays(60)).refreshFrom(willNotRefresh).currentKey(2))
+                .build();
+
+        saltRotation.rotateSalts(lastSnapshot, minAges, 1, targetDate());
+
+        var expected = Set.of(
+                "[INFO] salt_bucket_count_type=new-key-buckets target_date=2025-01-01 bucket_count=" + expectedNewKeyBuckets,
+                "[INFO] salt_bucket_count_type=total-key-buckets target_date=2025-01-01 bucket_count=" + expectedTotalKeyBuckets
+        );
+        var actual = appender.list.stream().map(Object::toString).filter(s -> s.contains("salt_bucket_count_type")).collect(Collectors.toSet());
+        assertThat(actual).isEqualTo(expected);
+    }
 }