[server][dvc][cc] Read checkpoint state from PubSubPosition field with offset fallback (linkedin#2123)

Dual-write to both offset and PubSubPosition fields was already in place. This  
change switches checkpoint reads in OffsetRecord to use the PubSubPosition field  
by default, with a fallback to the legacy offset field for backward compatibility.

Author: sushantmane

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

@@ -342,7 +342,7 @@ public BlobTransferPartitionMetadata prepareMetadata(BlobTransferPayload blobTra
 
     if (storageMetadataService.getStoreVersionState(blobTransferRequest.getTopicName()) == null
         || storageMetadataService
-            .getLastOffset(blobTransferRequest.getTopicName(), blobTransferRequest.getPartition()) == null) {
+            .getLastOffset(blobTransferRequest.getTopicName(), blobTransferRequest.getPartition(), null) == null) {
       throw new VeniceException("Cannot get store version state or offset record from storage metadata service.");
     }
 
@@ -352,8 +352,8 @@ public BlobTransferPartitionMetadata prepareMetadata(BlobTransferPayload blobTra
     java.nio.ByteBuffer storeVersionStateByte =
         ByteBuffer.wrap(storeVersionStateSerializer.serialize(blobTransferRequest.getTopicName(), storeVersionState));
 
-    OffsetRecord offsetRecord =
-        storageMetadataService.getLastOffset(blobTransferRequest.getTopicName(), blobTransferRequest.getPartition());
+    OffsetRecord offsetRecord = storageMetadataService
+        .getLastOffset(blobTransferRequest.getTopicName(), blobTransferRequest.getPartition(), null);
     java.nio.ByteBuffer offsetRecordByte = ByteBuffer.wrap(offsetRecord.toBytes());
 
     return new BlobTransferPartitionMetadata(

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

@@ -107,7 +107,7 @@ public void updateStorePartitionMetadata(
     storageMetadataService.put(
         transferredPartitionMetadata.topicName,
         transferredPartitionMetadata.partitionId,
-        new OffsetRecord(transferredPartitionMetadata.offsetRecord.array(), partitionStateSerializer));
+        new OffsetRecord(transferredPartitionMetadata.offsetRecord.array(), partitionStateSerializer, null));
     // update the metadata SVS
     updateStorageVersionState(storageMetadataService, transferredPartitionMetadata);
   }

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

@@ -6,6 +6,7 @@
 import com.linkedin.venice.compression.VeniceCompressor;
 import com.linkedin.venice.exceptions.VeniceException;
 import com.linkedin.venice.kafka.protocol.state.PartitionState;
+import com.linkedin.venice.pubsub.PubSubContext;
 import com.linkedin.venice.serialization.avro.InternalAvroSpecificSerializer;
 import com.linkedin.venice.utils.lazy.Lazy;
 import java.io.Closeable;
@@ -312,8 +313,10 @@ public final void onRecovery(
       StorageEngine storageEngine,
       int partitionId,
       InternalAvroSpecificSerializer<PartitionState> partitionStateSerializer,
-      Lazy<VeniceCompressor> compressor) {
-    recordTransformerUtility.onRecovery(storageEngine, partitionId, partitionStateSerializer, compressor);
+      Lazy<VeniceCompressor> compressor,
+      PubSubContext pubSubContext) {
+    recordTransformerUtility
+        .onRecovery(storageEngine, partitionId, partitionStateSerializer, compressor, pubSubContext);
   }
 
   /**

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

@@ -7,6 +7,7 @@
 import com.linkedin.venice.compression.VeniceCompressor;
 import com.linkedin.venice.kafka.protocol.state.PartitionState;
 import com.linkedin.venice.offsets.OffsetRecord;
+import com.linkedin.venice.pubsub.PubSubContext;
 import com.linkedin.venice.serialization.avro.InternalAvroSpecificSerializer;
 import com.linkedin.venice.serializer.FastSerializerDeserializerFactory;
 import com.linkedin.venice.serializer.RecordDeserializer;
@@ -128,10 +129,12 @@ public final void onRecovery(
       StorageEngine storageEngine,
       int partitionId,
       InternalAvroSpecificSerializer<PartitionState> partitionStateSerializer,
-      Lazy<VeniceCompressor> compressor) {
+      Lazy<VeniceCompressor> compressor,
+      PubSubContext pubSubContext) {
     int classHash = recordTransformer.getClassHash();
-    Optional<OffsetRecord> optionalOffsetRecord = storageEngine.getPartitionOffset(partitionId);
-    OffsetRecord offsetRecord = optionalOffsetRecord.orElseGet(() -> new OffsetRecord(partitionStateSerializer));
+    Optional<OffsetRecord> optionalOffsetRecord = storageEngine.getPartitionOffset(partitionId, pubSubContext);
+    OffsetRecord offsetRecord =
+        optionalOffsetRecord.orElseGet(() -> new OffsetRecord(partitionStateSerializer, pubSubContext));
 
     boolean transformerLogicChanged = hasTransformerLogicChanged(classHash, offsetRecord);
 
@@ -142,7 +145,7 @@ public final void onRecovery(
       storageEngine.clearPartitionOffset(partitionId);
 
       // Offset record is deleted, so create a new one and persist it
-      offsetRecord = new OffsetRecord(partitionStateSerializer);
+      offsetRecord = new OffsetRecord(partitionStateSerializer, pubSubContext);
       offsetRecord.setRecordTransformerClassHash(classHash);
       storageEngine.putPartitionOffset(partitionId, offsetRecord);
     } else {

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

@@ -5,6 +5,7 @@
 import com.linkedin.venice.annotation.Experimental;
 import com.linkedin.venice.compression.VeniceCompressor;
 import com.linkedin.venice.kafka.protocol.state.PartitionState;
+import com.linkedin.venice.pubsub.PubSubContext;
 import com.linkedin.venice.serialization.avro.InternalAvroSpecificSerializer;
 import com.linkedin.venice.utils.lazy.Lazy;
 import java.io.IOException;
@@ -110,8 +111,9 @@ public void internalOnRecovery(
       StorageEngine storageEngine,
       int partitionId,
       InternalAvroSpecificSerializer<PartitionState> partitionStateSerializer,
-      Lazy<VeniceCompressor> compressor) {
-    this.recordTransformer.onRecovery(storageEngine, partitionId, partitionStateSerializer, compressor);
+      Lazy<VeniceCompressor> compressor,
+      PubSubContext pubSubContext) {
+    this.recordTransformer.onRecovery(storageEngine, partitionId, partitionStateSerializer, compressor, pubSubContext);
   }
 
   public long getCountDownStartConsumptionLatchCount() {

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

@@ -5,6 +5,7 @@
 import com.linkedin.venice.controllerapi.D2ControllerClient;
 import com.linkedin.venice.pubsub.PubSubClientsFactory;
 import com.linkedin.venice.pubsub.PubSubConsumerAdapterFactory;
+import com.linkedin.venice.pubsub.PubSubContext;
 import com.linkedin.venice.pubsub.PubSubPositionDeserializer;
 import com.linkedin.venice.pubsub.PubSubPositionTypeRegistry;
 import com.linkedin.venice.pubsub.PubSubTopicRepository;
@@ -67,11 +68,8 @@ public class ChangelogClientConfig<T extends SpecificRecord> {
    * These are refreshed each time a new set of consumer properties is applied.
    */
   private PubSubConsumerAdapterFactory<? extends PubSubConsumerAdapter> pubSubConsumerAdapterFactory;
-  private PubSubPositionDeserializer pubSubPositionDeserializer;
   private PubSubMessageDeserializer pubSubMessageDeserializer;
-  private PubSubPositionTypeRegistry pubSubPositionTypeRegistry;
-
-  private final PubSubTopicRepository pubSubTopicRepository = new PubSubTopicRepository();
+  private PubSubContext pubSubContext;
 
   public ChangelogClientConfig(String storeName) {
     this.innerClientConfig = new ClientConfig<>(storeName);
@@ -358,10 +356,10 @@ public ChangelogClientConfig setMaxBufferSize(int maxBufferSize) {
    *
    * <p>This method sets up:
    * <ul>
-   *   <li>{@link #pubSubPositionTypeRegistry} – derived from consumer properties</li>
-   *   <li>{@link #pubSubPositionDeserializer} – uses the initialized position type registry</li>
-   *   <li>{@link #pubSubConsumerAdapterFactory} – created based on resolved consumer configuration</li>
-   *   <li>{@link #pubSubMessageDeserializer} – stateless shared instance</li>
+   *   <li>{@link PubSubPositionTypeRegistry} – derived from consumer properties</li>
+   *   <li>{@link PubSubPositionDeserializer} – uses the initialized position type registry</li>
+   *   <li>{@link PubSubConsumerAdapterFactory} – created based on resolved consumer configuration</li>
+   *   <li>{@link PubSubMessageDeserializer} – stateless shared instance</li>
    * </ul>
    *
    * <p><strong>Note:</strong> These fields are derived from the {@link #consumerProperties} and should
@@ -371,8 +369,14 @@ public ChangelogClientConfig setMaxBufferSize(int maxBufferSize) {
    */
   private void initializePubSubInternals() {
     VeniceProperties pubSubProperties = new VeniceProperties(this.consumerProperties);
-    this.pubSubPositionTypeRegistry = PubSubPositionTypeRegistry.fromPropertiesOrDefault(pubSubProperties);
-    this.pubSubPositionDeserializer = new PubSubPositionDeserializer(pubSubPositionTypeRegistry);
+    PubSubPositionTypeRegistry typeRegistry = PubSubPositionTypeRegistry.fromPropertiesOrDefault(pubSubProperties);
+    PubSubPositionDeserializer pubSubPositionDeserializer = new PubSubPositionDeserializer(typeRegistry);
+    PubSubTopicRepository pubSubTopicRepository = new PubSubTopicRepository();
+    // todo(sushantmane): Consider passing TopicManagerRepository from outside if required.
+    this.pubSubContext = new PubSubContext.Builder().setPubSubPositionDeserializer(pubSubPositionDeserializer)
+        .setPubSubPositionTypeRegistry(typeRegistry)
+        .setPubSubTopicRepository(pubSubTopicRepository)
+        .build();
     this.pubSubConsumerAdapterFactory = PubSubClientsFactory.createConsumerFactory(pubSubProperties);
     this.pubSubMessageDeserializer = PubSubMessageDeserializer.createOptimizedDeserializer();
   }
@@ -381,20 +385,12 @@ protected PubSubConsumerAdapterFactory<? extends PubSubConsumerAdapter> getPubSu
     return pubSubConsumerAdapterFactory;
   }
 
-  protected PubSubPositionDeserializer getPubSubPositionDeserializer() {
-    return pubSubPositionDeserializer;
-  }
-
   protected PubSubMessageDeserializer getPubSubMessageDeserializer() {
     return pubSubMessageDeserializer;
   }
 
-  protected PubSubPositionTypeRegistry getPubSubPositionTypeRegistry() {
-    return pubSubPositionTypeRegistry;
-  }
-
-  protected PubSubTopicRepository getPubSubTopicRepository() {
-    return pubSubTopicRepository;
+  protected PubSubContext getPubSubContext() {
+    return pubSubContext;
   }
 
   private ChangelogClientConfig setInnerClientConfig(ClientConfig<T> innerClientConfig) {

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

@@ -150,7 +150,7 @@ private Function<String, Boolean> functionToCheckWhetherStorageEngineShouldBeKep
       // subscribe to that position. If it's not able to, that means that the local state is off Venice retention,
       // and therefore should be completely re-bootstrapped.
       for (Integer partition: bootstrapStateMap.keySet()) {
-        OffsetRecord offsetRecord = storageMetadataService.getLastOffset(localStateTopicName, partition);
+        OffsetRecord offsetRecord = storageMetadataService.getLastOffset(localStateTopicName, partition, pubSubContext);
         if (offsetRecord == null) {
           // No offset info in local, need to bootstrap from beginning.
           return false;
@@ -304,7 +304,7 @@ public void onCompletion() {
    * This method flushes data partition on disk and syncs the underlying database with {@link OffsetRecord}.
    */
   private void syncOffset(int partitionId, BootstrapState bootstrapState) {
-    OffsetRecord lastOffset = storageMetadataService.getLastOffset(localStateTopicName, partitionId);
+    OffsetRecord lastOffset = storageMetadataService.getLastOffset(localStateTopicName, partitionId, pubSubContext);
     StorageEngine storageEngineReloadedFromRepo = getStorageEngine(localStateTopicName);
     if (storageEngineReloadedFromRepo == null) {
       String replicaId = Utils.getReplicaId(localStateTopicName, partitionId);
@@ -486,7 +486,7 @@ public CompletableFuture<Void> seekWithBootStrap(Set<Integer> partitions) {
             partition,
             () -> null);
         // Get the last persisted Offset record from metadata service
-        OffsetRecord offsetRecord = storageMetadataService.getLastOffset(localStateTopicName, partition);
+        OffsetRecord offsetRecord = storageMetadataService.getLastOffset(localStateTopicName, partition, pubSubContext);
         // Where we're at now
         String offsetString = offsetRecord.getDatabaseInfo().get(CHANGE_CAPTURE_COORDINATE);
         VeniceChangeCoordinate localCheckpoint;

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

@@ -219,8 +219,8 @@ protected static PubSubConsumerAdapter getPubSubConsumer(
     PubSubConsumerAdapterContext context = new PubSubConsumerAdapterContext.Builder().setConsumerName(consumerName)
         .setVeniceProperties(new VeniceProperties(changelogClientConfig.getConsumerProperties()))
         .setPubSubMessageDeserializer(changelogClientConfig.getPubSubMessageDeserializer())
-        .setPubSubTopicRepository(changelogClientConfig.getPubSubTopicRepository())
-        .setPubSubPositionTypeRegistry(changelogClientConfig.getPubSubPositionTypeRegistry())
+        .setPubSubTopicRepository(changelogClientConfig.getPubSubContext().getPubSubTopicRepository())
+        .setPubSubPositionTypeRegistry(changelogClientConfig.getPubSubContext().getPubSubPositionTypeRegistry())
         .build();
     return changelogClientConfig.getPubSubConsumerAdapterFactory().create(context);
   }

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

@@ -46,6 +46,7 @@
 import com.linkedin.venice.meta.PersistenceType;
 import com.linkedin.venice.meta.Store;
 import com.linkedin.venice.meta.Version;
+import com.linkedin.venice.pubsub.PubSubContext;
 import com.linkedin.venice.pubsub.PubSubPositionDeserializer;
 import com.linkedin.venice.pubsub.PubSubTopicPartitionImpl;
 import com.linkedin.venice.pubsub.PubSubTopicRepository;
@@ -150,6 +151,7 @@ public class VeniceChangelogConsumerImpl<K, V> implements VeniceChangelogConsume
   protected final PubSubConsumerAdapter pubSubConsumer;
   protected final PubSubTopicRepository pubSubTopicRepository;
   protected final PubSubPositionDeserializer pubSubPositionDeserializer;
+  protected final PubSubContext pubSubContext;
   protected final ExecutorService seekExecutorService;
 
   // This member is a map of maps in order to accommodate view topics. If the message we consume has the appropriate
@@ -181,8 +183,9 @@ public VeniceChangelogConsumerImpl(
       long consumerSequenceIdStartingValue) {
     Objects.requireNonNull(changelogClientConfig, "ChangelogClientConfig cannot be null");
     this.pubSubConsumer = pubSubConsumer;
-    this.pubSubTopicRepository = changelogClientConfig.getPubSubTopicRepository();
-    this.pubSubPositionDeserializer = changelogClientConfig.getPubSubPositionDeserializer();
+    this.pubSubContext = changelogClientConfig.getPubSubContext();
+    this.pubSubTopicRepository = pubSubContext.getPubSubTopicRepository();
+    this.pubSubPositionDeserializer = pubSubContext.getPubSubPositionDeserializer();
 
     seekExecutorService = Executors.newFixedThreadPool(10);
 

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

@@ -326,7 +326,8 @@ public boolean isOffsetLagged(
       long blobTransferDisabledOffsetLagThreshold,
       boolean hybridStore) {
     String topicName = Version.composeKafkaTopic(store, versionNumber);
-    OffsetRecord offsetRecord = storageMetadataService.getLastOffset(topicName, partition);
+    OffsetRecord offsetRecord =
+        storageMetadataService.getLastOffset(topicName, partition, storeIngestionService.getPubSubContext());
 
     if (offsetRecord == null) {
       return true;