move active key count initialization from first record to SOP

- initializeActiveKeyCount() at SOP sets -1→0, ensuring mid-batch config
  enablement (config OFF at SOP, ON after restart) stays at -1 (not tracked)
  rather than producing an inaccurate partial count
- incrementActiveKeyCountForBatchRecord skips when count < 0 (not initialized)
- Renamed finalizeActiveKeyCountForBatchPush → cleanupBatchKeyCountState
  (only releases dedup state, no longer manipulates count)
- Added tests: mid-batch config enablement, initializeActiveKeyCount
idempotency
- Fixed SpotBugs ICAST_INTEGER_MULTIPLY_CAST_TO_LONG in testConcurrency

Author: m-nagarajan

clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/PartitionConsumptionState.java

@@ -481,6 +481,17 @@ public void incrementActiveKeyCount() {
     activeKeyCount.incrementAndGet();
   }
 
+  /**
+   * Called at SOP to mark this partition as actively counting. Without this, a mid-batch
+   * restart with a newly enabled config would start counting partway through, producing
+   * an inaccurate partial count. On restart, SOP is not re-processed (already past the
+   * checkpoint), so the count stays at {@link #ACTIVE_KEY_COUNT_NOT_TRACKED} and batch
+   * records are skipped by {@link #incrementActiveKeyCountForBatchRecord}.
+   */
+  public void initializeActiveKeyCount() {
+    activeKeyCount.compareAndSet(ACTIVE_KEY_COUNT_NOT_TRACKED, 0);
+  }
+
   /**
    * Decrements activeKeyCount. If the count is already at 0, this indicates drift
    * (more deletes than creates), so we invalidate to {@link #ACTIVE_KEY_COUNT_NOT_TRACKED}
@@ -496,22 +507,26 @@ public boolean decrementActiveKeyCount() {
   /**
    * Increments activeKeyCount for a batch PUT, skipping speculative execution duplicates
    * (key <= last key). Only PUTs call this — DELETEs are tombstones and excluded.
-   * First call initializes count from {@link #ACTIVE_KEY_COUNT_NOT_TRACKED} to 1. Checkpoint-safe: partial counts are persisted.
+   * Requires prior {@link #initializeActiveKeyCount()} at SOP; skips if not initialized
+   * (mid-batch config enablement). Checkpoint-safe: partial counts are persisted.
    */
   public void incrementActiveKeyCountForBatchRecord(byte[] keyBytes) {
+    if (activeKeyCount.get() < 0) {
+      return; // Not initialized at SOP — mid-batch config enablement, skip
+    }
     if (lastBatchKeyForDedup != null && ArrayUtils.compareUnsigned(keyBytes, lastBatchKeyForDedup) <= 0) {
       return; // Speculative execution duplicate — key order went backwards
     }
     lastBatchKeyForDedup = keyBytes;
-    activeKeyCount.compareAndSet(ACTIVE_KEY_COUNT_NOT_TRACKED, 0);
     activeKeyCount.incrementAndGet();
   }
 
