feat(tiering): Serialized map (#5896)

Author: dranikpg

src/server/CMakeLists.txt

@@ -25,15 +25,16 @@ set_property(SOURCE dfly_main.cc APPEND PROPERTY COMPILE_DEFINITIONS
              SOURCE_PATH_FROM_BUILD_ENV=${CMAKE_SOURCE_DIR})
 
 if (WITH_TIERING)
-    SET(TX_LINUX_SRCS tiering/disk_storage.cc tiering/op_manager.cc tiering/small_bins.cc
-      tiering/external_alloc.cc tiering/decoders.cc)
+    SET(TX_LINUX_SRCS tiering/disk_storage.cc tiering/op_manager.cc
+      tiering/small_bins.cc tiering/external_alloc.cc tiering/serialized_map.cc tiering/decoders.cc)
 
     add_executable(dfly_bench dfly_bench.cc)
     cxx_link(dfly_bench dfly_parser_lib fibers2 absl::random_random redis_lib)
     cxx_test(tiering/disk_storage_test dfly_test_lib LABELS DFLY)
     cxx_test(tiering/op_manager_test dfly_test_lib LABELS DFLY)
     cxx_test(tiering/small_bins_test dfly_test_lib LABELS DFLY)
     cxx_test(tiering/external_alloc_test dfly_test_lib LABELS DFLY)
+    cxx_test(tiering/serialized_map_test dfly_test_lib LABELS DFLY)
 endif()
 
 # Optionally include command families as needed

src/server/tiering/serialized_map.cc

@@ -0,0 +1,80 @@
+#include "server/tiering/serialized_map.h"
+
+#include <absl/base/internal/endian.h>
+
+#include "base/logging.h"
+
+namespace dfly::tiering {
+
+SerializedMap::Iterator& SerializedMap::Iterator::operator++() {
+  slice_.remove_prefix(8 + key_.size() + value_.size());
+  Read();
+  return *this;
+}
+
+SerializedMap::Iterator::Iterator(std::string_view buffer) : slice_{buffer} {
+  Read();
+}
+
+void SerializedMap::Iterator::Read() {
+  if (slice_.empty())
+    return;
+
+  uint32_t key_len = absl::little_endian::Load32(slice_.data());
+  uint32_t value_len = absl::little_endian::Load32(slice_.data() + 4);
+  key_ = {slice_.data() + 8, key_len};
+  value_ = {slice_.data() + 8 + key_len, value_len};
+}
+
+SerializedMap::SerializedMap(std::string_view slice) {
+  size_ = absl::little_endian::Load32(slice.data());
+  DCHECK_GT(size_, 0u);
+  slice_ = slice;
+}
+
+SerializedMap::Iterator SerializedMap::Find(std::string_view key) const {
+  return std::find_if(begin(), end(), [key](auto p) { return p.first == key; });
+}
+
+SerializedMap::Iterator SerializedMap::begin() const {
+  return Iterator{slice_.substr(4)};
+}
+
+SerializedMap::Iterator SerializedMap::end() const {
+  return Iterator{slice_.substr(slice_.size(), 0)};
+}
+
+size_t SerializedMap::size() const {
+  return size_;
+}
+
+constexpr size_t kLenBytes = 4;
+
+size_t SerializedMap::SerializeSize(Input input) {
+  size_t out = kLenBytes;  // number of entries
+  for (const auto& [key, value] : input)
+    out += kLenBytes * 2 /* string lengts */ + key.size() + value.size();
+  return out;
+}
+
+size_t SerializedMap::Serialize(Input input, absl::Span<char> buffer) {
+  DCHECK_GE(buffer.size(), SerializeSize(input));
+  char* ptr = buffer.data();
+  absl::little_endian::Store32(ptr, input.size());
+  ptr += kLenBytes;
+
+  for (const auto& [key, value] : input) {
+    absl::little_endian::Store32(ptr, key.length());
+    ptr += kLenBytes;
+    absl::little_endian::Store32(ptr, value.length());
+    ptr += kLenBytes;
+    memcpy(ptr, key.data(), key.length());
+    ptr += key.length();
+    memcpy(ptr, value.data(), value.length());
+    ptr += value.length();
+  }
+
+  return ptr - buffer.data();
+}
+
+}  // namespace dfly::tiering

src/server/tiering/serialized_map.h

@@ -0,0 +1,64 @@
+#pragma once
+
+#include <absl/types/span.h>
+
+#include <string_view>
+
+namespace dfly::tiering {
+
+// Map built over single continuous byte slice to allow easy read operations.
+struct SerializedMap {
+  struct Iterator {
+    using iterator_category = std::forward_iterator_tag;
+    using difference_type = std::ptrdiff_t;
+    using value_type = std::pair<std::string_view, std::string_view>;
+    using reference = value_type;
+    using pointer = value_type*;
+
+    Iterator& operator++();
+
+    bool operator==(const Iterator& other) const {
+      return slice_.data() == other.slice_.data() && slice_.size() == other.slice_.size();
+    }
+
+    bool operator!=(const Iterator& other) const {
+      return !operator==(other);
+    }
+
+    std::pair<std::string_view, std::string_view> operator*() const {
+      return {key_, value_};
+    }
+
+   private:
+    friend struct SerializedMap;
+
+    explicit Iterator(std::string_view buffer);
+    void Read();
+
+    std::string_view slice_;  // the part left
+    std::string_view key_, value_;
+  };
+
+  explicit SerializedMap(std::string_view slice);
+
+  Iterator Find(std::string_view key) const;  // Linear search
+  Iterator begin() const;
+  Iterator end() const;
+  size_t size() const;
+
+  // Input for serialization
+  using Input = const absl::Span<const std::pair<std::string_view, std::string_view>>;
+
+  // Buffer size required for serialization
+  static size_t SerializeSize(Input);
+
+  // Write a slice that can be used to a SerializedMap on top of it.
+  // Returns number of bytes written
+  static size_t Serialize(Input, absl::Span<char> buffer);
+
+ private:
+  size_t size_;
+  std::string_view slice_;
+};
+
+}  // namespace dfly::tiering

src/server/tiering/serialized_map_test.cc

@@ -0,0 +1,38 @@
+#include "server/tiering/serialized_map.h"
+
+#include <map>
+
+#include "base/logging.h"
+#include "gmock/gmock.h"
+
+namespace dfly::tiering {
+
+using namespace std;
+
+class SerializedMapTest : public ::testing::Test {};
+
+TEST_F(SerializedMapTest, TestBasic) {
+  const vector<std::pair<string_view, string_view>> kBase = {{"first key", "first value"},
+                                                             {"second key", "second value"},
+                                                             {"third key", "third value"},
+                                                             {"fourth key", "fourth value"},
+                                                             {"fifth key", "fifth value"}};
+
+  // Serialize kBase to buffer
+  std::string buffer;
+  buffer.resize(SerializedMap::SerializeSize(kBase));
+  size_t written = SerializedMap::Serialize(kBase, absl::MakeSpan(buffer));
+  EXPECT_GT(written, 0u);
+
+  // Build map over buffer and check size
+  SerializedMap map{buffer};
+  EXPECT_EQ(map.size(), kBase.size());
+
+  // Check entries
+  size_t idx = 0;
+  for (auto it = map.begin(); it != map.end(); ++it, ++idx) {
+    EXPECT_EQ(*it, kBase[idx]);
+  }
+}
+
+}  // namespace dfly::tiering