address review comments

Author: m-nagarajan

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

@@ -937,7 +937,7 @@ private PubSubMessageHeaders computeActiveKeyCountSignal(
       if (partitionConsumptionState.decrementActiveKeyCount()) {
         return new PubSubMessageHeaders().add(KEY_DELETED_SIGNAL);
       }
-      invalidateActiveKeyCount(partitionConsumptionState, "Decrement underflow on leader during DCR");
+      invalidateActiveKeyCount(partitionConsumptionState, ActiveKeyCountInvalidationReason.LEADER_DCR_UNDERFLOW);
       return new PubSubMessageHeaders().add(KEY_COUNT_INVALIDATE_SIGNAL);
     }
     return EmptyPubSubMessageHeaders.SINGLETON;
@@ -1050,9 +1050,7 @@ private boolean isValuePresentForKey(
     try {
       return storageEngine.keyExists(partitionConsumptionState.getPartition(), storageKey);
     } catch (VeniceException e) {
-      // Transient RocksDB I/O failure: invalidate the count and assume key absent so ingestion
-      // doesn't halt and we stop publishing a wrong number.
-      invalidateActiveKeyCount(partitionConsumptionState, "keyExists failed", e);
+      invalidateActiveKeyCount(partitionConsumptionState, ActiveKeyCountInvalidationReason.KEY_EXISTS_FAILURE, e);
       return false;
     }
   }

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

@@ -0,0 +1,41 @@
+package com.linkedin.davinci.kafka.consumer;
+
+/**
+ * Reasons for invalidating a partition's active-key-count tracking. Each value carries a message
+ * template (some with a {@code %d} placeholder for runtime detail like the offending signal value
+ * or header length) that becomes the prefix of the operator-visible ERROR log emitted by
+ * {@link StoreIngestionTask#invalidateActiveKeyCount}.
+ */
+public enum ActiveKeyCountInvalidationReason {
+  /** Follower received {@code kcs=-1} but its count was already zero (count drifted). */
+  FOLLOWER_DECREMENT_UNDERFLOW("Decrement underflow on follower from kcs=-1"),
+  /** Follower received {@code kcs=0} (leader-propagated invalidation signal). */
+  LEADER_PROPAGATED_INVALIDATION("Leader propagated invalidation signal"),
+  /** Follower received a single-byte {@code kcs} value outside the {-1, 0, +1} contract. */
+  CORRUPT_KCS_SIGNAL_VALUE("Unexpected kcs signal value %d"),
+  /** Follower received a multi-byte {@code kcs} header (corrupt or future producer). */
+  CORRUPT_MULTI_BYTE_KCS_SIGNAL("Unexpected multi-byte kcs signal (length=%d)"),
+  /** Leader detected an underflow during DCR (count was zero but a delete was processed). */
+  LEADER_DCR_UNDERFLOW("Decrement underflow on leader during DCR"),
+  /**
+   * Leader's {@link com.linkedin.davinci.store.StorageEngine#keyExists} call (a RocksDB
+   * value-column-family lookup used to determine whether a key currently has a live value) threw.
+   * The transient I/O failure must not stop ingestion or leave the active count in a wrong state —
+   * invalidate so we stop publishing a stale value.
+   */
+  KEY_EXISTS_FAILURE("RocksDB value column family lookup failed");
+
+  private final String messageTemplate;
+
+  ActiveKeyCountInvalidationReason(String messageTemplate) {
+    this.messageTemplate = messageTemplate;
+  }
+
+  String getMessage() {
+    return messageTemplate;
+  }
+
+  String getMessage(int detail) {
+    return String.format(messageTemplate, detail);
+  }
+}

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

