[dvc][doc] Create MVP for DaVinciRecordTransformer

build.gradle

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

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

@@ -8,6 +8,7 @@
 import static java.lang.Thread.currentThread;
 
 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;
@@ -109,6 +110,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;
   private final boolean writeBatchingPushStatus;
@@ -119,7 +121,7 @@ public DaVinciBackend(
       Optional<Set<String>> managedClients,
       ICProvider icProvider,
       Optional<ObjectCacheConfig> cacheConfig,
-      Function<Integer, DaVinciRecordTransformer> getRecordTransformer) {
+      DaVinciRecordTransformerFunctionalInterface recordTransformerFunction) {
     LOGGER.info("Creating Da Vinci backend with managed clients: {}", managedClients);
     try {
       VeniceServerConfig backendConfig = configLoader.getVeniceServerConfig();
@@ -269,7 +271,7 @@ public DaVinciBackend(
           false,
           compressorFactory,
           cacheBackend,
-          getRecordTransformer,
+          recordTransformerFunction,
           true,
           // TODO: consider how/if a repair task would be valid for Davinci users?
           null,
@@ -279,6 +281,7 @@ public DaVinciBackend(
           null);
 
       ingestionService.start();
+
       ingestionService.addIngestionNotifier(ingestionListener);
 
       if (isIsolatedIngestion() && cacheConfig.isPresent()) {
@@ -292,6 +295,10 @@ public DaVinciBackend(
       }
 
       if (backendConfig.isBlobTransferManagerEnabled()) {
+        if (recordTransformerFunction != null) {
+          throw new VeniceException("DaVinciRecordTransformer doesn't support blob transfer.");
+        }
+
         blobTransferManager = BlobTransferUtil.getP2PBlobTransferManagerAndStart(
             configLoader.getVeniceServerConfig().getDvcP2pBlobTransferServerPort(),
             configLoader.getVeniceServerConfig().getDvcP2pBlobTransferClientPort(),
@@ -462,7 +469,8 @@ 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));
     });

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,9 @@ protected AvroGenericDaVinciClient(
     this.managedClients = managedClients;
     this.icProvider = icProvider;
     this.chunkingAdapter = chunkingAdapter;
+
+    this.recordTransformer = daVinciConfig.isRecordTransformerEnabled() ? daVinciConfig.getRecordTransformer(0) : null;
+
     preValidation.run();
   }
 
@@ -646,7 +650,7 @@ protected GenericRecordChunkingAdapter getGenericRecordChunkingAdapter() {
     return GenericRecordChunkingAdapter.INSTANCE;
   }
 
-  private D2ServiceDiscoveryResponse discoverService() {
+  D2ServiceDiscoveryResponse discoverService() {
     try (TransportClient client = getTransportClient(clientConfig)) {
       if (!(client instanceof D2TransportClient)) {
         throw new VeniceClientException(
@@ -677,6 +681,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 +696,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 +705,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 recordTransformerFunction) {
     synchronized (AvroGenericDaVinciClient.class) {
       if (daVinciBackend == null) {
         logger
@@ -721,7 +724,7 @@ private void initBackend(
                 managedClients,
                 icProvider,
                 cacheConfig,
-                getRecordTransformer),
+                recordTransformerFunction),
             backend -> {
               // Ensure that existing backend is fully closed before a new one can be created.
               synchronized (AvroGenericDaVinciClient.class) {
@@ -765,7 +768,7 @@ public synchronized void start() {
         managedClients,
         icProvider,
         cacheConfig,
-        daVinciConfig::getRecordTransformer);
+        daVinciConfig.getRecordTransformerFunction());
 
     try {
       getBackend().verifyCacheConfigEquality(daVinciConfig.getCacheConfig(), getStoreName());
@@ -784,12 +787,27 @@ 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) {
+          if (recordTransformer.getOutputValueClass() != clientConfig.getSpecificValueClass()) {
+            throw new VeniceClientException(
+                "Specific value class mismatch between ClientConfig and DaVinciRecordTransformer, expected="
+                    + clientConfig.getSpecificValueClass() + ", actual=" + recordTransformer.getOutputValueClass());
+          }
+
+          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.
@@ -126,8 +125,12 @@ public DaVinciRecordTransformer getRecordTransformer(Integer storeVersion) {
     return null;
   }
 
+  public DaVinciRecordTransformerFunctionalInterface getRecordTransformerFunction() {
+    return recordTransformerFunction;
+  }
+
   public DaVinciConfig setRecordTransformerFunction(
-      Function<Integer, DaVinciRecordTransformer> recordTransformerFunction) {
+      DaVinciRecordTransformerFunctionalInterface recordTransformerFunction) {
     this.recordTransformerFunction = recordTransformerFunction;
     return this;
   }