Charge blob cache usage against the global memory limit

Summary:

To help service owners to manage their memory budget effectively, we have been working towards counting all major memory users inside RocksDB towards a single global memory limit (see e.g. https://github.com/facebook/rocksdb/wiki/Write-Buffer-Manager#cost-memory-used-in-memtable-to-block-cache). The global limit is specified by the capacity of the block-based table's block cache, and is technically implemented by inserting dummy entries ("reservations") into the block cache. The goal of this task is to support charging the memory usage of the new blob cache against this global memory limit when the backing cache of the blob cache and the block cache are different.

This PR is a part of facebook#10156

cache/cache_reservation_manager.cc

@@ -181,4 +181,5 @@ template class CacheReservationManagerImpl<CacheEntryRole::kFilterConstruction>;
 template class CacheReservationManagerImpl<CacheEntryRole::kMisc>;
 template class CacheReservationManagerImpl<CacheEntryRole::kWriteBuffer>;
 template class CacheReservationManagerImpl<CacheEntryRole::kFileMetadata>;
+template class CacheReservationManagerImpl<CacheEntryRole::kBlobCache>;
 }  // namespace ROCKSDB_NAMESPACE

cache/lru_cache.h

@@ -484,9 +484,9 @@ class LRUCache
   virtual void WaitAll(std::vector<Handle*>& handles) override;
   std::string GetPrintableOptions() const override;
 
-  //  Retrieves number of elements in LRU, for unit test purpose only.
+  // Retrieves number of elements in LRU, for unit test purpose only.
   size_t TEST_GetLRUSize();
-  //  Retrieves high pri pool ratio.
+  // Retrieves high pri pool ratio.
   double GetHighPriPoolRatio();
 
  private:

cache/sharded_cache.cc

@@ -13,6 +13,7 @@
 #include <cstdint>
 #include <memory>
 
+#include "cache/cache_reservation_manager.h"
 #include "util/hash.h"
 #include "util/math.h"
 #include "util/mutexlock.h"
