feat(server): add lru data structure (#831)

Part of the heavy keeper algo, required for #257.
Also see #446 for the initial (abandoned) PR.

Signed-off-by: Roman Gershman <[email protected]>

src/core/CMakeLists.txt

@@ -1,6 +1,6 @@
 add_library(dfly_core compact_object.cc dragonfly_core.cc extent_tree.cc
     external_alloc.cc interpreter.cc json_object.cc mi_memory_resource.cc sds_utils.cc
-    segment_allocator.cc small_string.cc tx_queue.cc dense_set.cc
+    segment_allocator.cc simple_lru_counter.cc small_string.cc tx_queue.cc dense_set.cc
     string_set.cc string_map.cc detail/bitpacking.cc)
 
 cxx_link(dfly_core base absl::flat_hash_map absl::str_format redis_lib TRDP::lua lua_modules
@@ -16,5 +16,6 @@ cxx_test(external_alloc_test dfly_core LABELS DFLY)
 cxx_test(dash_test dfly_core file DATA testdata/ids.txt LABELS DFLY)
 cxx_test(interpreter_test dfly_core LABELS DFLY)
 cxx_test(json_test dfly_core TRDP::jsoncons LABELS DFLY)
+cxx_test(simple_lru_counter_test dfly_core LABELS DFLY)
 cxx_test(string_set_test dfly_core LABELS DFLY)
 cxx_test(string_map_test dfly_core LABELS DFLY)

src/core/simple_lru_counter.cc

@@ -0,0 +1,90 @@
+// Copyright 2023, DragonflyDB authors.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#include "core/simple_lru_counter.h"
+
+#include "base/logging.h"
+
+namespace dfly {
+
+using namespace std;
+
+SimpleLruCounter::SimpleLruCounter(size_t capacity) : head_(0) {
+  CHECK_GT(capacity, 1u);
+  node_arr_.resize(capacity);
+}
+
+SimpleLruCounter::~SimpleLruCounter() {
+}
+
+optional<uint64_t> SimpleLruCounter::Get(string_view key) const {
+  auto it = table_.find(key);
+  if (it == table_.end()) {
+    return nullopt;
+  }
+  const auto& node = node_arr_[it->second];
+
+  DCHECK_EQ(node.key, key);
+
+  return node.count;
+}
+
+void SimpleLruCounter::Put(string_view key, uint64_t value) {
+  auto [it, inserted] = table_.emplace(key, table_.size());
+
+  if (inserted) {
+    unsigned tail = node_arr_[head_].prev;  // 0 if we had 1 or 0 elements.
+
+    if (it->second < node_arr_.size()) {
+      auto& node = node_arr_[it->second];
+      // add new head.
+      node.prev = tail;
+      node.next = head_;
+      node_arr_[tail].next = it->second;
+      node_arr_[head_].prev = it->second;
+      head_ = it->second;
+    } else {
+      // Cache is full, remove the tail.
+      size_t res = table_.erase(string_view(node_arr_[tail].key));
+      DCHECK(res == 1);
+
+      it->second = tail;
+
+      DCHECK_EQ(table_.size(), node_arr_.size());
+    }
+
+    auto& node = node_arr_[it->second];
+    node.key = it->first;  // reference the key. We need it to erase the key referencing tail above.
+    node.count = value;
+  } else {  // not inserted.
+    auto& node = node_arr_[it->second];
+    node.count = value;
+  }
+
+  if (it->second != head_) {  // bump up to head.
+    BumpToHead(it->second);
+  }
+}
+
+void SimpleLruCounter::BumpToHead(uint32_t index) {
+  DCHECK_LT(index, node_arr_.size());
+  DCHECK_NE(index, head_);
+
+  unsigned tail = node_arr_[head_].prev;
+  if (index == tail) {
+    head_ = index;  // just shift the whole cycle.
+    return;
+  }
+
+  auto& node = node_arr_[index];
+
+  DCHECK(node.prev != node.next);
+
+  node_arr_[node.prev].next = node.next;
+  node_arr_[node.next].prev = node.prev;
+  node.prev = tail;
+  node.prev = head_;
+  head_ = index;
+}
+};  // namespace dfly

src/core/simple_lru_counter.h

@@ -0,0 +1,44 @@
+// Copyright 2023, DragonflyDB authors.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+#pragma once
+
+#include <absl/container/flat_hash_map.h>
+
+#include "base/string_view_sso.h"
+
+namespace dfly {
+
+class SimpleLruCounter {
+  struct Node {
+    base::string_view_sso key;  // key to the table.
+
+    uint32_t prev;
+    uint32_t next;
+
+    uint64_t count;
+
+    Node() : prev(0), next(0), count(0) {
+    }
+  };
+
+ public:
+  explicit SimpleLruCounter(size_t capacity);
+  ~SimpleLruCounter();
+
+  std::optional<uint64_t> Get(std::string_view key) const;
+  void Put(std::string_view key, uint64_t count);
+
+  size_t Size() const {
+    return table_.size();
+  }
+
+ private:
+  void BumpToHead(uint32_t index);
+
+  absl::flat_hash_map<std::string, uint32_t> table_;
+  std::vector<Node> node_arr_;
+  uint32_t head_;
+};
+
+};  // namespace dfly

src/core/simple_lru_counter_test.cc

@@ -0,0 +1,48 @@
+// Copyright 2023, DragonflyDB authors.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#include "core/simple_lru_counter.h"
+
+#include "base/gtest.h"
+#include "base/logging.h"
+
+using namespace std;
+
+namespace dfly {
+
+class SimpleLruTest : public ::testing::Test {
+ protected:
+  SimpleLruTest() : cache_(4) {
+  }
+
+  SimpleLruCounter cache_;
+};
+
+TEST_F(SimpleLruTest, Basic) {
+  cache_.Put("a", 1);
+  cache_.Put("b", 2);
+  cache_.Put("c", 3);
+  cache_.Put("d", 4);
+  cache_.Put("a", 1);
+
+  ASSERT_EQ(1, cache_.Get("a"));
+  ASSERT_EQ(2, cache_.Get("b"));
+  ASSERT_EQ(3, cache_.Get("c"));
+  ASSERT_EQ(4, cache_.Get("d"));
+
+  ASSERT_EQ(nullopt, cache_.Get("e"));
+  cache_.Put("e", 5);
+
+  ASSERT_EQ(nullopt, cache_.Get("b"));
+  ASSERT_EQ(3, cache_.Get("c"));
+  ASSERT_EQ(4, cache_.Get("d"));
+  ASSERT_EQ(5, cache_.Get("e"));
+
+  cache_.Put("f", 6);
+  ASSERT_EQ(nullopt, cache_.Get("c"));
+  ASSERT_EQ(5, cache_.Get("e"));
+  ASSERT_EQ(6, cache_.Get("f"));
+}
+
+}  // namespace dfly