Unified code paths into Segment.toProducerPartitionState(). 🤠

Author: KaiSernLim

clients/da-vinci-client/src/main/java/com/linkedin/davinci/validation/PartitionTracker.java

@@ -244,35 +244,8 @@ private void updateOffsetRecord(TopicType type, GUID guid, Segment segment, Offs
     }
 
     if (state == null) {
-      state = new ProducerPartitionState();
-
-      /**
-       * The aggregates and debugInfo being stored in the {@link ProducerPartitionState} will add a bit
-       * of overhead when we checkpoint this metadata to disk, so we should be careful not to add a very
-       * large number of elements to these arbitrary collections.
-       * <p>
-       * In the case of the debugInfo, it is expected (at the time of writing this comment) that all
-       * partitions produced by the same producer GUID would have the same debug values (though nothing
-       * precludes us from having per-partition debug values in the future if there is a use case for
-       * that). It is redundant that we store the same debug values once per partition. In the future,
-       * if we want to eliminate this redundancy, we could move the per-producer debug info to another
-       * data structure, though that would increase bookkeeping complexity. This is expected to be a
-       * minor overhead, and therefore it appears to be premature to optimize this now.
-       */
-      state.aggregates = CollectionUtils.substituteEmptyMap(segment.getAggregates());
-      state.debugInfo = CollectionUtils.substituteEmptyMap(segment.getDebugInfo());
+      state = segment.toProducerPartitionState();
     }
-    state.checksumType = segment.getCheckSumType().getValue();
-    /**
-     * {@link MD5Digest#getEncodedState()} is allocating a byte array to contain the intermediate state,
-     * which is expensive. We should only invoke this closure when necessary.
-     */
-    state.checksumState = ByteBuffer.wrap(segment.getCheckSumState());
-    state.segmentNumber = segment.getSegmentNumber();
-    state.messageSequenceNumber = segment.getSequenceNumber();
-    state.messageTimestamp = segment.getLastRecordProducerTimestamp();
-    state.segmentStatus = segment.getStatus().getValue();
-    state.isRegistered = segment.isRegistered();
 
     setProducerState(offsetRecord, type, guid, state);
   }

internal/venice-common/src/main/java/com/linkedin/venice/kafka/validation/Segment.java

@@ -413,15 +413,33 @@ private Map<CharSequence, CharSequence> getDedupedDebugInfo(Map<CharSequence, Ch
 
   public ProducerPartitionState toProducerPartitionState() {
     ProducerPartitionState pps = new ProducerPartitionState();
-    pps.segmentNumber = segmentNumber;
-    pps.segmentStatus = getStatus().getValue();
-    pps.messageSequenceNumber = sequenceNumber;
-    pps.checksumState = ByteBuffer.wrap(getCheckSumState());
+    /**
+     * The aggregates and debugInfo being stored in the {@link ProducerPartitionState} will add a bit
+     * of overhead when we checkpoint this metadata to disk, so we should be careful not to add a very
+     * large number of elements to these arbitrary collections.
+     * <p>
+     * In the case of the debugInfo, it is expected (at the time of writing this comment) that all
+     * partitions produced by the same producer GUID would have the same debug values (though nothing
+     * precludes us from having per-partition debug values in the future if there is a use case for
+     * that). It is redundant that we store the same debug values once per partition. In the future,
+     * if we want to eliminate this redundancy, we could move the per-producer debug info to another
+     * data structure, though that would increase bookkeeping complexity. This is expected to be a
+     * minor overhead, and therefore it appears to be premature to optimize this now.
+     */
+    pps.aggregates = CollectionUtils.substituteEmptyMap(getAggregates());
+    pps.debugInfo = CollectionUtils.substituteEmptyMap(getDebugInfo());
     pps.checksumType = getCheckSumType().getValue();
-    pps.aggregates = aggregates;
-    pps.debugInfo = debugInfo;
-    pps.messageTimestamp = lastRecordProducerTimestamp;
-    pps.isRegistered = registered;
+    /**
+     * {@link MD5Digest#getEncodedState()} is allocating a byte array to contain the intermediate,
+     * which is expensive. We should only invoke this closure when necessary.
+     */
+    pps.checksumState = ByteBuffer.wrap(getCheckSumState());
+    pps.segmentNumber = getSegmentNumber();
+    pps.messageSequenceNumber = getSequenceNumber();
+    pps.messageTimestamp = getLastRecordProducerTimestamp();
+    pps.segmentStatus = getStatus().getValue();
+    pps.isRegistered = isRegistered();
+
     return pps;
   }
 

internal/venice-common/src/test/java/com/linkedin/venice/kafka/validation/SegmentTest.java

@@ -206,6 +206,6 @@ public void testToProducerPartitionState(CheckSumType checkSumType) {
   @Test
   public void testToProducerPartitionState_NoCheckSum() {
     Segment segment = new Segment(1, 1, 1, CheckSumType.NONE, new HashMap<>(), new HashMap<>());
-    ProducerPartitionState pps = segment.toProducerPartitionState(); // assert does not throw with checksumtype NONE
+    segment.toProducerPartitionState(); // assert that a NullPointerException is not thrown with checksumtype NONE
   }
 }