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

6 files changed, 82 insertions, 0 deletions

diff --git a/cache/cache_test.cc b/cache/cache_test.cc
index 462c2ec74..12bcfe6cd 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<std::string> 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 078b922dd..090213cb0 100644
--- a/cache/clock_cache.cc
+++ b/cache/clock_cache.cc
@@ -1444,6 +1444,22 @@ const Cache::CacheItemHelper* BaseHyperClockCache<Table>::GetCacheItemHelper(
   return h->helper;
 }
 
+template <class Table>
+void BaseHyperClockCache<Table>::ApplyToHandle(
+    Cache* cache, Handle* handle,
+    const std::function<void(const Slice& key, Cache::ObjectPtr value,
+                             size_t charge, const CacheItemHelper* helper)>&
+        callback) {
+  BaseHyperClockCache<Table>* cache_ptr =
+      static_cast<BaseHyperClockCache<Table>*>(cache);
+  auto h = static_cast<const typename Table::HandleImpl*>(handle);
+  UniqueId64x2 unhashed;
+  auto hash_seed = cache_ptr->GetShard(h->GetHash()).GetTable().GetHashSeed();
+  callback(
+      ClockCacheShard<Table>::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 7423fa1f4..2d5d0d9ee 100644
--- a/cache/clock_cache.h
+++ b/cache/clock_cache.h
@@ -1128,6 +1128,12 @@ class BaseHyperClockCache : public ShardedCache<ClockCacheShard<Table>> {
 
   const CacheItemHelper* GetCacheItemHelper(Handle* handle) const override;
 
+  void ApplyToHandle(
+      Cache* cache, Handle* handle,
+      const std::function<void(const Slice& key, Cache::ObjectPtr obj,
+                               size_t charge, const CacheItemHelper* helper)>&
+          callback) override;
+
   void ReportProblems(
       const std::shared_ptr<Logger>& /*info_log*/) const override;
 };
diff --git a/cache/lru_cache.cc b/cache/lru_cache.cc
index 79c46bcc5..230a6726c 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<void(const Slice& key, ObjectPtr value, size_t charge,
+                             const CacheItemHelper* helper)>& callback) {
+  auto cache_ptr = static_cast<LRUCache*>(cache);
+  auto h = static_cast<const LRUHandle*>(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 045480fbc..7fb2a88a0 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<void(const Slice& key, ObjectPtr obj, size_t charge,
+                               const CacheItemHelper* helper)>& 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 e2aefdd01..ab9f722a0 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<void(const Slice& key, ObjectPtr obj, size_t charge,
+                               const CacheItemHelper* helper)>& 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<void(const Slice& key, ObjectPtr obj, size_t charge,
+                               const CacheItemHelper* helper)>& callback)
+      override {
+    auto cache_ptr = static_cast<CacheWrapper*>(cache);
+    target_->ApplyToHandle(cache_ptr->target_.get(), handle, callback);
+  }
+
   void EraseUnRefEntries() override { target_->EraseUnRefEntries(); }
 
   void StartAsyncLookup(AsyncLookupHandle& async_handle) override {