Support custom allocators for the blob cache (#10628)

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

Test Plan: `make check`

Reviewed By: akankshamahajan15

Differential Revision: D39228165

Pulled By: ltamasi

fbshipit-source-id: 591fdff08db400b170b26f0165551f86d33c1dbf

db/blob/blob_contents.cc

@@ -71,9 +71,11 @@ Cache::CacheItemHelper* BlobContents::GetCacheItemHelper() {
   return &cache_helper;
 }
 
-Status BlobContents::CreateCallback(const void* buf, size_t size,
+Status BlobContents::CreateCallback(CacheAllocationPtr&& allocation,
+                                    const void* buf, size_t size,
                                     void** out_obj, size_t* charge) {
-  CacheAllocationPtr allocation(new char[size]);
+  assert(allocation);
+
   memcpy(allocation.get(), buf, size);
 
   std::unique_ptr<BlobContents> obj = Create(std::move(allocation), size);

db/blob/blob_contents.h

@@ -42,8 +42,8 @@ class BlobContents {
 
   static Cache::CacheItemHelper* GetCacheItemHelper();
 
-  static Status CreateCallback(const void* buf, size_t size, void** out_obj,
-                               size_t* charge);
+  static Status CreateCallback(CacheAllocationPtr&& allocation, const void* buf,
+                               size_t size, void** out_obj, size_t* charge);
 
  private:
   BlobContents(CacheAllocationPtr&& allocation, size_t size)

db/blob/blob_file_builder.cc

@@ -410,8 +410,8 @@ Status BlobFileBuilder::PutBlobIntoCacheIfNeeded(const Slice& blob,
     // Objects to be put into the cache have to be heap-allocated and
     // self-contained, i.e. own their contents. The Cache has to be able to
     // take unique ownership of them.
-    // TODO: support custom allocators
-    CacheAllocationPtr allocation(new char[blob.size()]);
+    CacheAllocationPtr allocation =
+        AllocateBlock(blob.size(), blob_cache->memory_allocator());
     memcpy(allocation.get(), blob.data(), blob.size());
     std::unique_ptr<BlobContents> buf =
         BlobContents::Create(std::move(allocation), blob.size());

db/blob/blob_source.cc

@@ -83,8 +83,8 @@ Status BlobSource::PutBlobIntoCache(const Slice& cache_key,
   // Objects to be put into the cache have to be heap-allocated and
   // self-contained, i.e. own their contents. The Cache has to be able to take
   // unique ownership of them.
-  // TODO: support custom allocators
-  CacheAllocationPtr allocation(new char[blob->size()]);
+  CacheAllocationPtr allocation =
+      AllocateBlock(blob->size(), blob_cache_->memory_allocator());
   memcpy(allocation.get(), blob->data(), blob->size());
   std::unique_ptr<BlobContents> buf =
       BlobContents::Create(std::move(allocation), blob->size());
@@ -112,9 +112,17 @@ Cache::Handle* BlobSource::GetEntryFromCache(const Slice& key) const {
   Cache::Handle* cache_handle = nullptr;
 
   if (lowest_used_cache_tier_ == CacheTier::kNonVolatileBlockTier) {
-    cache_handle = blob_cache_->Lookup(
-        key, BlobContents::GetCacheItemHelper(), &BlobContents::CreateCallback,
-        Cache::Priority::BOTTOM, true /* wait_for_cache */, statistics_);
+    Cache::CreateCallback create_cb =
+        [allocator = blob_cache_->memory_allocator()](
+            const void* buf, size_t size, void** out_obj,
+            size_t* charge) -> Status {
+      return BlobContents::CreateCallback(AllocateBlock(size, allocator), buf,
+                                          size, out_obj, charge);
+    };
+
+    cache_handle = blob_cache_->Lookup(key, BlobContents::GetCacheItemHelper(),
+                                       create_cb, Cache::Priority::BOTTOM,
+                                       true /* wait_for_cache */, statistics_);
   } else {
     cache_handle = blob_cache_->Lookup(key, statistics_);
   }

db/blob/blob_source_test.cc

@@ -1149,6 +1149,15 @@ TEST_F(BlobSecondaryCacheTest, GetBlobsFromSecondaryCache) {
   auto blob_cache = options_.blob_cache;
   auto secondary_cache = lru_cache_opts_.secondary_cache;
 
+  Cache::CreateCallback create_cb = [](const void* buf, size_t size,
+                                       void** out_obj,
+                                       size_t* charge) -> Status {
+    CacheAllocationPtr allocation(new char[size]);
+
+    return BlobContents::CreateCallback(std::move(allocation), buf, size,
+                                        out_obj, charge);
+  };
+
   {
     // GetBlob
     std::vector<PinnableSlice> values(keys.size());
@@ -1196,8 +1205,8 @@ TEST_F(BlobSecondaryCacheTest, GetBlobsFromSecondaryCache) {
       // key0 should be in the secondary cache. After looking up key0 in the
       // secondary cache, it will be erased from the secondary cache.
       bool is_in_sec_cache = false;
-      auto sec_handle0 = secondary_cache->Lookup(
-          key0, &BlobContents::CreateCallback, true, is_in_sec_cache);
+      auto sec_handle0 =
+          secondary_cache->Lookup(key0, create_cb, true, is_in_sec_cache);
       ASSERT_FALSE(is_in_sec_cache);
       ASSERT_NE(sec_handle0, nullptr);
       ASSERT_TRUE(sec_handle0->IsReady());
@@ -1220,8 +1229,8 @@ TEST_F(BlobSecondaryCacheTest, GetBlobsFromSecondaryCache) {
       blob_cache->Release(handle1);
 
       bool is_in_sec_cache = false;
-      auto sec_handle1 = secondary_cache->Lookup(
-          key1, &BlobContents::CreateCallback, true, is_in_sec_cache);
+      auto sec_handle1 =
+          secondary_cache->Lookup(key1, create_cb, true, is_in_sec_cache);
       ASSERT_FALSE(is_in_sec_cache);
       ASSERT_EQ(sec_handle1, nullptr);