Merge remote-tracking branch 'upstream/main' into mnagaraj/addKeyCountMetric

Author: m-nagarajan

.gitignore

@@ -41,3 +41,4 @@ Gemfile.lock
 docs/vendor/
 docs/_site/
 
+.worktrees/

build.gradle

@@ -329,8 +329,6 @@ subprojects {
       // when actually using the new protocol. Example to pin KME to v12 when introducing v13:
       // project(':internal:venice-common').file('src/main/resources/avro/KafkaMessageEnvelope/v12', PathValidation.DIRECTORY)
       def versionOverrides = [
-        project(':internal:venice-common').file('src/main/resources/avro/StoreMetaValue/v40', PathValidation.DIRECTORY),
-        project(':services:venice-controller').file('src/main/resources/avro/AdminOperation/v95', PathValidation.DIRECTORY)
       ]
 
       def schemaDirs = [sourceDir]
@@ -645,7 +643,10 @@ spotless {
   }
   format 'markdown', {
     target '**/*.md'
-    targetExclude '**/build/**', '**/target/**', '**/.gradle/**', 'tests/venice-feature-matrix-test/**'
+    // docs/operations/** and docs/contributing/** are MkDocs website content formatted by Prettier
+    // via the MkDocs build pipeline, not by Spotless. Excluding them avoids a known issue where
+    // `clean check` deletes the spotless npm working directory, causing npm to fail on re-install.
+    targetExclude '**/build/**', '**/target/**', '**/.gradle/**', 'tests/venice-feature-matrix-test/**', 'docs/**'
     prettier(['prettier': '2.8.8', 'prettier-plugin-java': '2.2.0'])
         .config(['proseWrap': 'always', 'printWidth': 120])
   }

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

@@ -11,7 +11,6 @@
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Supplier;
-import org.apache.http.HttpStatus;
 
 
 /**
@@ -65,9 +64,7 @@ private static <T> CompletableFuture<T> trackRequest(
         }
       });
     } catch (Exception e) {
-      stats.emitUnhealthyRequestMetrics(
-          LatencyUtils.getElapsedTimeFromNSToMS(startTimeInNS),
-          HttpStatus.SC_INTERNAL_SERVER_ERROR);
+      stats.emitUnhealthyRequestMetricsForDavinciClient(LatencyUtils.getElapsedTimeFromNSToMS(startTimeInNS));
       throw e;
     }
   }

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

@@ -211,6 +211,7 @@
 import static com.linkedin.venice.ConfigKeys.SERVER_SSL_HANDSHAKE_THREAD_POOL_SIZE;
 import static com.linkedin.venice.ConfigKeys.SERVER_STOP_CONSUMPTION_TIMEOUT_IN_SECONDS;
 import static com.linkedin.venice.ConfigKeys.SERVER_STORE_TO_EARLY_TERMINATION_THRESHOLD_MS_MAP;
+import static com.linkedin.venice.ConfigKeys.SERVER_STORE_VERSION_METADATA_WAIT_DURING_STATE_TRANSITION_TIME_MS;
 import static com.linkedin.venice.ConfigKeys.SERVER_STUCK_CONSUMER_REPAIR_ENABLED;
 import static com.linkedin.venice.ConfigKeys.SERVER_STUCK_CONSUMER_REPAIR_INTERVAL_SECOND;
 import static com.linkedin.venice.ConfigKeys.SERVER_STUCK_CONSUMER_REPAIR_THRESHOLD_SECOND;
@@ -227,6 +228,7 @@
 import static com.linkedin.venice.ConfigKeys.SERVER_USE_HEARTBEAT_LAG_FOR_READY_TO_SERVE_CHECK_ENABLED;
 import static com.linkedin.venice.ConfigKeys.SERVER_USE_METRICS_BASED_POSITION_IN_LAG_COMPUTATION;
 import static com.linkedin.venice.ConfigKeys.SERVER_USE_UPSTREAM_PUBSUB_POSITIONS;
+import static com.linkedin.venice.ConfigKeys.SERVER_VERSION_SWAP_DISK_SIZE_DROP_ALERT_THRESHOLD;
 import static com.linkedin.venice.ConfigKeys.SERVER_ZSTD_DICT_COMPRESSION_LEVEL;
 import static com.linkedin.venice.ConfigKeys.SEVER_CALCULATE_QUOTA_USAGE_BASED_ON_PARTITIONS_ASSIGNMENT_ENABLED;
 import static com.linkedin.venice.ConfigKeys.SORTED_INPUT_DRAINER_SIZE;
@@ -432,6 +434,8 @@ public class VeniceServerConfig extends VeniceClusterConfig {
 
   private final double diskFullThreshold;
 
+  private final double versionSwapDiskSizeDropAlertThreshold;
+
   private final int partitionGracefulDropDelaySeconds;
 
   private final int stopConsumptionTimeoutInSeconds;
@@ -468,6 +472,8 @@ public class VeniceServerConfig extends VeniceClusterConfig {
 
   private final Duration serverMaxWaitForVersionInfo;
 
+  private final long storeVersionMetadataWaitDuringStateTransitionTimeMs;
+
   private final boolean computeFastAvroEnabled;
 
   private final long participantMessageConsumptionDelayMs;
@@ -871,6 +877,8 @@ public VeniceServerConfig(VeniceProperties serverProperties, Map<String, Map<Str
     databaseSyncBytesIntervalForDeferredWriteMode =
         serverProperties.getSizeInBytes(SERVER_DATABASE_SYNC_BYTES_INTERNAL_FOR_DEFERRED_WRITE_MODE, 60 * 1024 * 1024);
     diskFullThreshold = serverProperties.getDouble(SERVER_DISK_FULL_THRESHOLD, 0.95);
+    versionSwapDiskSizeDropAlertThreshold =
+        serverProperties.getDouble(SERVER_VERSION_SWAP_DISK_SIZE_DROP_ALERT_THRESHOLD, 0.5);
     partitionGracefulDropDelaySeconds = serverProperties.getInt(SERVER_PARTITION_GRACEFUL_DROP_DELAY_IN_SECONDS, 30);
     stopConsumptionTimeoutInSeconds = serverProperties.getInt(SERVER_STOP_CONSUMPTION_TIMEOUT_IN_SECONDS, 180);
     leakedResourceCleanUpIntervalInMS =
@@ -890,6 +898,8 @@ public VeniceServerConfig(VeniceProperties serverProperties, Map<String, Map<Str
     diskHealthCheckServiceEnabled = serverProperties.getBoolean(SERVER_DISK_HEALTH_CHECK_SERVICE_ENABLED, true);
     serverMaxWaitForVersionInfo =
         Duration.ofMillis(serverProperties.getLong(SERVER_MAX_WAIT_FOR_VERSION_INFO_MS_CONFIG, 5000));
+    storeVersionMetadataWaitDuringStateTransitionTimeMs =
+        serverProperties.getLong(SERVER_STORE_VERSION_METADATA_WAIT_DURING_STATE_TRANSITION_TIME_MS, 300_000);
     computeFastAvroEnabled = serverProperties.getBoolean(SERVER_COMPUTE_FAST_AVRO_ENABLED, true);
     participantMessageConsumptionDelayMs = serverProperties.getLong(PARTICIPANT_MESSAGE_CONSUMPTION_DELAY_MS, 60000);
     serverPromotionToLeaderReplicaDelayMs =
@@ -1455,6 +1465,10 @@ public double getDiskFullThreshold() {
     return diskFullThreshold;
   }
 
+  public double getVersionSwapDiskSizeDropAlertThreshold() {
+    return versionSwapDiskSizeDropAlertThreshold;
+  }
+
   public int getPartitionGracefulDropDelaySeconds() {
     return partitionGracefulDropDelaySeconds;
   }
@@ -1507,6 +1521,10 @@ public Duration getServerMaxWaitForVersionInfo() {
     return serverMaxWaitForVersionInfo;
   }
 
+  public long getStoreVersionMetadataWaitDuringStateTransitionTimeMs() {
+    return storeVersionMetadataWaitDuringStateTransitionTimeMs;
+  }
+
   public BlockingQueueType getBlockingQueueType() {
     return blockingQueueType;
   }

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

@@ -15,8 +15,11 @@
 import com.linkedin.venice.pushmonitor.ExecutionStatus;
 import com.linkedin.venice.pushmonitor.HybridStoreQuotaStatus;
 import com.linkedin.venice.utils.LogContext;
+import com.linkedin.venice.utils.RetryUtils;
 import com.linkedin.venice.utils.Timer;
 import com.linkedin.venice.utils.Utils;
+import java.time.Duration;
+import java.util.Collections;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
@@ -223,10 +226,45 @@ public void reset() {
     }
   }
 
+  /**
+   * Waits for the version metadata to become available in the store repository using exponential backoff.
+   * During OFFLINE->STANDBY transitions, version info may not have propagated from ZK yet.
+   * This must be called before the storage engine is created so that the correct partition type
+   * (with or without replication metadata) is used.
+   */
+  protected void waitForVersionToBeAvailable() {
+    String storeVersion = Version.composeKafkaTopic(storeName, versionNumber);
+    long waitTimeMs = storeAndServerConfigs.getStoreVersionMetadataWaitDuringStateTransitionTimeMs();
+    try {
+      RetryUtils.executeWithMaxAttemptAndExponentialBackoff(() -> {
+        Version version = storeRepository.getStoreOrThrow(storeName).getVersion(versionNumber);
+        if (version == null) {
+          throw new VeniceException(storeVersion + " not yet available in store repository");
+        }
+      },
+          Integer.MAX_VALUE,
+          Duration.ofMillis(10),
+          Duration.ofMillis(200),
+          Duration.ofMillis(waitTimeMs),
+          Collections.singletonList(VeniceException.class));
+    } catch (Exception e) {
+      String errorMsg = storeVersion + " did not become available in store repository within " + waitTimeMs + " ms.";
+      logger.error(errorMsg, e);
+      throw new VeniceException(errorMsg, e);
+    }
+  }
+
   /**
    * set up a new store partition and start the ingestion
    */
   protected void setupNewStorePartition() {
+    /**
+     * Wait for version metadata to become available in the store repository before opening the
+     * storage engine. This ensures the correct partition type (with or without replication metadata)
+     * is determined based on version-level config rather than falling back to store-level config.
+     */
+    waitForVersionToBeAvailable();
+
     /**
      * Waiting for push accessor to get initialized before starting ingestion.
      * Otherwise, it's possible that store ingestion starts without having the

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

@@ -2069,11 +2069,10 @@ protected void produceToLocalKafka(
             partition,
             kafkaUrl,
             beforeProcessingRecordTimestampNs,
-            leaderMetadataWrapper,
-            leaderProducedRecordContext);
+            kafkaClusterId);
       }
     } catch (Exception e) {
-      LOGGER.error("Failed to send Global RT DIV message", e); // don't fail ingestion if sending Global RT DIV fails
+      LOGGER.error("Failed to send Global RT DIV message", e);
     }
   }
 
@@ -2732,9 +2731,13 @@ protected Iterable<DefaultPubSubMessage> validateAndFilterOutDuplicateMessagesFr
         /**
          * TODO: An improvement can be made to fail all future versions for fatal DIV exceptions after EOP.
          */
+        // shouldProduceToVersionTopic() means this leader is consuming from a non-local-VT source
+        // (e.g. RT, remote VT in NR mode, or stream-reprocessing topic) and producing to local VT.
+        // For RT, use REALTIME_TOPIC_TYPE so segments are tracked per broker URL in rtSegments.
+        // For all other cases keep VERSION_TOPIC so segments are tracked in vtSegments and synced to OffsetRecord.
         TopicType topicType = PartitionTracker.VERSION_TOPIC;
-        // shouldProduceToVersionTopic() ensures this is a LEADER that is consuming from RT or remote VT
-        if (isGlobalRtDivEnabled() && shouldProduceToVersionTopic(pcs)) {
+        if (isGlobalRtDivEnabled() && shouldProduceToVersionTopic(pcs)
+            && topicPartition.getPubSubTopic().isRealTime()) {
           topicType = TopicType.of(REALTIME_TOPIC_TYPE, kafkaUrl);
         }
         validateMessage(topicType, consumerDiv, record, pcs, false);
@@ -3158,12 +3161,20 @@ void syncOffsetFromSnapshotIfNeeded(DefaultPubSubMessage record, PubSubTopicPart
 
     try {
       // VT DIV contains the latest consumed VT position (LCVP)
-      PartitionTracker vtDiv = consumerDiv.cloneVtProducerStates(partition, true);
+      long latestMessageTimeInMs = pcs.getLatestMessageTimeInMs();
+      PartitionTracker vtDiv = getConsumerDiv().cloneVtProducerStates(partition, true, latestMessageTimeInMs);
+
+      // Skip sync if no real VT progress has been made yet. Syncing with EARLIEST would persist
+      // EARLIEST to the OffsetRecord, causing the consumer to re-subscribe from EARLIEST on retry.
+      if (PubSubSymbolicPosition.EARLIEST.equals(vtDiv.getLatestConsumedVtPosition())) {
+        return;
+      }
+
       CompletableFuture<Void> lastFuture = pcs.getLastQueuedRecordPersistedFuture();
       storeBufferService.execSyncOffsetFromSnapshotAsync(topicPartition, vtDiv, lastFuture, this);
       // Reset consumer-side VT bytes so the size-based condition in shouldSyncOffsetFromSnapshot does not keep
       // firing for every subsequent record.
-      getConsumedBytesSinceLastSync().put(getVersionTopic().getName(), 0L);
+      pcs.resetConsumedBytesSinceLastGlobalRtDivSync(getVersionTopic().getName());
 
       // TODO: remove. this is a temporary log for debugging while the feature is in its infancy
       LOGGER.info(
@@ -3199,7 +3210,7 @@ boolean shouldSyncOffsetFromSnapshot(DefaultPubSubMessage consumerRecord, Partit
 
     // must be greater than the interval in shouldSendGlobalRtDiv() to not interfere
     final long syncBytesInterval = getSyncBytesInterval(pcs); // size-based sync condition
-    long vtConsumedBytesSinceLastSync = getConsumedBytesSinceLastSync().getOrDefault(getVersionTopic().getName(), 0L);
+    long vtConsumedBytesSinceLastSync = pcs.getConsumedBytesSinceLastGlobalRtDivSync(getVersionTopic().getName());
     return syncBytesInterval > 0 && (vtConsumedBytesSinceLastSync >= 2 * syncBytesInterval);
   }
 
@@ -3897,24 +3908,27 @@ public byte[] getGlobalRtDivKeyBytes(int partitionId, String brokerUrl) {
    * Upon completion, the {@link LeaderProducerCallback} will write the {@link GlobalRtDivState} to the StorageEngine.
    * When the drainer receives a Global RT DIV, that is the signal to sync the VT DIV to the OffsetRecord.
    * NOTE: This method is called per-broker. The broker url is included in the key.
-   * @param previousMessage the last message validated and produced to kafka before this GlobalRtDiv will be produced
+   * @param previousMessage the last RT message that was validated and produced to kafka before this GlobalRtDiv
+   *                        will be produced.
    */
   void sendGlobalRtDivMessage(
       DefaultPubSubMessage previousMessage,
       PartitionConsumptionState pcs,
       int partition,
       String brokerUrl,
       long beforeProcessingRecordTimestampNs,
-      LeaderMetadataWrapper leaderMetadataWrapper,
-      LeaderProducedRecordContext context) {
+      int kafkaClusterId) {
     final byte[] keyBytes = getGlobalRtDivKeyBytes(partition, brokerUrl);
     final PubSubTopicPartition topicPartition = previousMessage.getTopicPartition();
     TopicType realTimeTopicType = TopicType.of(REALTIME_TOPIC_TYPE, brokerUrl);
+    LeaderMetadataWrapper leaderMetadataWrapper =
+        new LeaderMetadataWrapper(previousMessage.getPosition(), kafkaClusterId, DEFAULT_TERM_ID);
 
     // Snapshot the RT DIV (single broker URL) in preparation to be produced
     // VT DIV contains the latest consumed VT position (LCVP)
-    PartitionTracker vtDiv = consumerDiv.cloneVtProducerStates(partition, true);
-    PartitionTracker rtDiv = consumerDiv.cloneRtProducerStates(partition, brokerUrl);
+    long latestMessageTimeInMs = pcs.getLatestMessageTimeInMs();
+    PartitionTracker vtDiv = consumerDiv.cloneVtProducerStates(partition, true, latestMessageTimeInMs);
+    PartitionTracker rtDiv = consumerDiv.cloneRtProducerStates(partition, brokerUrl, latestMessageTimeInMs);
     Map<CharSequence, ProducerPartitionState> rtDivPartitionStates = rtDiv.getPartitionStates(realTimeTopicType);
 
     // Create GlobalRtDivState (RT DIV + LCRP) and serialize into a byte array. Try compression.
@@ -3927,7 +3941,7 @@ void sendGlobalRtDivMessage(
         partition,
         brokerUrl,
         beforeProcessingRecordTimestampNs,
-        context,
+        kafkaClusterId,
         keyBytes,
         valueBytes,
         topicPartition,
@@ -3971,7 +3985,7 @@ void sendGlobalRtDivMessage(
             null,
             true);
 
-    consumedBytesSinceLastSync.put(brokerUrl, 0L); // reset the timer for the next sync, since RT DIV was just synced
+    pcs.resetConsumedBytesSinceLastGlobalRtDivSync(brokerUrl);
   }
 
   private byte[] createGlobalRtDivValueBytes(
@@ -4000,7 +4014,7 @@ private LeaderProducerCallback createGlobalRtDivCallback(
       int partition,
       String brokerUrl,
       long beforeProcessingRecordTimestampNs,
-      LeaderProducedRecordContext prevContext,
+      int kafkaClusterId,
       byte[] keyBytes,
       byte[] valueBytes,
       PubSubTopicPartition topicPartition,
@@ -4027,8 +4041,8 @@ private LeaderProducerCallback createGlobalRtDivCallback(
         prevMessage.getPosition(),
         System.currentTimeMillis(),
         divKey.getKeyLength() + valueBytes.length);
-    LeaderProducedRecordContext context = LeaderProducedRecordContext
-        .newPutRecord(prevContext.getConsumedKafkaClusterId(), prevContext.getConsumedPosition(), keyBytes, put);
+    LeaderProducedRecordContext context =
+        LeaderProducedRecordContext.newPutRecord(kafkaClusterId, prevMessage.getPosition(), keyBytes, put);
     LeaderProducerCallback divCallback =
         createProducerCallback(divMessage, pcs, context, partition, brokerUrl, beforeProcessingRecordTimestampNs);
 

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

@@ -164,6 +164,12 @@ enum LatchStatus {
    */
   private long processedRecordSizeSinceLastSync;
 
+  /**
+   * Tracks bytes consumed per source key (VT name or RT broker URL) since the last Global RT DIV sync.
+   * Stored per-partition so that each partition's sync cadence is independent.
+   */
+  private final Map<String, Long> consumedBytesSinceLastGlobalRtDivSync = new VeniceConcurrentHashMap<>();
+
   /** Minimum lgK supported by DataSketches HllSketch (mirrors package-private HllUtil.MIN_LOG_K). */
   static final int HLL_MIN_LOG_K = 4;
   /** Maximum lgK supported by DataSketches HllSketch (mirrors package-private HllUtil.MAX_LOG_K). */
@@ -629,6 +635,10 @@ public OffsetRecord getOffsetRecord() {
     return this.offsetRecord;
   }
 
+  public long getLatestMessageTimeInMs() {
+    return this.offsetRecord.calculateLatestMessageTimeInMs();
+  }
+
   public void setDeferredWrite(boolean deferredWrite) {
     this.deferredWrite = deferredWrite;
   }
@@ -760,6 +770,8 @@ public String toString() {
         .append(leaderFollowerState)
         .append(", leaderCompleteState=")
         .append(leaderCompleteState)
+        .append(", lastLeaderCompleteStateUpdateInMs=")
+        .append(lastLeaderCompleteStateUpdateInMs)
         .append(", consumeRemotely=")
         .append(consumeRemotely)
         .append(", latestMessageConsumedTimestampInMs=")
@@ -782,6 +794,21 @@ public void resetProcessedRecordSizeSinceLastSync() {
     this.processedRecordSizeSinceLastSync = 0;
   }
 
+  public long getConsumedBytesSinceLastGlobalRtDivSync(String key) {
+    return consumedBytesSinceLastGlobalRtDivSync.getOrDefault(key, 0L);
+  }
+
+  public void addConsumedBytesSinceLastGlobalRtDivSync(String key, long bytes) {
+    if (bytes <= 0) {
+      return;
+    }
+    consumedBytesSinceLastGlobalRtDivSync.merge(key, bytes, Long::sum);
+  }
+
+  public void resetConsumedBytesSinceLastGlobalRtDivSync(String key) {
+    consumedBytesSinceLastGlobalRtDivSync.put(key, 0L);
+  }
+
   public void setLeaderFollowerState(LeaderFollowerStateType state) {
     this.leaderFollowerState = state;
   }