ApplyToAllCacheEntries

Summary: Added a method that executes a callback on every cache entry.

Test Plan: added a unit test

Reviewers: haobo

Reviewed By: haobo

CC: leveldb, dhruba

Differential Revision: https://reviews.facebook.net/D18441

include/rocksdb/cache.h

@@ -116,6 +116,12 @@ class Cache {
     // default implementation is noop
   };
 
+  // Apply callback to all entries in the cache
+  // If thread_safe is true, it will also lock the accesses. Otherwise, it will
+  // access the cache without the lock held
+  virtual void ApplyToAllCacheEntries(void (*callback)(void*, size_t),
+                                      bool thread_safe) = 0;
+
  private:
   void LRU_Remove(Handle* e);
   void LRU_Append(Handle* e);

util/cache.cc

@@ -164,6 +164,9 @@ class LRUCache {
     return usage_;
   }
 
+  void ApplyToAllCacheEntries(void (*callback)(void*, size_t),
+                              bool thread_safe);
+
  private:
   void LRU_Remove(LRUHandle* e);
   void LRU_Append(LRUHandle* e);
@@ -220,6 +223,19 @@ void LRUCache::FreeEntry(LRUHandle* e) {
   free(e);
 }
 
+void LRUCache::ApplyToAllCacheEntries(void (*callback)(void*, size_t),
+                                      bool thread_safe) {
+  if (thread_safe) {
+    mutex_.Lock();
+  }
+  for (auto e = lru_.next; e != &lru_; e = e->next) {
+    callback(e->value, e->charge);
+  }
+  if (thread_safe) {
+    mutex_.Unlock();
+  }
+}
+
 void LRUCache::LRU_Remove(LRUHandle* e) {
   e->next->prev = e->prev;
   e->prev->next = e->next;
@@ -432,6 +448,14 @@ class ShardedLRUCache : public Cache {
   virtual void DisownData() {
     shards_ = nullptr;
   }
+
+  virtual void ApplyToAllCacheEntries(void (*callback)(void*, size_t),
+                                      bool thread_safe) override {
+    int num_shards = 1 << num_shard_bits_;
+    for (int s = 0; s < num_shards; s++) {
+      shards_[s].ApplyToAllCacheEntries(callback, thread_safe);
+    }
+  }
 };
 
 }  // end anonymous namespace

util/cache_test.cc

@@ -420,6 +420,28 @@ TEST(CacheTest, BadEviction) {
   std::cout << "Poor entries\n";
 }
 
+namespace {
+std::vector<std::pair<int, int>> callback_state;
+void callback(void* entry, size_t charge) {
+  callback_state.push_back({DecodeValue(entry), static_cast<int>(charge)});
+}
+};
+
+TEST(CacheTest, ApplyToAllCacheEntiresTest) {
+  std::vector<std::pair<int, int>> inserted;
+  callback_state.clear();
+
+  for (int i = 0; i < 10; ++i) {
+    Insert(i, i * 2, i + 1);
+    inserted.push_back({i * 2, i + 1});
+  }
+  cache_->ApplyToAllCacheEntries(callback, true);
+
+  sort(inserted.begin(), inserted.end());
+  sort(callback_state.begin(), callback_state.end());
+  ASSERT_TRUE(inserted == callback_state);
+}
+
 }  // namespace rocksdb
 
 int main(int argc, char** argv) {