RAII support for per cache reservation through handle (#9130)

Summary:
Note: This PR is the 3rd PR of a bigger PR stack (https://github.com/facebook/rocksdb/issues/9073) and depends on the second PR (https://github.com/facebook/rocksdb/pull/9071). **See changes from this PR only 00447324d0**

Context:
pdillinger brought up a good [point](https://github.com/facebook/rocksdb/pull/9073#discussion_r741478309) about lacking RAII support for per cache reservation in `CacheReservationManager`  when reviewing https://github.com/facebook/rocksdb/pull/9073.

To summarize the discussion, the current API `CacheReservationManager::UpdateCacheReservation()` requires callers to explicitly calculate and pass in a correct`new_mem_used` to release a cache reservation (if they don't want to rely on the clean-up during `CacheReservationManager`'s destruction - such as they want to release it earlier).

While this implementation has convenience in some use-case such as `WriteBufferManager`, where [reservation](https://github.com/facebook/rocksdb/blob/main/memtable/write_buffer_manager.cc#L69-L91) and [release](https://github.com/facebook/rocksdb/blob/main/memtable/write_buffer_manager.cc#L109-L129) amounts do not necessarily correspond symmetrically and thus a flexible `new_mem_used` inputing is needed, it can be prone to caller's calculation error as well as cause a mass of codes in releasing cache in other use-case such as filter construction, where reservation and release amounts do correspond symmetrically and many code paths requiring a cache release, as [pointed](https://github.com/facebook/rocksdb/pull/9073#discussion_r741478309) out by pdillinger.

Therefore we decided to provide a new API in `CacheReservationManager` to update reservation with better RAII support for per cache reservation, using a handle to manage the life time of that particular cache reservation.
- Added a new class `CacheReservationHandle`
- Added a new API `CacheReservationManager::MakeCacheReservation()` that outputs a `CacheReservationHandle` for managing the reservation
- Updated class comments to clarify two different cache reservation methods

Tests:
- Passing new tests

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

Reviewed By: pdillinger

Differential Revision: D32199446

Pulled By: hx235

fbshipit-source-id: 1cba7c636e5ecfb55b0c1e0c2d218cc9b5b30b4e

cache/cache_reservation_manager.cc

@@ -79,6 +79,23 @@ template Status CacheReservationManager::UpdateCacheReservation<
 template Status CacheReservationManager::UpdateCacheReservation<
     CacheEntryRole::kMisc>(std::size_t new_mem_used);
 
+template <CacheEntryRole R>
+Status CacheReservationManager::MakeCacheReservation(
+    std::size_t incremental_memory_used,
+    std::unique_ptr<CacheReservationHandle<R>>* handle) {
+  assert(handle != nullptr);
+  Status s =
+      UpdateCacheReservation<R>(GetTotalMemoryUsed() + incremental_memory_used);
+  (*handle).reset(new CacheReservationHandle<R>(incremental_memory_used,
+                                                shared_from_this()));
+  return s;
+}
+
+template Status
+CacheReservationManager::MakeCacheReservation<CacheEntryRole::kMisc>(
+    std::size_t incremental_memory_used,
+    std::unique_ptr<CacheReservationHandle<CacheEntryRole::kMisc>>* handle);
+
 template <CacheEntryRole R>
 Status CacheReservationManager::IncreaseCacheReservation(
     std::size_t new_mem_used) {
@@ -135,4 +152,27 @@ Slice CacheReservationManager::GetNextCacheKey() {
       EncodeVarint64(cache_key_ + kCacheKeyPrefixSize, next_cache_key_id_++);
   return Slice(cache_key_, static_cast<std::size_t>(end - cache_key_));
 }
+
+template <CacheEntryRole R>
+CacheReservationHandle<R>::CacheReservationHandle(
+    std::size_t incremental_memory_used,
+    std::shared_ptr<CacheReservationManager> cache_res_mgr)
+    : incremental_memory_used_(incremental_memory_used) {
+  assert(cache_res_mgr != nullptr);
+  cache_res_mgr_ = cache_res_mgr;
+}
+
+template <CacheEntryRole R>
+CacheReservationHandle<R>::~CacheReservationHandle() {
+  assert(cache_res_mgr_ != nullptr);
+  assert(cache_res_mgr_->GetTotalMemoryUsed() >= incremental_memory_used_);
+
+  Status s = cache_res_mgr_->UpdateCacheReservation<R>(
+      cache_res_mgr_->GetTotalMemoryUsed() - incremental_memory_used_);
+  s.PermitUncheckedError();
+}
+
+// Explicitly instantiate templates for "CacheEntryRole" values we use.
+// This makes it possible to keep the template definitions in the .cc file.
+template class CacheReservationHandle<CacheEntryRole::kMisc>;
 }  // namespace ROCKSDB_NAMESPACE

cache/cache_reservation_manager.h

@@ -24,12 +24,16 @@
 
 namespace ROCKSDB_NAMESPACE {
 
+template <CacheEntryRole R>
+class CacheReservationHandle;
+
 // CacheReservationManager is for reserving cache space for the memory used
 // through inserting/releasing dummy entries in the cache.
 //
 // This class is NOT thread-safe, except that GetTotalReservedCacheSize()
 // can be called without external synchronization.
-class CacheReservationManager {
+class CacheReservationManager
+    : public std::enable_shared_from_this<CacheReservationManager> {
  public:
   // Construct a CacheReservationManager
   // @param cache The cache where dummy entries are inserted and released for
@@ -54,6 +58,10 @@ class CacheReservationManager {
 
   template <CacheEntryRole R>
 
+  // One of the two ways of reserving/releasing cache,
+  // see CacheReservationManager::MakeCacheReservation() for the other.
+  // Use ONLY one of them to prevent unexpected behavior.
+  //
   // Insert and release dummy entries in the cache to
   // match the size of total dummy entries with the least multiple of
   // kSizeDummyEntry greater than or equal to new_mem_used
@@ -84,6 +92,48 @@ class CacheReservationManager {
   //         On keeping dummy entries the same, it always returns Status::OK().
   Status UpdateCacheReservation(std::size_t new_memory_used);
 
+  // One of the two ways of reserving/releasing cache,
+  // see CacheReservationManager::UpdateCacheReservation() for the other.
+  // Use ONLY one of them to prevent unexpected behavior.
+  //
+  // Insert dummy entries in the cache for the incremental memory usage
+  // to match the size of total dummy entries with the least multiple of
+  // kSizeDummyEntry greater than or equal to the total memory used.
+  //
+  // A CacheReservationHandle is returned as an output parameter.
+  // The reserved dummy entries are automatically released on the destruction of
+  // this handle, which achieves better RAII per cache reservation.
+  //
+  // WARNING: Deallocate all the handles of the CacheReservationManager object
+  //          before deallocating the object to prevent unexpected behavior.
+  //
+  // @param incremental_memory_used The number of bytes increased in memory
+  //        usage.
+  //
+  //        Calling GetTotalMemoryUsed() afterward will return the total memory
+  //        increased by this number, even when calling MakeCacheReservation()
+  //        returns non-ok status.
+  //
+  //        Since the class is NOT thread-safe, external synchronization in
+  //        calling MakeCacheReservation() is needed if you want
+  //        GetTotalMemoryUsed() indeed returns the latest memory used.
+  //
+  // @param handle An pointer to std::unique_ptr<CacheReservationHandle<R>> that
+  //        manages the lifetime of the handle and its cache reservation.
+  //
+  // @return It returns Status::OK() if all dummy
+  //         entry insertions succeed.
+  //         Otherwise, it returns the first non-ok status;
+  //
+  // REQUIRES: handle != nullptr
+  // REQUIRES: The CacheReservationManager object is NOT managed by
+  //           std::unique_ptr as CacheReservationHandle needs to
+  //           shares ownership to the CacheReservationManager object.
+  template <CacheEntryRole R>
+  Status MakeCacheReservation(
+      std::size_t incremental_memory_used,
+      std::unique_ptr<CacheReservationHandle<R>> *handle);
+
   // Return the size of the cache (which is a multiple of kSizeDummyEntry)
   // successfully reserved by calling UpdateCacheReservation().
   //
@@ -121,4 +171,21 @@ class CacheReservationManager {
   // The non-prefix part will be updated according to the ID to use.
   char cache_key_[kCacheKeyPrefixSize + kMaxVarint64Length];
 };
+
+// CacheReservationHandle is for managing the lifetime of a cache reservation
+// This class is NOT thread-safe
+template <CacheEntryRole R>
+class CacheReservationHandle {
+ public:
+  // REQUIRES: cache_res_mgr != nullptr
+  explicit CacheReservationHandle(
+      std::size_t incremental_memory_used,
+      std::shared_ptr<CacheReservationManager> cache_res_mgr);
+
+  ~CacheReservationHandle();
+
+ private:
+  std::size_t incremental_memory_used_;
+  std::shared_ptr<CacheReservationManager> cache_res_mgr_;
+};
 }  // namespace ROCKSDB_NAMESPACE

cache/cache_reservation_manager_test.cc

@@ -438,6 +438,74 @@ TEST(CacheReservationManagerDestructorTest,
       << "Failed to release remaining underlying dummy entries in cache in "
          "CacheReservationManager's destructor";
 }
+
+TEST(CacheReservationHandleTest, HandleTest) {
+  constexpr std::size_t kOneGigabyte = 1024 * 1024 * 1024;
+  constexpr std::size_t kSizeDummyEntry = 256 * 1024;
+  constexpr std::size_t kMetaDataChargeOverhead = 10000;
+
+  LRUCacheOptions lo;
+  lo.capacity = kOneGigabyte;
+  lo.num_shard_bits = 0;
+  std::shared_ptr<Cache> cache = NewLRUCache(lo);
+
+  std::shared_ptr<CacheReservationManager> test_cache_rev_mng(
+      std::make_shared<CacheReservationManager>(cache));
+
+  std::size_t mem_used = 0;
+  const std::size_t incremental_mem_used_handle_1 = 1 * kSizeDummyEntry;
+  const std::size_t incremental_mem_used_handle_2 = 2 * kSizeDummyEntry;
+  std::unique_ptr<CacheReservationHandle<CacheEntryRole::kMisc>> handle_1,
+      handle_2;
+
+  // To test consecutive CacheReservationManager::MakeCacheReservation works
+  // correctly in terms of returning the handle as well as updating cache
+  // reservation and the latest total memory used
+  Status s = test_cache_rev_mng->MakeCacheReservation<CacheEntryRole::kMisc>(
+      incremental_mem_used_handle_1, &handle_1);
+  mem_used = mem_used + incremental_mem_used_handle_1;
+  ASSERT_EQ(s, Status::OK());
+  EXPECT_TRUE(handle_1 != nullptr);
+  EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), mem_used);
+  EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), mem_used);
+  EXPECT_GE(cache->GetPinnedUsage(), mem_used);
+  EXPECT_LT(cache->GetPinnedUsage(), mem_used + kMetaDataChargeOverhead);
+
+  s = test_cache_rev_mng->MakeCacheReservation<CacheEntryRole::kMisc>(
+      incremental_mem_used_handle_2, &handle_2);
+  mem_used = mem_used + incremental_mem_used_handle_2;
+  ASSERT_EQ(s, Status::OK());
+  EXPECT_TRUE(handle_2 != nullptr);
+  EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), mem_used);
+  EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), mem_used);
+  EXPECT_GE(cache->GetPinnedUsage(), mem_used);
+  EXPECT_LT(cache->GetPinnedUsage(), mem_used + kMetaDataChargeOverhead);
+
+  // To test CacheReservationHandle::~CacheReservationHandle() works correctly
+  // in releasing the cache reserved for the handle
+  handle_1.reset();
+  EXPECT_TRUE(handle_1 == nullptr);
+  mem_used = mem_used - incremental_mem_used_handle_1;
+  EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), mem_used);
+  EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), mem_used);
+  EXPECT_GE(cache->GetPinnedUsage(), mem_used);
+  EXPECT_LT(cache->GetPinnedUsage(), mem_used + kMetaDataChargeOverhead);
+
+  // To test the actual CacheReservationManager object won't be deallocated
+  // as long as there remain handles pointing to it.
+  // We strongly recommend deallocating CacheReservationManager object only
+  // after all its handles are deallocated to keep things easy to reasonate
+  test_cache_rev_mng.reset();
+  EXPECT_GE(cache->GetPinnedUsage(), mem_used);
+  EXPECT_LT(cache->GetPinnedUsage(), mem_used + kMetaDataChargeOverhead);
+
+  handle_2.reset();
+  // The CacheReservationManager object is now deallocated since all the handles
+  // and its original pointer is gone
+  mem_used = mem_used - incremental_mem_used_handle_2;
+  EXPECT_EQ(mem_used, 0);
+  EXPECT_EQ(cache->GetPinnedUsage(), mem_used);
+}
 }  // namespace ROCKSDB_NAMESPACE
 
 int main(int argc, char** argv) {