Bug fixes

Works well

Full working copy

Close parition in VW before sending DOL

add latency

Delete bogus test

Fix flaky test

Improve logs

Fix flaky test

Author: sushantmane

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

@@ -988,7 +988,7 @@ protected Map<String, PubSubPosition> calculateRtConsumptionStartPositions(
               pubSubAddress,
               sourceTopicPartition);
           PubSubTopicPartition newSourceTopicPartition =
-              resolveRtTopicPartitionWithPubSubBrokerAddress(newSourceTopic, pcs, pubSubAddress);
+              resolveTopicPartitionWithPubSubBrokerAddress(newSourceTopic, pcs, pubSubAddress);
           try {
             rtStartPosition =
                 getRewindStartPositionForRealTimeTopic(pubSubAddress, newSourceTopicPartition, rewindStartTimestamp);
@@ -1372,7 +1372,7 @@ Runnable buildRepairTask(
     return () -> {
       PubSubTopic pubSubTopic = sourceTopicPartition.getPubSubTopic();
       PubSubTopicPartition resolvedTopicPartition =
-          resolveRtTopicPartitionWithPubSubBrokerAddress(pubSubTopic, pcs, sourceKafkaUrl);
+          resolveTopicPartitionWithPubSubBrokerAddress(pubSubTopic, pcs, sourceKafkaUrl);
       // Calculate upstream offset
       PubSubPosition upstreamOffset =
           getRewindStartPositionForRealTimeTopic(sourceKafkaUrl, resolvedTopicPartition, rewindStartTimestamp);

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

@@ -11,13 +11,15 @@
 public class DolStamp {
   private final long leadershipTerm;
   private final String hostId;
+  private final long produceStartTimeMs; // Timestamp when DoL production started
   private volatile boolean dolProduced; // DoL message was acked by broker
   private volatile boolean dolConsumed; // DoL message was consumed back by this replica
   private volatile CompletableFuture<PubSubProduceResult> dolProduceFuture; // Future tracking DoL produce result
 
   public DolStamp(long leadershipTerm, String hostId) {
     this.leadershipTerm = leadershipTerm;
     this.hostId = hostId;
+    this.produceStartTimeMs = System.currentTimeMillis();
     this.dolProduced = false;
     this.dolConsumed = false;
     this.dolProduceFuture = null;
@@ -59,9 +61,29 @@ public boolean isReady() {
     return dolProduced && dolConsumed;
   }
 
+  public long getProduceStartTimeMs() {
+    return produceStartTimeMs;
+  }
+
+  /**
+   * Calculate latency from DoL production start to now.
+   * @return latency in milliseconds
+   */
+  public long getLatencyMs() {
+    return System.currentTimeMillis() - produceStartTimeMs;
+  }
+
   @Override
   public String toString() {
+    String produceResult = "";
+    if (dolProduceFuture != null && dolProduceFuture.isDone()) {
+      try {
+        produceResult = ", offset=" + dolProduceFuture.get().getPubSubPosition();
+      } catch (Exception e) {
+        // Ignore, keep empty
+      }
+    }
     return "DolStamp{term=" + leadershipTerm + ", host=" + hostId + ", produced=" + dolProduced + ", consumed="
-        + dolConsumed + ", futureSet=" + (dolProduceFuture != null) + "}";
+        + dolConsumed + produceResult + ", latencyMs=" + getLatencyMs() + "}";
   }
 }

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

@@ -63,7 +63,7 @@
  *    c) {@link ConsumerSubscriptionCleaner}
  * 2. Receive various calls to interrogate or mutate consumer state, and delegate them to the correct unit, by
  *    maintaining a mapping of which unit belongs to which version-topic and subscribed topic-partition. Notably,
- *    the {@link #startConsumptionIntoDataReceiver(PartitionReplicaIngestionContext, long, ConsumedDataReceiver)} function allows the
+ *    the {@link #startConsumptionIntoDataReceiver(PartitionReplicaIngestionContext, PubSubPosition, ConsumedDataReceiver)} function allows the
  *    caller to start funneling consumed data into a receiver (i.e. into another task).
  * 3. Provide a single abstract function that must be overridden by subclasses in order to implement a consumption
  *    load balancing strategy: {@link #pickConsumerForPartition(PubSubTopic, PubSubTopicPartition)}
@@ -118,7 +118,7 @@ protected KafkaConsumerService(
       final boolean isKafkaConsumerOffsetCollectionEnabled,
       final ReadOnlyStoreRepository metadataRepository,
       final boolean isUnregisterMetricForDeletedStoreEnabled,
-      VeniceServerConfig serverConfig) {
+      final VeniceServerConfig serverConfig) {
     this.kafkaUrl = consumerProperties.getProperty(KAFKA_BOOTSTRAP_SERVERS);
     this.kafkaUrlForLogger = Utils.getSanitizedStringForLogger(kafkaUrl);
     this.LOGGER = LogManager.getLogger(
@@ -155,7 +155,9 @@ protected KafkaConsumerService(
           pubSubConsumerAdapterFactory.create(contextBuilder.build()),
           aggStats,
           this::recordPartitionsPerConsumerSensor,
-          this::handleUnsubscription);
+          this::handleUnsubscription,
+          serverConfig.getRegionName(),
+          i);
 
       Supplier<Map<PubSubTopicPartition, List<DefaultPubSubMessage>>> pollFunction =
           liveConfigBasedKafkaThrottlingEnabled

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

@@ -865,6 +865,8 @@ private void initializeAndSendDoLStamp(PartitionConsumptionState partitionConsum
           partitionConsumptionState.getReplicaId());
       return;
     }
+    // Close any existing VeniceWriter partition session
+    veniceWriter.get().closePartition(partitionConsumptionState.getPartition());
 
     DolStamp dolStamp = new DolStamp(leadershipTerm, veniceWriter.get().getWriterId());
     partitionConsumptionState.setDolState(dolStamp);
@@ -878,26 +880,13 @@ private void initializeAndSendDoLStamp(PartitionConsumptionState partitionConsum
     // Send DolStamp to local VT
     PubSubProducerCallback dolCallback = new DolStampProduceCallback(partitionConsumptionState, leadershipTerm);
     CompletableFuture<PubSubProduceResult> dolProduceFuture = veniceWriter.get()
-        .sendDoLStamp(partitionConsumptionState.getReplicaTopicPartition(), dolCallback, leadershipTerm);
-
+        .sendDoLStamp(
+            partitionConsumptionState.getReplicaTopicPartition(),
+            dolCallback,
+            leadershipTerm,
+            localKafkaClusterId);
     // Store the produce future in DolStamp
     dolStamp.setDolProduceFuture(dolProduceFuture);
-
-    // Chain logging for produce completion
-    dolProduceFuture.thenAccept((result) -> {
-      LOGGER.info(
-          "DoL stamp produced for replica: {} at offset: {} for term: {}",
-          partitionConsumptionState.getReplicaId(),
-          result.getPubSubPosition(),
-          leadershipTerm);
-    }).exceptionally((ex) -> {
-      LOGGER.error(
-          "Failed to produce DoL stamp for replica: {} for term: {}",
-          partitionConsumptionState.getReplicaId(),
-          leadershipTerm,
-          ex);
-      return null;
-    });
   }
 
   /**
@@ -916,8 +905,8 @@ public DolStampProduceCallback(PartitionConsumptionState pcs, long leadershipTer
     public void onCompletion(PubSubProduceResult produceResult, Exception exception) {
       if (exception != null) {
         LOGGER.error(
-            "Failed to produce DoL message for partition {} with term {}",
-            pcs.getPartition(),
+            "Failed to produce DoL message for replica: {} with term: {}",
+            pcs.getReplicaId(),
             leadershipTerm,
             exception);
         // Clear DoL state on failure
@@ -930,15 +919,15 @@ public void onCompletion(PubSubProduceResult produceResult, Exception exception)
       if (dolStamp != null && dolStamp.getLeadershipTerm() == leadershipTerm) {
         dolStamp.setDolProduced(true);
         LOGGER.info(
-            "DoL message produce confirmed for partition {} at position {} (term: {}) - dolStamp: {}",
-            pcs.getPartition(),
+            "DoL message produce confirmed for replica: {} at position: {} (term: {}) - dolStamp: {}",
+            pcs.getReplicaId(),
             produceResult.getPubSubPosition(),
             leadershipTerm,
             dolStamp);
       } else {
         LOGGER.warn(
-            "DoL state mismatch or null for partition {} - expected term: {}, dolStamp: {}",
-            pcs.getPartition(),
+            "DoL state mismatch or null for replica: {} - expected term: {}, dolStamp: {}",
+            pcs.getReplicaId(),
             leadershipTerm,
             dolStamp);
       }
@@ -958,17 +947,19 @@ private boolean canSwitchToLeaderTopic(PartitionConsumptionState pcs) {
     if (shouldUseDolMechanism() && dolStamp != null) {
       // Check if DoL state is ready (both produced and consumed)
       if (dolStamp.isReady()) {
+        long dolLatencyMs = dolStamp.getLatencyMs();
         LOGGER.info(
-            "DoL mechanism complete for replica: {} - switching to leader topic. DolStamp: {}",
+            "DoL mechanism complete for replica: {} - unblocking switch to the leader topic. Total DoL latency: {} ms. DolStamp: {}",
             pcs.getReplicaId(),
+            dolLatencyMs,
             dolStamp);
         // Clear DoL state as we're done with this transition
         pcs.clearDolState();
         return true;
       }
 
       // DoL not ready yet, stay on local VT
-      LOGGER.info("DoL mechanism not ready for partition {} - DolStamp: {}", pcs.getReplicaId(), dolStamp);
+      LOGGER.debug("DoL mechanism not ready for replica: {} - DolStamp: {}", pcs.getReplicaId(), dolStamp);
       return false;
     } else {
       // Use legacy time-based mechanism
@@ -1230,7 +1221,7 @@ void preparePositionCheckpointAndStartConsumptionAsLeader(
     syncConsumedUpstreamRTOffsetMapIfNeeded(pcs, Collections.singletonMap(pubSubAddress, startPos));
     LOGGER.info(
         "Leader replica: {} started consuming: {} from: {}",
-        pcs.getReplicaId(),
+        pcs,
         Utils.getReplicaId(leaderTopic, pcs.getPartition()),
         startPos);
   }
@@ -1336,6 +1327,14 @@ protected boolean shouldNewLeaderSwitchToRemoteConsumption(PartitionConsumptionS
   }
 
   private boolean isConsumingFromRemoteVersionTopic(PartitionConsumptionState partitionConsumptionState) {
+    // orint all three conditions for easier debugging
+    LOGGER.info(
+        "Checking remote VT consumption for replica: {}. EOP received: {}, isCurrentVersion: {}, nativeReplicationSourceVersionTopicKafkaURL: {}, localKafkaServer: {}",
+        partitionConsumptionState.getReplicaId(),
+        partitionConsumptionState.isEndOfPushReceived(),
+        isCurrentVersion.getAsBoolean(),
+        nativeReplicationSourceVersionTopicKafkaURL,
+        localKafkaServer);
     return !partitionConsumptionState.isEndOfPushReceived() && !isCurrentVersion.getAsBoolean()
     // Do not enable remote consumption for the source fabric leader. Otherwise, it will produce extra messages.
         && !Objects.equals(nativeReplicationSourceVersionTopicKafkaURL, localKafkaServer);

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

@@ -904,6 +904,8 @@ public void setLatestProcessedRemoteVtPosition(PubSubPosition upstreamVtPosition
    * @return the current upstream version topic position
    */
   public PubSubPosition getLatestProcessedRemoteVtPosition() {
+    // TODO: Ideally, we should get this from offset record to ensure durability
+    // return this.offsetRecord.getCheckpointedRemoteVtPosition();
     return this.latestProcessedRemoteVtPosition;
   }
 
@@ -988,6 +990,13 @@ public PubSubPosition getLeaderPosition(String pubSubBrokerAddress, boolean useC
           ? getDivRtCheckpointPosition(pubSubBrokerAddress)
           : getLatestProcessedRtPosition(pubSubBrokerAddress);
     } else {
+      LOGGER.info(
+          "### Getting leader position for replica: {}. Remote consumption enabled: {} remote: {} local: {}",
+          getReplicaTopicPartition(),
+          consumeRemotely(),
+          getLatestProcessedRemoteVtPosition(),
+          getLatestProcessedVtPosition());
+
       return consumeRemotely() ? getLatestProcessedRemoteVtPosition() : getLatestProcessedVtPosition();
     }
   }

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

@@ -177,7 +177,8 @@ protected synchronized SharedKafkaConsumer pickConsumerForPartition(
           "Did not find a suitable consumer after checking " + consumersChecked + " instances.");
     }
     LOGGER.info(
-        "Get shared consumer for: {} from the ingestion task belonging to version topic: {} with index: {}",
+        "Get shared consumer: {} for: {} from the ingestion task belonging to version topic: {} with index: {}",
+        consumer,
         topicPartition,
         versionTopic,
         consumerIndex);

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

@@ -72,6 +72,8 @@ class SharedKafkaConsumer implements PubSubConsumerAdapter {
 
   private final Time time;
 
+  private final String toString;
+
   /**
    * Used to keep track of which version-topic is intended to use a given subscription, in order to detect
    * regressions where we would end up using this consumer to subscribe to a given topic-partition on behalf
@@ -95,23 +97,28 @@ public SharedKafkaConsumer(
       PubSubConsumerAdapter delegate,
       AggKafkaConsumerServiceStats stats,
       Runnable assignmentChangeListener,
-      UnsubscriptionListener unsubscriptionListener) {
-    this(delegate, stats, assignmentChangeListener, unsubscriptionListener, new SystemTime());
+      UnsubscriptionListener unsubscriptionListener,
+      String regionName,
+      int idx) {
+    this(delegate, stats, assignmentChangeListener, unsubscriptionListener, new SystemTime(), regionName, idx);
   }
 
   SharedKafkaConsumer(
       PubSubConsumerAdapter delegate,
       AggKafkaConsumerServiceStats stats,
       Runnable assignmentChangeListener,
       UnsubscriptionListener unsubscriptionListener,
-      Time time) {
+      Time time,
+      String regionName,
+      int idx) {
     this.delegate = delegate;
     this.stats = stats;
     this.assignmentChangeListener = assignmentChangeListener;
     this.unsubscriptionListener = unsubscriptionListener;
     this.time = time;
     this.currentAssignment = Collections.emptySet();
     this.currentAssignmentSize = new AtomicInteger(0);
+    this.toString = String.format("SharedKafkaConsumer-%s:%s ", idx, regionName);
   }
 
   /**
@@ -132,17 +139,15 @@ protected synchronized void updateCurrentAssignment(Set<PubSubTopicPartition> ne
   @UnderDevelopment(value = "This API may not be implemented in all PubSubConsumerAdapter implementations.")
   @Override
   public synchronized void subscribe(PubSubTopicPartition pubSubTopicPartition, PubSubPosition lastReadPubSubPosition) {
-    throw new VeniceException(
-        this.getClass().getSimpleName() + " does not support subscribe without specifying a version-topic.");
+    throw new VeniceException(this + " does not support subscribe without specifying a version-topic.");
   }
 
   @Override
   public void subscribe(
       @Nonnull PubSubTopicPartition pubSubTopicPartition,
       @Nonnull PubSubPosition position,
       boolean isInclusive) {
-    throw new VeniceException(
-        this.getClass().getSimpleName() + " does not support subscribe without specifying a version-topic.");
+    throw new VeniceException(this + " does not support subscribe without specifying a version-topic.");
   }
 
   synchronized void subscribe(
@@ -209,7 +214,7 @@ protected synchronized void unSubscribeAction(Supplier<Set<PubSubTopicPartition>
     long elapsedTime = System.currentTimeMillis() - startTime;
     LOGGER.info(
         "Shared consumer {} unsubscribed {} partition(s): ({}) in {} ms",
-        this.getClass().getSimpleName(),
+        this,
         topicPartitions.size(),
         topicPartitions,
         elapsedTime);
@@ -419,4 +424,9 @@ public synchronized PubSubPosition decodePosition(
       ByteBuffer buffer) {
     return delegate.decodePosition(partition, positionTypeId, buffer);
   }
+
+  @Override
+  public String toString() {
+    return toString;
+  }
 }