@@ -59,8 +60,18 @@ Status ShardedCache::Insert(const Slice& key, void* value, size_t charge,
                             DeleterFn deleter, Handle** handle,
                             Priority priority) {
   uint32_t hash = HashSlice(key);
-  return GetShard(Shard(hash))
-      ->Insert(key, hash, value, charge, deleter, handle, priority);
+  Status s = GetShard(Shard(hash))
+                 ->Insert(key, hash, value, charge, deleter, handle, priority);
+  if (s.ok()) {
+    auto cache_res_mgr = cache_reservation_manager();
+    if (cache_res_mgr) {
+      // Insert may cause the cache entry eviction if the cache is full. So we
+      // directly call the reservation manager to update the total memory used
+      // in the cache.
+      s = cache_res_mgr->UpdateCacheReservation(GetUsage());
+    }
+  }
+  return s;
 }
 
 Status ShardedCache::Insert(const Slice& key, void* value,
@@ -106,18 +117,44 @@ bool ShardedCache::Ref(Handle* handle) {
 
 bool ShardedCache::Release(Handle* handle, bool erase_if_last_ref) {
   uint32_t hash = GetHash(handle);
-  return GetShard(Shard(hash))->Release(handle, erase_if_last_ref);
+  size_t memory_used_delta = GetUsage(handle);
+  bool erased = GetShard(Shard(hash))->Release(handle, erase_if_last_ref);
+  if (erased) {
+    auto cache_res_mgr = cache_reservation_manager();
+    if (cache_res_mgr) {
+      const Status s =
+          cache_res_mgr->UpdateCacheReservation(memory_used_delta, false);
+      s.PermitUncheckedError();
+    }
+  }
+  return erased;
 }
 
 bool ShardedCache::Release(Handle* handle, bool useful,
                            bool erase_if_last_ref) {
   uint32_t hash = GetHash(handle);
-  return GetShard(Shard(hash))->Release(handle, useful, erase_if_last_ref);
+  size_t memory_used_delta = GetUsage(handle);
+  bool erased =
+      GetShard(Shard(hash))->Release(handle, useful, erase_if_last_ref);
+  if (erased) {
+    auto cache_res_mgr = cache_reservation_manager();
+    if (cache_res_mgr) {
+      const Status s =
+          cache_res_mgr->UpdateCacheReservation(memory_used_delta, false);
+      s.PermitUncheckedError();
+    }
+  }
+  return erased;
 }
 
 void ShardedCache::Erase(const Slice& key) {
   uint32_t hash = HashSlice(key);
   GetShard(Shard(hash))->Erase(key, hash);
+  auto cache_res_mgr = cache_reservation_manager();
+  if (cache_res_mgr) {
+    const Status s = cache_res_mgr->UpdateCacheReservation(GetUsage());
+    s.PermitUncheckedError();
+  }
 }
 
 uint64_t ShardedCache::NewId() {
@@ -188,6 +225,11 @@ void ShardedCache::EraseUnRefEntries() {
   for (uint32_t s = 0; s < num_shards; s++) {
     GetShard(s)->EraseUnRefEntries();
   }
+  auto cache_res_mgr = cache_reservation_manager();
+  if (cache_res_mgr) {
+    const Status s = cache_res_mgr->UpdateCacheReservation(GetUsage());
+    s.PermitUncheckedError();
+  }
 }
 
 std::string ShardedCache::GetPrintableOptions() const {

db/blob/blob_source.cc

@@ -8,6 +8,7 @@
 #include <cassert>
 #include <string>
 
+#include "cache/cache_reservation_manager.h"
 #include "db/blob/blob_file_reader.h"
 #include "db/blob/blob_log_format.h"
 #include "monitoring/statistics.h"
@@ -17,16 +18,31 @@
 
 namespace ROCKSDB_NAMESPACE {
 
-BlobSource::BlobSource(const ImmutableOptions* immutable_options,
+BlobSource::BlobSource(const ImmutableOptions* ioptions,
                        const std::string& db_id,
                        const std::string& db_session_id,
                        BlobFileCache* blob_file_cache)
     : db_id_(db_id),
       db_session_id_(db_session_id),
-      statistics_(immutable_options->statistics.get()),
+      statistics_(ioptions->statistics.get()),
       blob_file_cache_(blob_file_cache),
-      blob_cache_(immutable_options->blob_cache),
-      lowest_used_cache_tier_(immutable_options->lowest_used_cache_tier) {}
+      blob_cache_(ioptions->blob_cache),
+      lowest_used_cache_tier_(ioptions->lowest_used_cache_tier) {
+#ifndef ROCKSDB_LITE
+  auto bbto = ioptions->table_factory->GetOptions<BlockBasedTableOptions>();
+  if (bbto && bbto->block_cache && blob_cache_ &&
+      bbto->block_cache != blob_cache_ &&
+      bbto->cache_usage_options.options_overrides.at(CacheEntryRole::kBlobCache)
+              .charged == CacheEntryRoleOptions::Decision::kEnabled) {
+    std::shared_ptr<ConcurrentCacheReservationManager> cache_res_mgr(
+        new ConcurrentCacheReservationManager(
+            std::make_shared<
+                CacheReservationManagerImpl<CacheEntryRole::kBlobCache>>(
+                bbto->block_cache)));
+    blob_cache_->SetCacheReservationManager(cache_res_mgr);
+  }
+#endif  // ROCKSDB_LITE
+}
 
 BlobSource::~BlobSource() = default;
 

db/blob/blob_source.h

@@ -29,9 +29,8 @@ class Slice;
 // storage with minimal cost.
 class BlobSource {
  public:
-  BlobSource(const ImmutableOptions* immutable_options,
-             const std::string& db_id, const std::string& db_session_id,
-             BlobFileCache* blob_file_cache);
+  BlobSource(const ImmutableOptions* ioptions, const std::string& db_id,
+             const std::string& db_session_id, BlobFileCache* blob_file_cache);
 
   BlobSource(const BlobSource&) = delete;
   BlobSource& operator=(const BlobSource&) = delete;

include/rocksdb/cache.h

@@ -38,6 +38,7 @@ namespace ROCKSDB_NAMESPACE {
 class Cache;
 struct ConfigOptions;
 class SecondaryCache;
+class ConcurrentCacheReservationManager;
 
 extern const bool kDefaultToAdaptiveMutex;
 
@@ -540,8 +541,22 @@ class Cache {
   // to each of the handles.
   virtual void WaitAll(std::vector<Handle*>& /*handles*/) {}
 
+  // Reserves (block) cache space for memory used in this cache.
+  void SetCacheReservationManager(
+      std::shared_ptr<ConcurrentCacheReservationManager> cache_res_mgr) {
+    assert(cache_res_mgr == nullptr);
+    cache_res_mgr_ = cache_res_mgr;
+  }
+
+  ConcurrentCacheReservationManager* cache_reservation_manager() const {
+    return cache_res_mgr_.get();
+  }
+
  private:
   std::shared_ptr<MemoryAllocator> memory_allocator_;
+
+  // ONLY USED for charging blob cache usage
+  std::shared_ptr<ConcurrentCacheReservationManager> cache_res_mgr_;
 };
 
 // Classifications of block cache entries.
@@ -571,12 +586,12 @@ enum class CacheEntryRole {
   // Filter's charge to account for
   // (new) bloom and ribbon filter construction's memory usage
   kFilterConstruction,
-  // BlockBasedTableReader's charge to account for
-  // its memory usage
+  // BlockBasedTableReader's charge to account for its memory usage
   kBlockBasedTableReader,
-  // FileMetadata's charge to account for
-  // its memory usage
+  // FileMetadata's charge to account for its memory usage
   kFileMetadata,
+  // Blob cache's charge to account for its memory usage
+  kBlobCache,
   // Default bucket, for miscellaneous cache entries. Do not use for
   // entries that could potentially add up to large usage.
   kMisc,