Addressed more review comments

Author: xunyin8

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

@@ -91,7 +91,7 @@ public class ChangelogClientConfig<T extends SpecificRecord> {
    * Total region count used for version swap in A/A setup. Each subscribed partition need to receive this many
    * corresponding version swap messages before it can safely go to the new version to ensure data completeness.
    */
-  private int totalRegionCount = 1;
+  private int totalRegionCount = 0;
   /**
    * 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.

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

@@ -7,8 +7,6 @@
 import static com.linkedin.venice.ConfigKeys.DATA_BASE_PATH;
 import static com.linkedin.venice.ConfigKeys.KAFKA_BOOTSTRAP_SERVERS;
 import static com.linkedin.venice.ConfigKeys.ZOOKEEPER_ADDRESS;
-import static com.linkedin.venice.VeniceConstants.ENVIRONMENT_CONFIG_KEY_FOR_REGION_NAME;
-import static com.linkedin.venice.VeniceConstants.SYSTEM_PROPERTY_FOR_APP_RUNNING_REGION;
 import static com.linkedin.venice.kafka.protocol.enums.ControlMessageType.START_OF_SEGMENT;
 import static com.linkedin.venice.schema.rmd.RmdConstants.REPLICATION_CHECKPOINT_VECTOR_FIELD_POS;
 import static com.linkedin.venice.stats.dimensions.VeniceResponseStatusCategory.FAIL;
@@ -98,13 +96,15 @@
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.locks.ReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.stream.Collectors;
@@ -185,6 +185,7 @@ public class VeniceChangelogConsumerImpl<K, V> implements VeniceChangelogConsume
   protected final String clientRegionName;
   protected final int totalRegionCount;
   protected final long versionSwapTimeoutInMs;
+  protected final AtomicReference<CountDownLatch> onGoingVersionSwapSignal = new AtomicReference<>();
   /**
    * Interaction of this field should acquire the subscriptionLock.readLock()
    */
@@ -219,29 +220,24 @@ public VeniceChangelogConsumerImpl(
     this.pubSubMessageDeserializer = pubSubMessageDeserializer;
     this.versionSwapByControlMessage = changelogClientConfig.isVersionSwapByControlMessageEnabled();
     this.totalRegionCount = changelogClientConfig.getTotalRegionCount();
+    this.clientRegionName = changelogClientConfig.getClientRegionName();
     this.versionSwapTimeoutInMs = changelogClientConfig.getVersionSwapTimeoutInMs();
     this.time = new SystemTime();
+    this.onGoingVersionSwapSignal.set(new CountDownLatch(0));
     if (versionSwapByControlMessage) {
-      String clientRegionNameFromConfig = changelogClientConfig.getClientRegionName();
-      if (clientRegionNameFromConfig.isEmpty()) {
-        String regionFromEnv = System.getenv(ENVIRONMENT_CONFIG_KEY_FOR_REGION_NAME);
-        if (regionFromEnv == null) {
-          regionFromEnv = System.getProperty(SYSTEM_PROPERTY_FOR_APP_RUNNING_REGION);
-        }
-        if (regionFromEnv == null) {
-          throw new VeniceException(
-              "Failed to enable version swap by control message because cannot resolve client region name from config, environment or system property");
-        }
-        clientRegionName = regionFromEnv;
-      } else {
-        clientRegionName = clientRegionNameFromConfig;
+      // Version swap related configs should all be resolved or explicitly set at this point.
+      if (this.clientRegionName.isEmpty()) {
+        throw new VeniceException(
+            "Failed to enable version swap by control message because client region name is missing");
+      }
+      if (this.totalRegionCount <= 0) {
+        throw new VeniceException(
+            "Failed to enable version swap by control message because total region count is not set");
       }
       LOGGER.info(
           "VeniceChangelogConsumer version swap by control message is enabled. Client region name: {}, total region count: {}",
           clientRegionName,
           totalRegionCount);
-    } else {
-      clientRegionName = "";
     }
 
     seekExecutorService = Executors.newFixedThreadPool(10, new DaemonThreadFactory(getClass().getSimpleName()));
@@ -382,21 +378,39 @@ protected CompletableFuture<Void> internalSubscribe(Set<Integer> partitions, Pub
         throw new RuntimeException(e);
       }
 
-      PubSubTopic topicToSubscribe;
-      if (topic == null) {
-        topicToSubscribe = getCurrentServingVersionTopic();
+      if (versionSwapByControlMessage) {
+        boolean lockAcquiredAndNoOngoingVersionSwap = false;
+        for (int i = 0; i <= MAX_SUBSCRIBE_RETRIES; i++) {
+          // If version swap is in progress, wait for it to finish
+          try {
+            onGoingVersionSwapSignal.get().await();
+          } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+          }
+          subscriptionLock.writeLock().lock();
+          if (versionSwapMessageState != null) {
+            // A new version swap is in progress, wait for it to finish again
+            subscriptionLock.writeLock().unlock();
+          } else {
+            // No version swap is in progress, proceed with subscription
+            lockAcquiredAndNoOngoingVersionSwap = true;
+            break;
+          }
+        }
+        if (!lockAcquiredAndNoOngoingVersionSwap) {
+          // This should be extremely rare where the subscribe request is constantly conflicting with new version swaps
+          throw new VeniceException("Unable to subscribe to new partitions due to conflicting version swaps");
+        }
       } else {
-        topicToSubscribe = topic;
+        subscriptionLock.writeLock().lock();
       }
 
-      subscriptionLock.writeLock().lock();
       try {
-        if (versionSwapByControlMessage && versionSwapMessageState != null) {
-          throw new VeniceException(
-              String.format(
-                  "Unable to subscribe to new partitions while the changelog consumer is undergoing version swap from topic %s to topic %s",
-                  versionSwapMessageState.getOldVersionTopic(),
-                  versionSwapMessageState.getNewVersionTopic()));
+        PubSubTopic topicToSubscribe;
+        if (topic == null) {
+          topicToSubscribe = getCurrentServingVersionTopic();
+        } else {
+          topicToSubscribe = topic;
         }
         Set<PubSubTopicPartition> topicPartitionSet = getTopicAssignment();
         for (PubSubTopicPartition topicPartition: topicPartitionSet) {
@@ -841,6 +855,7 @@ protected Collection<PubSubMessage<K, ChangeEvent<V>, VeniceChangeCoordinate>> i
               changeCaptureStats.emitVersionSwapCountMetrics(SUCCESS);
               changeCaptureStats.setUndergoingVersionSwap(0);
               versionSwapMessageState = null;
+              onGoingVersionSwapSignal.get().countDown();
             } else {
               return Collections.emptyList();
             }
@@ -860,7 +875,13 @@ protected Collection<PubSubMessage<K, ChangeEvent<V>, VeniceChangeCoordinate>> i
               changeCaptureStats.emitVersionSwapCountMetrics(SUCCESS);
               changeCaptureStats.setUndergoingVersionSwap(0);
               versionSwapMessageState = null;
+              onGoingVersionSwapSignal.get().countDown();
             } else {
+              LOGGER.warn(
+                  "Version swap from topic: {} to topic: {}, generation id: {} already timed out but still unable to find new topic checkpoints to go to.",
+                  versionSwapMessageState.getOldVersionTopic(),
+                  versionSwapMessageState.getNewVersionTopic(),
+                  versionSwapMessageState.getVersionSwapGenerationId());
               return Collections.emptyList();
             }
           }
@@ -974,6 +995,7 @@ private boolean isNewVersionCheckpointsReady(long timeoutInMs) throws Interrupte
               versionSwapMessageState.getNewVersionTopic(),
               versionSwapMessageState.getVersionSwapGenerationId(),
               versionSwapMessageState.getAssignedPartitions()));
+      return false;
     }
     return true;
   }
