[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

@@ -9,7 +9,7 @@
 
 import com.linkedin.davinci.blobtransfer.BlobTransferManager;
 import com.linkedin.davinci.blobtransfer.BlobTransferUtil;
-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;
@@ -119,7 +119,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();
@@ -268,7 +268,7 @@ public DaVinciBackend(
           false,
           compressorFactory,
           cacheBackend,
-          getRecordTransformer,
+          recordTransformerFunction,
           true,
           // TODO: consider how/if a repair task would be valid for Davinci users?
           null,
@@ -290,6 +290,10 @@ public DaVinciBackend(
       }
 
       if (backendConfig.isBlobTransferManagerEnabled()) {
+        if (recordTransformerFunction != null) {
+          throw new VeniceException("DaVinciRecordTransformer doesn't support blob transfer.");
+        }
+
         blobTransferManager = BlobTransferUtil.getP2PBlobTransferManagerForDVCAndStart(
             configLoader.getVeniceServerConfig().getDvcP2pBlobTransferServerPort(),
             configLoader.getVeniceServerConfig().getDvcP2pBlobTransferClientPort(),
@@ -462,7 +466,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

@@ -78,7 +78,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;
@@ -138,6 +137,8 @@ protected boolean removeEldestEntry(Map.Entry<Integer, GenericRecord> eldest) {
   private final AbstractAvroChunkingAdapter<V> chunkingAdapter;
   private final Executor readChunkExecutorForLargeRequest;
 
+  private final DaVinciRecordTransformerConfig recordTransformerConfig;
+
   public AvroGenericDaVinciClient(
       DaVinciConfig daVinciConfig,
       ClientConfig clientConfig,
@@ -180,8 +181,17 @@ protected AvroGenericDaVinciClient(
     this.managedClients = managedClients;
     this.icProvider = icProvider;
     this.chunkingAdapter = chunkingAdapter;
+    this.recordTransformerConfig = daVinciConfig.getRecordTransformerConfig();
     this.readChunkExecutorForLargeRequest =
         readChunkExecutorForLargeRequest != null ? readChunkExecutorForLargeRequest : READ_CHUNK_EXECUTOR;
+
+    if (daVinciConfig.isIsolated() && recordTransformerConfig != null) {
+      // When both are enabled, this causes the storage engine to be deleted everytime the client starts,
+      // since the record transformer config is never persisted to disk. Additionally, this will spawn multiple
+      // transformers per version, and if the user's transformer is stateful this could cause issues.
+      throw new VeniceClientException("Ingestion Isolation is not supported with DaVinciRecordTransformer");
+    }
+
     preValidation.run();
   }
 
@@ -657,7 +667,8 @@ protected GenericRecordChunkingAdapter getGenericRecordChunkingAdapter() {
     return GenericRecordChunkingAdapter.INSTANCE;
   }
 
-  private D2ServiceDiscoveryResponse discoverService() {
+  // Visible for testing
+  protected D2ServiceDiscoveryResponse discoverService() {
     try (TransportClient client = getTransportClient(clientConfig)) {
       if (!(client instanceof D2TransportClient)) {
         throw new VeniceClientException(
@@ -688,6 +699,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
@@ -703,8 +715,7 @@ private VeniceConfigLoader buildVeniceConfig() {
         .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()
+                ? recordTransformerConfig.getOutputValueSchema().toString()
                 : "null")
         .put(INGESTION_ISOLATION_CONFIG_PREFIX + "." + INGESTION_MEMORY_LIMIT, -1) // Explicitly disable memory limiter
                                                                                    // in Isolated Process
@@ -714,13 +725,14 @@ private VeniceConfigLoader buildVeniceConfig() {
     return new VeniceConfigLoader(config, config);
   }
 
-  private void initBackend(
+  // Visible for testing
+  protected 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
@@ -732,7 +744,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) {
@@ -776,7 +788,7 @@ public synchronized void start() {
         managedClients,
         icProvider,
         cacheConfig,
-        daVinciConfig::getRecordTransformer);
+        daVinciConfig.getRecordTransformerFunction());
 
     try {
       getBackend().verifyCacheConfigEquality(daVinciConfig.getCacheConfig(), getStoreName());
@@ -795,12 +807,28 @@ 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 (daVinciConfig.isRecordTransformerEnabled()) {
+          if (recordTransformerConfig.getOutputValueClass() != clientConfig.getSpecificValueClass()) {
+            throw new VeniceClientException(
+                "Specific value class mismatch between ClientConfig and DaVinciRecordTransformer, expected="
+                    + clientConfig.getSpecificValueClass() + ", actual="
+                    + recordTransformerConfig.getOutputValueClass());
+          }
+
+          this.storeDeserializerCache = new AvroSpecificStoreDeserializerCache<>(
+              recordTransformerConfig.getOutputValueSchema(),
+              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,41 +19,44 @@ 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;
   }
 
-  public Schema getKeyOutputSchema() {
-    return this.recordTransformer.getKeyOutputSchema();
+  public Schema getKeySchema() {
+    return this.recordTransformer.getKeySchema();
   }
 
-  public Schema getValueOutputSchema() {
-    return this.recordTransformer.getValueOutputSchema();
+  public Schema getOutputValueSchema() {
+    return this.recordTransformer.getOutputValueSchema();
   }
 
-  public O put(Lazy<K> key, Lazy<V> value) {
+  public DaVinciRecordTransformerResult<O> transform(Lazy<K> key, Lazy<V> value) {
+    return 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 void processDelete(Lazy<K> key) {
+    this.recordTransformer.processDelete(key);
   }
 
-  public void onStartIngestionTask() {
-    this.recordTransformer.onStartIngestionTask();
+  public void onStartVersionIngestion() {
+    this.recordTransformer.onStartVersionIngestion();
     startLatch.countDown();
   }
 
-  public void onEndIngestionTask() {
-    this.recordTransformer.onEndIngestionTask();
+  public void onEndVersionIngestion() {
+    this.recordTransformer.onEndVersionIngestion();
   }
 }

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 DaVinciRecordTransformerConfig recordTransformerConfig;
 
   /**
    * Whether to enable read-path metrics.
@@ -107,7 +106,7 @@ public boolean isCacheEnabled() {
   }
 
   public boolean isRecordTransformerEnabled() {
-    return recordTransformerFunction != null;
+    return recordTransformerConfig != null;
   }
 
   public ObjectCacheConfig getCacheConfig() {
@@ -119,17 +118,27 @@ public DaVinciConfig setCacheConfig(ObjectCacheConfig cacheConfig) {
     return this;
   }
 
+  public DaVinciConfig setRecordTransformerConfig(DaVinciRecordTransformerConfig recordTransformerConfig) {
+    this.recordTransformerConfig = recordTransformerConfig;
+    return this;
+  }
+
+  public DaVinciRecordTransformerConfig getRecordTransformerConfig() {
+    return recordTransformerConfig;
+  }
+
   public DaVinciRecordTransformer getRecordTransformer(Integer storeVersion) {
-    if (recordTransformerFunction != null) {
-      return recordTransformerFunction.apply(storeVersion);
+    if (recordTransformerConfig == null) {
+      return null;
     }
-    return null;
+    return recordTransformerConfig.getRecordTransformer(storeVersion);
   }
 
-  public DaVinciConfig setRecordTransformerFunction(
-      Function<Integer, DaVinciRecordTransformer> recordTransformerFunction) {
-    this.recordTransformerFunction = recordTransformerFunction;
-    return this;
+  public DaVinciRecordTransformerFunctionalInterface getRecordTransformerFunction() {
+    if (recordTransformerConfig == null) {
+      return null;
+    }
+    return recordTransformerConfig.getRecordTransformerFunction();
   }
 
   public boolean isReadMetricsEnabled() {