From f48b64460e90e679cb349aa820d73e66482e2f93 Mon Sep 17 00:00:00 2001 From: anand76 Date: Fri, 6 Sep 2024 14:54:09 -0700 Subject: [PATCH] Provide a way to invoke a callback for a Cache handle (#12987) Summary: Add the `ApplyToHandle` method to the `Cache` interface to allow a caller to request the invocation of a callback on the given cache handle. The goal here is to allow a cache that manages multiple cache instances to use a callback on a handle to determine which instance it belongs to. For example, the callback can hash the key and use that to pick the correct target instance. This is useful to redirect methods like `Ref` and `Release`, which don't know the cache key. Pull Request resolved: https://github.com/facebook/rocksdb/pull/12987 Reviewed By: pdillinger Differential Revision: D62151907 Pulled By: anand1976 fbshipit-source-id: e4ffbbb96eac9061d2ab0e7e1739eea5ebb1cd58 --- cache/cache_test.cc | 26 ++++++++++++++++++++++++++ cache/clock_cache.cc | 16 ++++++++++++++++ cache/clock_cache.h | 6 ++++++ cache/lru_cache.cc | 11 +++++++++++ cache/lru_cache.h | 6 ++++++ include/rocksdb/advanced_cache.h | 17 +++++++++++++++++ 6 files changed, 82 insertions(+) diff --git a/cache/cache_test.cc b/cache/cache_test.cc index 462c2ec741..12bcfe6cd4 100644 --- a/cache/cache_test.cc +++ b/cache/cache_test.cc @@ -886,6 +886,32 @@ TEST_P(CacheTest, ApplyToAllEntriesDuringResize) { ASSERT_EQ(special_count, kSpecialCount); } +TEST_P(CacheTest, ApplyToHandleTest) { + std::string callback_state; + const auto callback = [&](const Slice& key, Cache::ObjectPtr value, + size_t charge, + const Cache::CacheItemHelper* helper) { + callback_state = std::to_string(DecodeKey(key)) + "," + + std::to_string(DecodeValue(value)) + "," + + std::to_string(charge); + assert(helper == &CacheTest::kHelper); + }; + + std::vector inserted; + + for (int i = 0; i < 10; ++i) { + Insert(i, i * 2, i + 1); + inserted.push_back(std::to_string(i) + "," + std::to_string(i * 2) + "," + + std::to_string(i + 1)); + } + for (int i = 0; i < 10; ++i) { + Cache::Handle* handle = cache_->Lookup(EncodeKey(i)); + cache_->ApplyToHandle(cache_.get(), handle, callback); + EXPECT_EQ(inserted[i], callback_state); + cache_->Release(handle); + } +} + TEST_P(CacheTest, DefaultShardBits) { // Prevent excessive allocation (to save time & space) estimated_value_size_ = 100000; diff --git a/cache/clock_cache.cc b/cache/clock_cache.cc index 078b922dd3..090213cb0d 100644 --- a/cache/clock_cache.cc +++ b/cache/clock_cache.cc @@ -1444,6 +1444,22 @@ const Cache::CacheItemHelper* BaseHyperClockCache::GetCacheItemHelper( return h->helper; } +template +void BaseHyperClockCache
::ApplyToHandle( + Cache* cache, Handle* handle, + const std::function& + callback) { + BaseHyperClockCache
* cache_ptr = + static_cast*>(cache); + auto h = static_cast(handle); + UniqueId64x2 unhashed; + auto hash_seed = cache_ptr->GetShard(h->GetHash()).GetTable().GetHashSeed(); + callback( + ClockCacheShard
::ReverseHash(h->hashed_key, &unhashed, hash_seed), + h->value, h->GetTotalCharge(), h->helper); +} + namespace { // For each cache shard, estimate what the table load factor would be if diff --git a/cache/clock_cache.h b/cache/clock_cache.h index 7423fa1f41..2d5d0d9eef 100644 --- a/cache/clock_cache.h +++ b/cache/clock_cache.h @@ -1128,6 +1128,12 @@ class BaseHyperClockCache : public ShardedCache> { const CacheItemHelper* GetCacheItemHelper(Handle* handle) const override; + void ApplyToHandle( + Cache* cache, Handle* handle, + const std::function& + callback) override; + void ReportProblems( const std::shared_ptr& /*info_log*/) const override; }; diff --git a/cache/lru_cache.cc b/cache/lru_cache.cc index 79c46bcc5c..230a6726ca 100644 --- a/cache/lru_cache.cc +++ b/cache/lru_cache.cc @@ -677,6 +677,17 @@ const Cache::CacheItemHelper* LRUCache::GetCacheItemHelper( return h->helper; } +void LRUCache::ApplyToHandle( + Cache* cache, Handle* handle, + const std::function& callback) { + auto cache_ptr = static_cast(cache); + auto h = static_cast(handle); + callback(h->key(), h->value, + h->GetCharge(cache_ptr->GetShard(0).metadata_charge_policy_), + h->helper); +} + size_t LRUCache::TEST_GetLRUSize() { return SumOverShards([](LRUCacheShard& cs) { return cs.TEST_GetLRUSize(); }); } diff --git a/cache/lru_cache.h b/cache/lru_cache.h index 045480fbcf..7fb2a88a00 100644 --- a/cache/lru_cache.h +++ b/cache/lru_cache.h @@ -452,6 +452,12 @@ class LRUCache size_t GetCharge(Handle* handle) const override; const CacheItemHelper* GetCacheItemHelper(Handle* handle) const override; + void ApplyToHandle( + Cache* cache, Handle* handle, + const std::function& callback) + override; + // Retrieves number of elements in LRU, for unit test purpose only. size_t TEST_GetLRUSize(); // Retrieves high pri pool ratio. diff --git a/include/rocksdb/advanced_cache.h b/include/rocksdb/advanced_cache.h index e2aefdd011..ab9f722a05 100644 --- a/include/rocksdb/advanced_cache.h +++ b/include/rocksdb/advanced_cache.h @@ -411,6 +411,14 @@ class Cache { const CacheItemHelper* helper)>& callback, const ApplyToAllEntriesOptions& opts) = 0; + // Apply a callback to a cache handle. The Cache must ensure the lifetime + // of the key passed to the callback is valid for the duration of the + // callback. The handle may not belong to the cache, but is guaranteed to + // be type compatible. + virtual void ApplyToHandle( + Cache* cache, Handle* handle, + const std::function& callback) = 0; // Remove all entries. // Prerequisite: no entry is referenced. virtual void EraseUnRefEntries() = 0; @@ -636,6 +644,15 @@ class CacheWrapper : public Cache { target_->ApplyToAllEntries(callback, opts); } + virtual void ApplyToHandle( + Cache* cache, Handle* handle, + const std::function& callback) + override { + auto cache_ptr = static_cast(cache); + target_->ApplyToHandle(cache_ptr->target_.get(), handle, callback); + } + void EraseUnRefEntries() override { target_->EraseUnRefEntries(); } void StartAsyncLookup(AsyncLookupHandle& async_handle) override {