diff --git a/src/core/CMakeLists.txt b/src/core/CMakeLists.txt index c9930b8abca2..4a6cfe6a0d4d 100644 --- a/src/core/CMakeLists.txt +++ b/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) diff --git a/src/core/simple_lru_counter.cc b/src/core/simple_lru_counter.cc new file mode 100644 index 000000000000..9856df9e2619 --- /dev/null +++ b/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 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 diff --git a/src/core/simple_lru_counter.h b/src/core/simple_lru_counter.h new file mode 100644 index 000000000000..75b7676f6590 --- /dev/null +++ b/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 + +#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 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 table_; + std::vector node_arr_; + uint32_t head_; +}; + +}; // namespace dfly diff --git a/src/core/simple_lru_counter_test.cc b/src/core/simple_lru_counter_test.cc new file mode 100644 index 000000000000..719c95312762 --- /dev/null +++ b/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