From c153a83ea3e5ff56f679ef52c0ac5dfa63e73584 Mon Sep 17 00:00:00 2001
From: Sourav Maji <tester@linkedin.com>
Date: Fri, 31 Jan 2025 15:40:17 -0800
Subject: [PATCH] add new API to access CF metric value

---
 .../davinci/stats/RocksDBMemoryStats.java         | 15 +++++++--------
 .../store/rocksdb/RocksDBStoragePartition.java    | 15 ++++++++++++++-
 2 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/stats/RocksDBMemoryStats.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/stats/RocksDBMemoryStats.java
index ed8690356f..ef869797fe 100644
--- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/stats/RocksDBMemoryStats.java
+++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/stats/RocksDBMemoryStats.java
@@ -112,14 +112,13 @@ public RocksDBMemoryStats(MetricsRepository metricsRepository, String name, bool
         // Lock down the list of RocksDB interfaces while the collection is ongoing
         synchronized (hostedRocksDBPartitions) {
           for (RocksDBStoragePartition dbPartition: hostedRocksDBPartitions.values()) {
-            if (!(dbPartition instanceof ReplicationMetadataRocksDBStoragePartition)) {
-              continue;
-            }
-            try {
-              total += dbPartition.getRocksDBStatValue(metric);
-            } catch (VeniceException e) {
-              LOGGER.warn("Could not get rocksDB metric {} with error:", metric, e);
-              continue;
+            if (dbPartition instanceof ReplicationMetadataRocksDBStoragePartition) {
+              try {
+                total += dbPartition.getRocksDBCFStatValue(metric);
+              } catch (VeniceException e) {
+                LOGGER.warn("Could not get rocksDB CF metric {} with error:", metric, e);
+                continue;
+              }
             }
             if (INSTANCE_METRIC_DOMAINS.contains(metric)) {
               // Collect this metric once from any available partition and move on
diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/store/rocksdb/RocksDBStoragePartition.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/store/rocksdb/RocksDBStoragePartition.java
index 3ec74037c2..04808a97fc 100644
--- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/store/rocksdb/RocksDBStoragePartition.java
+++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/store/rocksdb/RocksDBStoragePartition.java
@@ -1,6 +1,7 @@
 package com.linkedin.davinci.store.rocksdb;
 
 import static com.linkedin.davinci.store.AbstractStorageEngine.METADATA_PARTITION_ID;
+import static com.linkedin.davinci.store.rocksdb.RocksDBSstFileWriter.REPLICATION_METADATA_COLUMN_FAMILY_INDEX;
 
 import com.linkedin.davinci.blobtransfer.BlobSnapshotManager;
 import com.linkedin.davinci.callback.BytesStreamingCallback;
@@ -924,10 +925,22 @@ private void deRegisterDBStats() {
   }
 
   public long getRocksDBStatValue(String statName) {
+    return getRocksDBStatValue(statName, 0);
+  }
+
+  public long getRocksDBCFStatValue(String statName) {
+    return getRocksDBStatValue(statName, REPLICATION_METADATA_COLUMN_FAMILY_INDEX);
+  }
+
+  private long getRocksDBStatValue(String statName, int cfIndex) {
     readCloseRWLock.readLock().lock();
     try {
       makeSureRocksDBIsStillOpen();
-      return rocksDB.getLongProperty(statName);
+      if (cfIndex == REPLICATION_METADATA_COLUMN_FAMILY_INDEX) {
+        return rocksDB.getLongProperty(columnFamilyHandleList.get(REPLICATION_METADATA_COLUMN_FAMILY_INDEX), statName);
+      } else {
+        return rocksDB.getLongProperty(statName);
+      }
     } catch (RocksDBException e) {
       throw new VeniceException(
           "Failed to get property value from RocksDB: " + replicaId + " for property: " + statName,