[da-vinci] Global RT DIV: Per-Partition Consumed Bytes Counter (linkedin#2739)

consumedBytesSinceLastSync was a Map<String, Long> on StoreIngestionTask, shared across all partitions of the same version. This meant the sync cadence for shouldSendGlobalRtDiv and shouldSyncOffsetFromSnapshot was driven by aggregate bytes across all partitions rather than per-partition progress.

Fix: Remove the SIT-level field entirely. Add consumedBytesSinceLastGlobalRtDivSync: Map<String, Long> to PartitionConsumptionState with get/add/reset methods, keyed by VT name (for local VT tracking) or broker URL (for RT tracking).

---------

Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

Author: KaiSernLim

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

@@ -3174,7 +3174,7 @@ void syncOffsetFromSnapshotIfNeeded(DefaultPubSubMessage record, PubSubTopicPart
       storeBufferService.execSyncOffsetFromSnapshotAsync(topicPartition, vtDiv, lastFuture, this);
       // Reset consumer-side VT bytes so the size-based condition in shouldSyncOffsetFromSnapshot does not keep
       // firing for every subsequent record.
-      getConsumedBytesSinceLastSync().put(getVersionTopic().getName(), 0L);
+      pcs.resetConsumedBytesSinceLastGlobalRtDivSync(getVersionTopic().getName());
 
       // TODO: remove. this is a temporary log for debugging while the feature is in its infancy
       LOGGER.info(
@@ -3210,7 +3210,7 @@ boolean shouldSyncOffsetFromSnapshot(DefaultPubSubMessage consumerRecord, Partit
 
     // must be greater than the interval in shouldSendGlobalRtDiv() to not interfere
     final long syncBytesInterval = getSyncBytesInterval(pcs); // size-based sync condition
-    long vtConsumedBytesSinceLastSync = getConsumedBytesSinceLastSync().getOrDefault(getVersionTopic().getName(), 0L);
+    long vtConsumedBytesSinceLastSync = pcs.getConsumedBytesSinceLastGlobalRtDivSync(getVersionTopic().getName());
     return syncBytesInterval > 0 && (vtConsumedBytesSinceLastSync >= 2 * syncBytesInterval);
   }
 
@@ -3985,7 +3985,7 @@ void sendGlobalRtDivMessage(
             null,
             true);
 
-    consumedBytesSinceLastSync.put(brokerUrl, 0L); // reset the timer for the next sync, since RT DIV was just synced
+    pcs.resetConsumedBytesSinceLastGlobalRtDivSync(brokerUrl);
   }
 
   private byte[] createGlobalRtDivValueBytes(

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

@@ -162,6 +162,12 @@ enum LatchStatus {
    */
   private long processedRecordSizeSinceLastSync;
 
+  /**
+   * Tracks bytes consumed per source key (VT name or RT broker URL) since the last Global RT DIV sync.
+   * Stored per-partition so that each partition's sync cadence is independent.
+   */
+  private final Map<String, Long> consumedBytesSinceLastGlobalRtDivSync = new VeniceConcurrentHashMap<>();
+
   /** Minimum lgK supported by DataSketches HllSketch (mirrors package-private HllUtil.MIN_LOG_K). */
   static final int HLL_MIN_LOG_K = 4;
   /** Maximum lgK supported by DataSketches HllSketch (mirrors package-private HllUtil.MAX_LOG_K). */
@@ -718,6 +724,21 @@ public void resetProcessedRecordSizeSinceLastSync() {
     this.processedRecordSizeSinceLastSync = 0;
   }
 
+  public long getConsumedBytesSinceLastGlobalRtDivSync(String key) {
+    return consumedBytesSinceLastGlobalRtDivSync.getOrDefault(key, 0L);
+  }
+
+  public void addConsumedBytesSinceLastGlobalRtDivSync(String key, long bytes) {
+    if (bytes <= 0) {
+      return;
+    }
+    consumedBytesSinceLastGlobalRtDivSync.merge(key, bytes, Long::sum);
+  }
+
+  public void resetConsumedBytesSinceLastGlobalRtDivSync(String key) {
+    consumedBytesSinceLastGlobalRtDivSync.put(key, 0L);
+  }
+
   public void setLeaderFollowerState(LeaderFollowerStateType state) {
     this.leaderFollowerState = state;
   }

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

@@ -311,9 +311,6 @@ public abstract class StoreIngestionTask implements Runnable, Closeable {
    * to local VT. This will also be used to send DIV snapshots to the drainer to persist the VT + RT DIV on-disk.
    */
   protected final DataIntegrityValidator consumerDiv;
-  /** Map of (RT broker URL | VT name) to the total bytes consumed by ConsumptionTask since the last Global RT DIV sync */
-  // TODO: clear it out when the sync is done
-  protected final VeniceConcurrentHashMap<String, Long> consumedBytesSinceLastSync;
   protected final HostLevelIngestionStats hostLevelIngestionStats;
   protected final AggVersionedDIVStats versionedDIVStats;
   protected final AggVersionedIngestionStats versionedIngestionStats;
@@ -535,7 +532,6 @@ public StoreIngestionTask(
         pubSubContext.getPubSubPositionDeserializer(),
         DISABLED,
         producerStateMaxAgeMs);
-    this.consumedBytesSinceLastSync = new VeniceConcurrentHashMap<>();
     this.ingestionTaskName = String.format(CONSUMER_TASK_ID_FORMAT, kafkaVersionTopic);
     this.readOnlyForBatchOnlyStoreEnabled = storeVersionConfig.isReadOnlyForBatchOnlyStoreEnabled();
     this.hostLevelIngestionStats = builder.getIngestionStats().getStoreStats(storeName);
@@ -1494,7 +1490,7 @@ protected void produceToStoreBufferServiceOrKafka(
       PubSubTopic topic = topicPartition.getPubSubTopic();
       if (isGlobalRtDivEnabled() && (versionTopic.equals(topic) || topic.isRealTime())) {
         String consumedBytesKey = versionTopic.equals(topic) ? versionTopic.getName() : kafkaUrl;
-        consumedBytesSinceLastSync.compute(consumedBytesKey, (k, v) -> (v == null) ? recordSize : v + recordSize);
+        partitionConsumptionState.addConsumedBytesSinceLastGlobalRtDivSync(consumedBytesKey, recordSize);
       }
     }
 
@@ -1582,7 +1578,7 @@ protected void produceToStoreBufferServiceOrKafkaInBatch(
             linkBackManifestFromTransientRecord(processedRecord, partitionConsumptionState);
           }
 
-          totalBytesRead += handleSingleMessage(
+          int recordSize = handleSingleMessage(
               processedRecord,
               topicPartition,
               partitionConsumptionState,
@@ -1591,6 +1587,11 @@ protected void produceToStoreBufferServiceOrKafkaInBatch(
               beforeProcessingPerRecordTimestampNs,
               beforeProcessingBatchRecordsTimestampMs,
               elapsedTimeForPuttingIntoQueue);
+          totalBytesRead += recordSize;
+          // Batch path only handles RT messages (guaranteed by isAllMessagesFromRTTopic), so key by kafkaUrl.
+          if (isGlobalRtDivEnabled()) {
+            partitionConsumptionState.addConsumedBytesSinceLastGlobalRtDivSync(kafkaUrl, recordSize);
+          }
 
           // Only track keys that were actually produced (not ignored by DCR).
           // Ignored records don't call setChunkingInfo, so the transient record's
@@ -3372,7 +3373,7 @@ boolean shouldSendGlobalRtDiv(DefaultPubSubMessage record, PartitionConsumptionS
       return false;
     }
     final long syncBytesInterval = getSyncBytesInterval(pcs);
-    return syncBytesInterval > 0 && (getConsumedBytesSinceLastSync().getOrDefault(brokerUrl, 0L) >= syncBytesInterval);
+    return syncBytesInterval > 0 && (pcs.getConsumedBytesSinceLastGlobalRtDivSync(brokerUrl) >= syncBytesInterval);
   }
 
   abstract void syncOffsetFromSnapshotIfNeeded(DefaultPubSubMessage record, PubSubTopicPartition topicPartition);
@@ -5672,10 +5673,6 @@ boolean isDaVinciClient() {
     return isDaVinciClient;
   }
 
-  VeniceConcurrentHashMap<String, Long> getConsumedBytesSinceLastSync() {
-    return consumedBytesSinceLastSync; // mainly for unit test mocks
-  }
-
   boolean isGlobalRtDivEnabled() {
     return isGlobalRtDivEnabled;
   }

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

@@ -859,8 +859,6 @@ public void testShouldSyncOffsetFromSnapshot() throws InterruptedException {
     // Stub early so size-based branch can call getVersionTopic().getName()
     PubSubTopic versionTopic = TOPIC_REPOSITORY.getTopic("test-topic_v1");
     doReturn(versionTopic).when(mockIngestionTask).getVersionTopic();
-    VeniceConcurrentHashMap<String, Long> consumedBytesSinceLastSync = new VeniceConcurrentHashMap<>();
-    doReturn(consumedBytesSinceLastSync).when(mockIngestionTask).getConsumedBytesSinceLastSync();
 
     // Set up Global RT DIV message
     final DefaultPubSubMessage globalRtDivMessage = getMockMessage(1).getMessage();
@@ -913,15 +911,18 @@ public void testShouldSyncOffsetFromSnapshot() throws InterruptedException {
     doReturn(false).when(regularMockKey).isControlMessage();
 
     // Test case 1: When VT consumed bytes since last sync is less than 2*syncBytesInterval
-    consumedBytesSinceLastSync.put(versionTopic.getName(), 1500L);
+    doReturn(1500L).when(mockPartitionConsumptionState)
+        .getConsumedBytesSinceLastGlobalRtDivSync(versionTopic.getName());
     assertFalse(mockIngestionTask.shouldSyncOffsetFromSnapshot(regularMessage, mockPartitionConsumptionState));
 
     // Test case 2: When VT consumed bytes since last sync is equal to 2*syncBytesInterval
-    consumedBytesSinceLastSync.put(versionTopic.getName(), 2000L);
+    doReturn(2000L).when(mockPartitionConsumptionState)
+        .getConsumedBytesSinceLastGlobalRtDivSync(versionTopic.getName());
     assertTrue(mockIngestionTask.shouldSyncOffsetFromSnapshot(regularMessage, mockPartitionConsumptionState));
 
     // Test case 3: When VT consumed bytes since last sync is greater than 2*syncBytesInterval
-    consumedBytesSinceLastSync.put(versionTopic.getName(), 2500L);
+    doReturn(2500L).when(mockPartitionConsumptionState)
+        .getConsumedBytesSinceLastGlobalRtDivSync(versionTopic.getName());
     assertTrue(mockIngestionTask.shouldSyncOffsetFromSnapshot(regularMessage, mockPartitionConsumptionState));
 
     // Test case 4: When syncBytesInterval is 0 (disabled)
@@ -973,8 +974,6 @@ public void testSyncOffsetFromSnapshotIfNeededSkipsWhenLcvpIsEarliest() throws I
     doReturn(ApacheKafkaOffsetPosition.of(10L)).when(snapshotReady).getLatestConsumedVtPosition();
     doReturn(Collections.singletonMap("producerGuid", new Object())).when(snapshotReady).getPartitionStates(any());
     doReturn(snapshotReady).when(mockConsumerDiv).cloneVtProducerStates(anyInt(), anyBoolean(), anyLong());
-    VeniceConcurrentHashMap<String, Long> consumedBytes = new VeniceConcurrentHashMap<>();
-    doReturn(consumedBytes).when(leaderFollowerStoreIngestionTask).getConsumedBytesSinceLastSync();
 
     leaderFollowerStoreIngestionTask.syncOffsetFromSnapshotIfNeeded(mockRecord, versionTopicPartition);
     verify(mockStoreBufferService, times(1)).execSyncOffsetFromSnapshotAsync(any(), any(), any(), any());

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

@@ -60,6 +60,7 @@
 import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.atMost;
+import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doCallRealMethod;
 import static org.mockito.Mockito.doNothing;
@@ -6101,14 +6102,16 @@ public void testShouldSendGlobalRtDiv(boolean isGlobalRtDivEnabled) {
     doReturn(key).when(message).getKey();
     PartitionConsumptionState pcs = mock(PartitionConsumptionState.class);
     doReturn(100L).when(pcs).getProcessedRecordSizeSinceLastSync(); // just needs to be greater than syncBytesInterval
-    VeniceConcurrentHashMap<String, Long> lastProcessedMap = new VeniceConcurrentHashMap<>();
-    doReturn(lastProcessedMap).when(storeIngestionTask).getConsumedBytesSinceLastSync();
+    // Default: 0 bytes consumed, so shouldSendGlobalRtDiv returns false regardless of host
+    doReturn(0L).when(pcs).getConsumedBytesSinceLastGlobalRtDivSync(any());
 
-    // Two sanity tests: empty map should not cause divide by zero, and host not present in map should return false
+    // Two sanity tests: 0 bytes should not trigger sync, and an untracked host key returning 0 bytes should return
+    // false
     storeIngestionTask.shouldSendGlobalRtDiv(message, pcs, brokerUrl);
     assertFalse(storeIngestionTask.shouldSendGlobalRtDiv(message, pcs, "fakehost:5678"));
 
-    lastProcessedMap.put(brokerUrl, 100L); // just needs to be greater than syncBytesInterval
+    doReturn(100L).when(pcs).getConsumedBytesSinceLastGlobalRtDivSync(brokerUrl); // just needs to be >=
+                                                                                  // syncBytesInterval
     boolean shouldSendGlobalRtDiv = storeIngestionTask.shouldSendGlobalRtDiv(message, pcs, brokerUrl);
     boolean shouldSyncOffset = storeIngestionTask.shouldSyncOffset(pcs, message, null);
 
@@ -6121,12 +6124,12 @@ public void testShouldSendGlobalRtDiv(boolean isGlobalRtDivEnabled) {
   }
 
   /**
-   * Verifies that {@link StoreIngestionTask#produceToStoreBufferServiceOrKafka} populates
-   * {@link StoreIngestionTask#consumedBytesSinceLastSync} only for the local VT (keyed by VT name)
-   * and RT topics (keyed by broker URL). Remote VTs must be excluded from the map entirely.
+   * Verifies that {@link StoreIngestionTask#produceToStoreBufferServiceOrKafka} calls
+   * {@link PartitionConsumptionState#addConsumedBytesSinceLastGlobalRtDivSync} only for the local VT
+   * (keyed by VT name) and RT topics (keyed by broker URL). Remote VTs must be excluded entirely.
    */
   @Test
-  public void testConsumedBytesSinceLastSyncTracking() throws Exception {
+  public void testConsumedBytesSinceLastGlobalRtDivSyncTracking() throws Exception {
     String storeName = "test-store";
     PubSubTopic localVt = pubSubTopicRepository.getTopic(Version.composeKafkaTopic(storeName, 1));
     PubSubTopic rtTopic = pubSubTopicRepository.getTopic(storeName + "_rt");
@@ -6143,17 +6146,16 @@ public void testConsumedBytesSinceLastSyncTracking() throws Exception {
     doReturn(StoreIngestionTask.DelegateConsumerRecordResult.SKIPPED_MESSAGE).when(sit)
         .delegateConsumerRecord(any(), anyInt(), any(), anyInt(), anyLong(), anyLong());
 
-    VeniceConcurrentHashMap<String, Long> consumedBytesMap = new VeniceConcurrentHashMap<>();
+    PartitionConsumptionState pcs = mock(PartitionConsumptionState.class);
     VeniceConcurrentHashMap<Integer, PartitionConsumptionState> pcsMap = new VeniceConcurrentHashMap<>();
-    pcsMap.put(partition, mock(PartitionConsumptionState.class));
+    pcsMap.put(partition, pcs);
 
     for (String fieldName: new String[] { "isActiveActiveReplicationEnabled", "isWriteComputationEnabled" }) {
       Field f = StoreIngestionTask.class.getDeclaredField(fieldName);
       f.setAccessible(true);
       f.set(sit, false);
     }
     for (Object[] entry: new Object[][] { { "versionTopic", localVt },
-        { "consumedBytesSinceLastSync", consumedBytesMap },
         { "storageUtilizationManager", mock(StorageUtilizationManager.class) },
         { "partitionConsumptionStateMap", pcsMap } }) {
       Field f = StoreIngestionTask.class.getDeclaredField((String) entry[0]);
@@ -6170,19 +6172,19 @@ public void testConsumedBytesSinceLastSyncTracking() throws Exception {
 
     // Local VT: keyed by VT name
     sit.produceToStoreBufferServiceOrKafka(records, new PubSubTopicPartitionImpl(localVt, partition), kafkaUrl, 0);
-    assertTrue(consumedBytesMap.containsKey(localVt.getName()), "Local VT should be tracked by VT name");
-    assertFalse(consumedBytesMap.containsKey(kafkaUrl));
+    verify(pcs).addConsumedBytesSinceLastGlobalRtDivSync(eq(localVt.getName()), anyLong());
+    verify(pcs, never()).addConsumedBytesSinceLastGlobalRtDivSync(eq(kafkaUrl), anyLong());
 
     // RT topic: keyed by kafkaUrl
-    consumedBytesMap.clear();
+    clearInvocations(pcs);
     sit.produceToStoreBufferServiceOrKafka(records, new PubSubTopicPartitionImpl(rtTopic, partition), kafkaUrl, 0);
-    assertTrue(consumedBytesMap.containsKey(kafkaUrl), "RT topic should be tracked by kafkaUrl");
-    assertFalse(consumedBytesMap.containsKey(rtTopic.getName()));
+    verify(pcs).addConsumedBytesSinceLastGlobalRtDivSync(eq(kafkaUrl), anyLong());
+    verify(pcs, never()).addConsumedBytesSinceLastGlobalRtDivSync(eq(rtTopic.getName()), anyLong());
 
     // Remote VT: excluded entirely
-    consumedBytesMap.clear();
+    clearInvocations(pcs);
     sit.produceToStoreBufferServiceOrKafka(records, new PubSubTopicPartitionImpl(remoteVt, partition), kafkaUrl, 0);
-    assertTrue(consumedBytesMap.isEmpty(), "Remote VT should not be tracked");
+    verify(pcs, never()).addConsumedBytesSinceLastGlobalRtDivSync(any(), anyLong());
   }
 
   /**