Add new API CacheReservationManager::GetTotalMemoryUsage() (#9071)

Summary:
Note: This PR is the 2nd PR of a bigger PR stack (https://github.com/facebook/rocksdb/pull/9073).

Context:
`CacheReservationManager::UpdateCacheReservation(std::size_t new_memory_used)` accepts an accumulated total memory used (e.g, used 10MB so far) instead of usage change (e.g, increase by 5 MB, decrease by 5 MB). It has benefits including consolidating API for increase and decrease as described in https://github.com/facebook/rocksdb/pull/8506.

However, not every `CacheReservationManager` user keeps track of this accumulated total memory usage. For example, Bloom/Ribbon Filter construction (e.g, [here](822d729fcd/table/block_based/filter_policy.cc (L587)) in https://github.com/facebook/rocksdb/pull/9073) does not  while WriteBufferManager and compression dictionary buffering do.

Considering future users might or might not keep track of this counter and implementing this counter within `CacheReservationManager` is easy due to the passed-in `std::size_t new_memory_used` in calling `CacheReservationManager::UpdateCacheReservation(std::size_t new_memory_used)`, it is proposed to add a new API `CacheReservationManager::GetTotalMemoryUsage()`.

As noted in the API comments,   since `CacheReservationManager` is NOT thread-safe, external synchronization is
 needed in calling `UpdateCacheReservation()` if you want `GetTotalMemoryUsed()` returns the indeed latest memory used.
- Added and updated private counter `memory_used_` every time `CacheReservationManager::UpdateCacheReservation(std::size_t new_memory_used)` is called regardless if the call returns non-okay status
- Added `CacheReservationManager::GetTotalMemoryUsage()` to return `memory_used_`

Pull Request resolved: https://github.com/facebook/rocksdb/pull/9071

Test Plan:
- Passing new tests
- Passing existing tests

Reviewed By: ajkr

Differential Revision: D31887813

Pulled By: hx235

fbshipit-source-id: 9a09f0c8683822673260362894c878b61ee60ceb

cache/cache_reservation_manager.cc

@@ -23,7 +23,9 @@
 namespace ROCKSDB_NAMESPACE {
 CacheReservationManager::CacheReservationManager(std::shared_ptr<Cache> cache,
                                                  bool delayed_decrease)
-    : delayed_decrease_(delayed_decrease), cache_allocated_size_(0) {
+    : delayed_decrease_(delayed_decrease),
+      cache_allocated_size_(0),
+      memory_used_(0) {
   assert(cache != nullptr);
   cache_ = cache;
   std::memset(cache_key_, 0, kCacheKeyPrefixSize + kMaxVarint64Length);
@@ -39,6 +41,7 @@ CacheReservationManager::~CacheReservationManager() {
 template <CacheEntryRole R>
 Status CacheReservationManager::UpdateCacheReservation(
     std::size_t new_mem_used) {
+  memory_used_ = new_mem_used;
   std::size_t cur_cache_allocated_size =
       cache_allocated_size_.load(std::memory_order_relaxed);
   if (new_mem_used == cur_cache_allocated_size) {
@@ -118,6 +121,10 @@ std::size_t CacheReservationManager::GetTotalReservedCacheSize() {
   return cache_allocated_size_.load(std::memory_order_relaxed);
 }
 
+std::size_t CacheReservationManager::GetTotalMemoryUsed() {
+  return memory_used_;
+}
+
 Slice CacheReservationManager::GetNextCacheKey() {
   // Calling this function will have the side-effect of changing the
   // underlying cache_key_ that is shared among other keys generated from this
@@ -128,4 +135,4 @@ Slice CacheReservationManager::GetNextCacheKey() {
       EncodeVarint64(cache_key_ + kCacheKeyPrefixSize, next_cache_key_id_++);
   return Slice(cache_key_, static_cast<std::size_t>(end - cache_key_));
 }
-}  // namespace ROCKSDB_NAMESPACE
+}  // namespace ROCKSDB_NAMESPACE

cache/cache_reservation_manager.h

@@ -26,7 +26,9 @@ namespace ROCKSDB_NAMESPACE {
 
 // CacheReservationManager is for reserving cache space for the memory used
 // through inserting/releasing dummy entries in the cache.
-// This class is not thread-safe.
+//
+// This class is NOT thread-safe, except that GetTotalReservedCacheSize()
+// can be called without external synchronization.
 class CacheReservationManager {
  public:
   // Construct a CacheReservationManager
@@ -53,8 +55,8 @@ class CacheReservationManager {
   template <CacheEntryRole R>
 
   // Insert and release dummy entries in the cache to
-  // match the size of total dummy entries with the smallest multiple of
+  // match the size of total dummy entries with the least multiple of
-  // kSizeDummyEntry that is greater than or equal to new_mem_used
+  // kSizeDummyEntry greater than or equal to new_mem_used
   //
   // Insert dummy entries if new_memory_used > cache_allocated_size_;
   //
@@ -68,13 +70,34 @@ class CacheReservationManager {
   // is set true.
   //
   // @param new_memory_used The number of bytes used by new memory
+  //        The most recent new_memoy_used passed in will be returned
+  //        in GetTotalMemoryUsed() even when the call return non-ok status.
+  //
+  //        Since the class is NOT thread-safe, external synchronization on the
+  //        order of calling UpdateCacheReservation() is needed if you want
+  //        GetTotalMemoryUsed() indeed returns the latest memory used.
+  //
   // @return On inserting dummy entries, it returns Status::OK() if all dummy
-  // entry insertions succeed. Otherwise, it returns the first non-ok status;
+  //         entry insertions succeed.
-  // On releasing dummy entries, it always returns Status::OK().
+  //         Otherwise, it returns the first non-ok status;
-  // On keeping dummy entries the same, it always returns Status::OK().
+  //         On releasing dummy entries, it always returns Status::OK().
+  //         On keeping dummy entries the same, it always returns Status::OK().
   Status UpdateCacheReservation(std::size_t new_memory_used);
+
+  // Return the size of the cache (which is a multiple of kSizeDummyEntry)
+  // successfully reserved by calling UpdateCacheReservation().
+  //
+  // When UpdateCacheReservation() returns non-ok status,
+  // calling GetTotalReservedCacheSize() after that might return a slightly
+  // smaller number than the actual reserved cache size due to
+  // the returned number will always be a multiple of kSizeDummyEntry
+  // and cache full might happen in the middle of inserting a dummy entry.
   std::size_t GetTotalReservedCacheSize();
 
+  // Return the latest total memory used indicated by the most recent call of
+  // UpdateCacheReservation(std::size_t new_memory_used);
+  std::size_t GetTotalMemoryUsed();
+
   static constexpr std::size_t GetDummyEntrySize() { return kSizeDummyEntry; }
 
  private:
@@ -92,6 +115,7 @@ class CacheReservationManager {
   std::shared_ptr<Cache> cache_;
   bool delayed_decrease_;
   std::atomic<std::size_t> cache_allocated_size_;
+  std::size_t memory_used_;
   std::vector<Cache::Handle *> dummy_handles_;
   std::uint64_t next_cache_key_id_ = 0;
   // The non-prefix part will be updated according to the ID to use.

cache/cache_reservation_manager_test.cc

@@ -81,6 +81,7 @@ TEST_F(CacheReservationManagerTest, KeepCacheReservationTheSame) {
   ASSERT_EQ(s, Status::OK());
   ASSERT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(),
             1 * kSizeDummyEntry);
+  ASSERT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used);
   std::size_t initial_pinned_usage = cache->GetPinnedUsage();
   ASSERT_GE(initial_pinned_usage, 1 * kSizeDummyEntry);
   ASSERT_LT(initial_pinned_usage,
@@ -96,6 +97,9 @@ TEST_F(CacheReservationManagerTest, KeepCacheReservationTheSame) {
             1 * kSizeDummyEntry)
       << "Failed to bookkeep correctly when new_mem_used equals to current "
          "cache reservation";
+  EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
+      << "Failed to bookkeep the used memory correctly when new_mem_used "
+         "equals to current cache reservation";
   EXPECT_EQ(cache->GetPinnedUsage(), initial_pinned_usage)
       << "Failed to keep underlying dummy entries the same when new_mem_used "
          "equals to current cache reservation";
@@ -113,6 +117,8 @@ TEST_F(CacheReservationManagerTest,
   EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(),
             2 * kSizeDummyEntry)
       << "Failed to bookkeep cache reservation increase correctly";
+  EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
+      << "Failed to bookkeep the used memory correctly";
   EXPECT_GE(cache->GetPinnedUsage(), 2 * kSizeDummyEntry)
       << "Failed to increase underlying dummy entries in cache correctly";
   EXPECT_LT(cache->GetPinnedUsage(),
@@ -132,6 +138,8 @@ TEST_F(CacheReservationManagerTest,
   EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(),
             3 * kSizeDummyEntry)
       << "Failed to bookkeep cache reservation increase correctly";
+  EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
+      << "Failed to bookkeep the used memory correctly";
   EXPECT_GE(cache->GetPinnedUsage(), 3 * kSizeDummyEntry)
       << "Failed to increase underlying dummy entries in cache correctly";
   EXPECT_LT(cache->GetPinnedUsage(),
@@ -173,6 +181,8 @@ TEST(CacheReservationManagerIncreaseReservcationOnFullCacheTest,
       << "Failed to bookkeep correctly (i.e, bookkeep only successful dummy "
          "entry insertions) when encountering cache resevation failure due to "
          "full cache";
+  EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
+      << "Failed to bookkeep the used memory correctly";
   EXPECT_GE(cache->GetPinnedUsage(), 1 * kSizeDummyEntry)
       << "Failed to insert underlying dummy entries correctly when "
          "encountering cache resevation failure due to full cache";
@@ -191,6 +201,8 @@ TEST(CacheReservationManagerIncreaseReservcationOnFullCacheTest,
             2 * kSizeDummyEntry)
       << "Failed to bookkeep cache reservation decrease correctly after "
          "encountering cache reservation due to full cache";
+  EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
+      << "Failed to bookkeep the used memory correctly";
   EXPECT_GE(cache->GetPinnedUsage(), 2 * kSizeDummyEntry)
       << "Failed to release underlying dummy entries correctly on cache "
          "reservation decrease after encountering cache resevation failure due "
@@ -218,6 +230,8 @@ TEST(CacheReservationManagerIncreaseReservcationOnFullCacheTest,
       << "Failed to bookkeep correctly (i.e, bookkeep only successful dummy "
          "entry insertions) when encountering cache resevation failure due to "
          "full cache";
+  EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
+      << "Failed to bookkeep the used memory correctly";
   EXPECT_GE(cache->GetPinnedUsage(), 1 * kSizeDummyEntry)
       << "Failed to insert underlying dummy entries correctly when "
          "encountering cache resevation failure due to full cache";
@@ -239,6 +253,8 @@ TEST(CacheReservationManagerIncreaseReservcationOnFullCacheTest,
             5 * kSizeDummyEntry)
       << "Failed to bookkeep cache reservation increase correctly after "
          "increasing cache capacity and mitigating cache full error";
+  EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
+      << "Failed to bookkeep the used memory correctly";
   EXPECT_GE(cache->GetPinnedUsage(), 5 * kSizeDummyEntry)
       << "Failed to insert underlying dummy entries correctly after increasing "
          "cache capacity and mitigating cache full error";
@@ -258,6 +274,7 @@ TEST_F(CacheReservationManagerTest,
   ASSERT_EQ(s, Status::OK());
   ASSERT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(),
             2 * kSizeDummyEntry);
+  ASSERT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used);
   ASSERT_GE(cache->GetPinnedUsage(), 2 * kSizeDummyEntry);
   ASSERT_LT(cache->GetPinnedUsage(),
             2 * kSizeDummyEntry + kMetaDataChargeOverhead);
@@ -271,6 +288,8 @@ TEST_F(CacheReservationManagerTest,
   EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(),
             1 * kSizeDummyEntry)
       << "Failed to bookkeep cache reservation decrease correctly";
+  EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
+      << "Failed to bookkeep the used memory correctly";
   EXPECT_GE(cache->GetPinnedUsage(), 1 * kSizeDummyEntry)
       << "Failed to decrease underlying dummy entries in cache correctly";
   EXPECT_LT(cache->GetPinnedUsage(),
@@ -288,6 +307,7 @@ TEST_F(CacheReservationManagerTest,
   ASSERT_EQ(s, Status::OK());
   ASSERT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(),
             2 * kSizeDummyEntry);
+  ASSERT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used);
   ASSERT_GE(cache->GetPinnedUsage(), 2 * kSizeDummyEntry);
   ASSERT_LT(cache->GetPinnedUsage(),
             2 * kSizeDummyEntry + kMetaDataChargeOverhead);
@@ -301,6 +321,8 @@ TEST_F(CacheReservationManagerTest,
   EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(),
             1 * kSizeDummyEntry)
       << "Failed to bookkeep cache reservation decrease correctly";
+  EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
+      << "Failed to bookkeep the used memory correctly";
   EXPECT_GE(cache->GetPinnedUsage(), 1 * kSizeDummyEntry)
       << "Failed to decrease underlying dummy entries in cache correctly";
   EXPECT_LT(cache->GetPinnedUsage(),
@@ -330,6 +352,7 @@ TEST(CacheReservationManagerWithDelayedDecreaseTest,
   ASSERT_EQ(s, Status::OK());
   ASSERT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(),
             8 * kSizeDummyEntry);
+  ASSERT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used);
   std::size_t initial_pinned_usage = cache->GetPinnedUsage();
   ASSERT_GE(initial_pinned_usage, 8 * kSizeDummyEntry);
   ASSERT_LT(initial_pinned_usage,
@@ -344,6 +367,8 @@ TEST(CacheReservationManagerWithDelayedDecreaseTest,
             8 * kSizeDummyEntry)
       << "Failed to bookkeep correctly when delaying cache reservation "
          "decrease";
+  EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
+      << "Failed to bookkeep the used memory correctly";
   EXPECT_EQ(cache->GetPinnedUsage(), initial_pinned_usage)
       << "Failed to delay decreasing underlying dummy entries in cache";
 
@@ -356,6 +381,8 @@ TEST(CacheReservationManagerWithDelayedDecreaseTest,
             8 * kSizeDummyEntry)
       << "Failed to bookkeep correctly when delaying cache reservation "
          "decrease";
+  EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
+      << "Failed to bookkeep the used memory correctly";
   EXPECT_EQ(cache->GetPinnedUsage(), initial_pinned_usage)
       << "Failed to delay decreasing underlying dummy entries in cache";
 
@@ -370,6 +397,8 @@ TEST(CacheReservationManagerWithDelayedDecreaseTest,
             6 * kSizeDummyEntry)
       << "Failed to bookkeep correctly when new_mem_used < "
          "GetTotalReservedCacheSize() * 3 / 4 on delayed decrease mode";
+  EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
+      << "Failed to bookkeep the used memory correctly";
   EXPECT_GE(cache->GetPinnedUsage(), 6 * kSizeDummyEntry)
       << "Failed to decrease underlying dummy entries in cache when "
          "new_mem_used < GetTotalReservedCacheSize() * 3 / 4 on delayed "