Fix deadlock with CC view topic switch logic and adding more tests

Author: xunyin8

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

@@ -227,7 +227,7 @@ protected CompletableFuture<Void> internalFindNewVersionCheckpoints(
               newVersionTopic,
               generationId,
               partitions);
-          while (!versionSwapConsumedPerPartitionMap.values().stream().allMatch(x -> x)) {
+          while (!areAllTrue(versionSwapConsumedPerPartitionMap.values())) {
             polledResults = consumerAdapter.poll(5000L);
             for (Map.Entry<PubSubTopicPartition, List<DefaultPubSubMessage>> entry: polledResults.entrySet()) {
               PubSubTopicPartition pubSubTopicPartition = entry.getKey();
@@ -396,6 +396,15 @@ protected CompletableFuture<Void> internalSeekToEndOfPush(
     }, seekExecutorService);
   }
 
+  private boolean areAllTrue(Collection<Boolean> booleanCollections) {
+    for (Boolean b: booleanCollections) {
+      if (!b) {
+        return false;
+      }
+    }
+    return true;
+  }
+
   @Override
   public CompletableFuture<Void> seekToEndOfPush(Set<Integer> partitions) {
     if (partitions.isEmpty()) {

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

@@ -187,9 +187,9 @@ public class VeniceChangelogConsumerImpl<K, V> implements VeniceChangelogConsume
   protected final long versionSwapTimeoutInMs;
   protected final AtomicReference<CountDownLatch> onGoingVersionSwapSignal = new AtomicReference<>();
   /**
-   * Interaction of this field should acquire the subscriptionLock.readLock()
+   * Interaction of this field should acquire the subscriptionLock.writeLock()
    */
-  protected VersionSwapMessageState versionSwapMessageState = null;
+  protected volatile VersionSwapMessageState versionSwapMessageState = null;
   protected Time time;
 
   public VeniceChangelogConsumerImpl(
@@ -343,6 +343,11 @@ public VeniceChangelogConsumerImpl(
     LOGGER.info("Start a change log consumer client for store: {}", storeName);
   }
 
+  // Unit test only and read only
+  VersionSwapMessageState getVersionSwapMessageState() {
+    return this.versionSwapMessageState;
+  }
+
   @Override
   public int getPartitionCount() {
     Store store = getStore();
@@ -818,7 +823,7 @@ protected Collection<PubSubMessage<K, ChangeEvent<V>, VeniceChangeCoordinate>> i
       String topicSuffix,
       boolean includeControlMessage) {
     Collection<PubSubMessage<K, ChangeEvent<V>, VeniceChangeCoordinate>> pubSubMessages = new ArrayList<>();
-    Map<PubSubTopicPartition, List<DefaultPubSubMessage>> messagesMap = Collections.EMPTY_MAP;
+    Map<PubSubTopicPartition, List<DefaultPubSubMessage>> messagesMap;
     boolean lockAcquired = false;
 
     try {
@@ -853,7 +858,6 @@ protected Collection<PubSubMessage<K, ChangeEvent<V>, VeniceChangeCoordinate>> i
                   versionSwapMessageState.getNewVersionTopic(),
                   versionSwapMessageState.getVersionSwapGenerationId());
               changeCaptureStats.emitVersionSwapCountMetrics(SUCCESS);
-              changeCaptureStats.setUndergoingVersionSwap(0);
               versionSwapMessageState = null;
               onGoingVersionSwapSignal.get().countDown();
             } else {
@@ -873,7 +877,6 @@ protected Collection<PubSubMessage<K, ChangeEvent<V>, VeniceChangeCoordinate>> i
                   versionSwapMessageState.getVersionSwapGenerationId(),
                   versionSwapMessageState.getIncompletePartitions());
               changeCaptureStats.emitVersionSwapCountMetrics(SUCCESS);
-              changeCaptureStats.setUndergoingVersionSwap(0);
               versionSwapMessageState = null;
               onGoingVersionSwapSignal.get().countDown();
             } else {
@@ -1109,8 +1112,6 @@ protected Optional<PubSubMessage<K, ChangeEvent<V>, VeniceChangeCoordinate>> con
       }
     }
     if (messageType.equals(MessageType.DELETE)) {
-      Delete delete = (Delete) message.getValue().payloadUnion;
-
       // Deletes have a previous and current value of null. So just fill it in!
       ChangeEvent<V> changeEvent = new ChangeEvent<>(null, null);
       pubSubChangeEventMessage = Optional.of(
@@ -1464,7 +1465,6 @@ protected boolean handleVersionSwapMessageInVT(
         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: {}",
             storeName,
@@ -1650,7 +1650,12 @@ protected boolean switchToNewTopic(PubSubTopic newTopic, String topicSuffix, Int
     }
     unsubscribe(partitions);
     try {
-      internalSubscribe(partitions, mergedTopicName).get();
+      Set<VeniceChangeCoordinate> beginningOfNewTopic = new HashSet<>(partitions.size());
+      for (Integer p: partitions) {
+        beginningOfNewTopic
+            .add(new VeniceChangeCoordinate(mergedTopicName.getName(), PubSubSymbolicPosition.EARLIEST, p));
+      }
+      synchronousSeekToCheckpoint(beginningOfNewTopic);
     } catch (Exception e) {
       throw new VeniceException("Subscribe to new topic:" + mergedTopicName + " is not successful, error: " + e);
     }

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

@@ -138,7 +138,7 @@ public synchronized Set<VeniceChangeCoordinate> getNewTopicCheckpointsWithEOPAsB
     return checkpoints;
   }
 
-  public Set<Integer> getAssignedPartitions() {
+  public synchronized Set<Integer> getAssignedPartitions() {
     return Collections.unmodifiableSet(assignedPartitions);
   }
 

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

@@ -13,7 +13,6 @@
 import com.linkedin.venice.stats.VeniceOpenTelemetryMetricsRepository;
 import com.linkedin.venice.stats.dimensions.VeniceMetricsDimensions;
 import com.linkedin.venice.stats.dimensions.VeniceResponseStatusCategory;
-import com.linkedin.venice.stats.metrics.AsyncMetricEntityStateBase;
 import com.linkedin.venice.stats.metrics.MetricEntity;
 import com.linkedin.venice.stats.metrics.MetricEntityStateBase;
 import com.linkedin.venice.stats.metrics.MetricEntityStateOneEnum;
@@ -23,7 +22,6 @@
 import com.linkedin.venice.stats.metrics.TehutiMetricNameEnum;
 import io.opentelemetry.api.common.Attributes;
 import io.tehuti.metrics.MetricsRepository;
-import io.tehuti.metrics.stats.AsyncGauge;
 import io.tehuti.metrics.stats.Avg;
 import io.tehuti.metrics.stats.Gauge;
 import io.tehuti.metrics.stats.Max;
@@ -34,7 +32,6 @@
 import java.util.Collections;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.atomic.AtomicInteger;
 
 
 /**
@@ -56,8 +53,6 @@ public class BasicConsumerStats extends AbstractVeniceStats {
   private final MetricEntityStateOneEnum<VeniceResponseStatusCategory> versionSwapFailCountMetric;
   private final MetricEntityStateOneEnum<VeniceResponseStatusCategory> chunkedRecordSuccessCountMetric;
   private final MetricEntityStateOneEnum<VeniceResponseStatusCategory> chunkedRecordFailCountMetric;
-  private final AsyncMetricEntityStateBase undergoingVersionSwapMetric;
-  private final AtomicInteger undergoingVersionSwap = new AtomicInteger(0);
 
   public BasicConsumerStats(MetricsRepository metricsRepository, String consumerName, String storeName) {
     super(metricsRepository, consumerName);
@@ -162,19 +157,6 @@ public BasicConsumerStats(MetricsRepository metricsRepository, String consumerNa
         baseDimensionsMap,
         VeniceResponseStatusCategory.class);
 
-    undergoingVersionSwapMetric = AsyncMetricEntityStateBase.create(
-        BasicConsumerMetricEntity.UNDERGOING_VERSION_SWAP.getMetricEntity(),
-        otelRepository,
-        this::registerSensor,
-        BasicConsumerTehutiMetricName.UNDERGOING_VERSION_SWAP,
-        Collections.singletonList(
-            new AsyncGauge(
-                (ignored, ignored2) -> this.undergoingVersionSwap.get(),
-                BasicConsumerTehutiMetricName.UNDERGOING_VERSION_SWAP.getMetricName())),
-        baseDimensionsMap,
-        baseAttributes,
-        this.undergoingVersionSwap::get);
-
     /*
      * Record default value for version swap metrics so the UP_DOWN_COUNTER in OTEL will emit a default 0.
      * If you don't do this, the OTEL metric will return no data upon query time until a value is recorded.
@@ -223,10 +205,6 @@ public void emitChunkedRecordCountMetrics(VeniceResponseStatusCategory responseS
     }
   }
 
-  public void setUndergoingVersionSwap(int value) {
-    undergoingVersionSwap.set(value);
-  }
-
   @VisibleForTesting
   public Attributes getBaseAttributes() {
     return baseAttributes;
@@ -238,7 +216,7 @@ public Attributes getBaseAttributes() {
   public enum BasicConsumerTehutiMetricName implements TehutiMetricNameEnum {
     MAX_PARTITION_LAG, RECORDS_CONSUMED, MINIMUM_CONSUMING_VERSION, MAXIMUM_CONSUMING_VERSION, POLL_SUCCESS_COUNT,
     POLL_FAIL_COUNT, VERSION_SWAP_SUCCESS_COUNT, VERSION_SWAP_FAIL_COUNT, CHUNKED_RECORD_SUCCESS_COUNT,
-    CHUNKED_RECORD_FAIL_COUNT, UNDERGOING_VERSION_SWAP;
+    CHUNKED_RECORD_FAIL_COUNT;
 
     private final String metricName;
 
@@ -293,13 +271,6 @@ public enum BasicConsumerMetricEntity implements ModuleMetricEntityInterface {
     CHUNKED_RECORD_COUNT(
         MetricType.COUNTER, MetricUnit.NUMBER, "Measures the count of chunked records consumed",
         setOf(VENICE_STORE_NAME, VENICE_RESPONSE_STATUS_CODE_CATEGORY)
-    ),
-    /**
-     * Indicate if the changelog consumer is undergoing version swap triggered by version swap messages
-     */
-    UNDERGOING_VERSION_SWAP(
-        MetricType.ASYNC_GAUGE, MetricUnit.NUMBER, "Indicate if the changelog consumer is undergoing version swap",
-        setOf(VENICE_STORE_NAME)
     );
 
     private final MetricEntity metricEntity;