[cc][dvc][test] Add support for SpecificRecord deserialization in DVRT, and add isCaughtUp API for DVRT CDC (#1790)

Problem Statement
1. Currently, DVRT only supports deserializing keys and values into Avro GenericRecords. This is a problem as for complex schemas this can lead to poor deserialization performance. Additionally, regular CDC client supports deserializing values to SpecificRecords, but DVRT CDC doesn't.
2. In DVRT CDC, there is no way for the user to tell whether the client has caught up.
3. During the start phase of DVRT CDC, subscription to the DVC can die silently.
4. We allow users to call start multiple times in DVRT CDC.
5. The javadoc for BootstrappingVeniceChangelogConsumer hasn't been updated since the introduction of DVRT CDC.
6. DaVinciClientRecordTransformerTest::testRecordTransformer has been failing consistently in the CI.

Solution
1. Add support for SpecificRecord deserialization for keys and values in DVRT to improve performance, and so DVRT CDC can benefit from it. Additionally, I've moved deserialization and serialization to FastAvro as it can perform 90% better when deserializing complex schemas. Please note that the regular CDC client doesn't support SpecificRecord for keys, but we are adding it for DVRT CDC since a user is requesting it.
2. To provide context to the user on whether they're caught up in DVRT CDC, the isCaughtUp API is added to the BootstrappingVeniceChangelogConsumer interface. Since the original BootstrappingVeniceChangelogConsumer implementation is extended from VeniceAfterImageConsumerImpl, it already supports isCaughtUp.
3. To prevent subscription to DVC in DVRT CDC from dying silently, I re-organized the futures and if subscription fails to DVC we will complete the future returned to the user exceptionally.
4. If a user calls start multiple times on DVRT CDC, we now throw an exception. Additionally, if a user passes in an empty set to start we will subscribe to all partitions. I've also added sychronized to start.
5. Updated the javadoc for BootstrappingVeniceChangelogConsumer, explaining how it behaves differently compared to the regular CDC client.
6. To make DaVinciClientRecordTransformerTest::testRecordTransformer pass consistently, we need to ensure that this test runs first.

Author: kvargha

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

@@ -4,30 +4,39 @@
 import com.linkedin.venice.utils.lazy.Lazy;
 import java.util.Optional;
 import org.apache.avro.Schema;
+import org.apache.avro.specific.SpecificRecord;
 
 
 /**
  * Configuration class for {@link DaVinciRecordTransformer}, which is passed into {@link DaVinciConfig}.
  */
 public class DaVinciRecordTransformerConfig {
   private final DaVinciRecordTransformerFunctionalInterface recordTransformerFunction;
+  private final Class keyClass;
   private final Class outputValueClass;
   private final Schema outputValueSchema;
   private final boolean storeRecordsInDaVinci;
   private final boolean alwaysBootstrapFromVersionTopic;
   private final boolean skipCompatibilityChecks;
+  private final boolean useSpecificRecordKeyDeserializer;
+  private final boolean useSpecificRecordValueDeserializer;
 
   public DaVinciRecordTransformerConfig(Builder builder) {
     this.recordTransformerFunction = Optional.ofNullable(builder.recordTransformerFunction)
         .orElseThrow(() -> new VeniceException("recordTransformerFunction cannot be null"));
 
+    this.keyClass = builder.keyClass;
     this.outputValueClass = builder.outputValueClass;
     this.outputValueSchema = builder.outputValueSchema;
     if ((this.outputValueClass != null && this.outputValueSchema == null)
         || (this.outputValueClass == null && this.outputValueSchema != null)) {
       throw new VeniceException("outputValueClass and outputValueSchema must be defined together");
     }
 
+    this.useSpecificRecordKeyDeserializer = keyClass != null && SpecificRecord.class.isAssignableFrom(keyClass);
+    this.useSpecificRecordValueDeserializer =
+        outputValueClass != null && SpecificRecord.class.isAssignableFrom(outputValueClass);
+
     this.storeRecordsInDaVinci = builder.storeRecordsInDaVinci;
     this.alwaysBootstrapFromVersionTopic = builder.alwaysBootstrapFromVersionTopic;
     this.skipCompatibilityChecks = builder.skipCompatibilityChecks;
@@ -40,13 +49,34 @@ public DaVinciRecordTransformerFunctionalInterface getRecordTransformerFunction(
     return recordTransformerFunction;
   }
 
+  /**
+   * @return {@link #keyClass}
+   */
+  public Class getKeyClass() {
+    return keyClass;
+  }
+
+  /**
+   * @return Whether the {@link SpecificRecord} deserializer should be used for keys
+   */
+  public boolean useSpecificRecordKeyDeserializer() {
+    return useSpecificRecordKeyDeserializer;
+  }
+
   /**
    * @return {@link #outputValueClass}
    */
   public Class getOutputValueClass() {
     return outputValueClass;
   }
 
+  /**
+   * @return Whether the {@link SpecificRecord} deserializer should be used for values
+   */
+  public boolean useSpecificRecordValueDeserializer() {
+    return useSpecificRecordValueDeserializer;
+  }
+
   /**
    * @return {@link #outputValueSchema}
    */
@@ -77,6 +107,7 @@ public boolean shouldSkipCompatibilityChecks() {
 
   public static class Builder {
     private DaVinciRecordTransformerFunctionalInterface recordTransformerFunction;
+    private Class keyClass;
     private Class outputValueClass;
     private Schema outputValueSchema;
     private Boolean storeRecordsInDaVinci = true;
@@ -92,7 +123,18 @@ public Builder setRecordTransformerFunction(DaVinciRecordTransformerFunctionalIn
     }
 
     /**
-     * Set this if you modify the schema during transformation. Must be used in conjunction with {@link #setOutputValueSchema(Schema)}
+     * Set this if you want to deserialize keys into {@link org.apache.avro.specific.SpecificRecord}.
+     * @param keyClass the class of the key
+     */
+    public Builder setKeyClass(Class keyClass) {
+      this.keyClass = keyClass;
+      return this;
+    }
+
+    /**
+     * Set this if you modify the schema during transformation, or you want to deserialize values
+     * into {@link org.apache.avro.specific.SpecificRecord}.
+     * Must be used in conjunction with {@link #setOutputValueSchema(Schema)}
      * @param outputValueClass the class of the output value
      */
     public Builder setOutputValueClass(Class outputValueClass) {
@@ -101,7 +143,9 @@ public Builder setOutputValueClass(Class outputValueClass) {
     }
 
     /**
-     * Set this if you modify the schema during transformation. Must be used in conjunction with {@link #setOutputValueClass(Class)}
+     * Set this if you modify the schema during transformation, or you want to deserialize values
+     * into {@link org.apache.avro.specific.SpecificRecord}.
+     * Must be used in conjunction with {@link #setOutputValueClass(Class)}
      * @param outputValueSchema the schema of the output value
      */
     public Builder setOutputValueSchema(Schema outputValueSchema) {

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

@@ -6,8 +6,9 @@
 import com.linkedin.venice.kafka.protocol.state.PartitionState;
 import com.linkedin.venice.offsets.OffsetRecord;
 import com.linkedin.venice.serialization.avro.InternalAvroSpecificSerializer;
-import com.linkedin.venice.serializer.AvroGenericDeserializer;
-import com.linkedin.venice.serializer.AvroSerializer;
+import com.linkedin.venice.serializer.FastSerializerDeserializerFactory;
+import com.linkedin.venice.serializer.RecordDeserializer;
+import com.linkedin.venice.serializer.RecordSerializer;
 import com.linkedin.venice.utils.lazy.Lazy;
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -27,9 +28,9 @@ public class DaVinciRecordTransformerUtility<K, O> {
   private static final Logger LOGGER = LogManager.getLogger(DaVinciRecordTransformerUtility.class);
   private final DaVinciRecordTransformer recordTransformer;
   private final DaVinciRecordTransformerConfig recordTransformerConfig;
-  private final AvroGenericDeserializer<K> keyDeserializer;
-  private final AvroGenericDeserializer<O> outputValueDeserializer;
-  private final AvroSerializer<O> outputValueSerializer;
+  private final RecordDeserializer<K> keyDeserializer;
+  private final RecordDeserializer<O> outputValueDeserializer;
+  private final RecordSerializer<O> outputValueSerializer;
 
   public DaVinciRecordTransformerUtility(
       DaVinciRecordTransformer recordTransformer,
@@ -38,10 +39,23 @@ public DaVinciRecordTransformerUtility(
     this.recordTransformerConfig = recordTransformerConfig;
 
     Schema keySchema = recordTransformer.getKeySchema();
+    if (recordTransformerConfig.useSpecificRecordKeyDeserializer()) {
+      this.keyDeserializer = FastSerializerDeserializerFactory
+          .getFastAvroSpecificDeserializer(keySchema, recordTransformerConfig.getKeyClass());
+    } else {
+      this.keyDeserializer = FastSerializerDeserializerFactory.getFastAvroGenericDeserializer(keySchema, keySchema);
+    }
+
     Schema outputValueSchema = recordTransformer.getOutputValueSchema();
-    this.keyDeserializer = new AvroGenericDeserializer<>(keySchema, keySchema);
-    this.outputValueDeserializer = new AvroGenericDeserializer<>(outputValueSchema, outputValueSchema);
-    this.outputValueSerializer = new AvroSerializer<>(outputValueSchema);
+    if (recordTransformerConfig.useSpecificRecordValueDeserializer()) {
+      this.outputValueDeserializer = FastSerializerDeserializerFactory
+          .getFastAvroSpecificDeserializer(outputValueSchema, recordTransformerConfig.getOutputValueClass());
+    } else {
+      this.outputValueDeserializer =
+          FastSerializerDeserializerFactory.getFastAvroGenericDeserializer(outputValueSchema, outputValueSchema);
+    }
+
+    this.outputValueSerializer = FastSerializerDeserializerFactory.getFastAvroGenericSerializer(outputValueSchema);
   }
 
   /**

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

@@ -29,6 +29,13 @@ public interface BootstrappingVeniceChangelogConsumer<K, V> {
    * NOTE: This future may take some time to complete depending on how much data needs to be ingested in order to catch
    * up with the time that this client started.
    *
+   * NOTE: In the experimental client, the future will complete when there is at least one message to be polled.
+   * We don't wait for all partitions to catch up, as loading every message into a buffer will result in an
+   * Out Of Memory error. Instead, use the {@link #isCaughtUp()} method to determine once all subscribed partitions have
+   * caught up.
+   *
+   * NOTE: In the experimental client, if you pass in an empty set, it will subscribe to all partitions for the store
+   *
    * @param partitions which partition id's to catch up with
    * @return a future that completes once catch up is complete for all passed in partitions.
    */
@@ -41,9 +48,21 @@ public interface BootstrappingVeniceChangelogConsumer<K, V> {
   /**
    * polls for the next batch of change events. The first records returned following calling 'start()' will be from the bootstrap state.
    * Once this state is consumed, subsequent calls to poll will be based off of recent updates to the Venice store.
+   *
+   * In the experimental client, records will be returned in batches configured to the MAX_BUFFER_SIZE. So the initial
+   * calls to poll will be from records from the bootstrap state, until the partitions have caught up.
+   * Additionally, if the buffer hits the MAX_BUFFER_SIZE before the timeout is hit, poll will return immediately.
+   *
    * @param timeoutInMs
    * @return
    */
   Collection<PubSubMessage<K, ChangeEvent<V>, VeniceChangeCoordinate>> poll(long timeoutInMs);
 
+  /**
+   * In the experimental client, once this becomes true it will stay true even if we start to lag after the
+   * bootstrapping phase.
+   * @return True if all subscribed partitions have caught up.
+   */
+  boolean isCaughtUp();
+
 }

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

@@ -32,6 +32,7 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
@@ -46,6 +47,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.ReentrantLock;
 import org.apache.avro.Schema;
@@ -65,12 +67,12 @@ public class BootstrappingVeniceChangelogConsumerDaVinciRecordTransformerImpl<K,
   // A buffer of messages that will be returned to the user
   private final BlockingQueue<PubSubMessage<K, ChangeEvent<V>, VeniceChangeCoordinate>> pubSubMessages;
   // Determines what version per partition is currently serving
-  private final ConcurrentHashMap<Integer, Integer> partitionToVersionToServe;
+  private final Map<Integer, Integer> partitionToVersionToServe;
   private final DaVinciRecordTransformerConfig recordTransformerConfig;
   // CachingDaVinciClientFactory used instead of DaVinciClientFactory, so we have the ability to close down the client
   private final CachingDaVinciClientFactory daVinciClientFactory;
   private final DaVinciClient<Object, Object> daVinciClient;
-  private boolean isStarted = false;
+  private AtomicBoolean isStarted = new AtomicBoolean(false);
   private final CountDownLatch startLatch = new CountDownLatch(1);
   // Using a dedicated thread pool for CompletableFutures created by this class to avoid potential thread starvation
   // issues in the default ForkJoinPool
@@ -80,6 +82,7 @@ public class BootstrappingVeniceChangelogConsumerDaVinciRecordTransformerImpl<K,
   private final ApacheKafkaOffsetPosition placeHolderOffset = ApacheKafkaOffsetPosition.of(0);
   private final ReentrantLock bufferLock = new ReentrantLock();
   private final Condition bufferIsFullCondition = bufferLock.newCondition();
+  private final AtomicBoolean isCaughtUp = new AtomicBoolean(false);
 
   public BootstrappingVeniceChangelogConsumerDaVinciRecordTransformerImpl(ChangelogClientConfig changelogClientConfig) {
     this.changelogClientConfig = changelogClientConfig;
@@ -96,6 +99,9 @@ public BootstrappingVeniceChangelogConsumerDaVinciRecordTransformerImpl(Changelo
          * DVRT implmentation. This is to prevent the local state from being wiped everytime a change is deployed
          */
         .setSkipCompatibilityChecks(true)
+        .setKeyClass(innerClientConfig.getSpecificKeyClass())
+        .setOutputValueClass(innerClientConfig.getSpecificValueClass())
+        .setOutputValueSchema(innerClientConfig.getSpecificValueSchema())
         .build();
     daVinciConfig.setRecordTransformerConfig(recordTransformerConfig);
 
@@ -114,25 +120,24 @@ public BootstrappingVeniceChangelogConsumerDaVinciRecordTransformerImpl(Changelo
   }
 
   @Override
-  public CompletableFuture<Void> start(Set<Integer> partitions) {
-    internalStart();
-    subscribedPartitions.addAll(partitions);
+  public synchronized CompletableFuture<Void> start(Set<Integer> partitions) {
+    if (isStarted.get()) {
+      throw new VeniceException("BootstrappingVeniceChangelogConsumer is already started!");
+    }
 
-    /*
-     * Avoid waiting on the CompletableFuture to prevent a circular dependency.
-     * When subscribe is called, DVRT scans the entire storage engine and fills pubSubMessages.
-     * Because pubSubMessages has limited capacity, blocking on the CompletableFuture
-     * prevents the user from calling poll to drain pubSubMessages, so the threads populating pubSubMessages
-     * will wait forever for capacity to become available. This leads to a deadlock.
-     */
-    daVinciClient.subscribe(partitions).whenComplete((result, error) -> {
-      if (error != null) {
-        LOGGER.error("Failed to subscribe to partitions: {} for store: {}", partitions, storeName, error);
-        throw new VeniceException(error);
+    daVinciClient.start();
+    isStarted.set(true);
+
+    // If a user passes in empty partitions set, we subscribe to all partitions
+    if (partitions.isEmpty()) {
+      for (int i = 0; i < daVinciClient.getPartitionCount(); i++) {
+        subscribedPartitions.add(i);
       }
-    });
+    } else {
+      subscribedPartitions.addAll(partitions);
+    }
 
-    return CompletableFuture.supplyAsync(() -> {
+    CompletableFuture<Void> startFuture = CompletableFuture.supplyAsync(() -> {
       try {
         /*
          * When this latch gets released, this means there's at least one message in pubSubMessages. So when the user
@@ -147,24 +152,40 @@ public CompletableFuture<Void> start(Set<Integer> partitions) {
       }
       return null;
     }, completableFutureThreadPool);
+
+    /*
+     * Avoid waiting on the CompletableFuture to prevent a circular dependency.
+     * When subscribe is called, DVRT scans the entire storage engine and fills pubSubMessages.
+     * Because pubSubMessages has limited capacity, blocking on the CompletableFuture
+     * prevents the user from calling poll to drain pubSubMessages, so the threads populating pubSubMessages
+     * will wait forever for capacity to become available. This leads to a deadlock.
+     */
+    daVinciClient.subscribe(subscribedPartitions).whenComplete((result, error) -> {
+      if (error != null) {
+        LOGGER.error("Failed to subscribe to partitions: {} for store: {}", subscribedPartitions, storeName, error);
+        startFuture.completeExceptionally(new VeniceException(error));
+        return;
+      }
+
+      isCaughtUp.set(true);
+      LOGGER.info(
+          "BootstrappingVeniceChangelogConsumer is caught up for store: {} for partitions: {}",
+          storeName,
+          subscribedPartitions);
+    });
+
+    return startFuture;
   }
 
   @Override
   public CompletableFuture<Void> start() {
-    internalStart();
-
-    Set<Integer> allPartitions = new HashSet<>();
-    for (int i = 0; i < daVinciClient.getPartitionCount(); i++) {
-      allPartitions.add(i);
-    }
-
-    return this.start(allPartitions);
+    return this.start(Collections.emptySet());
   }
 
   @Override
   public void stop() throws Exception {
     daVinciClientFactory.close();
-    isStarted = false;
+    isStarted.set(false);
   }
 
   @Override
@@ -204,13 +225,9 @@ public Collection<PubSubMessage<K, ChangeEvent<V>, VeniceChangeCoordinate>> poll
     return drainedPubSubMessages;
   }
 
-  private void internalStart() {
-    if (isStarted) {
-      return;
-    }
-
-    daVinciClient.start();
-    isStarted = true;
+  @Override
+  public boolean isCaughtUp() {
+    return isCaughtUp.get();
   }
 
   private VeniceProperties buildVeniceConfig() {