From 4f0e887702ba059b1cb93f520e71f6c7c178dbb8 Mon Sep 17 00:00:00 2001
From: Yueh-Hsuan Chiang <93241502+yhchiang-sol@users.noreply.github.com>
Date: Thu, 7 Apr 2022 00:15:00 -0700
Subject: [PATCH] (LedgerStore) Report RocksDB perf metrics for Protobuf
 Columns (#24065)

This PR enables the reporting of both RocksDB read and write perf metrics for ProtobufColumns,
including TransactionStatus and Rewards.
---
 ledger/src/blockstore_db.rs | 25 ++++++++++++++++++++++---
 1 file changed, 22 insertions(+), 3 deletions(-)

diff --git a/ledger/src/blockstore_db.rs b/ledger/src/blockstore_db.rs
index b97080d53..14c03882d 100644
--- a/ledger/src/blockstore_db.rs
+++ b/ledger/src/blockstore_db.rs
@@ -2399,7 +2399,13 @@ where
         &self,
         key: C::Index,
     ) -> Result<Option<C::Type>> {
-        if let Some(serialized_value) = self.backend.get_cf(self.handle(), &C::key(key))? {
+        let is_perf_context_enabled = maybe_collect_perf_context();
+        let result = self.backend.get_cf(self.handle(), &C::key(key));
+        if is_perf_context_enabled {
+            report_read_perf_context(C::rocksdb_get_perf_metric_header(&self.column_options));
+        }
+
+        if let Some(serialized_value) = result? {
             let value = match C::Type::decode(&serialized_value[..]) {
                 Ok(value) => value,
                 Err(_) => deserialize::<T>(&serialized_value)?.into(),
@@ -2411,7 +2417,13 @@ where
     }
 
     pub fn get_protobuf(&self, key: C::Index) -> Result<Option<C::Type>> {
-        if let Some(serialized_value) = self.backend.get_cf(self.handle(), &C::key(key))? {
+        let is_perf_context_enabled = maybe_collect_perf_context();
+        let result = self.backend.get_cf(self.handle(), &C::key(key));
+        if is_perf_context_enabled {
+            report_read_perf_context(C::rocksdb_get_perf_metric_header(&self.column_options));
+        }
+
+        if let Some(serialized_value) = result? {
             Ok(Some(C::Type::decode(&serialized_value[..])?))
         } else {
             Ok(None)
@@ -2421,7 +2433,14 @@ where
     pub fn put_protobuf(&self, key: C::Index, value: &C::Type) -> Result<()> {
         let mut buf = Vec::with_capacity(value.encoded_len());
         value.encode(&mut buf)?;
-        self.backend.put_cf(self.handle(), &C::key(key), &buf)
+
+        let is_perf_context_enabled = maybe_collect_perf_context();
+        let result = self.backend.put_cf(self.handle(), &C::key(key), &buf);
+        if is_perf_context_enabled {
+            report_write_perf_context(C::rocksdb_put_perf_metric_header(&self.column_options));
+        }
+
+        result
     }
 }