[server][protocol] Add active unique key count tracking for store versions (linkedin#2676)

Introduces an exact per-partition active key count and exposed via OTel and Tehuti 
metrics. The count is built in three stages: a batch baseline that counts all valid PUTs 
during ingestion, an optional RMD for batch push for hybrid stores to distinguish 
batch keys without extra reads, and hybrid ingestion continuing to build on top of 
batch baseline. This design ensures accuracy while minimizing I/O by avoiding 
redundant value column family reads and delegating computation primarily to the leader.

- Exact per-partition count via `AtomicLong` in  `PartitionConsumptionState`  
- Persisted with offsets in `PartitionState v22` (PR `linkedin#2761`)  
- Metrics:  
  - OTel `ASYNC_GAUGE` → `ingestion.key.active_count`  
  - Tehuti `AsyncGauge` → `active_key_count`  

# Approach  
- Controlled by 3 configs
## Batch RMD (Hybrid)  
- Config: `server.add.rmd.to.batch.push.for.hybrid.stores`  
- Add `ts=0` RMD for batch `PUT`s only  
- Helps with zero-I/O detection of batch keys  
## Batch Baseline  
- Config: `server.active.key.count.for.all.batch.push.enabled`  
- Count `PUT`s during SOP→EOP  
- Filters: chunk fragments + duplicates  
- SOP: `-1 → 0`; missed SOP ⇒ stays `-1`  
- Persist via `syncOffset()`  
## RT Signals  
- Config: `server.active.key.count.for.hybrid.store.enabled`  
- Leader computes alive state change  
- Emits `+1/-1/0` via `"kcs"` header  
- Followers mirror without DCR or value lookup

# Key Decisions  
- `ACTIVE_KEY_COUNT_NOT_TRACKED = -1` sentinel  
- Any inconsistency ⇒ invalidate the counting for that version  
- `decrementActiveKeyCount()` detects underflow  
- `keyExists` failure ⇒ invalidate, continue ingestion  
- Batch `DELETE` has no `ts=0` RMD  
- 3-tier lookup: cache → transient → `keyExists`

Author: m-nagarajan

build.gradle

@@ -329,7 +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/PartitionState/v21', PathValidation.DIRECTORY)
       ]
 
       def schemaDirs = [sourceDir]

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

@@ -69,12 +69,15 @@
 import static com.linkedin.venice.ConfigKeys.SERVER_AA_WC_INGESTION_STORAGE_LOOKUP_THREAD_POOL_SIZE;
 import static com.linkedin.venice.ConfigKeys.SERVER_AA_WC_WORKLOAD_PARALLEL_PROCESSING_ENABLED;
 import static com.linkedin.venice.ConfigKeys.SERVER_AA_WC_WORKLOAD_PARALLEL_PROCESSING_THREAD_POOL_SIZE;
+import static com.linkedin.venice.ConfigKeys.SERVER_ACTIVE_KEY_COUNT_FOR_ALL_BATCH_PUSH_ENABLED;
+import static com.linkedin.venice.ConfigKeys.SERVER_ACTIVE_KEY_COUNT_FOR_HYBRID_STORE_ENABLED;
 import static com.linkedin.venice.ConfigKeys.SERVER_ADAPTIVE_THROTTLER_ENABLED;
 import static com.linkedin.venice.ConfigKeys.SERVER_ADAPTIVE_THROTTLER_MULTI_GET_LATENCY_THRESHOLD;
 import static com.linkedin.venice.ConfigKeys.SERVER_ADAPTIVE_THROTTLER_READ_COMPUTE_GET_LATENCY_THRESHOLD;
 import static com.linkedin.venice.ConfigKeys.SERVER_ADAPTIVE_THROTTLER_SIGNAL_IDLE_THRESHOLD;
 import static com.linkedin.venice.ConfigKeys.SERVER_ADAPTIVE_THROTTLER_SIGNAL_REFRESH_INTERVAL_IN_SECONDS;
 import static com.linkedin.venice.ConfigKeys.SERVER_ADAPTIVE_THROTTLER_SINGLE_GET_LATENCY_THRESHOLD;
+import static com.linkedin.venice.ConfigKeys.SERVER_ADD_RMD_TO_BATCH_PUSH_FOR_HYBRID_STORES;
 import static com.linkedin.venice.ConfigKeys.SERVER_BATCH_REPORT_END_OF_INCREMENTAL_PUSH_STATUS_ENABLED;
 import static com.linkedin.venice.ConfigKeys.SERVER_BLOB_TRANSFER_ADAPTIVE_THROTTLER_ENABLED;
 import static com.linkedin.venice.ConfigKeys.SERVER_BLOB_TRANSFER_ADAPTIVE_THROTTLER_UPDATE_PERCENTAGE;
