Add version swap timeout to deal with edge cases

Author: xunyin8

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

@@ -92,6 +92,11 @@ public class ChangelogClientConfig<T extends SpecificRecord> {
    * corresponding version swap messages before it can safely go to the new version to ensure data completeness.
    */
   private int totalRegionCount = 1;
+  /**
+   * Version swap timeout in milliseconds. If the version swap is not completed within this time, the consumer will swap
+   * to the new version and resume normal consumption from EOP for any incomplete partitions. Default is 30 minutes.
+   */
+  private long versionSwapTimeoutInMs = 30 * 60 * 1000;
 
   public ChangelogClientConfig(String storeName) {
     this.innerClientConfig = new ClientConfig<>(storeName);
@@ -365,6 +370,15 @@ public ChangelogClientConfig setTotalRegionCount(int totalRegionCount) {
     return this;
   }
 
+  public long getVersionSwapTimeoutInMs() {
+    return this.versionSwapTimeoutInMs;
+  }
+
+  public ChangelogClientConfig setVersionSwapTimeoutInMs(long versionSwapTimeoutInMs) {
+    this.versionSwapTimeoutInMs = versionSwapTimeoutInMs;
+    return this;
+  }
+
   public static <V extends SpecificRecord> ChangelogClientConfig<V> cloneConfig(ChangelogClientConfig<V> config) {
     ChangelogClientConfig<V> newConfig = new ChangelogClientConfig<V>().setStoreName(config.getStoreName())
         .setLocalD2ZkHosts(config.getLocalD2ZkHosts())
@@ -395,7 +409,8 @@ public static <V extends SpecificRecord> ChangelogClientConfig<V> cloneConfig(Ch
         .setIsStateful(false)
         .setVersionSwapByControlMessageEnabled(config.isVersionSwapByControlMessageEnabled())
         .setClientRegionName(config.getClientRegionName())
-        .setTotalRegionCount(config.getTotalRegionCount());
+        .setTotalRegionCount(config.getTotalRegionCount())
+        .setVersionSwapTimeoutInMs(config.getVersionSwapTimeoutInMs());
     return newConfig;
   }
 

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

@@ -13,6 +13,7 @@
 import com.linkedin.venice.pubsub.api.PubSubSymbolicPosition;
 import com.linkedin.venice.pubsub.api.PubSubTopic;
 import com.linkedin.venice.pubsub.api.PubSubTopicPartition;
+import com.linkedin.venice.utils.Time;
 import com.linkedin.venice.utils.lazy.Lazy;
 import java.util.Collection;
 import java.util.Collections;
@@ -78,6 +79,11 @@ protected VeniceAfterImageConsumerImpl(
         changelogClientConfig.isVersionSwapByControlMessageEnabled());
   }
 
+  // Intended for unit test only.
+  void setTime(Time time) {
+    this.time = time;
+  }
+
   @Override
   public Collection<PubSubMessage<K, ChangeEvent<V>, VeniceChangeCoordinate>> poll(long timeoutInMs) {
     try {
@@ -181,7 +187,8 @@ protected static void adjustSeekCheckPointsBasedOnHeartbeats(
   }
 
   /**
-   * Similar to {@link #internalSeekToEndOfPush} this can also be optimized later for a faster find.
+   * Similar to {@link #internalSeekToEndOfPush} exception in addition to finding the EOP of each partition we will also
+   * be looking for the first relevant version swap. This can also be optimized later for a faster find.
    */
   @Override
   protected CompletableFuture<Void> internalFindNewVersionCheckpoints(
@@ -195,6 +202,7 @@ protected CompletableFuture<Void> internalFindNewVersionCheckpoints(
     return CompletableFuture.supplyAsync(() -> {
       boolean lockAcquired = false;
       Map<Integer, VeniceChangeCoordinate> checkpoints = new HashMap<>();
+      Map<Integer, VeniceChangeCoordinate> eopCheckpoints = new HashMap<>();
       try {
         synchronized (internalSeekConsumer) {
           PubSubConsumerAdapter consumerAdapter = internalSeekConsumer.get();
@@ -228,17 +236,25 @@ protected CompletableFuture<Void> internalFindNewVersionCheckpoints(
                 if (message.getKey().isControlMessage()) {
                   ControlMessage controlMessage = (ControlMessage) message.getValue().getPayloadUnion();
                   ControlMessageType controlMessageType = ControlMessageType.valueOf(controlMessage);
-                  if (controlMessageType.equals(ControlMessageType.VERSION_SWAP)) {
+                  if (controlMessageType.equals(ControlMessageType.END_OF_PUSH)) {
+                    VeniceChangeCoordinate eopCoordinate = new VeniceChangeCoordinate(
+                        pubSubTopicPartition.getPubSubTopic().getName(),
+                        message.getPosition(),
+                        pubSubTopicPartition.getPartitionNumber());
+                    eopCheckpoints.put(pubSubTopicPartition.getPartitionNumber(), eopCoordinate);
+                    LOGGER.info(
+                        "Found EOP for version swap message with generation id: {} for partition: {}",
+                        generationId,
+                        pubSubTopicPartition.getPartitionNumber());
+                    // We continue to poll until we find the corresponding version swap which should be after EOP
+                  } else if (controlMessageType.equals(ControlMessageType.VERSION_SWAP)) {
                     VersionSwap versionSwap = (VersionSwap) controlMessage.getControlMessageUnion();
                     // In theory just matching the generation id and source region should be sufficient but just to be
                     // safe we will match all fields
                     if (versionSwap.getGenerationId() == generationId
                         && versionSwap.getSourceRegion().toString().equals(clientRegionName)
                         && oldVersionTopic.equals(versionSwap.getOldServingVersionTopic().toString())
                         && newVersionTopic.equals(versionSwap.getNewServingVersionTopic().toString())) {
-                      LOGGER.info(
-                          "Found corresponding version swap message for partition: {}",
-                          pubSubTopicPartition.getPartitionNumber());
                       versionSwapConsumedPerPartitionMap.put(pubSubTopicPartition.getPartitionNumber(), true);
                       VeniceChangeCoordinate coordinate = new VeniceChangeCoordinate(
                           pubSubTopicPartition.getPubSubTopic().getName(),
@@ -247,6 +263,10 @@ protected CompletableFuture<Void> internalFindNewVersionCheckpoints(
                       checkpoints.put(pubSubTopicPartition.getPartitionNumber(), coordinate);
                       // We are done with this partition
                       consumerAdapter.unSubscribe(pubSubTopicPartition);
+                      LOGGER.info(
+                          "Found corresponding version swap message with generation id: {} for partition: {}",
+                          generationId,
+                          pubSubTopicPartition.getPartitionNumber());
                       break;
                     }
                   }
@@ -264,7 +284,8 @@ protected CompletableFuture<Void> internalFindNewVersionCheckpoints(
         // old version topic yet. We can acquire the lock and update the VersionSwapMessageState.
         subscriptionLock.writeLock().lock();
         lockAcquired = true;
-        versionSwapMessageState.setNewTopicCheckpoints(checkpoints);
+        versionSwapMessageState.setNewTopicVersionSwapCheckpoints(checkpoints);
+        versionSwapMessageState.setNewTopicEOPCheckpoints(eopCheckpoints);
       } finally {
         if (lockAcquired) {
           subscriptionLock.writeLock().unlock();

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

@@ -73,6 +73,8 @@
 import com.linkedin.venice.store.rocksdb.RocksDBUtils;
 import com.linkedin.venice.utils.DaemonThreadFactory;
 import com.linkedin.venice.utils.DictionaryUtils;
+import com.linkedin.venice.utils.SystemTime;
+import com.linkedin.venice.utils.Time;
 import com.linkedin.venice.utils.Utils;
 import com.linkedin.venice.utils.VeniceProperties;
 import com.linkedin.venice.utils.concurrent.VeniceConcurrentHashMap;
@@ -182,10 +184,12 @@ public class VeniceChangelogConsumerImpl<K, V> implements VeniceChangelogConsume
   protected final boolean versionSwapByControlMessage;
   protected final String clientRegionName;
   protected final int totalRegionCount;
+  protected final long versionSwapTimeoutInMs;
   /**
    * Interaction of this field should acquire the subscriptionLock.readLock()
    */
   protected VersionSwapMessageState versionSwapMessageState = null;
+  protected Time time;
 
   public VeniceChangelogConsumerImpl(
       ChangelogClientConfig changelogClientConfig,
@@ -215,6 +219,8 @@ public VeniceChangelogConsumerImpl(
     this.pubSubMessageDeserializer = pubSubMessageDeserializer;
     this.versionSwapByControlMessage = changelogClientConfig.isVersionSwapByControlMessageEnabled();
     this.totalRegionCount = changelogClientConfig.getTotalRegionCount();
+    this.versionSwapTimeoutInMs = changelogClientConfig.getVersionSwapTimeoutInMs();
+    this.time = new SystemTime();
     if (versionSwapByControlMessage) {
       String clientRegionNameFromConfig = changelogClientConfig.getClientRegionName();
       if (clientRegionNameFromConfig.isEmpty()) {
@@ -352,7 +358,7 @@ public CompletableFuture<Void> subscribe(Set<Integer> partitions) {
     for (int partition: partitions) {
       getPartitionToBootstrapState().put(partition, false);
     }
-    subscribeTime = System.currentTimeMillis();
+    subscribeTime = time.getMilliseconds();
     return internalSubscribe(partitions, null);
   }
 
@@ -815,12 +821,18 @@ protected Collection<PubSubMessage<K, ChangeEvent<V>, VeniceChangeCoordinate>> i
           return Collections.emptyList();
         }
 
-        if (versionSwapByControlMessage) {
-          if (versionSwapMessageState != null
-              && versionSwapMessageState.isVersionSwapMessagesReceivedForAllPartitions()) {
-            try {
-              versionSwapMessageState.getFindNewTopicCheckpointFuture().get(timeoutInMs, TimeUnit.MILLISECONDS);
-              synchronousSeekToCheckpoint(versionSwapMessageState.getNewTopicCheckpoints());
+        if (versionSwapByControlMessage && versionSwapMessageState != null) {
+          /*
+           * If version swap by control message is enabled and the consumer is undergoing version swap we need to check
+           * and act on two scenarios:
+           * 1. If all version swap messages have been received for all partitions, we need to seek to the new topic.
+           * 2. If we have reached the timeout for the version swap, we need to forcefully seek to the new topic using
+           * the EOP positions for any remaining partitions as our backup plan. See javadoc of
+           * VersionSwapMessageState.getVersionSwapStartTimestamp() for more details.
+           */
+          if (versionSwapMessageState.isVersionSwapMessagesReceivedForAllPartitions()) {
+            if (isNewVersionCheckpointsReady(timeoutInMs)) {
+              synchronousSeekToCheckpoint(versionSwapMessageState.getNewTopicVersionSwapCheckpoints());
               LOGGER.info(
                   "Version swap completed from topic: {} to topic: {}, generation id: {}",
                   versionSwapMessageState.getOldVersionTopic(),
@@ -829,23 +841,27 @@ protected Collection<PubSubMessage<K, ChangeEvent<V>, VeniceChangeCoordinate>> i
               changeCaptureStats.emitVersionSwapCountMetrics(SUCCESS);
               changeCaptureStats.setUndergoingVersionSwap(0);
               versionSwapMessageState = null;
-            } catch (TimeoutException timeoutException) {
-              // Still waiting for internalFindNewVersionCheckpoints to complete
+            } else {
               return Collections.emptyList();
-            } catch (ExecutionException e) {
-              // Re-attempt the seek but should report the error
-              LOGGER.warn(
-                  "Completed consuming old topic: {} for version swap but caught an exception when looking for corresponding checkpoint in new topic: {}. Retrying.",
+            }
+          } else if (time.getMilliseconds()
+              - versionSwapMessageState.getVersionSwapStartTimestamp() > versionSwapTimeoutInMs) {
+            if (!getTopicAssignment().isEmpty()) {
+              internalUnsubscribe(versionSwapMessageState.getIncompletePartitions(), true);
+            }
+            if (isNewVersionCheckpointsReady(timeoutInMs)) {
+              synchronousSeekToCheckpoint(versionSwapMessageState.getNewTopicCheckpointsWithEOPAsBackup());
+              LOGGER.info(
+                  "Version swap completed after timeout from topic: {} to topic: {}, generation id: {}. Partitions: {} are seeked to EOP positions.",
                   versionSwapMessageState.getOldVersionTopic(),
                   versionSwapMessageState.getNewVersionTopic(),
-                  e);
-              changeCaptureStats.emitVersionSwapCountMetrics(FAIL);
-              versionSwapMessageState.setFindNewTopicCheckpointFuture(
-                  internalFindNewVersionCheckpoints(
-                      versionSwapMessageState.getOldVersionTopic(),
-                      versionSwapMessageState.getNewVersionTopic(),
-                      versionSwapMessageState.getVersionSwapGenerationId(),
-                      versionSwapMessageState.getAssignedPartitions()));
+                  versionSwapMessageState.getVersionSwapGenerationId(),
+                  versionSwapMessageState.getIncompletePartitions());
+              changeCaptureStats.emitVersionSwapCountMetrics(SUCCESS);
+              changeCaptureStats.setUndergoingVersionSwap(0);
+              versionSwapMessageState = null;
+            } else {
+              return Collections.emptyList();
             }
           }
         }
@@ -935,6 +951,33 @@ protected Collection<PubSubMessage<K, ChangeEvent<V>, VeniceChangeCoordinate>> i
     }
   }
 
+  private boolean isNewVersionCheckpointsReady(long timeoutInMs) throws InterruptedException {
+    if (versionSwapMessageState == null) {
+      return false;
+    }
+    try {
+      versionSwapMessageState.getFindNewTopicCheckpointFuture().get(timeoutInMs, TimeUnit.MILLISECONDS);
+    } catch (TimeoutException timeoutException) {
+      // Still waiting for internalFindNewVersionCheckpoints to complete.
+      return false;
+    } catch (ExecutionException e) {
+      // Re-attempt the seek but should report the error.
+      LOGGER.warn(
+          "Caught an exception when looking for corresponding checkpoints with generation id: {} in new topic: {}. Retrying.",
+          versionSwapMessageState.getVersionSwapGenerationId(),
+          versionSwapMessageState.getNewVersionTopic(),
+          e);
+      changeCaptureStats.emitVersionSwapCountMetrics(FAIL);
+      versionSwapMessageState.setFindNewTopicCheckpointFuture(
+          internalFindNewVersionCheckpoints(
+              versionSwapMessageState.getOldVersionTopic(),
+              versionSwapMessageState.getNewVersionTopic(),
+              versionSwapMessageState.getVersionSwapGenerationId(),
+              versionSwapMessageState.getAssignedPartitions()));
+    }
+    return true;
+  }
+
   void maybeUpdatePartitionToBootstrapMap(DefaultPubSubMessage message, PubSubTopicPartition pubSubTopicPartition) {
     if (getSubscribeTime() - message.getValue().producerMetadata.messageTimestamp <= TimeUnit.MINUTES.toMillis(1)) {
       getPartitionToBootstrapState().put(pubSubTopicPartition.getPartitionNumber(), true);
@@ -1026,7 +1069,7 @@ protected <T> T processRecordBytes(
    * from seeking in between a sequence of related version swap messages with a partition. During a version swap we will
    * also use a low watermark approach for the {@link VeniceChangeCoordinate} returned. However, the changelog consumer
    * is still vulnerable to this edge case when seekToTimestamp and seekToTail is used. These edge cases will be rare
-   * so for now we will have a metric to detect it and restarting the changelog consumer and re-seek should fix it.
+   * and will be handled by the version swap timeout and backup strategy of seek to new version's EOP.
    */
   protected Optional<PubSubMessage<K, ChangeEvent<V>, VeniceChangeCoordinate>> convertPubSubMessageToPubSubChangeEventWithVersionSwapState(
       DefaultPubSubMessage message,
@@ -1396,7 +1439,8 @@ protected boolean handleVersionSwapMessageInVT(
         .isVersionSwapRelevant(pubSubTopicPartition.getTopicName(), clientRegionName, versionSwap)) {
       if (versionSwapMessageState == null) {
         Set<PubSubTopicPartition> currentAssignment = getTopicAssignment();
-        versionSwapMessageState = new VersionSwapMessageState(versionSwap, totalRegionCount, currentAssignment);
+        versionSwapMessageState =
+            new VersionSwapMessageState(versionSwap, totalRegionCount, currentAssignment, time.getMilliseconds());
         changeCaptureStats.setUndergoingVersionSwap(1);
         LOGGER.info(
             "New version detected for store: {} through version swap messages. Performing version swap from topic: {} to topic: {}, generation id: {}",