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
main
Levi Tamasi 2 years ago committed by Facebook GitHub Bot
parent 5a97e6b1d2
commit c8543296ca
  1. 6
      db/blob/blob_contents.cc
  2. 4
      db/blob/blob_contents.h
  3. 4
      db/blob/blob_file_builder.cc
  4. 18
      db/blob/blob_source.cc
  5. 17
      db/blob/blob_source_test.cc

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

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

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

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

@ -1149,6 +1149,15 @@ TEST_F(BlobSecondaryCacheTest, GetBlobsFromSecondaryCache) {
auto blob_cache = options_.blob_cache; auto blob_cache = options_.blob_cache;
auto secondary_cache = lru_cache_opts_.secondary_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 // GetBlob
std::vector<PinnableSlice> values(keys.size()); 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 // key0 should be in the secondary cache. After looking up key0 in the
// secondary cache, it will be erased from the secondary cache. // secondary cache, it will be erased from the secondary cache.
bool is_in_sec_cache = false; bool is_in_sec_cache = false;
auto sec_handle0 = secondary_cache->Lookup( auto sec_handle0 =
key0, &BlobContents::CreateCallback, true, is_in_sec_cache); secondary_cache->Lookup(key0, create_cb, true, is_in_sec_cache);
ASSERT_FALSE(is_in_sec_cache); ASSERT_FALSE(is_in_sec_cache);
ASSERT_NE(sec_handle0, nullptr); ASSERT_NE(sec_handle0, nullptr);
ASSERT_TRUE(sec_handle0->IsReady()); ASSERT_TRUE(sec_handle0->IsReady());
@ -1220,8 +1229,8 @@ TEST_F(BlobSecondaryCacheTest, GetBlobsFromSecondaryCache) {
blob_cache->Release(handle1); blob_cache->Release(handle1);
bool is_in_sec_cache = false; bool is_in_sec_cache = false;
auto sec_handle1 = secondary_cache->Lookup( auto sec_handle1 =
key1, &BlobContents::CreateCallback, true, is_in_sec_cache); secondary_cache->Lookup(key1, create_cb, true, is_in_sec_cache);
ASSERT_FALSE(is_in_sec_cache); ASSERT_FALSE(is_in_sec_cache);
ASSERT_EQ(sec_handle1, nullptr); ASSERT_EQ(sec_handle1, nullptr);

Loading…
Cancel
Save