From fe7543c31a273d9a07b44787db389c50b68e9ad6 Mon Sep 17 00:00:00 2001
From: Yueh-Hsuan Chiang <93241502+yhchiang-sol@users.noreply.github.com>
Date: Wed, 19 Jan 2022 19:31:19 -0800
Subject: [PATCH] (Ledger Store) APIs for obtaining physical size of all data
 and coding shreds (#22443)

---
 ledger/src/blockstore.rs    | 19 +++++++++++++++++++
 ledger/src/blockstore_db.rs | 22 ++++++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/ledger/src/blockstore.rs b/ledger/src/blockstore.rs
index f998a0e690..ec09242abe 100644
--- a/ledger/src/blockstore.rs
+++ b/ledger/src/blockstore.rs
@@ -70,6 +70,7 @@ use {
 pub mod blockstore_purge;
 
 pub const BLOCKSTORE_DIRECTORY: &str = "rocksdb";
+pub const ROCKSDB_TOTAL_SST_FILES_SIZE: &str = "rocksdb.total-sst-files-size";
 
 thread_local!(static PAR_THREAD_POOL: RefCell<ThreadPool> = RefCell::new(rayon::ThreadPoolBuilder::new()
                     .num_threads(get_thread_count())
@@ -3167,6 +3168,24 @@ impl Blockstore {
         self.db.storage_size()
     }
 
+    /// Returns the total physical storage size contributed by all data shreds.
+    ///
+    /// Note that the reported size does not include those recently inserted
+    /// shreds that are still in memory.
+    pub fn total_data_shred_storage_size(&self) -> Result<u64> {
+        let shred_data_cf = self.db.column::<cf::ShredData>();
+        shred_data_cf.get_int_property(ROCKSDB_TOTAL_SST_FILES_SIZE)
+    }
+
+    /// Returns the total physical storage size contributed by all coding shreds.
+    ///
+    /// Note that the reported size does not include those recently inserted
+    /// shreds that are still in memory.
+    pub fn total_coding_shred_storage_size(&self) -> Result<u64> {
+        let shred_code_cf = self.db.column::<cf::ShredCode>();
+        shred_code_cf.get_int_property(ROCKSDB_TOTAL_SST_FILES_SIZE)
+    }
+
     pub fn is_primary_access(&self) -> bool {
         self.db.is_primary_access()
     }
diff --git a/ledger/src/blockstore_db.rs b/ledger/src/blockstore_db.rs
index 9e868a7f59..3a157fbcc9 100644
--- a/ledger/src/blockstore_db.rs
+++ b/ledger/src/blockstore_db.rs
@@ -501,6 +501,19 @@ impl Rocks {
     fn is_primary_access(&self) -> bool {
         self.1 == ActualAccessType::Primary
     }
+
+    /// Retrieves the specified RocksDB integer property of the current
+    /// column family.
+    ///
+    /// Full list of properties that return int values could be found
+    /// [here](https://github.com/facebook/rocksdb/blob/08809f5e6cd9cc4bc3958dd4d59457ae78c76660/include/rocksdb/db.h#L654-L689).
+    fn get_int_property_cf(&self, cf: &ColumnFamily, name: &str) -> Result<u64> {
+        match self.0.property_int_value_cf(cf, name) {
+            Ok(Some(value)) => Ok(value),
+            Ok(None) => Ok(0),
+            Err(e) => Err(BlockstoreError::RocksDb(e)),
+        }
+    }
 }
 
 pub trait Column {
@@ -1138,6 +1151,15 @@ where
     pub fn put_bytes(&self, key: C::Index, value: &[u8]) -> Result<()> {
         self.backend.put_cf(self.handle(), &C::key(key), value)
     }
+
+    /// Retrieves the specified RocksDB integer property of the current
+    /// column family.
+    ///
+    /// Full list of properties that return int values could be found
+    /// [here](https://github.com/facebook/rocksdb/blob/08809f5e6cd9cc4bc3958dd4d59457ae78c76660/include/rocksdb/db.h#L654-L689).
+    pub fn get_int_property(&self, name: &str) -> Result<u64> {
+        self.backend.get_int_property_cf(self.handle(), name)
+    }
 }
 
 impl<C> LedgerColumn<C>