[dvc][common][test] leader complete status to follower part 1: leader produce HB and SOS to local VT with new header

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

@@ -61,6 +61,7 @@
 import com.linkedin.venice.utils.Utils;
 import com.linkedin.venice.utils.lazy.Lazy;
 import com.linkedin.venice.writer.ChunkAwareCallback;
+import com.linkedin.venice.writer.LeaderCompleteState;
 import com.linkedin.venice.writer.LeaderMetadataWrapper;
 import com.linkedin.venice.writer.VeniceWriter;
 import com.linkedin.venice.writer.VeniceWriterFactory;
@@ -1676,7 +1677,7 @@ protected void reportIfCatchUpVersionTopicOffset(PartitionConsumptionState pcs)
         if (isCurrentVersion.getAsBoolean()) {
           amplificationFactorAdapter
               .executePartitionConsumptionState(pcs.getUserPartition(), PartitionConsumptionState::lagHasCaughtUp);
-          statusReportAdapter.reportCompleted(pcs, true);
+          reportCompleted(pcs, true);
         }
       }
     }
@@ -2087,7 +2088,15 @@ protected DelegateConsumerRecordResult delegateConsumerRecord(
                 LeaderMetadataWrapper leaderMetadataWrapper =
                     new LeaderMetadataWrapper(consumerRecord.getOffset(), kafkaClusterId);
                 PubSubTopicPartition topicPartition = new PubSubTopicPartitionImpl(getVersionTopic(), subPartition);
-                veniceWriter.get().sendHeartbeat(topicPartition, callback, leaderMetadataWrapper);
+                // Leaders forward HB SOS message to local VT with updated LeaderCompleteState header
+                veniceWriter.get()
+                    .sendHeartbeat(
+                        topicPartition,
+                        callback,
+                        leaderMetadataWrapper,
+                        true,
+                        LeaderCompleteState.getLeaderCompleteState(partitionConsumptionState.isCompletionReported()),
+                        consumerRecord.getValue().producerMetadata.messageTimestamp); // original producers timestamp
               } else {
                 /**
                  * Based on current design handling this case (specially EOS) is tricky as we don't produce the SOS/EOS
@@ -3194,4 +3203,21 @@ protected void maybeSendIngestionHeartbeat() {
     }
     lastSendIngestionHeartbeatTimestamp = currentTimestamp;
   }
+
+  /**
+   * Once leader is marked completed, immediately send a heart beat message to the local VT such that
+   * followers don't have to wait till the periodic heartbeat to know that the leader is completed
+   */
+  void reportCompleted(PartitionConsumptionState partitionConsumptionState, boolean forceCompletion) {
+    super.reportCompleted(partitionConsumptionState, forceCompletion);
+    if (partitionConsumptionState.getLeaderFollowerState().equals(LeaderFollowerStateType.LEADER)) {
+      veniceWriter.get()
+          .sendHeartbeat(
+              new PubSubTopicPartitionImpl(versionTopic, partitionConsumptionState.getPartition()),
+              null,
+              DEFAULT_LEADER_METADATA_WRAPPER,
+              true,
+              LeaderCompleteState.LEADER_COMPLETED);
+    }
+  }
 }

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

@@ -1694,7 +1694,7 @@ private void checkConsumptionStateWhenStart(
               amplificationFactorAdapter.executePartitionConsumptionState(
                   newPartitionConsumptionState.getUserPartition(),
                   PartitionConsumptionState::lagHasCaughtUp);
-              statusReportAdapter.reportCompleted(newPartitionConsumptionState, true);
+              reportCompleted(newPartitionConsumptionState, true);
               isCompletedReport = true;
             }
             // Clear offset lag in metadata, it is only used in restart.
@@ -3473,9 +3473,7 @@ private ReadyToServeCheck getDefaultReadyToServeChecker() {
                 partition,
                 partitionConsumptionState.getLatestProcessedLocalVersionTopicOffset());
           } else {
-            statusReportAdapter.reportCompleted(partitionConsumptionState);
-            LOGGER.info("{} Partition {} is ready to serve", consumerTaskId, partition);
-
+            reportCompleted(partitionConsumptionState);
             warmupSchemaCache(store);
           }
           if (suppressLiveUpdates) {
@@ -3495,6 +3493,15 @@ private ReadyToServeCheck getDefaultReadyToServeChecker() {
     };
   }
 
