address review comments: invalidation, metrics, and correctness fixes

- Corrupt/multi-byte kcs signal invalidates count instead of silently ignoring
- Strict length==1 check on follower signal header (rejects multi-byte)
- Decrement underflow invalidates to ACTIVE_KEY_COUNT_NOT_TRACKED (not clamp
to 0)
- keyExists failure in Tier 3 catches VeniceException, invalidates, returns
false
- Leader propagates KEY_COUNT_INVALIDATE_SIGNAL (value=0) to followers on
invalidation
- Follower handles invalidate signal, corrupt signal, and underflow
consistently
- Added ACTIVE_KEY_COUNT_NOT_TRACKED constant, replaced all inline -1
sentinels
- Added OTel COUNTER metric (ACTIVE_KEY_COUNT_INVALIDATION) via
AggVersionedIngestionStats
- Added Tehuti AsyncGauge (active_key_count) and LongAdderRateGauge
(active_key_count_invalidation)
- Added activeKeyCountBeforeAliveCheck snapshot for mid-record invalidation
detection
- Rate-limited all error logs via REDUNDANT_LOGGING_FILTER
- Updated all Javadocs to reflect 3 signal types (+1/-1/0) and invalidation
semantics
- Added tests: underflow, invalidate signal, multi-byte, keyExists failure,
OTel gauge after invalidation

Author: m-nagarajan

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

@@ -2,6 +2,7 @@
 
 import static com.linkedin.davinci.kafka.consumer.AggKafkaConsumerService.getKeyLevelLockMaxPoolSizeBasedOnServerConfig;
 import static com.linkedin.davinci.kafka.consumer.LeaderFollowerStateType.LEADER;
+import static com.linkedin.davinci.kafka.consumer.PartitionConsumptionState.ACTIVE_KEY_COUNT_NOT_TRACKED;
 import static com.linkedin.venice.VeniceConstants.REWIND_TIME_DECIDED_BY_SERVER;
 import static com.linkedin.venice.writer.VeniceWriter.APP_DEFAULT_LOGICAL_TS;
 
@@ -110,14 +111,19 @@ public class ActiveActiveStoreIngestionTask extends LeaderFollowerStoreIngestion
   private final boolean addRmdToBatchPushForHybridStores;
   /** @see ConfigKeys#SERVER_ACTIVE_KEY_COUNT_FOR_HYBRID_STORE_ENABLED */
   private final boolean activeKeyCountForHybridStoreEnabled;
+  /** @see #computeActiveKeyCountSignal for signal semantics. */
   static final byte KEY_CREATED_SIGNAL_VALUE = 1;
   static final byte KEY_DELETED_SIGNAL_VALUE = -1;
+  static final byte KEY_COUNT_INVALIDATE_SIGNAL_VALUE = 0;
   static final PubSubMessageHeader KEY_CREATED_SIGNAL =
       new PubSubMessageHeader(StoreIngestionTask.KEY_COUNT_SIGNAL_HEADER, new byte[] { KEY_CREATED_SIGNAL_VALUE });
   static final PubSubMessageHeader KEY_DELETED_SIGNAL =
       new PubSubMessageHeader(StoreIngestionTask.KEY_COUNT_SIGNAL_HEADER, new byte[] { KEY_DELETED_SIGNAL_VALUE });
+  static final PubSubMessageHeader KEY_COUNT_INVALIDATE_SIGNAL = new PubSubMessageHeader(
+      StoreIngestionTask.KEY_COUNT_SIGNAL_HEADER,
+      new byte[] { KEY_COUNT_INVALIDATE_SIGNAL_VALUE });
 
-  /** RMD bytes (ts=0) with superset schema ID prepended. For chunk manifests and DELETEs. */
+  /** RMD bytes (ts=0) with superset schema ID prepended. For chunk manifests (superset schema ID known at construction). */
   private final byte[] defaultBatchRmdWithSchemaIdPrefix;
   /** RMD bytes (ts=0) without schema ID prefix. For non-chunked PUTs (schema ID varies). */
   private final byte[] defaultBatchRmdBytes;