@@ -4740,48 +4740,76 @@ private void trackActiveKeyCount(
           partitionConsumptionState.incrementActiveKeyCount();
         } else if (signal == ActiveActiveStoreIngestionTask.KEY_DELETED_SIGNAL_VALUE) {
           if (!partitionConsumptionState.decrementActiveKeyCount()) {
-            invalidateActiveKeyCount(partitionConsumptionState, "Decrement underflow on follower from kcs=-1");
+            invalidateActiveKeyCount(
+                partitionConsumptionState,
+                ActiveKeyCountInvalidationReason.FOLLOWER_DECREMENT_UNDERFLOW);
           }
         } else if (signal == ActiveActiveStoreIngestionTask.KEY_COUNT_INVALIDATE_SIGNAL_VALUE) {
-          invalidateActiveKeyCount(partitionConsumptionState, "Leader propagated invalidation signal");
+          invalidateActiveKeyCount(
+              partitionConsumptionState,
+              ActiveKeyCountInvalidationReason.LEADER_PROPAGATED_INVALIDATION);
         } else {
-          invalidateActiveKeyCount(partitionConsumptionState, "Unexpected kcs signal value " + signal);
+          invalidateActiveKeyCount(
+              partitionConsumptionState,
+              ActiveKeyCountInvalidationReason.CORRUPT_KCS_SIGNAL_VALUE,
+              signal);
         }
       } else if (signalHeader != null && signalHeader.value() != null && signalHeader.value().length > 1) {
         invalidateActiveKeyCount(
             partitionConsumptionState,
-            "Unexpected multi-byte kcs signal (length=" + signalHeader.value().length + ")");
+            ActiveKeyCountInvalidationReason.CORRUPT_MULTI_BYTE_KCS_SIGNAL,
+            signalHeader.value().length);
       }
     }
   }
 
   /**
    * Sets the active key count to {@link OffsetRecord#ACTIVE_KEY_COUNT_NOT_TRACKED}, records the
-   * invalidation metric, and emits a rate-limited ERROR log. The state mutation and metric are
-   * wrapped in try/finally so the log fires even if either throws.
+   * invalidation metric on both the OTel and Tehuti paths, and emits a rate-limited ERROR log.
+   * The state mutation and metric are wrapped in try/finally so the log fires even if either throws.
+   */
+  final void invalidateActiveKeyCount(
+      PartitionConsumptionState partitionConsumptionState,
+      ActiveKeyCountInvalidationReason reason) {
+    invalidateActiveKeyCountAndLog(partitionConsumptionState, reason.getMessage(), null);
+  }
+
+  /**
+   * @param cause attached to the ERROR log; {@code null} yields a no-stack-trace log
+   */
+  final void invalidateActiveKeyCount(
+      PartitionConsumptionState partitionConsumptionState,
+      ActiveKeyCountInvalidationReason reason,
+      Throwable cause) {
+    invalidateActiveKeyCountAndLog(partitionConsumptionState, reason.getMessage(), cause);
+  }
+
+  /**
+   * @param detail integer rendered into the reason's {@code %d} placeholder
    */
   final void invalidateActiveKeyCount(
       PartitionConsumptionState partitionConsumptionState,
-      String reason,
+      ActiveKeyCountInvalidationReason reason,
+      int detail) {
+    invalidateActiveKeyCountAndLog(partitionConsumptionState, reason.getMessage(detail), null);
+  }
+
+  private void invalidateActiveKeyCountAndLog(
+      PartitionConsumptionState partitionConsumptionState,
+      String reasonText,
       Throwable cause) {
     try {
       partitionConsumptionState.setActiveKeyCount(ACTIVE_KEY_COUNT_NOT_TRACKED);
       recordActiveKeyCountInvalidation();
     } finally {
       String msg =
-          reason + " for replica " + partitionConsumptionState.getReplicaId() + "; invalidating activeKeyCount.";
+          reasonText + " for replica " + partitionConsumptionState.getReplicaId() + "; invalidating activeKeyCount.";
       if (!REDUNDANT_LOGGING_FILTER.isRedundantException(msg)) {
-        // log4j2's error(String, Throwable) accepts a null throwable, so this single call covers
-        // both with-cause and no-cause invocations.
         LOGGER.error(msg, cause);
       }
     }
   }
 