-  /** Called at EOP. Sets {@link #ACTIVE_KEY_COUNT_NOT_TRACKED} to 0 for empty partitions so the metric reports 0. */
-  public void finalizeActiveKeyCountForBatchPush() {
-    if (activeKeyCount.get() == ACTIVE_KEY_COUNT_NOT_TRACKED) {
-      activeKeyCount.set(0);
-    }
+  /**
+   * Called at EOP. Releases dedup state. Empty partitions are already at 0 from
+   * {@link #initializeActiveKeyCount()} at SOP. If SOP was missed (mid-batch config enablement),
+   * the count stays at {@link #ACTIVE_KEY_COUNT_NOT_TRACKED} — no partial baseline is created.
+   */
+  public void cleanupBatchKeyCountState() {
     lastBatchKeyForDedup = null; // Release reference; dedup is only needed during batch
   }
 

clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTask.java

@@ -3793,6 +3793,10 @@ private void processStartOfPush(
     beginBatchWrite(persistedStoreVersionState.sorted, partitionConsumptionState);
     partitionConsumptionState.setStartOfPushTimestamp(startOfPushKME.producerMetadata.messageTimestamp);
 
+    // Initialize active key count at SOP
+    if (serverConfig.isActiveKeyCountForAllBatchPushEnabled()) {
+      partitionConsumptionState.initializeActiveKeyCount();
+    }
   }
 
   @VisibleForTesting
@@ -3872,11 +3876,11 @@ protected void processEndOfPush(
      */
     partitionConsumptionState.finalizeExpectedChecksum();
 
-    // Finalize batch key count (sets ACTIVE_KEY_COUNT_NOT_TRACKED→0 for empty partitions).
+    // Release batch key count dedup state (no longer needed after EOP).
     if (serverConfig.isActiveKeyCountForAllBatchPushEnabled()) {
-      partitionConsumptionState.finalizeActiveKeyCountForBatchPush();
+      partitionConsumptionState.cleanupBatchKeyCountState();
       LOGGER.info(
-          "Active key count finalized at EOP for replica: {}, activeKeyCount: {}",
+          "Active key count at EOP for replica: {}, activeKeyCount: {}",
           partitionConsumptionState.getReplicaId(),
           partitionConsumptionState.getActiveKeyCount());
     }

clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/ActiveKeyCountScenarioTest.java

@@ -133,6 +133,7 @@ public void testCrashRecoveryWithConfigCombinations(
       String desc) {
     PartitionConsumptionState pcs = freshPcs();
     if (batchCountingEnabled) {
+      pcs.initializeActiveKeyCount(); // SOP
       for (int i = 0; i < 30; i++) {
         pcs.incrementActiveKeyCountForBatchRecord(ActiveKeyCountTestUtils.sortedKeyBytes(i));
       }
@@ -143,7 +144,7 @@ public void testCrashRecoveryWithConfigCombinations(
       for (int i = 30; i < 50; i++) {
         restarted.incrementActiveKeyCountForBatchRecord(ActiveKeyCountTestUtils.sortedKeyBytes(i));
       }
-      restarted.finalizeActiveKeyCountForBatchPush();
+      restarted.cleanupBatchKeyCountState();
       assertEquals(restarted.getActiveKeyCount(), 50L, desc);
     } else {
       assertEquals(restarted.getActiveKeyCount(), -1L, desc);

clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/ActiveKeyCountTest.java

@@ -295,15 +295,19 @@ public Object[][] followerSignalSkippedCases() {
   public void testBatchKeyCountAndFinalize() {
     PartitionConsumptionState localPcs = freshPcs();
     assertEquals(localPcs.getActiveKeyCount(), -1L);
+    localPcs.initializeActiveKeyCount();
+    assertEquals(localPcs.getActiveKeyCount(), 0L);
     for (int i = 0; i < 5; i++) {
       localPcs.incrementActiveKeyCountForBatchRecord(ActiveKeyCountTestUtils.sortedKeyBytes(i));
     }
     assertEquals(localPcs.getActiveKeyCount(), 5L);
-    localPcs.finalizeActiveKeyCountForBatchPush();
+    localPcs.cleanupBatchKeyCountState();
     assertEquals(localPcs.getActiveKeyCount(), 5L);
-    // Empty batch finalize yields 0
+    // Empty batch: initializeActiveKeyCount at SOP sets -1→0, finalize is a no-op
     PartitionConsumptionState emptyPcs = freshPcs();
-    emptyPcs.finalizeActiveKeyCountForBatchPush();
+    emptyPcs.initializeActiveKeyCount();
+    assertEquals(emptyPcs.getActiveKeyCount(), 0L);
+    emptyPcs.cleanupBatchKeyCountState();
     assertEquals(emptyPcs.getActiveKeyCount(), 0L);
     // Verify increment/decrement works post-finalize (RT signals)
     localPcs.incrementActiveKeyCount();
@@ -315,6 +319,7 @@ public void testBatchKeyCountAndFinalize() {
   @Test
   public void testBatchDedupSkipsDuplicateKeys() {
     PartitionConsumptionState localPcs = freshPcs();
+    localPcs.initializeActiveKeyCount();
     for (int i = 0; i < 5; i++) {
       localPcs.incrementActiveKeyCountForBatchRecord(ActiveKeyCountTestUtils.sortedKeyBytes(i));
     }
@@ -335,14 +340,15 @@ public void testBatchDedupSkipsDuplicateKeys() {
   @Test
   public void testFinalizeDoesNotOverwriteRTSignals() {
     PartitionConsumptionState localPcs = freshPcs();
+    localPcs.initializeActiveKeyCount();
     for (int i = 0; i < 10; i++) {
       localPcs.incrementActiveKeyCountForBatchRecord(ActiveKeyCountTestUtils.sortedKeyBytes(i));
     }
-    localPcs.finalizeActiveKeyCountForBatchPush();
+    localPcs.cleanupBatchKeyCountState();
     localPcs.incrementActiveKeyCount(); // RT adjustment -> 11
     assertEquals(localPcs.getActiveKeyCount(), 11L);
     // Second finalize (e.g., from duplicate EOP) does NOT overwrite RT signals
-    localPcs.finalizeActiveKeyCountForBatchPush();
+    localPcs.cleanupBatchKeyCountState();
     assertEquals(localPcs.getActiveKeyCount(), 11L);
   }
 
@@ -790,9 +796,9 @@ public void testProcessEndOfPush() throws Exception {
       kme.producerMetadata.messageTimestamp = System.currentTimeMillis();
       sitMock.processEndOfPush(kme, mock(PubSubPosition.class), mockPcs, new EndOfPush());
       if (enabled) {
-        verify(mockPcs).finalizeActiveKeyCountForBatchPush();
+        verify(mockPcs).cleanupBatchKeyCountState();
       } else {
-        verify(mockPcs, never()).finalizeActiveKeyCountForBatchPush();
+        verify(mockPcs, never()).cleanupBatchKeyCountState();
       }
     }
   }
@@ -991,12 +997,51 @@ public void testDecrementUnderflowInvalidatesAndPreservesInvalidState() {
   public void testDecrementUnderflowAfterEmptyBatch() {
     PartitionConsumptionState localPcs = freshPcs();
 
-    // Simulate empty batch push → finalize sets -1 to 0
-    localPcs.finalizeActiveKeyCountForBatchPush();
+    // Simulate empty batch push: SOP initializes -1→0, finalize is no-op
+    localPcs.initializeActiveKeyCount();
     assertEquals(localPcs.getActiveKeyCount(), 0);
 
     // RT DELETE signal on empty batch: underflow invalidates to -1 (drift)
     Assert.assertFalse(localPcs.decrementActiveKeyCount());
     assertEquals(localPcs.getActiveKeyCount(), -1);
   }
+
+  @Test
+  public void testMidBatchConfigEnablementSkipsCounting() {
+    // Simulate: PCS restored from checkpoint mid-batch (SOP already processed without config ON)
+    // activeKeyCount is ACTIVE_KEY_COUNT_NOT_TRACKED (-1) because SOP didn't initialize it
+    PartitionConsumptionState localPcs = freshPcs();
+    assertEquals(localPcs.getActiveKeyCount(), -1);
+
+    // Batch records arrive with config now ON, but initializeActiveKeyCount was never called (SOP missed)
+    // incrementActiveKeyCountForBatchRecord should skip (count stays -1)
+    localPcs.incrementActiveKeyCountForBatchRecord(ActiveKeyCountTestUtils.sortedKeyBytes(0));
+    localPcs.incrementActiveKeyCountForBatchRecord(ActiveKeyCountTestUtils.sortedKeyBytes(1));
+    localPcs.incrementActiveKeyCountForBatchRecord(ActiveKeyCountTestUtils.sortedKeyBytes(2));
+    assertEquals(localPcs.getActiveKeyCount(), -1);
+
+    // EOP finalize: count stays -1 (no partial baseline created). Finalize only releases
+    // dedup state; it does not set -1→0 since initializeActiveKeyCount was never called at SOP.
+    localPcs.cleanupBatchKeyCountState();
+    assertEquals(localPcs.getActiveKeyCount(), -1);
+  }
+
+  @Test
+  public void testInitializeActiveKeyCountIdempotent() {
+    PartitionConsumptionState localPcs = freshPcs();
+    assertEquals(localPcs.getActiveKeyCount(), -1);
+
+    // First init: -1 → 0
+    localPcs.initializeActiveKeyCount();
+    assertEquals(localPcs.getActiveKeyCount(), 0);
+
+    // Count some records
+    localPcs.incrementActiveKeyCountForBatchRecord(ActiveKeyCountTestUtils.sortedKeyBytes(0));
+    localPcs.incrementActiveKeyCountForBatchRecord(ActiveKeyCountTestUtils.sortedKeyBytes(1));
+    assertEquals(localPcs.getActiveKeyCount(), 2);
+
+    // Second init (e.g., duplicate SOP): no-op because compareAndSet(-1, 0) fails (count is 2)
+    localPcs.initializeActiveKeyCount();
+    assertEquals(localPcs.getActiveKeyCount(), 2);
+  }
 }

clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/ActiveKeyCountTestUtils.java

@@ -79,11 +79,12 @@ static OffsetRecord checkpoint(PartitionConsumptionState pcs) {
 
   /** Simulates a non-chunked batch push of {@code count} records followed by finalization. */
   static void doBatch(PartitionConsumptionState pcs, int count) {
+    pcs.initializeActiveKeyCount(); // SOP initialization
     for (int i = 0; i < count; i++) {
       // Synthetic sorted key bytes (ascending order) to pass dedup check
       pcs.incrementActiveKeyCountForBatchRecord(sortedKeyBytes(i));
     }
-    pcs.finalizeActiveKeyCountForBatchPush();
+    pcs.cleanupBatchKeyCountState();
   }
 
   /**
@@ -92,12 +93,13 @@ static void doBatch(PartitionConsumptionState pcs, int count) {
    * Chunk fragment filtering itself is tested via reflection in ActiveKeyCountTest.
    */
   static void doBatchChunked(PartitionConsumptionState pcs, int logicalKeyCount) {
+    pcs.initializeActiveKeyCount(); // SOP initialization
     for (int key = 0; key < logicalKeyCount; key++) {
       // Chunk fragments: schemaId == CHUNK -> filtered out by processKafkaDataMessage, NOT counted
       // Manifest: schemaId == CHUNKED_VALUE_MANIFEST -> passes filter, counted
       pcs.incrementActiveKeyCountForBatchRecord(sortedKeyBytes(key));
     }
-    pcs.finalizeActiveKeyCountForBatchPush();
+    pcs.cleanupBatchKeyCountState();
   }
 
   /** Generates key bytes in ascending unsigned byte order for test dedup checks. */

internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestActiveKeyCount.java

@@ -671,8 +671,8 @@ public void testBothExactAndHllCountsWorkSimultaneously() throws Exception {
 
   /**
    * Validates that an empty batch push (0 user records) results in an active key count of 0,
-   * not -1 (untracked). The finalizeActiveKeyCountForBatchPush at EOP transitions -1 to 0
-   * for empty partitions.
+   * not -1 (untracked). The initializeActiveKeyCount at SOP transitions -1 to 0,
+   * and empty partitions stay at 0 through EOP.
    */
   @Test(timeOut = TEST_TIMEOUT)
   public void testEmptyBatchPushGetsZeroKeyCount() throws Exception {