[dvc][duckdb] API change proposal to support DVRT checkpointing

build.gradle

@@ -82,7 +82,7 @@ ext.libraries = [
     commonsLang: 'commons-lang:commons-lang:2.6',
     conscrypt: 'org.conscrypt:conscrypt-openjdk-uber:2.5.2',
     d2: "com.linkedin.pegasus:d2:${pegasusVersion}",
-    duckdbJdbc: "org.duckdb:duckdb_jdbc:1.2.0-20250124.011319-133", // TODO: Remove SNAPSHOT when the real release is published!
+    duckdbJdbc: "org.duckdb:duckdb_jdbc:1.2.0-20250127.011437-139", // TODO: Remove SNAPSHOT when the real release is published!
     failsafe: 'net.jodah:failsafe:2.4.0',
     fastUtil: 'it.unimi.dsi:fastutil:8.3.0',
     grpcNettyShaded: "io.grpc:grpc-netty-shaded:${grpcVersion}",

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

@@ -246,7 +246,9 @@ public DaVinciBackend(
               new MetadataUpdateStats(metricsRepository),
               configLoader,
               storageService.getStoreVersionStateSyncer())
-          : new StorageEngineMetadataService(storageService.getStorageEngineRepository(), partitionStateSerializer);
+          : new StorageEngineMetadataService(
+              storageService.getStorageEngineRepository()::getCheckpointStorageEngine,
+              partitionStateSerializer);
       // Start storage metadata service
       ((AbstractVeniceService) storageMetadataService).start();
       compressorFactory = new StorageEngineBackedCompressorFactory(storageMetadataService);

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

@@ -1,5 +1,6 @@
 package com.linkedin.davinci.client;
 
+import com.linkedin.davinci.storage.StorageMetadataService;
 import com.linkedin.venice.annotation.Experimental;
 import com.linkedin.venice.utils.lazy.Lazy;
 import java.io.IOException;
@@ -58,6 +59,16 @@ public void onEndVersionIngestion(int currentVersion) {
     this.recordTransformer.onEndVersionIngestion(currentVersion);
   }
 
+  @Override
+  public boolean useUniformInputValueSchema() {
+    return this.recordTransformer.useUniformInputValueSchema();
+  }
+
+  @Override
+  public StorageMetadataService getStorageMetadataService() {
+    return this.recordTransformer.getStorageMetadataService();
+  }
+
   @Override
   public void close() throws IOException {
     this.recordTransformer.close();

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

@@ -1,5 +1,6 @@
 package com.linkedin.davinci.client;
 
+import com.linkedin.davinci.storage.StorageMetadataService;
 import com.linkedin.davinci.store.AbstractStorageEngine;
 import com.linkedin.venice.annotation.Experimental;
 import com.linkedin.venice.compression.VeniceCompressor;
@@ -135,6 +136,11 @@ public boolean useUniformInputValueSchema() {
     return false;
   }
 
+  /** A subclass can return a concrete implementation of {@link StorageMetadataService} to control checkpointing. */
+  public StorageMetadataService getStorageMetadataService() {
+    return null;
+  }
+
   // Final methods below
 
   /**

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

@@ -131,8 +131,9 @@ public InternalLocalBootstrappingVeniceChangelogConsumer(
         true,
         true,
         functionToCheckWhetherStorageEngineShouldBeKeptOrNot());
-    storageMetadataService =
-        new StorageEngineMetadataService(storageService.getStorageEngineRepository(), partitionStateSerializer);
+    storageMetadataService = new StorageEngineMetadataService(
+        storageService.getStorageEngineRepository()::getCheckpointStorageEngine,
+        partitionStateSerializer);
   }
 
   private Function<String, Boolean> functionToCheckWhetherStorageEngineShouldBeKeptOrNot() {

clients/da-vinci-client/src/main/java/com/linkedin/davinci/ingestion/isolated/IsolatedIngestionServer.java

@@ -685,8 +685,9 @@ private void initializeIsolatedIngestionServer() {
             .setStoreName(AvroProtocolDefinition.KAFKA_MESSAGE_ENVELOPE.getSystemStoreName()),
         null);
 
-    storageMetadataService =
-        new StorageEngineMetadataService(storageService.getStorageEngineRepository(), partitionStateSerializer);
+    storageMetadataService = new StorageEngineMetadataService(
+        storageService.getStorageEngineRepository()::getCheckpointStorageEngine,
+        partitionStateSerializer);
 
     StorageEngineBackedCompressorFactory compressorFactory =
         new StorageEngineBackedCompressorFactory(storageMetadataService);

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

@@ -386,7 +386,6 @@ public StoreIngestionTask(
     this.kafkaProps = kafkaConsumerProperties;
     this.storageService = storageService;
     this.storageEngineRepository = builder.getStorageEngineRepository();
-    this.storageMetadataService = builder.getStorageMetadataService();
     this.storeRepository = builder.getMetadataRepo();
     this.schemaRepository = builder.getSchemaRepo();
     this.kafkaVersionTopic = storeVersionConfig.getStoreVersionName();
@@ -475,6 +474,7 @@ public StoreIngestionTask(
     this.chunkAssembler = new ChunkAssembler(storeName);
     this.cacheBackend = cacheBackend;
 
+    StorageMetadataService dvrtCheckpointer = null;
     if (recordTransformerConfig != null && recordTransformerConfig.getRecordTransformerFunction() != null) {
       Schema keySchema = schemaRepository.getKeySchema(storeName).getSchema();
       this.recordTransformerKeyDeserializer = new AvroGenericDeserializer(keySchema, keySchema);
@@ -483,6 +483,7 @@ public StoreIngestionTask(
 
       DaVinciRecordTransformer clientRecordTransformer = recordTransformerConfig.getRecordTransformerFunction()
           .apply(versionNumber, keySchema, this.recordTransformerInputValueSchema, outputValueSchema);
+      dvrtCheckpointer = clientRecordTransformer.getStorageMetadataService();
 
       this.recordTransformer = new BlockingDaVinciRecordTransformer(
           clientRecordTransformer,
@@ -512,6 +513,7 @@ public StoreIngestionTask(
       this.recordTransformerInputValueSchema = null;
       this.recordTransformerDeserializersByPutSchemaId = null;
     }
+    this.storageMetadataService = dvrtCheckpointer == null ? builder.getStorageMetadataService() : dvrtCheckpointer;
 
     this.localKafkaServer = this.kafkaProps.getProperty(KAFKA_BOOTSTRAP_SERVERS);
     this.localKafkaServerSingletonSet = Collections.singleton(localKafkaServer);

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

@@ -1,14 +1,12 @@
 package com.linkedin.davinci.storage;
 
-import com.linkedin.davinci.store.AbstractStorageEngine;
-import com.linkedin.davinci.store.AbstractStoragePartition;
+import com.linkedin.davinci.store.CheckpointStorageEngine;
 import com.linkedin.venice.exceptions.VeniceException;
 import com.linkedin.venice.kafka.protocol.state.PartitionState;
 import com.linkedin.venice.kafka.protocol.state.StoreVersionState;
 import com.linkedin.venice.offsets.OffsetRecord;
 import com.linkedin.venice.serialization.avro.InternalAvroSpecificSerializer;
 import com.linkedin.venice.service.AbstractVeniceService;
-import java.util.Optional;
 import java.util.function.Function;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -21,13 +19,13 @@
 public class StorageEngineMetadataService extends AbstractVeniceService implements StorageMetadataService {
   private static final Logger LOGGER = LogManager.getLogger(StorageEngineMetadataService.class);
 
-  private final StorageEngineRepository storageEngineRepository;
+  private final Function<String, CheckpointStorageEngine> checkpointStorageEngineGetter;
   private final InternalAvroSpecificSerializer<PartitionState> partitionStateSerializer;
 
   public StorageEngineMetadataService(
-      StorageEngineRepository storageEngineRepository,
+      Function<String, CheckpointStorageEngine> checkpointStorageEngineGetter,
       InternalAvroSpecificSerializer<PartitionState> serializer) {
-    this.storageEngineRepository = storageEngineRepository;
+    this.checkpointStorageEngineGetter = checkpointStorageEngineGetter;
     this.partitionStateSerializer = serializer;
   }
 
@@ -38,7 +36,7 @@ public void put(String topicName, int partitionId, OffsetRecord record) throws V
 
   @Override
   public void clearOffset(String topicName, int partitionId) {
-    AbstractStorageEngine<?> storageEngine = this.storageEngineRepository.getLocalStorageEngine(topicName);
+    CheckpointStorageEngine storageEngine = this.checkpointStorageEngineGetter.apply(topicName);
     if (storageEngine == null) {
       LOGGER.info("Store: {} could not be located, ignoring the reset partition message.", topicName);
       return;
@@ -48,8 +46,11 @@ public void clearOffset(String topicName, int partitionId) {
 
   @Override
   public OffsetRecord getLastOffset(String topicName, int partitionId) throws VeniceException {
-    Optional<OffsetRecord> record = getStorageEngineOrThrow(topicName).getPartitionOffset(partitionId);
-    return record.orElseGet(() -> new OffsetRecord(partitionStateSerializer));
+    OffsetRecord record = getStorageEngineOrThrow(topicName).getPartitionOffset(partitionId);
+    if (record == null) {
+      return new OffsetRecord(this.partitionStateSerializer);
+    }
+    return record;
   }
 
   @Override
@@ -64,7 +65,7 @@ public void stopInner() throws Exception {
   @Override
   public void computeStoreVersionState(String topicName, Function<StoreVersionState, StoreVersionState> mapFunction)
       throws VeniceException {
-    AbstractStorageEngine engine = getStorageEngineOrThrow(topicName);
+    CheckpointStorageEngine engine = getStorageEngineOrThrow(topicName);
     synchronized (engine) {
       StoreVersionState previousSVS = engine.getStoreVersionState();
       StoreVersionState newSVS = mapFunction.apply(previousSVS);
@@ -86,8 +87,8 @@ public StoreVersionState getStoreVersionState(String topicName) throws VeniceExc
     }
   }
 
-  private AbstractStorageEngine<? extends AbstractStoragePartition> getStorageEngineOrThrow(String topicName) {
-    AbstractStorageEngine<?> storageEngine = this.storageEngineRepository.getLocalStorageEngine(topicName);
+  private CheckpointStorageEngine getStorageEngineOrThrow(String topicName) {
+    CheckpointStorageEngine storageEngine = this.checkpointStorageEngineGetter.apply(topicName);
     if (storageEngine == null) {
       throw new VeniceException("Topic " + topicName + " not found in storageEngineRepository");
     }

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

@@ -1,6 +1,7 @@
 package com.linkedin.davinci.storage;
 
 import com.linkedin.davinci.store.AbstractStorageEngine;
+import com.linkedin.davinci.store.CheckpointStorageEngine;
 import com.linkedin.venice.exceptions.VeniceException;
 import java.util.ArrayList;
 import java.util.List;
@@ -25,6 +26,10 @@ public AbstractStorageEngine getLocalStorageEngine(String storeName) {
     return localStorageEngines.get(storeName);
   }
 
+  public CheckpointStorageEngine getCheckpointStorageEngine(String storeName) {
+    return getLocalStorageEngine(storeName);
+  }
+
   public AbstractStorageEngine removeLocalStorageEngine(String storeName) {
     AbstractStorageEngine engine = localStorageEngines.remove(storeName);
     return engine;

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

@@ -51,7 +51,8 @@
  * The point of having one storage engine(environment) or one database for one partition, is to simplify the complexity of rebalancing/partition migration/host swap.
  * The team agreed to take (2.2) as default storage-partition model for now, and run performance tests to see if it goes well.
  */
-public abstract class AbstractStorageEngine<Partition extends AbstractStoragePartition> implements Closeable {
+public abstract class AbstractStorageEngine<Partition extends AbstractStoragePartition>
+    implements CheckpointStorageEngine, Closeable {
   public enum StoragePartitionAdjustmentTrigger {
     CHECK_DATABASE_INTEGRITY, BEGIN_BATCH_PUSH, END_BATCH_PUSH, PREPARE_FOR_READ, PROMOTE_TO_LEADER, DEMOTE_TO_FOLLOWER
   }
@@ -573,7 +574,7 @@ public synchronized void putPartitionOffset(int partitionId, OffsetRecord offset
   /**
    * Retrieve the offset associated with the partitionId from the metadata partition.
    */
-  public synchronized Optional<OffsetRecord> getPartitionOffset(int partitionId) {
+  public synchronized OffsetRecord getPartitionOffset(int partitionId) {
     if (!metadataPartitionCreated()) {
       throw new StorageInitializationException("Metadata partition not created!");
     }
@@ -585,9 +586,9 @@ public synchronized Optional<OffsetRecord> getPartitionOffset(int partitionId) {
     }
     byte[] value = metadataPartition.get(getPartitionMetadataKey(partitionId));
     if (value == null) {
-      return Optional.empty();
+      return null;
     }
-    return Optional.of(new OffsetRecord(value, partitionStateSerializer));
+    return new OffsetRecord(value, partitionStateSerializer);
   }
 
   /**

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

@@ -0,0 +1,25 @@
+package com.linkedin.davinci.store;
+
+import com.linkedin.venice.kafka.protocol.state.StoreVersionState;
+import com.linkedin.venice.offsets.OffsetRecord;
+
+
+/**
+ * Interface to read and write the metadata needed for checkpointing, including:
+ *
+ * - {@link com.linkedin.venice.offsets.OffsetRecord}
+ * - {@link com.linkedin.venice.kafka.protocol.state.StoreVersionState}
+ */
+public interface CheckpointStorageEngine {
+  void putPartitionOffset(int partitionId, OffsetRecord offsetRecord);
+
+  OffsetRecord getPartitionOffset(int partitionId);
+
+  void clearPartitionOffset(int partitionId);
+
+  void putStoreVersionState(StoreVersionState versionState);
+
+  StoreVersionState getStoreVersionState();
+
+  void clearStoreVersionState();
+}

clients/da-vinci-client/src/test/java/com/linkedin/davinci/store/rocksdb/RocksDBStorageEngineTest.java

@@ -3,6 +3,7 @@
 import static com.linkedin.davinci.store.AbstractStorageEngine.METADATA_PARTITION_ID;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertNull;
 
 import com.linkedin.davinci.config.VeniceStoreVersionConfig;
 import com.linkedin.davinci.stats.AggVersionedStorageEngineStats;
@@ -132,13 +133,13 @@ public void testGetAndPutPartitionOffset() {
     offsetRecord.setRealtimeTopicProducerState(kafkaUrl, guid, ppState);
     offsetRecord.setCheckpointLocalVersionTopicOffset(666L);
     rocksDBStorageEngine.putPartitionOffset(PARTITION_ID, offsetRecord);
-    Assert.assertEquals(rocksDBStorageEngine.getPartitionOffset(PARTITION_ID).get().getLocalVersionTopicOffset(), 666L);
+    Assert.assertEquals(rocksDBStorageEngine.getPartitionOffset(PARTITION_ID).getLocalVersionTopicOffset(), 666L);
     ProducerPartitionState ppStateFromRocksDB =
-        rocksDBStorageEngine.getPartitionOffset(PARTITION_ID).get().getRealTimeProducerState(kafkaUrl, guid);
+        rocksDBStorageEngine.getPartitionOffset(PARTITION_ID).getRealTimeProducerState(kafkaUrl, guid);
     Assert.assertEquals(ppStateFromRocksDB.getSegmentNumber(), segment);
     Assert.assertEquals(ppStateFromRocksDB.getMessageSequenceNumber(), sequence);
     rocksDBStorageEngine.clearPartitionOffset(PARTITION_ID);
-    Assert.assertEquals(rocksDBStorageEngine.getPartitionOffset(PARTITION_ID).isPresent(), false);
+    assertNull(rocksDBStorageEngine.getPartitionOffset(PARTITION_ID));
   }
 
   private ProducerPartitionState createProducerPartitionState(int segment, int sequence) {

services/venice-server/src/main/java/com/linkedin/venice/server/VeniceServer.java

@@ -327,8 +327,9 @@ private List<AbstractVeniceService> createServices() {
         storeVersionStateSerializer,
         partitionStateSerializer,
         metadataRepo);
-    storageEngineMetadataService =
-        new StorageEngineMetadataService(storageService.getStorageEngineRepository(), partitionStateSerializer);
+    storageEngineMetadataService = new StorageEngineMetadataService(
+        storageService.getStorageEngineRepository()::getCheckpointStorageEngine,
+        partitionStateSerializer);
     services.add(storageEngineMetadataService);
     storageMetadataService = storageEngineMetadataService;
     services.add(storageService);