1 files changed, 197 insertions, 0 deletions
diff --git a/src/rocksdb/cache/lru_cache_test.cc b/src/rocksdb/cache/lru_cache_test.cc
new file mode 100644
index 00000000..9980dd72
--- /dev/null
+++ b/src/rocksdb/cache/lru_cache_test.cc
@@ -0,0 +1,197 @@
+//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#include "cache/lru_cache.h"
+
+#include <string>
+#include <vector>
+#include "port/port.h"
+#include "util/testharness.h"
+
+namespace rocksdb {
+
+class LRUCacheTest : public testing::Test {
+ public:
+  LRUCacheTest() {}
+  ~LRUCacheTest() override { DeleteCache(); }
+
+  void DeleteCache() {
+    if (cache_ != nullptr) {
+      cache_->~LRUCacheShard();
+      port::cacheline_aligned_free(cache_);
+      cache_ = nullptr;
+    }
+  }
+
+  void NewCache(size_t capacity, double high_pri_pool_ratio = 0.0,
+                bool use_adaptive_mutex = kDefaultToAdaptiveMutex) {
+    DeleteCache();
+    cache_ = reinterpret_cast<LRUCacheShard*>(
+        port::cacheline_aligned_alloc(sizeof(LRUCacheShard)));
+    new (cache_) LRUCacheShard(capacity, false /*strict_capcity_limit*/,
+                               high_pri_pool_ratio, use_adaptive_mutex);
+  }
+
+  void Insert(const std::string& key,
+              Cache::Priority priority = Cache::Priority::LOW) {
+    cache_->Insert(key, 0 /*hash*/, nullptr /*value*/, 1 /*charge*/,
+                   nullptr /*deleter*/, nullptr /*handle*/, priority);
+  }
+
+  void Insert(char key, Cache::Priority priority = Cache::Priority::LOW) {
+    Insert(std::string(1, key), priority);
+  }
+
+  bool Lookup(const std::string& key) {
+    auto handle = cache_->Lookup(key, 0 /*hash*/);
+    if (handle) {
+      cache_->Release(handle);
+      return true;
+    }
+    return false;
+  }
+
+  bool Lookup(char key) { return Lookup(std::string(1, key)); }
+
+  void Erase(const std::string& key) { cache_->Erase(key, 0 /*hash*/); }
+
+  void ValidateLRUList(std::vector<std::string> keys,
+                       size_t num_high_pri_pool_keys = 0) {
+    LRUHandle* lru;
+    LRUHandle* lru_low_pri;
+    cache_->TEST_GetLRUList(&lru, &lru_low_pri);
+    LRUHandle* iter = lru;
+    bool in_high_pri_pool = false;
+    size_t high_pri_pool_keys = 0;
+    if (iter == lru_low_pri) {
+      in_high_pri_pool = true;
+    }
+    for (const auto& key : keys) {
+      iter = iter->next;
+      ASSERT_NE(lru, iter);
+      ASSERT_EQ(key, iter->key().ToString());
+      ASSERT_EQ(in_high_pri_pool, iter->InHighPriPool());
+      if (in_high_pri_pool) {
+        high_pri_pool_keys++;
+      }
+      if (iter == lru_low_pri) {
+        ASSERT_FALSE(in_high_pri_pool);
+        in_high_pri_pool = true;
+      }
+    }
+    ASSERT_EQ(lru, iter->next);
+    ASSERT_TRUE(in_high_pri_pool);
+    ASSERT_EQ(num_high_pri_pool_keys, high_pri_pool_keys);
+  }
+
+ private:
+  LRUCacheShard* cache_ = nullptr;
+};
+
+TEST_F(LRUCacheTest, BasicLRU) {
+  NewCache(5);
+  for (char ch = 'a'; ch <= 'e'; ch++) {
+    Insert(ch);
+  }
+  ValidateLRUList({"a", "b", "c", "d", "e"});
+  for (char ch = 'x'; ch <= 'z'; ch++) {
+    Insert(ch);
+  }
+  ValidateLRUList({"d", "e", "x", "y", "z"});
+  ASSERT_FALSE(Lookup("b"));
+  ValidateLRUList({"d", "e", "x", "y", "z"});
+  ASSERT_TRUE(Lookup("e"));
+  ValidateLRUList({"d", "x", "y", "z", "e"});
+  ASSERT_TRUE(Lookup("z"));
+  ValidateLRUList({"d", "x", "y", "e", "z"});
+  Erase("x");
+  ValidateLRUList({"d", "y", "e", "z"});
+  ASSERT_TRUE(Lookup("d"));
+  ValidateLRUList({"y", "e", "z", "d"});
+  Insert("u");
+  ValidateLRUList({"y", "e", "z", "d", "u"});
+  Insert("v");
+  ValidateLRUList({"e", "z", "d", "u", "v"});
+}
+
+TEST_F(LRUCacheTest, MidpointInsertion) {
+  // Allocate 2 cache entries to high-pri pool.
+  NewCache(5, 0.45);
+
+  Insert("a", Cache::Priority::LOW);
+  Insert("b", Cache::Priority::LOW);
+  Insert("c", Cache::Priority::LOW);
+  Insert("x", Cache::Priority::HIGH);
+  Insert("y", Cache::Priority::HIGH);
+  ValidateLRUList({"a", "b", "c", "x", "y"}, 2);
+
+  // Low-pri entries inserted to the tail of low-pri list (the midpoint).
+  // After lookup, it will move to the tail of the full list.
+  Insert("d", Cache::Priority::LOW);
+  ValidateLRUList({"b", "c", "d", "x", "y"}, 2);
+  ASSERT_TRUE(Lookup("d"));
+  ValidateLRUList({"b", "c", "x", "y", "d"}, 2);
+
+  // High-pri entries will be inserted to the tail of full list.
+  Insert("z", Cache::Priority::HIGH);
+  ValidateLRUList({"c", "x", "y", "d", "z"}, 2);
+}
+
+TEST_F(LRUCacheTest, EntriesWithPriority) {
+  // Allocate 2 cache entries to high-pri pool.
+  NewCache(5, 0.45);
+
+  Insert("a", Cache::Priority::LOW);
+  Insert("b", Cache::Priority::LOW);
+  Insert("c", Cache::Priority::LOW);
+  ValidateLRUList({"a", "b", "c"}, 0);
+
+  // Low-pri entries can take high-pri pool capacity if available
+  Insert("u", Cache::Priority::LOW);
+  Insert("v", Cache::Priority::LOW);
+  ValidateLRUList({"a", "b", "c", "u", "v"}, 0);
+
+  Insert("X", Cache::Priority::HIGH);
+  Insert("Y", Cache::Priority::HIGH);
+  ValidateLRUList({"c", "u", "v", "X", "Y"}, 2);
+
+  // High-pri entries can overflow to low-pri pool.
+  Insert("Z", Cache::Priority::HIGH);
+  ValidateLRUList({"u", "v", "X", "Y", "Z"}, 2);
+
+  // Low-pri entries will be inserted to head of low-pri pool.
+  Insert("a", Cache::Priority::LOW);
+  ValidateLRUList({"v", "X", "a", "Y", "Z"}, 2);
+
+  // Low-pri entries will be inserted to head of high-pri pool after lookup.
+  ASSERT_TRUE(Lookup("v"));
+  ValidateLRUList({"X", "a", "Y", "Z", "v"}, 2);
+
+  // High-pri entries will be inserted to the head of the list after lookup.
+  ASSERT_TRUE(Lookup("X"));
+  ValidateLRUList({"a", "Y", "Z", "v", "X"}, 2);
+  ASSERT_TRUE(Lookup("Z"));
+  ValidateLRUList({"a", "Y", "v", "X", "Z"}, 2);
+
+  Erase("Y");
+  ValidateLRUList({"a", "v", "X", "Z"}, 2);
+  Erase("X");
+  ValidateLRUList({"a", "v", "Z"}, 1);
+  Insert("d", Cache::Priority::LOW);
+  Insert("e", Cache::Priority::LOW);
+  ValidateLRUList({"a", "v", "d", "e", "Z"}, 1);
+  Insert("f", Cache::Priority::LOW);
+  Insert("g", Cache::Priority::LOW);
+  ValidateLRUList({"d", "e", "f", "g", "Z"}, 1);
+  ASSERT_TRUE(Lookup("d"));
+  ValidateLRUList({"e", "f", "g", "Z", "d"}, 2);
+}
+
+}  // namespace rocksdb
+
+int main(int argc, char** argv) {
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}