@@ -1441,6 +1463,7 @@ protected boolean handleVersionSwapMessageInVT(
         Set<PubSubTopicPartition> currentAssignment = getTopicAssignment();
         versionSwapMessageState =
             new VersionSwapMessageState(versionSwap, totalRegionCount, currentAssignment, time.getMilliseconds());
+        onGoingVersionSwapSignal.set(new CountDownLatch(1));
         changeCaptureStats.setUndergoingVersionSwap(1);
         LOGGER.info(
             "New version detected for store: {} through version swap messages. Performing version swap from topic: {} to topic: {}, generation id: {}",

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

@@ -1,11 +1,11 @@
 package com.linkedin.davinci.consumer;
 
-import com.linkedin.venice.annotation.NotThreadsafe;
 import com.linkedin.venice.exceptions.VeniceException;
 import com.linkedin.venice.kafka.protocol.VersionSwap;
 import com.linkedin.venice.meta.Version;
 import com.linkedin.venice.pubsub.api.PubSubPosition;
 import com.linkedin.venice.pubsub.api.PubSubTopicPartition;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
@@ -15,10 +15,13 @@
 
 /**
  * A class initialized to indicate that the changelog consumer is undergoing version swap. The class is also keeping
- * various states about this version swap. This class is NOT thread safe and the intended caller should at least
- * acquire the subscriptionLock.readLock().
+ * various states about this version swap. This class is thread safe by having all methods that change or read a state
+ * that's mutable synchronized. However, it's important to keep in mind that race can still occur if the caller is
+ * trying to perform a sequence of events that depend on each other. In those scenarios an external lock is required.
+ * E.g. (1) getFindNewTopicCheckpointFuture() and if the future is complete call (2) getNewTopicVersionSwapCheckpoints()
+ * By the time (2) is called it's possible that the result from (1) is no longer valid. A different thread slipped in
+ * between and changed the state.
  */
-@NotThreadsafe
 public class VersionSwapMessageState {
   private final String oldVersionTopic;
   private final String newVersionTopic;
@@ -55,7 +58,17 @@ public VersionSwapMessageState(
     this.versionSwapStartTimestamp = versionSwapStartTimestamp;
   }
 
-  public PubSubPosition getVersionSwapLowWatermarkPosition(String topic, int partitionId) {
+  /**
+   * Get the pub sub position of the first relevant version swap message for the given partition. Null will be returned
+   * if the partition have not consumed its version swap yet. This is acceptable because different partitions could be
+   * making progress towards version swap at different pace. e.g. partition 0 consumed its version swap message already
+   * but partition 1 could still be consuming regular messages from the old version topic before encountering any
+   * version swap messages.
+   * @param topic where the version swap message originated from
+   * @param partitionId of the version topic
+   * @return the pub sub position or null
+   */
+  public synchronized PubSubPosition getVersionSwapLowWatermarkPosition(String topic, int partitionId) {
     if (oldVersionTopic.equals(topic)) {
       return partitionToVersionSwapLowWatermarkPositionMap.get(partitionId);
     } else {
@@ -75,23 +88,26 @@ public long getVersionSwapGenerationId() {
     return versionSwapGenerationId;
   }
 
-  public void setFindNewTopicCheckpointFuture(CompletableFuture<Void> findNewTopicCheckpointFuture) {
+  public synchronized void setFindNewTopicCheckpointFuture(CompletableFuture<Void> findNewTopicCheckpointFuture) {
     this.findNewTopicCheckpointFuture = findNewTopicCheckpointFuture;
+    this.newTopicEOPCheckpoints.clear();
+    this.newTopicVersionSwapCheckpoints.clear();
   }
 
-  public CompletableFuture<Void> getFindNewTopicCheckpointFuture() {
+  public synchronized CompletableFuture<Void> getFindNewTopicCheckpointFuture() {
     return findNewTopicCheckpointFuture;
   }
 
-  public void setNewTopicVersionSwapCheckpoints(Map<Integer, VeniceChangeCoordinate> newTopicVersionSwapCheckpoints) {
+  public synchronized void setNewTopicVersionSwapCheckpoints(
+      Map<Integer, VeniceChangeCoordinate> newTopicVersionSwapCheckpoints) {
     this.newTopicVersionSwapCheckpoints = newTopicVersionSwapCheckpoints;
   }
 
-  public void setNewTopicEOPCheckpoints(Map<Integer, VeniceChangeCoordinate> newTopicEOPCheckpoints) {
+  public synchronized void setNewTopicEOPCheckpoints(Map<Integer, VeniceChangeCoordinate> newTopicEOPCheckpoints) {
     this.newTopicEOPCheckpoints = newTopicEOPCheckpoints;
   }
 
-  public Set<VeniceChangeCoordinate> getNewTopicVersionSwapCheckpoints() {
+  public synchronized Set<VeniceChangeCoordinate> getNewTopicVersionSwapCheckpoints() {
     // Defensive coding
     if (findNewTopicCheckpointFuture == null || !findNewTopicCheckpointFuture.isDone()) {
       throw new VeniceException("New topic checkpoints are not available yet");
@@ -109,7 +125,7 @@ public Set<VeniceChangeCoordinate> getNewTopicVersionSwapCheckpoints() {
    * Intended to be used as a backup strategy if any partition still did not complete version swap within the timeout.
    * Remaining partitions will be resumed from EOP instead of first relevant version swap message in the new topic.
    */
-  public Set<VeniceChangeCoordinate> getNewTopicCheckpointsWithEOPAsBackup() {
+  public synchronized Set<VeniceChangeCoordinate> getNewTopicCheckpointsWithEOPAsBackup() {
     Set<VeniceChangeCoordinate> checkpoints = getNewTopicVersionSwapCheckpoints();
     for (Integer partition: getIncompletePartitions()) {
       if (newTopicEOPCheckpoints.containsKey(partition)) {
@@ -123,10 +139,10 @@ public Set<VeniceChangeCoordinate> getNewTopicCheckpointsWithEOPAsBackup() {
   }
 
   public Set<Integer> getAssignedPartitions() {
-    return assignedPartitions;
+    return Collections.unmodifiableSet(assignedPartitions);
   }
 
-  public Set<Integer> getIncompletePartitions() {
+  public synchronized Set<Integer> getIncompletePartitions() {
     Set<Integer> incompletePartitions = new HashSet<>(assignedPartitions);
     incompletePartitions.removeAll(completedPartitions);
     return incompletePartitions;
@@ -149,7 +165,7 @@ public long getVersionSwapStartTimestamp() {
    * This means we can subscribe to the new version topic and resume normal consumption from the first relevant version
    * swap message.
    */
-  public boolean isVersionSwapMessagesReceivedForAllPartitions() {
+  public synchronized boolean isVersionSwapMessagesReceivedForAllPartitions() {
     return completedPartitions.size() == receivedVersionSwapPartitionToRegionsMap.size();
   }
 
@@ -161,7 +177,7 @@ public boolean isVersionSwapMessagesReceivedForAllPartitions() {
    * @param position of the version swap message.
    * @return true if all version swap messages related to this version swap event have been received.
    */
-  public boolean handleVersionSwap(
+  public synchronized boolean handleVersionSwap(
       VersionSwap versionSwap,
       PubSubTopicPartition pubSubTopicPartition,
       PubSubPosition position) {
@@ -197,7 +213,7 @@ public boolean handleVersionSwap(
    * Remove unsubscribed partitions from the ongoing version swap states.
    * @param partitions to unsubscribe
    */
-  public void handleUnsubscribe(Set<Integer> partitions) {
+  public synchronized void handleUnsubscribe(Set<Integer> partitions) {
     for (Integer partition: partitions) {
       receivedVersionSwapPartitionToRegionsMap.remove(partition);
       partitionToVersionSwapLowWatermarkPositionMap.remove(partition);