[ResourceManagement] Refactor WriteBufferManager::CacheRep into CacheReservationManager

CMakeLists.txt

@@ -624,6 +624,7 @@ find_package(Threads REQUIRED)
 set(SOURCES
         cache/cache.cc
         cache/cache_entry_roles.cc
+        cache/cache_reservation_manager.cc
         cache/clock_cache.cc
         cache/lru_cache.cc
         cache/sharded_cache.cc
@@ -1123,6 +1124,7 @@ if(WITH_TESTS)
   )
   if(WITH_ALL_TESTS)
     list(APPEND TESTS
+        cache/cache_reservation_manager_test.cc
         cache/cache_test.cc
         cache/lru_cache_test.cc
         db/blob/blob_counting_iterator_test.cc

Makefile

@@ -1897,7 +1897,9 @@ clipping_iterator_test: $(OBJ_DIR)/db/compaction/clipping_iterator_test.o $(TEST
 
 ribbon_bench: $(OBJ_DIR)/microbench/ribbon_bench.o $(LIBRARY)
 	$(AM_LINK)
-
+
+cache_reservation_manager_test: $(OBJ_DIR)/cache/cache_reservation_manager_test.o $(TEST_LIBRARY) $(LIBRARY)
+	$(AM_LINK)
 #-------------------------------------------------
 # make install related stuff
 PREFIX ?= /usr/local

TARGETS

@@ -133,6 +133,7 @@ cpp_library(
     srcs = [
         "cache/cache.cc",
         "cache/cache_entry_roles.cc",
+        "cache/cache_reservation_manager.cc",
         "cache/clock_cache.cc",
         "cache/lru_cache.cc",
         "cache/sharded_cache.cc",
@@ -452,6 +453,7 @@ cpp_library(
     srcs = [
         "cache/cache.cc",
         "cache/cache_entry_roles.cc",
+        "cache/cache_reservation_manager.cc",
         "cache/clock_cache.cc",
         "cache/lru_cache.cc",
         "cache/sharded_cache.cc",
@@ -1025,6 +1027,13 @@ ROCKS_TESTS = [
         [],
         [],
     ],
+    [
+        "cache_reservation_manager_test",
+        "cache/cache_reservation_manager_test.cc",
+        "parallel",
+        [],
+        [],
+    ],
     [
         "cache_simulator_test",
         "utilities/simulator_cache/cache_simulator_test.cc",

cache/cache_reservation_manager.cc

@@ -0,0 +1,128 @@
+//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+//
+// Copyright (c) 2011 The LevelDB Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file. See the AUTHORS file for names of contributors.
+#include "cache/cache_reservation_manager.h"
+
+#include <cassert>
+#include <cstddef>
+#include <cstring>
+#include <memory>
+
+#include "cache/cache_entry_roles.h"
+#include "rocksdb/cache.h"
+#include "rocksdb/slice.h"
+#include "rocksdb/status.h"
+#include "table/block_based/block_based_table_reader.h"
+#include "util/coding.h"
+
+namespace ROCKSDB_NAMESPACE {
+CacheReservationManager::CacheReservationManager(std::shared_ptr<Cache> cache,
+                                                 bool delayed_decrease)
+    : delayed_decrease_(delayed_decrease), cache_allocated_size_(0) {
+  assert(cache != nullptr);
+  cache_ = cache;
+  std::memset(cache_key_, 0, kCacheKeyPrefixSize + kMaxVarint64Length);
+  EncodeVarint64(cache_key_, cache_->NewId());
+}
+
+CacheReservationManager::~CacheReservationManager() {
+  for (auto* handle : dummy_handles_) {
+    cache_->Release(handle, true);
+  }
+}
+
+template <CacheEntryRole R>
+Status CacheReservationManager::UpdateCacheReservation(
+    std::size_t 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) {
+    return Status::OK();
+  } else if (new_mem_used > cur_cache_allocated_size) {
+    Status s = IncreaseCacheReservation<R>(new_mem_used);
+    return s;
+  } else {
+    // In delayed decrease mode, we don't decrease cache reservation
+    // untill the memory usage is less than 3/4 of what we reserve
+    // in the cache.
+    // We do this because
+    // (1) Dummy entry insertion is expensive in block cache
+    // (2) Delayed releasing previously inserted dummy entries can save such
+    // expensive dummy entry insertion on memory increase in the near future,
+    // which is likely to happen when the memory usage is greater than or equal
+    // to 3/4 of what we reserve
+    if (delayed_decrease_ && new_mem_used >= cur_cache_allocated_size / 4 * 3) {
+      return Status::OK();
+    } else {
+      Status s = DecreaseCacheReservation(new_mem_used);
+      return s;
+    }
+  }
+}
+
+// Explicitly instantiate templates for "CacheEntryRole" values we use.
+// This makes it possible to keep the template definitions in the .cc file.
+template Status CacheReservationManager::UpdateCacheReservation<
+    CacheEntryRole::kWriteBuffer>(std::size_t new_mem_used);
+// For cache reservation manager unit tests
+template Status CacheReservationManager::UpdateCacheReservation<
+    CacheEntryRole::kMisc>(std::size_t new_mem_used);
+
+template <CacheEntryRole R>
+Status CacheReservationManager::IncreaseCacheReservation(
+    std::size_t new_mem_used) {
+  Status return_status = Status::OK();
+  while (new_mem_used > cache_allocated_size_.load(std::memory_order_relaxed)) {
+    Cache::Handle* handle = nullptr;
+    return_status = cache_->Insert(GetNextCacheKey(), nullptr, kSizeDummyEntry,
+                                   GetNoopDeleterForRole<R>(), &handle);
+
+    if (return_status != Status::OK()) {
+      return return_status;
+    }
+
+    dummy_handles_.push_back(handle);
+    cache_allocated_size_ += kSizeDummyEntry;
+  }
+  return return_status;
+}
+
+Status CacheReservationManager::DecreaseCacheReservation(
+    std::size_t new_mem_used) {
+  Status return_status = Status::OK();
+
+  // Decrease to the smallest multiple of kSizeDummyEntry that is greater than
+  // or equal to new_mem_used We do addition instead of new_mem_used <=
+  // cache_allocated_size_.load(std::memory_order_relaxed) - kSizeDummyEntry to
+  // avoid underflow of size_t when cache_allocated_size_ = 0
+  while (new_mem_used + kSizeDummyEntry <=
+         cache_allocated_size_.load(std::memory_order_relaxed)) {
+    assert(!dummy_handles_.empty());
+    auto* handle = dummy_handles_.back();
+    cache_->Release(handle, true);
+    dummy_handles_.pop_back();
+    cache_allocated_size_ -= kSizeDummyEntry;
+  }
+  return return_status;
+}
+
+std::size_t CacheReservationManager::GetTotalReservedCacheSize() {
+  return cache_allocated_size_.load(std::memory_order_relaxed);
+}
+
+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
+  // fucntion. Therefore please make sure the previous keys are saved/copied
+  // before calling this function.
+  std::memset(cache_key_ + kCacheKeyPrefixSize, 0, kMaxVarint64Length);
+  char* end =
+      EncodeVarint64(cache_key_ + kCacheKeyPrefixSize, next_cache_key_id_++);
+  return Slice(cache_key_, static_cast<std::size_t>(end - cache_key_));
+}
+}  // namespace ROCKSDB_NAMESPACE

cache/cache_reservation_manager.h

@@ -0,0 +1,97 @@
+//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+//
+// Copyright (c) 2011 The LevelDB Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file. See the AUTHORS file for names of contributors.
+
+#pragma once
+
+#include <atomic>
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include "cache/cache_entry_roles.h"
+#include "rocksdb/cache.h"
+#include "rocksdb/slice.h"
+#include "rocksdb/status.h"
+#include "table/block_based/block_based_table_reader.h"
+#include "util/coding.h"
+
+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.
+class CacheReservationManager {
+ public:
+  // Construct a CacheReservationManager
+  // @param cache The cache where dummy entries are inserted and released for
+  // reserving cache space
+  // @param delayed_decrease If set true, then dummy entries won't be released
+  // immediately when memory usage decreases.
+  //                         Instead, it will be released when the memory usage
+  //                         decreases to 3/4 of what we have reserved so far.
+  //                         This is for saving some future dummy entry
+  //                         insertion when memory usage increases are likely to
+  //                         happen in the near future.
+  explicit CacheReservationManager(std::shared_ptr<Cache> cache,
+                                   bool delayed_decrease = false);
+
+  // no copy constructor, copy assignment, move constructor, move assignment
+  CacheReservationManager(const CacheReservationManager &) = delete;
+  CacheReservationManager &operator=(const CacheReservationManager &) = delete;
+  CacheReservationManager(CacheReservationManager &&) = delete;
+  CacheReservationManager &operator=(CacheReservationManager &&) = delete;
+
+  ~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
+  // kSizeDummyEntry that is greater than or equal to new_mem_used
+  //
+  // Insert dummy entries if new_memory_used > cache_allocated_size_;
+  //
+  // Release dummy entries if new_memory_used < cache_allocated_size_
+  // (and new_memory_used < cache_allocated_size_ * 3/4
+  // when delayed_decrease is set true);
+  //
+  // Keey dummy entries the same if (1) new_memory_used == cache_allocated_size_
+  // or (2) new_memory_used is in the interval of
+  // [cache_allocated_size_ * 3/4, cache_allocated_size) when delayed_decrease
+  // is set true.
+  //
+  // On inserting dummy entries, it returns Status::OK() if all dummy entry
+  // insertions succeed. Otherwise, it returns the first non-ok status;
+  // 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);
+  std::size_t GetTotalReservedCacheSize();
+
+ private:
+  static constexpr std::size_t kSizeDummyEntry = 256 * 1024;
+  // The key will be longer than keys for blocks in SST files so they won't
+  // conflict.
+  static const std::size_t kCacheKeyPrefixSize =
+      BlockBasedTable::kMaxCacheKeyPrefixSize + kMaxVarint64Length;
+
+  Slice GetNextCacheKey();
+  template <CacheEntryRole R>
+  Status IncreaseCacheReservation(std::size_t new_mem_used);
+  Status DecreaseCacheReservation(std::size_t new_mem_used);
+
+  std::shared_ptr<Cache> cache_;
+  bool delayed_decrease_;
+  std::atomic<std::size_t> cache_allocated_size_;
+  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.
+  char cache_key_[kCacheKeyPrefixSize + kMaxVarint64Length];
+};
+}  // namespace ROCKSDB_NAMESPACE