From a6350796dfef212dda597b36178b2b9b082107e4 Mon Sep 17 00:00:00 2001
From: Dmitri Petrov <dmitri.petrov@rockset.com>
Date: Fri, 8 Dec 2023 19:45:56 +0000
Subject: [PATCH] Add counters for SST File reads/blocks in MultiGet.

---
 db/version_set.cc            | 30 ++++++++++++++++++++++++++++++
 include/rocksdb/statistics.h |  9 +++++++++
 monitoring/statistics.cc     |  8 +++++++-
 3 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/db/version_set.cc b/db/version_set.cc
index a6b772948ea..37ba0b73d9d 100644
--- a/db/version_set.cc
+++ b/db/version_set.cc
@@ -2521,9 +2521,17 @@ void Version::MultiGet(const ReadOptions& read_options, MultiGetRange* range,
           if (fp.GetHitFileLevel() == 0) {
             dump_stats_for_l0_file = true;
           }
+          // RocksDB-Cloud contribution begin
+          RecordTick(db_statistics_, MULTIGET_SST_FILE_READ_COUNT, 1);
+          // RocksDB-Cloud contribution end
         }
         if (s.ok()) {
           f = fp.GetNextFileInLevel();
+          // RocksDB-Cloud contribution begin
+          if (f) {
+            RecordTick(db_statistics_, MULTIGET_SST_SERIALIZED_FILE_READ_COUNT, 1);
+          }
+          // RocksDB-Cloud contribution end
         }
 #if USE_COROUTINES
       } else {
@@ -2567,6 +2575,9 @@ void Version::MultiGet(const ReadOptions& read_options, MultiGetRange* range,
         if (mget_tasks.size() > 0) {
           RecordTick(db_statistics_, MULTIGET_COROUTINE_COUNT,
                      mget_tasks.size());
+          // RocksDB-Cloud contribution begin
+          RecordTick(db_statistics_, MULTIGET_SST_FILE_READ_COUNT, mget_tasks.size());
+          // RocksDB-Cloud contribution end
           // Collect all results so far
           std::vector<Status> statuses = folly::coro::blockingWait(
               folly::coro::collectAllRange(std::move(mget_tasks))
@@ -2582,6 +2593,12 @@ void Version::MultiGet(const ReadOptions& read_options, MultiGetRange* range,
 
           if (s.ok() && fp.KeyMaySpanNextFile()) {
             f = fp.GetNextFileInLevel();
+            // RocksDB-Cloud contribution begin
+            if (f) {
+              RecordTick(db_statistics_, 
+                         MULTIGET_SST_SERIALIZED_FILE_READ_COUNT, 1);
+            }
+            // RocksDB-Cloud contribution end
           }
         }
 #endif  // USE_COROUTINES
@@ -2763,6 +2780,10 @@ Status Version::ProcessBatch(
                             table_handle, std::get<0>(stat->second),
                             std::get<1>(stat->second),
                             std::get<2>(stat->second));
+        // RocksDB-Cloud contribution begin
+        RecordTick(db_statistics_, 
+                   MULTIGET_SST_FILE_READ_COUNT, mget_tasks.size());
+        // RocksDB-Cloud contribution end
       } else {
         mget_tasks.emplace_back(MultiGetFromSSTCoroutine(
             read_options, file_range, fp.GetHitFileLevel(), skip_filters,
@@ -2853,6 +2874,10 @@ Status Version::MultiGetAsync(
       if (mget_tasks.size() > 0) {
         assert(waiting.size());
         RecordTick(db_statistics_, MULTIGET_COROUTINE_COUNT, mget_tasks.size());
+        // RocksDB-Cloud contribution begin
+        RecordTick(db_statistics_, 
+                   MULTIGET_SST_FILE_READ_COUNT, mget_tasks.size());
+        // RocksDB-Cloud contribution end
         // Collect all results so far
         std::vector<Status> statuses = folly::coro::blockingWait(
             folly::coro::collectAllRange(std::move(mget_tasks))
@@ -2879,6 +2904,11 @@ Status Version::MultiGetAsync(
           // and no need to prepare the next level.
           if (!fp.GetHitFile() && !fp.GetRange().empty()) {
             fp.PrepareNextLevelForSearch();
+          } else {
+            // RocksDB-Cloud contribution begin
+            RecordTick(db_statistics_, 
+                       MULTIGET_SST_SERIALIZED_FILE_READ_COUNT, 1);
+            // RocksDB-Cloud contribution end
           }
         }
         to_process.swap(waiting);
diff --git a/include/rocksdb/statistics.h b/include/rocksdb/statistics.h
index 42a938f30c4..cae36be882d 100644
--- a/include/rocksdb/statistics.h
+++ b/include/rocksdb/statistics.h
@@ -452,6 +452,15 @@ enum Tickers : uint32_t {
   READ_ASYNC_MICROS,
   // Number of errors returned to the async read callback
   ASYNC_READ_ERROR_COUNT,
+  
+  // RocksDB-Cloud contribution begin
+
+  // Total number of files read in MultiGet operations
+  MULTIGET_SST_FILE_READ_COUNT,
+  // Number of times file reads were serialized behind other reads
+  MULTIGET_SST_SERIALIZED_FILE_READ_COUNT,
+
+  // RocksDB-Cloud contribution end
 
   TICKER_ENUM_MAX
 };
diff --git a/monitoring/statistics.cc b/monitoring/statistics.cc
index e01eed3f381..4178a3072b6 100644
--- a/monitoring/statistics.cc
+++ b/monitoring/statistics.cc
@@ -235,7 +235,13 @@ const std::vector<std::pair<Tickers, std::string>> TickersNameMap = {
     {BLOB_DB_CACHE_BYTES_READ, "rocksdb.blobdb.cache.bytes.read"},
     {BLOB_DB_CACHE_BYTES_WRITE, "rocksdb.blobdb.cache.bytes.write"},
     {READ_ASYNC_MICROS, "rocksdb.read.async.micros"},
-    {ASYNC_READ_ERROR_COUNT, "rocksdb.async.read.error.count"}};
+    {ASYNC_READ_ERROR_COUNT, "rocksdb.async.read.error.count"},
+    // RocksDB-Cloud contribution begin
+    {MULTIGET_SST_FILE_READ_COUNT, "rodksdb.number.multiget.sst_file_reads"},
+    {MULTIGET_SST_SERIALIZED_FILE_READ_COUNT,
+        "rocksdb.number.multiget.sst_serialized_file_reads"}
+    // RocksDB-Cloud contribution end
+    };
 
 const std::vector<std::pair<Histograms, std::string>> HistogramsNameMap = {
     {DB_GET, "rocksdb.db.get.micros"},