Added assembled_record_rmd_size_in_bytes to track the RMD size and refined the unit test. 🖇️

Author: KaiSernLim

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

@@ -108,6 +108,7 @@
 import com.linkedin.venice.utils.VeniceProperties;
 import com.linkedin.venice.utils.concurrent.VeniceConcurrentHashMap;
 import com.linkedin.venice.utils.lazy.Lazy;
+import com.linkedin.venice.writer.VeniceWriter;
 import it.unimi.dsi.fastutil.objects.Object2IntMap;
 import java.io.Closeable;
 import java.io.IOException;
@@ -2922,13 +2923,11 @@ private int internalProcessConsumerRecord(
     return sizeOfPersistedData;
   }
 
-  protected void recordWriterStats(
+  protected abstract void recordWriterStats(
       long consumerTimestampMs,
       long producerBrokerLatencyMs,
       long brokerConsumerLatencyMs,
-      PartitionConsumptionState partitionConsumptionState) {
-
-  }
+      PartitionConsumptionState partitionConsumptionState);
 
   /**
    * Message validation using DIV. Leaders should pass in the validator instance from {@link LeaderFollowerStoreIngestionTask};
@@ -3205,10 +3204,14 @@ private int processKafkaDataMessage(
             ChunkedValueManifest chunkedValueManifest = manifestSerializer.deserialize(
                 ByteUtils.extractByteArray(put.getPutValue()), // must be done before recordTransformer changes putValue
                 AvroProtocolDefinition.CHUNKED_VALUE_MANIFEST.getCurrentProtocolVersion());
-            hostLevelIngestionStats.recordAssembledValueSize(chunkedValueManifest.getSize(), currentTimeMs);
-
-            // Ideally, the record size calculation should include rmd, but it's too difficult to keep track of it
-            recordAssembledRecordSizeRatio(keyLen + chunkedValueManifest.getSize(), currentTimeMs);
+            ByteBuffer rmdPayload = put.getReplicationMetadataPayload();
+            boolean isValueManifest = rmdPayload != null && rmdPayload.equals(VeniceWriter.EMPTY_BYTE_BUFFER);
+            if (isValueManifest) {
+              hostLevelIngestionStats.recordAssembledValueSize(chunkedValueManifest.getSize(), currentTimeMs);
+              recordAssembledRecordSizeRatio(keyLen + chunkedValueManifest.getSize(), currentTimeMs);
+            } else { // the manifest pertains to a chunked rmd
+              hostLevelIngestionStats.recordAssembledRmdSize(chunkedValueManifest.getSize(), currentTimeMs);
+            }
           } catch (VeniceException | IllegalArgumentException | AvroRuntimeException e) {
             LOGGER.error("Failed to deserialize ChunkedValueManifest to record assembled value size", e);
           }

clients/da-vinci-client/src/main/java/com/linkedin/davinci/stats/HostLevelIngestionStats.java

@@ -33,6 +33,7 @@
  */
 public class HostLevelIngestionStats extends AbstractVeniceStats {
   public static final String ASSEMBLED_RECORD_SIZE_RATIO = "assembled_record_size_ratio";
+  public static final String ASSEMBLED_RECORD_RMD_SIZE_IN_BYTES = "assembled_record_rmd_size_in_bytes";
   public static final String ASSEMBLED_RECORD_VALUE_SIZE_IN_BYTES = "assembled_record_value_size_in_bytes";
 
   // The aggregated bytes ingested rate for the entire host
@@ -51,6 +52,7 @@ public class HostLevelIngestionStats extends AbstractVeniceStats {
   private final Sensor keySizeSensor;
   private final Sensor valueSizeSensor;
   private final Sensor assembledValueSizeSensor;
+  private final Sensor assembledRmdSizeSensor;
   private final Sensor assembledRecordSizeRatioSensor;
   private final Sensor unexpectedMessageSensor;
   private final Sensor inconsistentStoreMetadataSensor;
@@ -282,6 +284,8 @@ public HostLevelIngestionStats(
 
     this.assembledValueSizeSensor = registerSensor(ASSEMBLED_RECORD_VALUE_SIZE_IN_BYTES, avgAndMax());
 
+    this.assembledRmdSizeSensor = registerSensor(ASSEMBLED_RECORD_RMD_SIZE_IN_BYTES, avgAndMax());
+
     this.assembledRecordSizeRatioSensor = registerSensor(ASSEMBLED_RECORD_SIZE_RATIO, new Max());
 
     String viewTimerSensorName = "total_view_writer_latency";
@@ -471,6 +475,10 @@ public void recordAssembledValueSize(long bytes, long currentTimeMs) {
     assembledValueSizeSensor.record(bytes, currentTimeMs);
   }
 
+  public void recordAssembledRmdSize(long bytes, long currentTimeMs) {
+    assembledRmdSizeSensor.record(bytes, currentTimeMs);
+  }
+
   public void recordAssembledRecordSizeRatio(double ratio, long currentTimeMs) {
     assembledRecordSizeRatioSensor.record(ratio, currentTimeMs);
   }

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

@@ -4368,11 +4368,10 @@ public void testStoreIngestionRecordTransformerError(AAConfig aaConfig) throws E
 
   @DataProvider
   public static Object[][] testAssembledValueSizeProvider() {
-    Set<Integer> schemaIdSet = new HashSet<>();
-    for (AvroProtocolDefinition avroProtocolDefinition: AvroProtocolDefinition.values()) {
-      avroProtocolDefinition.currentProtocolVersion.ifPresent(schemaIdSet::add);
-    }
-    return DataProviderUtils.allPermutationGenerator(AAConfig.values(), schemaIdSet.toArray());
+    Object[] testSchemaIds =
+        { VeniceWriter.VENICE_DEFAULT_VALUE_SCHEMA_ID, AvroProtocolDefinition.CHUNK.getCurrentProtocolVersion(),
+            AvroProtocolDefinition.CHUNKED_VALUE_MANIFEST.getCurrentProtocolVersion() };
+    return DataProviderUtils.allPermutationGenerator(AAConfig.values(), testSchemaIds, DataProviderUtils.BOOLEAN);
   }
 
   /**
@@ -4383,7 +4382,7 @@ public static Object[][] testAssembledValueSizeProvider() {
    * @throws Exception
    */
   @Test(dataProvider = "testAssembledValueSizeProvider")
-  public void testAssembledValueSizeSensor(AAConfig aaConfig, int testSchemaId) throws Exception {
+  public void testAssembledValueSizeSensor(AAConfig aaConfig, int testSchemaId, boolean hasRmd) throws Exception {
     int numChunks = 10;
     long expectedRecordSize = (long) numChunks * ChunkingTestUtils.CHUNK_LENGTH;
     PubSubTopicPartition tp = new PubSubTopicPartitionImpl(pubSubTopic, PARTITION_FOO);
@@ -4395,8 +4394,6 @@ public void testAssembledValueSizeSensor(AAConfig aaConfig, int testSchemaId) th
         ChunkingTestUtils.createChunkValueManifestRecord(putKeyFoo, messages.get(0), numChunks, tp);
     messages.add(manifestMessage);
 
-    // The only expected real-life case should be when schemaId values are chunk=-10 and manifest=-20
-    boolean useRealSchemaId = testSchemaId == AvroProtocolDefinition.CHUNKED_VALUE_MANIFEST.getCurrentProtocolVersion();
     runTest(Collections.singleton(PARTITION_FOO), () -> {
       TestUtils.waitForNonDeterministicAssertion(
           5,
@@ -4406,8 +4403,9 @@ public void testAssembledValueSizeSensor(AAConfig aaConfig, int testSchemaId) th
       for (PubSubMessage<KafkaKey, KafkaMessageEnvelope, Long> message: messages) {
         try {
           Put put = (Put) message.getValue().getPayloadUnion();
-          if (!useRealSchemaId) {
-            put.schemaId = testSchemaId;
+          if (put.schemaId == AvroProtocolDefinition.CHUNKED_VALUE_MANIFEST.getCurrentProtocolVersion()) {
+            put.schemaId = testSchemaId; // set manifest schemaId to testSchemaId to see if metrics are still recorded
+            put.replicationMetadataPayload = (hasRmd) ? ByteBuffer.allocate(10) : VeniceWriter.EMPTY_BYTE_BUFFER;
           }
           LeaderProducedRecordContext leaderProducedRecordContext = mock(LeaderProducedRecordContext.class);
           when(leaderProducedRecordContext.getMessageType()).thenReturn(MessageType.PUT);
@@ -4426,13 +4424,25 @@ public void testAssembledValueSizeSensor(AAConfig aaConfig, int testSchemaId) th
         }
       }
 
-      // Verify that the size of the assembled record was recorded in the metrics only if schemaId=-20
+      // Verify that the assembled record metrics are only recorded if schemaId=-20 which indicates a manifest
       HostLevelIngestionStats hostLevelIngestionStats = storeIngestionTaskUnderTest.hostLevelIngestionStats;
-      if (useRealSchemaId) {
-        verify(hostLevelIngestionStats, timeout(1000)).recordAssembledValueSize(eq(expectedRecordSize), anyLong());
-        verify(hostLevelIngestionStats, timeout(1000).times(1)).recordAssembledValueSize(anyLong(), anyLong());
+      if (testSchemaId == AvroProtocolDefinition.CHUNKED_VALUE_MANIFEST.getCurrentProtocolVersion()) {
+        if (hasRmd) { // should have metrics for rmd but not value size
+          verify(hostLevelIngestionStats, timeout(1000)).recordAssembledRmdSize(eq(expectedRecordSize), anyLong());
+          verify(hostLevelIngestionStats, timeout(1000).times(1)).recordAssembledRmdSize(anyLong(), anyLong());
+          verify(hostLevelIngestionStats, times(0)).recordAssembledValueSize(anyLong(), anyLong());
+          verify(hostLevelIngestionStats, times(0)).recordAssembledRecordSizeRatio(anyDouble(), anyLong());
+        } else { // should have metrics for value but not rmd size
+          verify(hostLevelIngestionStats, timeout(1000)).recordAssembledValueSize(eq(expectedRecordSize), anyLong());
+          verify(hostLevelIngestionStats, timeout(1000).times(1)).recordAssembledValueSize(anyLong(), anyLong());
+          verify(hostLevelIngestionStats, timeout(1000).times(1))
+              .recordAssembledRecordSizeRatio(anyDouble(), anyLong());
+          verify(hostLevelIngestionStats, times(0)).recordAssembledRmdSize(anyLong(), anyLong());
+        }
       } else {
         verify(hostLevelIngestionStats, times(0)).recordAssembledValueSize(anyLong(), anyLong());
+        verify(hostLevelIngestionStats, times(0)).recordAssembledRmdSize(anyLong(), anyLong());
+        verify(hostLevelIngestionStats, times(0)).recordAssembledRecordSizeRatio(anyDouble(), anyLong());
       }
     }, aaConfig);
   }

internal/venice-test-common/src/main/java/com/linkedin/venice/utils/ChunkingTestUtils.java

@@ -19,6 +19,7 @@
 import com.linkedin.venice.storage.protocol.ChunkId;
 import com.linkedin.venice.storage.protocol.ChunkedKeySuffix;
 import com.linkedin.venice.storage.protocol.ChunkedValueManifest;
+import com.linkedin.venice.writer.VeniceWriter;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 
@@ -94,7 +95,7 @@ public static PubSubMessage<KafkaKey, KafkaMessageEnvelope, Long> createChunkedR
     put.schemaId = AvroProtocolDefinition.CHUNK.getCurrentProtocolVersion();
     byte[] valueBytes = createChunkBytes(chunkIndex * CHUNK_LENGTH, CHUNK_LENGTH);
     put.putValue = prependSchemaId(valueBytes, put.schemaId);
-    put.replicationMetadataPayload = ByteBuffer.allocate(10);
+    put.replicationMetadataPayload = VeniceWriter.EMPTY_BYTE_BUFFER;
     messageEnvelope.payloadUnion = put;
     return new ImmutablePubSubMessage<>(kafkaKey, messageEnvelope, pubSubTopicPartition, newOffset, 0, 20);
   }
@@ -123,7 +124,7 @@ public static PubSubMessage<KafkaKey, KafkaMessageEnvelope, Long> createChunkVal
     Put put = new Put();
     put.schemaId = AvroProtocolDefinition.CHUNKED_VALUE_MANIFEST.getCurrentProtocolVersion();
     put.putValue = prependSchemaId(valueBytes, put.schemaId);
-    put.replicationMetadataPayload = ByteBuffer.allocate(10);
+    put.replicationMetadataPayload = VeniceWriter.EMPTY_BYTE_BUFFER;
     messageEnvelope.payloadUnion = put;
     return new ImmutablePubSubMessage<>(kafkaKey, messageEnvelope, pubSubTopicPartition, newOffset, 0, 20);
   }