From c8543296cac9beb9acfea7d7826c6236b218f62b Mon Sep 17 00:00:00 2001 From: Levi Tamasi Date: Tue, 6 Sep 2022 13:31:48 -0700 Subject: [PATCH] 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 | 6 ++++-- db/blob/blob_contents.h | 4 ++-- db/blob/blob_file_builder.cc | 4 ++-- db/blob/blob_source.cc | 18 +++++++++++++----- db/blob/blob_source_test.cc | 17 +++++++++++++---- 5 files changed, 34 insertions(+), 15 deletions(-) diff --git a/db/blob/blob_contents.cc b/db/blob/blob_contents.cc index 50f6132de..9015609e7 100644 --- a/db/blob/blob_contents.cc +++ b/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 obj = Create(std::move(allocation), size); diff --git a/db/blob/blob_contents.h b/db/blob/blob_contents.h index 02d0e3f18..9b7c5b969 100644 --- a/db/blob/blob_contents.h +++ b/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) diff --git a/db/blob/blob_file_builder.cc b/db/blob/blob_file_builder.cc index 84d00ab3b..5e0e7f6cb 100644 --- a/db/blob/blob_file_builder.cc +++ b/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 buf = BlobContents::Create(std::move(allocation), blob.size()); diff --git a/db/blob/blob_source.cc b/db/blob/blob_source.cc index b0cafb224..12bf2b01d 100644 --- a/db/blob/blob_source.cc +++ b/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 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_); } diff --git a/db/blob/blob_source_test.cc b/db/blob/blob_source_test.cc index 3ce9a1072..a8323bc28 100644 --- a/db/blob/blob_source_test.cc +++ b/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 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);