-  final void invalidateActiveKeyCount(PartitionConsumptionState partitionConsumptionState, String reason) {
-    invalidateActiveKeyCount(partitionConsumptionState, reason, null);
-  }
-
   /** Records active-key-count invalidation on both the OTel (per-version) and Tehuti (host-level) paths. */
   protected final void recordActiveKeyCountInvalidation() {
     versionedIngestionStats.recordActiveKeyCountInvalidation(storeName, versionNumber);

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

@@ -314,29 +314,7 @@ public HostLevelIngestionStats(
       registerSensor(new AsyncGauge((ignored, ignored2) -> ingestionTaskMap.size(), "ingestion_task_count"));
     }
 
-    // ACTIVE_KEY_COUNT_NOT_TRACKED = not tracked, 0 = tracked but empty.
-    // Cannot use measurable() because its 0-fallback conflates "untracked" with "empty".
-    if (activeKeyCountEnabled) {
-      if (isTotalStats) {
-        registerSensor(new AsyncGauge((ignored, ignored2) -> {
-          long total = 0;
-          boolean anyTracked = false;
-          for (StoreIngestionTask task: ingestionTaskMap.values()) {
-            long storeCount = task.getActiveKeyCount();
-            if (storeCount != ACTIVE_KEY_COUNT_NOT_TRACKED) {
-              anyTracked = true;
-              total += storeCount;
-            }
-          }
-          return anyTracked ? total : ACTIVE_KEY_COUNT_NOT_TRACKED;
-        }, "active_key_count"));
-      } else {
-        registerSensor(new AsyncGauge((ignored, ignored2) -> {
-          StoreIngestionTask sit = ingestionTaskMap.get(storeName);
-          return sit == null ? ACTIVE_KEY_COUNT_NOT_TRACKED : sit.getActiveKeyCount();
-        }, "active_key_count"));
-      }
-    }
+    registerActiveKeyCountGauge(activeKeyCountEnabled, isTotalStats, ingestionTaskMap, storeName);
 
     // Stats which are per-store only:
     String keySizeSensorName = "record_key_size_in_bytes";
@@ -760,6 +738,36 @@ public void recordActiveKeyCountInvalidation() {
     }
   }
 
+  /** Uses the ACTIVE_KEY_COUNT_NOT_TRACKED sentinel to distinguish untracked from empty (which {@code measurable()}'s 0-fallback would conflate). */
+  private void registerActiveKeyCountGauge(
+      boolean activeKeyCountEnabled,
+      boolean isTotalStats,
+      Map<String, StoreIngestionTask> ingestionTaskMap,
+      String storeName) {
+    if (!activeKeyCountEnabled) {
+      return;
+    }
+    if (isTotalStats) {
+      registerSensor(new AsyncGauge((ignored, ignored2) -> {
+        long total = 0;
+        boolean anyTracked = false;
+        for (StoreIngestionTask task: ingestionTaskMap.values()) {
+          long storeCount = task.getActiveKeyCount();
+          if (storeCount != ACTIVE_KEY_COUNT_NOT_TRACKED) {
+            anyTracked = true;
+            total += storeCount;
+          }
+        }
+        return anyTracked ? total : ACTIVE_KEY_COUNT_NOT_TRACKED;
+      }, "active_key_count"));
+      return;
+    }
+    registerSensor(new AsyncGauge((ignored, ignored2) -> {
+      StoreIngestionTask sit = ingestionTaskMap.get(storeName);
+      return sit == null ? ACTIVE_KEY_COUNT_NOT_TRACKED : sit.getActiveKeyCount();
+    }, "active_key_count"));
+  }
+
   public void recordTotalLeaderBytesConsumed(long bytes) {
     totalLeaderBytesConsumedRate.record(bytes);
   }

clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/ActiveKeyCountScenarioTest.java

@@ -338,7 +338,7 @@ public void testGaugeEmitsNegativeOneAfterFollowerInvalidateSignal() {
 
   /** A single invalidation event must increment both the Tehuti rate and the OTel counter exactly once. */
   @Test
-  public void testInvalidationParity_tehutiAndOtelBothRecord() {
+  public void testInvalidationParityRecordsToBothTehutiAndOtel() {
     InMemoryMetricReader reader = InMemoryMetricReader.create();
     VeniceMetricsRepository otelRepo = new VeniceMetricsRepository(
         new VeniceMetricsConfig.Builder().setMetricEntities(SERVER_METRIC_ENTITIES)