+  void reportCompleted(PartitionConsumptionState partitionConsumptionState) {
+    reportCompleted(partitionConsumptionState, false);
+  }
+
+  void reportCompleted(PartitionConsumptionState partitionConsumptionState, boolean forceCompletion) {
+    statusReportAdapter.reportCompleted(partitionConsumptionState, forceCompletion);
+    LOGGER.info("{} Partition {} is ready to serve", consumerTaskId, partitionConsumptionState.getPartition());
+  }
+
   /**
    * Try to warm-up the schema repo cache before reporting completion as new value schema could cause latency degradation
    * while trying to compile it in the read-path.

internal/venice-common/src/main/java/com/linkedin/venice/pubsub/ImmutablePubSubMessage.java

@@ -1,6 +1,7 @@
 package com.linkedin.venice.pubsub;
 
 import com.linkedin.venice.pubsub.api.PubSubMessage;
+import com.linkedin.venice.pubsub.api.PubSubMessageHeaders;
 import com.linkedin.venice.pubsub.api.PubSubTopicPartition;
 import java.util.Objects;
 
@@ -13,19 +14,33 @@ public class ImmutablePubSubMessage<K, V> implements PubSubMessage<K, V, Long> {
   private final long timestamp;
   private final int payloadSize;
 
+  private final PubSubMessageHeaders pubSubMessageHeaders;
+
   public ImmutablePubSubMessage(
       K key,
       V value,
       PubSubTopicPartition topicPartition,
       long offset,
       long timestamp,
       int payloadSize) {
+    this(key, value, topicPartition, offset, timestamp, payloadSize, null);
+  }
+
+  public ImmutablePubSubMessage(
+      K key,
+      V value,
+      PubSubTopicPartition topicPartition,
+      long offset,
+      long timestamp,
+      int payloadSize,
+      PubSubMessageHeaders pubSubMessageHeaders) {
     this.key = key;
     this.value = value;
     this.topicPartition = Objects.requireNonNull(topicPartition);
     this.offset = offset;
     this.timestamp = timestamp;
     this.payloadSize = payloadSize;
+    this.pubSubMessageHeaders = pubSubMessageHeaders;
   }
 
   @Override
@@ -62,4 +77,9 @@ public int getPayloadSize() {
   public boolean isEndOfBootstrap() {
     return false;
   }
+
+  @Override
+  public PubSubMessageHeaders getPubSubMessageHeaders() {
+    return pubSubMessageHeaders;
+  }
 }

internal/venice-common/src/main/java/com/linkedin/venice/pubsub/api/PubSubMessage.java

@@ -1,5 +1,8 @@
 package com.linkedin.venice.pubsub.api;
 
+import static com.linkedin.venice.writer.VeniceWriter.EMPTY_MSG_HEADERS;
+
+
 public interface PubSubMessage<K, V, OFFSET> {
   /**
    * @return the key part of this message
@@ -43,4 +46,8 @@ default int getPartition() {
    * @return whether this message marks the end of bootstrap.
    */
   boolean isEndOfBootstrap();
+
+  default PubSubMessageHeaders getPubSubMessageHeaders() {
+    return EMPTY_MSG_HEADERS;
+  }
 }

internal/venice-common/src/main/java/com/linkedin/venice/pubsub/api/PubSubMessageDeserializer.java

@@ -1,5 +1,7 @@
 package com.linkedin.venice.pubsub.api;
 
+import static com.linkedin.venice.pubsub.api.PubSubMessageHeaders.VENICE_TRANSPORT_PROTOCOL_HEADER;
+
 import com.linkedin.avroutil1.compatibility.AvroCompatibilityHelper;
 import com.linkedin.venice.kafka.protocol.KafkaMessageEnvelope;
 import com.linkedin.venice.kafka.protocol.enums.MessageType;
