[dvc][doc] Create MVP for DaVinciRecordTransformer

build.gradle

@@ -269,6 +269,7 @@ subprojects {
     implementation libraries.grpcProtobuf
     implementation libraries.grpcServices
     implementation libraries.grpcStub
+    compileOnly 'org.ow2.asm:asm:9.7'
     compileOnly libraries.tomcatAnnotations
   }
 

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

@@ -7,7 +7,9 @@
 import static com.linkedin.venice.pushmonitor.ExecutionStatus.DVC_INGESTION_ERROR_OTHER;
 import static java.lang.Thread.currentThread;
 
+import com.linkedin.davinci.client.BlockingDaVinciRecordTransformer;
 import com.linkedin.davinci.client.DaVinciRecordTransformer;
+import com.linkedin.davinci.client.DaVinciRecordTransformerFunctionalInterface;
 import com.linkedin.davinci.compression.StorageEngineBackedCompressorFactory;
 import com.linkedin.davinci.config.StoreBackendConfig;
 import com.linkedin.davinci.config.VeniceConfigLoader;
@@ -106,6 +108,7 @@ public class DaVinciBackend implements Closeable {
   private IngestionBackend ingestionBackend;
   private final AggVersionedStorageEngineStats aggVersionedStorageEngineStats;
   private final boolean useDaVinciSpecificExecutionStatusForError;
+  private DaVinciRecordTransformer recordTransformer;
   private final ClientConfig clientConfig;
   private BlobTransferManager<Void> blobTransferManager;
 
@@ -115,7 +118,7 @@ public DaVinciBackend(
       Optional<Set<String>> managedClients,
       ICProvider icProvider,
       Optional<ObjectCacheConfig> cacheConfig,
-      Function<Integer, DaVinciRecordTransformer> getRecordTransformer) {
+      DaVinciRecordTransformerFunctionalInterface getRecordTransformer) {
     LOGGER.info("Creating Da Vinci backend with managed clients: {}", managedClients);
     try {
       VeniceServerConfig backendConfig = configLoader.getVeniceServerConfig();
@@ -247,6 +250,19 @@ public DaVinciBackend(
       cacheBackend = cacheConfig
           .map(objectCacheConfig -> new ObjectCacheBackend(clientConfig, objectCacheConfig, schemaRepository));
 
+      String storeName = clientConfig.getStoreName();
+
+      if (getRecordTransformer != null && getRecordTransformer.apply(0) != null) {
+        storeRepository.refreshOneStore(storeName);
+        Store store = storeRepository.getStoreOrThrow(storeName);
+        int version = store.getCurrentVersion();
+
+        DaVinciRecordTransformer clientRecordTransformer = getRecordTransformer.apply(version);
+        recordTransformer = new BlockingDaVinciRecordTransformer(
+            clientRecordTransformer,
+            clientRecordTransformer.getStoreRecordsInDaVinci());
+      }
+
       ingestionService = new KafkaStoreIngestionService(
           storageService.getStorageEngineRepository(),
           configLoader,
@@ -264,7 +280,7 @@ public DaVinciBackend(
           false,
           compressorFactory,
           cacheBackend,
-          getRecordTransformer,
+          recordTransformer,
           true,
           // TODO: consider how/if a repair task would be valid for Davinci users?
           null,
@@ -274,6 +290,7 @@ public DaVinciBackend(
           null);
 
       ingestionService.start();
+
       ingestionService.addIngestionNotifier(ingestionListener);
 
       if (isIsolatedIngestion() && cacheConfig.isPresent()) {
@@ -457,9 +474,16 @@ private synchronized void bootstrap() {
       List<Integer> partitions = storeNameToPartitionListMap.get(storeName);
       String versionTopic = version.kafkaTopicName();
       LOGGER.info("Bootstrapping partitions {} for {}", partitions, versionTopic);
-      aggVersionedStorageEngineStats.setStorageEngine(versionTopic, storageService.getStorageEngine(versionTopic));
+      AbstractStorageEngine storageEngine = storageService.getStorageEngine(versionTopic);
+      aggVersionedStorageEngineStats.setStorageEngine(versionTopic, storageEngine);
       StoreBackend storeBackend = getStoreOrThrow(storeName);
-      storeBackend.subscribe(ComplementSet.newSet(partitions), Optional.of(version));
+
+      if (recordTransformer != null) {
+        recordTransformer
+            .onRecovery(storageEngine, storeBackend, blobTransferManager, partitions, Optional.of(version));
+      } else {
+        storeBackend.subscribe(ComplementSet.newSet(partitions), Optional.of(version));
+      }
     });
   }
 

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

@@ -128,7 +128,7 @@ public CompletableFuture<Void> subscribe(ComplementSet<Integer> partitions) {
     return subscribe(partitions, Optional.empty());
   }
 
-  synchronized CompletableFuture<Void> subscribe(
+  public synchronized CompletableFuture<Void> subscribe(
       ComplementSet<Integer> partitions,
       Optional<Version> bootstrapVersion) {
     if (daVinciCurrentVersion == null) {

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

@@ -77,7 +77,6 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
-import java.util.function.Function;
 import org.apache.avro.Schema;
 import org.apache.avro.generic.GenericData;
 import org.apache.avro.generic.GenericRecord;
@@ -136,6 +135,8 @@ protected boolean removeEldestEntry(Map.Entry<Integer, GenericRecord> eldest) {
 
   private final AbstractAvroChunkingAdapter<V> chunkingAdapter;
 
+  private final DaVinciRecordTransformer recordTransformer;
+
   public AvroGenericDaVinciClient(
       DaVinciConfig daVinciConfig,
       ClientConfig clientConfig,
@@ -175,6 +176,13 @@ protected AvroGenericDaVinciClient(
     this.managedClients = managedClients;
     this.icProvider = icProvider;
     this.chunkingAdapter = chunkingAdapter;
+
+    this.recordTransformer = daVinciConfig.isRecordTransformerEnabled() ? daVinciConfig.getRecordTransformer(0) : null;
+
+    if (this.recordTransformer != null) {
+      this.clientConfig.setSpecificValueClass(recordTransformer.getOutputValueClass());
+    }
+
     preValidation.run();
   }
 
@@ -646,7 +654,8 @@ protected GenericRecordChunkingAdapter getGenericRecordChunkingAdapter() {
     return GenericRecordChunkingAdapter.INSTANCE;
   }
 
-  private D2ServiceDiscoveryResponse discoverService() {
+  // Visible for testing
+  public D2ServiceDiscoveryResponse discoverService() {
     try (TransportClient client = getTransportClient(clientConfig)) {
       if (!(client instanceof D2TransportClient)) {
         throw new VeniceClientException(
@@ -677,6 +686,7 @@ private VeniceConfigLoader buildVeniceConfig() {
     if (kafkaBootstrapServers == null) {
       kafkaBootstrapServers = backendConfig.getString(KAFKA_BOOTSTRAP_SERVERS);
     }
+
     VeniceProperties config = new PropertyBuilder().put(KAFKA_ADMIN_CLASS, ApacheKafkaAdminAdapter.class.getName())
         .put(ROCKSDB_LEVEL0_FILE_NUM_COMPACTION_TRIGGER, 4) // RocksDB default config
         .put(ROCKSDB_LEVEL0_SLOWDOWN_WRITES_TRIGGER, 20) // RocksDB default config
@@ -691,10 +701,7 @@ private VeniceConfigLoader buildVeniceConfig() {
         .put(INGESTION_USE_DA_VINCI_CLIENT, true)
         .put(
             RECORD_TRANSFORMER_VALUE_SCHEMA,
-            daVinciConfig.isRecordTransformerEnabled()
-                // We're creating a new record transformer here just to get the schema
-                ? daVinciConfig.getRecordTransformer(0).getValueOutputSchema().toString()
-                : "null")
+            recordTransformer != null ? recordTransformer.getValueOutputSchema().toString() : "null")
         .put(INGESTION_ISOLATION_CONFIG_PREFIX + "." + INGESTION_MEMORY_LIMIT, -1) // Explicitly disable memory limiter
                                                                                    // in Isolated Process
         .put(backendConfig.toProperties())
@@ -703,13 +710,14 @@ private VeniceConfigLoader buildVeniceConfig() {
     return new VeniceConfigLoader(config, config);
   }
 
-  private void initBackend(
+  // Visible for testing
+  public void initBackend(
       ClientConfig clientConfig,
       VeniceConfigLoader configLoader,
       Optional<Set<String>> managedClients,
       ICProvider icProvider,
       Optional<ObjectCacheConfig> cacheConfig,
-      Function<Integer, DaVinciRecordTransformer> getRecordTransformer) {
+      DaVinciRecordTransformerFunctionalInterface getRecordTransformer) {
     synchronized (AvroGenericDaVinciClient.class) {
       if (daVinciBackend == null) {
         logger
@@ -784,12 +792,21 @@ public synchronized void start() {
       this.keyDeserializer = FastSerializerDeserializerFactory.getFastAvroGenericDeserializer(keySchema, keySchema);
       this.genericRecordStoreDeserializerCache =
           new AvroStoreDeserializerCache(daVinciBackend.get().getSchemaRepository(), getStoreName(), true);
-      this.storeDeserializerCache = clientConfig.isSpecificClient()
-          ? new AvroSpecificStoreDeserializerCache<>(
+
+      if (clientConfig.isSpecificClient()) {
+        if (recordTransformer != null) {
+          this.storeDeserializerCache = new AvroSpecificStoreDeserializerCache<>(
+              recordTransformer.getValueOutputSchema(),
+              clientConfig.getSpecificValueClass());
+        } else {
+          this.storeDeserializerCache = new AvroSpecificStoreDeserializerCache<>(
               daVinciBackend.get().getSchemaRepository(),
               getStoreName(),
-              clientConfig.getSpecificValueClass())
-          : (AvroStoreDeserializerCache<V>) this.genericRecordStoreDeserializerCache;
+              clientConfig.getSpecificValueClass());
+        }
+      } else {
+        this.storeDeserializerCache = (AvroStoreDeserializerCache<V>) this.genericRecordStoreDeserializerCache;
+      }
 
       ready.set(true);
       logger.info("Client is started successfully, storeName=" + getStoreName());

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

@@ -19,8 +19,8 @@ public class BlockingDaVinciRecordTransformer<K, V, O> extends DaVinciRecordTran
   private final DaVinciRecordTransformer recordTransformer;
   private final CountDownLatch startLatch = new CountDownLatch(1);
 
-  public BlockingDaVinciRecordTransformer(DaVinciRecordTransformer recordTransformer) {
-    super(recordTransformer.getStoreVersion());
+  public BlockingDaVinciRecordTransformer(DaVinciRecordTransformer recordTransformer, boolean storeRecordsInDaVinci) {
+    super(recordTransformer.getStoreVersion(), storeRecordsInDaVinci);
     this.recordTransformer = recordTransformer;
   }
 
@@ -32,28 +32,31 @@ public Schema getValueOutputSchema() {
     return this.recordTransformer.getValueOutputSchema();
   }
 
-  public O put(Lazy<K> key, Lazy<V> value) {
+  public O transform(Lazy<K> key, Lazy<V> value) {
+    return (O) this.recordTransformer.transform(key, value);
+  }
+
+  public void processPut(Lazy<K> key, Lazy<O> value) {
     try {
       // Waiting for onStartIngestionTask to complete before proceeding
       startLatch.await();
-      return (O) this.recordTransformer.put(key, value);
+      this.recordTransformer.processPut(key, value);
     } catch (InterruptedException e) {
       // Restore the interrupt status
       Thread.currentThread().interrupt();
-      return null;
     }
   }
 
-  public O delete(Lazy<K> key) {
-    return (O) this.recordTransformer.delete(key);
+  public O processDelete(Lazy<K> key) {
+    return (O) this.recordTransformer.processDelete(key);
   }
 
-  public void onStartIngestionTask() {
-    this.recordTransformer.onStartIngestionTask();
+  public void onStart() {
+    this.recordTransformer.onStart();
     startLatch.countDown();
   }
 
-  public void onEndIngestionTask() {
-    this.recordTransformer.onEndIngestionTask();
+  public void onEnd() {
+    this.recordTransformer.onEnd();
   }
 }

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

@@ -1,7 +1,6 @@
 package com.linkedin.davinci.client;
 
 import com.linkedin.davinci.store.cache.backend.ObjectCacheConfig;
-import java.util.function.Function;
 
 
 public class DaVinciConfig {
@@ -32,7 +31,7 @@ public class DaVinciConfig {
   /**
    * Record transformer reference
    */
-  private Function<Integer, DaVinciRecordTransformer> recordTransformerFunction;
+  private DaVinciRecordTransformerFunctionalInterface recordTransformerFunction;
 
   /**
    * Whether to enable read-path metrics.
@@ -127,7 +126,7 @@ public DaVinciRecordTransformer getRecordTransformer(Integer storeVersion) {
   }
 
   public DaVinciConfig setRecordTransformerFunction(
-      Function<Integer, DaVinciRecordTransformer> recordTransformerFunction) {
+      DaVinciRecordTransformerFunctionalInterface recordTransformerFunction) {
     this.recordTransformerFunction = recordTransformerFunction;
     return this;
   }