@@ -619,6 +625,7 @@ private PubSubMessageProcessedResult processActiveActiveMessage(
               oldValueProvider,
               oldValueByteBufferProvider,
               false,
+              ACTIVE_KEY_COUNT_NOT_TRACKED, // ignored update — no signal computed
               rmdWithValueSchemaID,
               valueManifestContainer,
               null,
@@ -653,6 +660,10 @@ private PubSubMessageProcessedResult processActiveActiveMessage(
       final ByteBuffer updatedRmdBytes =
           rmdSerDe.serializeRmdRecord(mergeConflictResult.getValueSchemaId(), mergeConflictResult.getRmdRecord());
 
+      // Snapshot count before wasOldValueAlive, which may invalidate it on keyExists failure.
+      // Used by computeActiveKeyCountSignal to detect mid-record invalidation and propagate to followers.
+      long activeKeyCountBeforeAliveCheck = partitionConsumptionState.getActiveKeyCount();
+
       // Must be captured before setTransientRecord() below, which overwrites the transient cache.
       boolean oldValueAlive =
           wasOldValueAlive(rmdWithValueSchemaID, oldValueByteBufferProvider, partitionConsumptionState, keyBytes);
@@ -680,6 +691,7 @@ private PubSubMessageProcessedResult processActiveActiveMessage(
               oldValueProvider,
               oldValueByteBufferProvider,
               oldValueAlive,
+              activeKeyCountBeforeAliveCheck,
               rmdWithValueSchemaID,
               valueManifestContainer,
               updatedValueBytes,
@@ -867,21 +879,40 @@ private ByteBufferValueRecord<ByteBuffer> getValueBytesForKey(
   }
 
   /**
-   * Computes the active key count signal (+1/-1/none) and returns the VT headers to propagate
-   * to followers. Updates the PCS active key count directly on the leader.
+   * Computes the active key count signal and returns the VT headers to propagate to followers.
+   * Updates the PCS active key count directly on the leader.
+   *
+   * <p>Three signal types:
+   * <ul>
+   *   <li>{@code +1} (KEY_CREATED_SIGNAL): dead→alive transition, count incremented.</li>
+   *   <li>{@code -1} (KEY_DELETED_SIGNAL): alive→dead transition, count decremented.</li>
+   *   <li>{@code 0} (KEY_COUNT_INVALIDATE_SIGNAL): count invalidated — emitted when decrement
+   *       underflows (count was 0, drift detected) or when the count was invalidated mid-record
+   *       (e.g., keyExists failure in {@link #isValuePresentForKey}). Tells followers to also
+   *       invalidate. We choose to report {@link PartitionConsumptionState#ACTIVE_KEY_COUNT_NOT_TRACKED}
+   *       rather than continue publishing an inaccurate count — a missing metric is better
+   *       than a wrong one.</li>
+   * </ul>
    *
    * <p>Returns {@link EmptyPubSubMessageHeaders#SINGLETON} when no signal is needed (feature
-   * disabled, no batch baseline, or alive-to-alive/dead-to-dead transition). Returns a freshly
-   * created {@link PubSubMessageHeaders} with the "kcs" signal header when the key's alive/dead
-   * state changes. The fresh-per-record creation is safe even though
-   * {@code VeniceWriter.sendMessage} may mutate the headers (e.g., adding a view-partition
-   * header) — each record gets its own instance.
+   * disabled, count already invalidated before this record, or alive-to-alive/dead-to-dead
+   * transition). Returns a freshly created {@link PubSubMessageHeaders} for each signal — safe
+   * even though {@code VeniceWriter.sendMessage} may mutate headers.
    */
   private PubSubMessageHeaders computeActiveKeyCountSignal(
       MergeConflictResultWrapper mergeConflictResultWrapper,
       MergeConflictResult mergeConflictResult,
       PartitionConsumptionState partitionConsumptionState) {
-    if (!activeKeyCountForHybridStoreEnabled || partitionConsumptionState.getActiveKeyCount() < 0) {
+    if (!activeKeyCountForHybridStoreEnabled) {
+      return EmptyPubSubMessageHeaders.SINGLETON;
+    }
+    long currentCount = partitionConsumptionState.getActiveKeyCount();
+    if (currentCount < 0) {
+      // Count was invalidated. If it was valid before this record's wasOldValueAlive call,
+      // the invalidation happened mid-record (e.g., keyExists failure). Propagate to followers.
+      if (mergeConflictResultWrapper.getActiveKeyCountBeforeAliveCheck() >= 0) {
+        return new PubSubMessageHeaders().add(KEY_COUNT_INVALIDATE_SIGNAL);
+      }
       return EmptyPubSubMessageHeaders.SINGLETON;
     }
     boolean wasAlive = mergeConflictResultWrapper.wasOldValueAlive();
@@ -890,8 +921,14 @@ private PubSubMessageHeaders computeActiveKeyCountSignal(
       partitionConsumptionState.incrementActiveKeyCount();
       return new PubSubMessageHeaders().add(KEY_CREATED_SIGNAL);
     } else if (wasAlive && !isAlive) {
-      partitionConsumptionState.decrementActiveKeyCount();
-      return new PubSubMessageHeaders().add(KEY_DELETED_SIGNAL);
+      if (partitionConsumptionState.decrementActiveKeyCount()) {
+        return new PubSubMessageHeaders().add(KEY_DELETED_SIGNAL);
+      }
+      // Underflow detected (count was 0 but got a delete) — count drifted, now invalidated to
+      // ACTIVE_KEY_COUNT_NOT_TRACKED.
+      aggVersionedIngestionStats.recordActiveKeyCountInvalidation(storeName, versionNumber);
+      getHostLevelIngestionStats().recordActiveKeyCountInvalidation();
+      return new PubSubMessageHeaders().add(KEY_COUNT_INVALIDATE_SIGNAL);
     }
     return EmptyPubSubMessageHeaders.SINGLETON;
   }
@@ -921,7 +958,9 @@ private PubSubMessageHeaders computeActiveKeyCountSignal(
    *       is the only branch that adds a net-new RocksDB read.</li>
    *   <li>{@code rmd.ts==0} (batch sentinel): batch PUT, never RT-written — alive.
    *       <b>Value CF lookup: NO.</b> Residual edge case: reprocessing PUT then DELETE for
-   *       same key during batch leaves stale ts=0 RMD (rare).</li>
+   *       same key during batch leaves stale ts=0 RMD (rare). Impact: at most one incorrect
+   *       +1 or -1 signal per affected key, bounded by the number of batch-PUT-then-DELETE
+   *       sequences for the same key in a single batch push.</li>
    *   <li>{@code rmd.ts>0}: previously RT-written, could be alive or dead — must check.
    *       <b>Value CF lookup:</b> Tier 1 (DCR cached) for field-level/tie/UPDATE — free.
    *       Tier 2 (transient cache) for value-level new-wins — typically free.
@@ -973,8 +1012,8 @@ private boolean wasOldValueAlive(
    * </ol>
    *
    * <p>In steady state with recommended config (addRmdToBatchPush ON), Tier 3 is rarely
-   * reached: branch 1 (new keys) and branch 3 (batch keys with ts=0) avoid this method
-   * entirely, and branch 4 (RT keys) typically hits Tier 1 or Tier 2.
+   * reached: new keys with rmd==null and batch keys with ts=0 sentinel avoid this method
+   * entirely, and RT keys (rmd.ts&gt;0) typically hit Tier 1 or Tier 2.
    */
   private boolean isValuePresentForKey(
       Lazy<ByteBuffer> oldValueByteBufferProvider,
@@ -989,11 +1028,25 @@ private boolean isValuePresentForKey(
     if (transientRecord != null) {
       return transientRecord.getValue() != null;
     }
-    // Tier 3: Storage engine existence check (bloom filter → disk if needed).
+    // Tier 3: Storage engine existence check (RocksDB disk read).
     // For chunked stores, the actual RocksDB key has a chunking suffix appended.
     byte[] storageKey =
         isChunked() ? ChunkingUtils.KEY_WITH_CHUNKING_SUFFIX_SERIALIZER.serializeNonChunkedKey(key) : key;
-    return storageEngine.keyExists(partitionConsumptionState.getPartition(), storageKey);
+    try {
+      return storageEngine.keyExists(partitionConsumptionState.getPartition(), storageKey);
+    } catch (VeniceException e) {
+      // A transient RocksDB I/O failure must not halt ingestion. Invalidate the count so we stop
+      // publishing a wrong number, and return false (assume key absent) to skip the signal.
+      partitionConsumptionState.setActiveKeyCount(ACTIVE_KEY_COUNT_NOT_TRACKED);
+      aggVersionedIngestionStats.recordActiveKeyCountInvalidation(storeName, versionNumber);
+      getHostLevelIngestionStats().recordActiveKeyCountInvalidation();
+      String msg =
+          "keyExists failed for replica " + partitionConsumptionState.getReplicaId() + "; invalidating activeKeyCount.";
+      if (!REDUNDANT_LOGGING_FILTER.isRedundantException(msg)) {
+        LOGGER.error(msg, e);
+      }
+      return false;
+    }
   }
 
   ByteBuffer getCurrentValueFromTransientRecord(PartitionConsumptionState.TransientRecord transientRecord) {

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

@@ -23,6 +23,8 @@ public class MergeConflictResultWrapper {
 
   /** Whether a live value existed before this write's conflict resolution (for key count signal). */
   private final boolean oldValueAlive;
+  /** Snapshot of activeKeyCount before wasOldValueAlive ran. Used to detect mid-record invalidation. */
+  private final long activeKeyCountBeforeAliveCheck;
 
   // Serialized and potentially compressed updated value bytes
   private final ByteBuffer updatedValueBytes;
@@ -39,6 +41,7 @@ public MergeConflictResultWrapper(
       Lazy<ByteBufferValueRecord<ByteBuffer>> oldValueProvider,
       Lazy<ByteBuffer> oldValueByteBufferProvider,
       boolean oldValueAlive,
+      long activeKeyCountBeforeAliveCheck,
       RmdWithValueSchemaId oldRmdWithValueSchemaId,
       ChunkedValueManifestContainer oldValueManifestContainer,
       ByteBuffer updatedValueBytes,
@@ -48,6 +51,7 @@ public MergeConflictResultWrapper(
     this.oldValueProvider = oldValueProvider;
     this.oldValueByteBufferProvider = oldValueByteBufferProvider;
     this.oldValueAlive = oldValueAlive;
+    this.activeKeyCountBeforeAliveCheck = activeKeyCountBeforeAliveCheck;
     this.oldRmdWithValueSchemaId = oldRmdWithValueSchemaId;
     this.oldValueManifestContainer = oldValueManifestContainer;
     this.updatedValueBytes = updatedValueBytes;
@@ -87,6 +91,11 @@ public boolean wasOldValueAlive() {
     return oldValueAlive;
   }
 
+  /** Returns the activeKeyCount snapshot taken before wasOldValueAlive ran (for invalidation detection). */
+  public long getActiveKeyCountBeforeAliveCheck() {
+    return activeKeyCountBeforeAliveCheck;
+  }
+
   public RmdWithValueSchemaId getOldRmdWithValueSchemaId() {
     return oldRmdWithValueSchemaId;
   }

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

@@ -344,13 +344,17 @@ enum LatchStatus {
    */
   private final Map<String, HeartbeatKey> cachedHeartbeatKeys;
 
+  /** Sentinel value indicating the active key count is not tracked or has been invalidated. */
+  public static final long ACTIVE_KEY_COUNT_NOT_TRACKED = -1;
+
   /**
    * Active logical key count. Batch phase: incremented per logical PUT. RT phase (hybrid A/A):
-   * adjusted by +1/-1 signals from conflict resolution. -1 = not tracked (no batch baseline);
-   * RT signals are skipped when -1. AtomicLong for AA/WC parallel processing.
+   * adjusted by +1/-1 signals from conflict resolution. {@link #ACTIVE_KEY_COUNT_NOT_TRACKED} = not tracked
+   * (no batch baseline); RT signals are skipped when not tracked. AtomicLong for AA/WC parallel processing.
    */
-  private final AtomicLong activeKeyCount = new AtomicLong(-1);
-  /** Last PUT key for batch dedup. Not persisted. @see #incrementActiveKeyCountForBatchRecord */
+  private final AtomicLong activeKeyCount = new AtomicLong(ACTIVE_KEY_COUNT_NOT_TRACKED);
+  /** Last PUT key for batch dedup. Not persisted; used by {@link #incrementActiveKeyCountForBatchRecord}.
+   *  Single-threaded: only accessed from the drainer thread during batch (pre-EOP). */
   private byte[] lastBatchKeyForDedup;
 
   /** Lazily allocated per-partition detector for partial-update amplification. */
@@ -471,33 +475,41 @@ public void incrementActiveKeyCount() {
     activeKeyCount.incrementAndGet();
   }
 
-  public void decrementActiveKeyCount() {
-    activeKeyCount.decrementAndGet();
+  /**
+   * Decrements activeKeyCount. If the count is already at 0, this indicates drift
+   * (more deletes than creates), so we invalidate to {@link #ACTIVE_KEY_COUNT_NOT_TRACKED}
+   * rather than continue tracking with a wrong baseline. If already not tracked, stays not tracked.
+   *
+   * @return true if the count was successfully decremented, false if invalidated or already invalid
+   */
+  public boolean decrementActiveKeyCount() {
+    long prev = activeKeyCount.getAndUpdate(v -> v > 0 ? v - 1 : ACTIVE_KEY_COUNT_NOT_TRACKED);
+    return prev > 0;
   }
 
   /**
    * Increments activeKeyCount for a batch PUT, skipping speculative execution duplicates
    * (key &lt;= last key). Only PUTs call this — DELETEs are tombstones and excluded.
-   * First call initializes count from -1 to 1. Checkpoint-safe: partial counts are persisted.
+   * First call initializes count from {@link #ACTIVE_KEY_COUNT_NOT_TRACKED} to 1. Checkpoint-safe: partial counts are persisted.
    */
   public void incrementActiveKeyCountForBatchRecord(byte[] keyBytes) {
     if (lastBatchKeyForDedup != null && ArrayUtils.compareUnsigned(keyBytes, lastBatchKeyForDedup) <= 0) {
       return; // Speculative execution duplicate — key order went backwards
     }
     lastBatchKeyForDedup = keyBytes;
-    activeKeyCount.compareAndSet(-1, 0);
+    activeKeyCount.compareAndSet(ACTIVE_KEY_COUNT_NOT_TRACKED, 0);
     activeKeyCount.incrementAndGet();
   }
 
-  /** Called at EOP. Sets -1 to 0 for empty partitions so the metric reports 0. */
+  /** Called at EOP. Sets {@link #ACTIVE_KEY_COUNT_NOT_TRACKED} to 0 for empty partitions so the metric reports 0. */
   public void finalizeActiveKeyCountForBatchPush() {
-    if (activeKeyCount.get() == -1) {
+    if (activeKeyCount.get() == ACTIVE_KEY_COUNT_NOT_TRACKED) {
       activeKeyCount.set(0);
     }
     lastBatchKeyForDedup = null; // Release reference; dedup is only needed during batch
   }
 
-  // --- HLL unique ingested key count methods (from main) ---
+  // --- HLL unique ingested key count methods ---
 
   private void restoreUniqueKeyCountHllFromCheckpoint(int lgK, ByteBuffer hllBytes) {
     byte[] bytes = new byte[hllBytes.remaining()];