@@ -722,6 +725,9 @@ public class VeniceServerConfig extends VeniceClusterConfig {
   private final int lagMonitorCleanupCycle;
   private final boolean readQuotaInitializationFallbackEnabled;
   private final boolean ingestionProgressLoggingEnabled;
+  private final boolean addRmdToBatchPushForHybridStores;
+  private final boolean activeKeyCountForAllBatchPushEnabled;
+  private final boolean activeKeyCountForHybridStoreEnabled;
   private final int partialUpdateLargeResultLogThresholdBytes;
   private final long partialUpdateAmplificationReportIntervalMs;
 
@@ -1246,6 +1252,12 @@ public VeniceServerConfig(VeniceProperties serverProperties, Map<String, Map<Str
     this.readQuotaInitializationFallbackEnabled =
         serverProperties.getBoolean(SERVER_READ_QUOTA_INITIALIZATION_FALLBACK_ENABLED, true);
     this.ingestionProgressLoggingEnabled = serverProperties.getBoolean(POSITIONAL_PROGRESS_LOGGING_ENABLED, false);
+    this.addRmdToBatchPushForHybridStores =
+        serverProperties.getBoolean(SERVER_ADD_RMD_TO_BATCH_PUSH_FOR_HYBRID_STORES, false);
+    this.activeKeyCountForAllBatchPushEnabled =
+        serverProperties.getBoolean(SERVER_ACTIVE_KEY_COUNT_FOR_ALL_BATCH_PUSH_ENABLED, false);
+    this.activeKeyCountForHybridStoreEnabled =
+        serverProperties.getBoolean(SERVER_ACTIVE_KEY_COUNT_FOR_HYBRID_STORE_ENABLED, false);
     this.partialUpdateLargeResultLogThresholdBytes =
         serverProperties.getInt(PARTIAL_UPDATE_LARGE_RESULT_LOG_THRESHOLD_BYTES, 100 * 1024);
     this.partialUpdateAmplificationReportIntervalMs =
@@ -2261,6 +2273,18 @@ public boolean isIngestionProgressLoggingEnabled() {
     return ingestionProgressLoggingEnabled;
   }
 
+  public boolean isAddRmdToBatchPushForHybridStoresEnabled() {
+    return addRmdToBatchPushForHybridStores;
+  }
+
+  public boolean isActiveKeyCountForAllBatchPushEnabled() {
+    return activeKeyCountForAllBatchPushEnabled;
+  }
+
+  public boolean isActiveKeyCountForHybridStoreEnabled() {
+    return activeKeyCountForHybridStoreEnabled;
+  }
+
   public int getPartialUpdateLargeResultLogThresholdBytes() {
     return partialUpdateLargeResultLogThresholdBytes;
   }

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

@@ -21,6 +21,11 @@ public class MergeConflictResultWrapper {
   private final RmdWithValueSchemaId oldRmdWithValueSchemaId;
   private final ChunkedValueManifestContainer oldValueManifestContainer;
 
+  /** 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;
   private final ByteBuffer updatedRmdBytes;
@@ -35,6 +40,8 @@ public MergeConflictResultWrapper(
       MergeConflictResult mergeConflictResult,
       Lazy<ByteBufferValueRecord<ByteBuffer>> oldValueProvider,
       Lazy<ByteBuffer> oldValueByteBufferProvider,
+      boolean oldValueAlive,
+      long activeKeyCountBeforeAliveCheck,
       RmdWithValueSchemaId oldRmdWithValueSchemaId,
       ChunkedValueManifestContainer oldValueManifestContainer,
       ByteBuffer updatedValueBytes,
@@ -43,6 +50,8 @@ public MergeConflictResultWrapper(
     this.mergeConflictResult = mergeConflictResult;
     this.oldValueProvider = oldValueProvider;
     this.oldValueByteBufferProvider = oldValueByteBufferProvider;
+    this.oldValueAlive = oldValueAlive;
+    this.activeKeyCountBeforeAliveCheck = activeKeyCountBeforeAliveCheck;
     this.oldRmdWithValueSchemaId = oldRmdWithValueSchemaId;
     this.oldValueManifestContainer = oldValueManifestContainer;
     this.updatedValueBytes = updatedValueBytes;
@@ -77,6 +86,16 @@ public Lazy<ByteBuffer> getOldValueByteBufferProvider() {
     return oldValueByteBufferProvider;
   }
 
+  /** Captured before the transient record cache is updated, so it reflects the true pre-DCR state. */
+  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/MergeConflictResultWrapper.java

@@ -21,6 +21,7 @@
 import com.linkedin.venice.pubsub.api.PubSubTopicPartition;
 import com.linkedin.venice.serialization.avro.AvroProtocolDefinition;
 import com.linkedin.venice.storage.protocol.ChunkedValueManifest;
+import com.linkedin.venice.utils.ArrayUtils;
 import com.linkedin.venice.utils.LatencyUtils;
 import com.linkedin.venice.utils.concurrent.VeniceConcurrentHashMap;
 import com.linkedin.venice.utils.lazy.Lazy;
@@ -34,6 +35,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BooleanSupplier;
 import org.apache.avro.generic.GenericRecord;
@@ -348,6 +350,16 @@ enum LatchStatus {
    */
   private final Map<String, HeartbeatKey> cachedHeartbeatKeys;
 
+  /**
+   * Active logical key count. Batch phase: incremented per logical PUT. RT phase (hybrid A/A):
+   * adjusted by +1/-1 signals from conflict resolution. {@link OffsetRecord#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(OffsetRecord.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. */
   private volatile PartialUpdateAmplificationDetector partialUpdateAmplificationDetector;
 
@@ -415,6 +427,7 @@ public PartitionConsumptionState(
     this.lastLeaderCompleteStateUpdateInMs = 0;
     this.pendingReportIncPushVersionList = offsetRecord.getPendingReportIncPushVersionList();
     this.hasResubscribedAfterBootstrapAsCurrentVersion = false;
+    this.activeKeyCount.set(offsetRecord.getActiveKeyCount());
   }
 
   /** Create a fresh HLL sketch with {@link #HLL_DEFAULT_LOG_K}. */
@@ -450,6 +463,73 @@ public void restoreUniqueKeyCountHll(int lgK) {
       LOGGER.warn("Partition {} has no HLL checkpoint data to restore.", getPartition());
       return;
     }
+    restoreUniqueKeyCountHllFromCheckpoint(lgK, hllBytes);
+  }
+
+  public long getActiveKeyCount() {
+    return activeKeyCount.get();
+  }
+
+  public void setActiveKeyCount(long count) {
+    activeKeyCount.set(count);
+  }
+
+  public void incrementActiveKeyCount() {
+    activeKeyCount.incrementAndGet();
+  }
+
+  /**
+   * Called at SOP to mark this partition as actively counting. Without this, a mid-batch
+   * restart with a newly enabled config would start counting partway through, producing
+   * an inaccurate partial count. On restart, SOP is not re-processed (already past the
+   * checkpoint), so the count stays at {@link OffsetRecord#ACTIVE_KEY_COUNT_NOT_TRACKED} and batch
+   * records are skipped by {@link #incrementActiveKeyCountForBatchRecord}.
+   */
+  public void initializeActiveKeyCount() {
+    activeKeyCount.compareAndSet(OffsetRecord.ACTIVE_KEY_COUNT_NOT_TRACKED, 0);
+  }
+
+  /**
+   * Decrements activeKeyCount. If the count is already at 0, this indicates drift
+   * (more deletes than creates), so we invalidate to {@link OffsetRecord#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 : OffsetRecord.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.
+   * Requires prior {@link #initializeActiveKeyCount()} at SOP; skips if not initialized
+   * (mid-batch config enablement). Checkpoint-safe: partial counts are persisted.
+   */
+  public void incrementActiveKeyCountForBatchRecord(byte[] keyBytes) {
+    if (activeKeyCount.get() < 0) {
+      return; // Not initialized at SOP — mid-batch config enablement, skip
+    }
+    if (lastBatchKeyForDedup != null && ArrayUtils.compareUnsigned(keyBytes, lastBatchKeyForDedup) <= 0) {
+      return; // Speculative execution duplicate — key order went backwards
+    }
+    lastBatchKeyForDedup = keyBytes;
+    activeKeyCount.incrementAndGet();
+  }
+
+  /**
+   * Called at EOP. Releases dedup state. Empty partitions are already at 0 from
+   * {@link #initializeActiveKeyCount()} at SOP. If SOP was missed (mid-batch config enablement),
+   * the count stays at {@link OffsetRecord#ACTIVE_KEY_COUNT_NOT_TRACKED} — no partial baseline is created.
+   */
+  public void cleanupBatchKeyCountState() {
+    lastBatchKeyForDedup = null; // Release reference; dedup is only needed during batch
+  }
+
+  // --- HLL unique ingested key count methods ---
+
+  private void restoreUniqueKeyCountHllFromCheckpoint(int lgK, ByteBuffer hllBytes) {
     byte[] bytes = new byte[hllBytes.remaining()];
     hllBytes.duplicate().get(bytes);
     this.uniqueIngestedKeyCountHll = HllSketch.heapify(Memory.wrap(bytes));