@@ -19,9 +21,6 @@
  */
 public class PubSubMessageDeserializer {
   private static final Logger LOGGER = LogManager.getLogger(PubSubMessageDeserializer.class);
-
-  public static final String VENICE_TRANSPORT_PROTOCOL_HEADER = "vtp";
-
   private final KafkaKeySerializer keySerializer = new KafkaKeySerializer();
   private final KafkaValueSerializer valueSerializer;
   private final ObjectPool<KafkaMessageEnvelope> putEnvelopePool;
@@ -59,6 +58,8 @@ public PubSubMessage<KafkaKey, KafkaMessageEnvelope, Long> deserialize(
     KafkaMessageEnvelope value = null;
     if (key.isControlMessage()) {
       for (PubSubMessageHeader header: headers.toList()) {
+        // only process VENICE_TRANSPORT_PROTOCOL_HEADER here. Other headers will be stored in
+        // ImmutablePubSubMessage and used down the ingestion path later
         if (header.key().equals(VENICE_TRANSPORT_PROTOCOL_HEADER)) {
           try {
             Schema providedProtocolSchema = AvroCompatibilityHelper.parse(new String(header.value()));
@@ -85,7 +86,8 @@ public PubSubMessage<KafkaKey, KafkaMessageEnvelope, Long> deserialize(
         topicPartition,
         position,
         timestamp,
-        keyBytes.length + valueBytes.length);
+        keyBytes.length + valueBytes.length,
+        headers);
   }
 
   private KafkaMessageEnvelope getEnvelope(byte keyHeaderByte) {

internal/venice-common/src/main/java/com/linkedin/venice/pubsub/api/PubSubMessageHeaders.java

@@ -17,6 +17,10 @@ public class PubSubMessageHeaders {
   // from the beginning.
   private final Map<String, PubSubMessageHeader> headers = new LinkedHashMap<>();
 
+  public static final String VENICE_TRANSPORT_PROTOCOL_HEADER = "vtp";
+  /** Header to denote whether the leader is completed or not */
+  public static final String VENICE_LEADER_COMPLETION_STATE_HEADER = "lcs";
+
   public PubSubMessageHeaders add(PubSubMessageHeader header) {
     headers.put(header.key(), header);
     return this;

internal/venice-common/src/main/java/com/linkedin/venice/writer/LeaderCompleteState.java

@@ -0,0 +1,53 @@
+package com.linkedin.venice.writer;
+
+import com.linkedin.venice.exceptions.VeniceException;
+import com.linkedin.venice.utils.EnumUtils;
+import com.linkedin.venice.utils.VeniceEnumValue;
+
+
+/**
+ * Completion state of the leader partition.
+ */
+public enum LeaderCompleteState implements VeniceEnumValue {
+  /**
+   * Leader partition is not marked completed yet
+   */
+  LEADER_NOT_COMPLETED(0),
+  /**
+   * Leader partition is marked completed
+   */
+  LEADER_COMPLETED(1),
+  /**
+   * Default state for the standby to know that the feature is not supported, to
+   * not consider leader's completion status for its completion.
+   */
+  NOT_SUPPORTED(2);
+
+  private final int value;
+  private static final LeaderCompleteState[] TYPES_ARRAY = EnumUtils.getEnumValuesArray(LeaderCompleteState.class);
+
+  LeaderCompleteState(int value) {
+    this.value = value;
+  }
+
+  public static LeaderCompleteState getLeaderCompleteState(boolean isLeaderCompleted) {
+    if (isLeaderCompleted) {
+      return LEADER_COMPLETED;
+    } else {
+      return LEADER_NOT_COMPLETED;
+    }
+  }
+
+  public static LeaderCompleteState valueOf(int value) {
+    try {
+      return TYPES_ARRAY[value];
+    } catch (IndexOutOfBoundsException e) {
+      throw new VeniceException("Invalid LeaderCompleteState: " + value);
+    }
+  }
+
+  @Override
+  public int getValue() {
+    return value;
+  }
+}