chore(tiering): Optimize serialization

dranikpg · dranikpg · commit 83e4b522661b · 2025-12-05T12:48:17.000+03:00
diff --git a/src/core/compact_object.cc b/src/core/compact_object.cc
@@ -7,6 +7,8 @@
 // #define XXH_INLINE_ALL
 #include <xxhash.h>
 
+#include <array>
+
 extern "C" {
 #include "redis/intset.h"
 #include "redis/listpack.h"
@@ -1609,19 +1611,22 @@ void CompactObj::EncodeString(string_view str, bool is_key) {
   u_.r_obj.SetString(encoded, tl.local_mr);
 }
 
-StringOrView CompactObj::GetRawString() const {
+std::array<std::string_view, 2> CompactObj::GetRawString() const {
   DCHECK(!IsExternal());
 
   if (taglen_ == ROBJ_TAG) {
     CHECK_EQ(OBJ_STRING, u_.r_obj.type());
     DCHECK_EQ(OBJ_ENCODING_RAW, u_.r_obj.encoding());
-    return StringOrView::FromView(u_.r_obj.AsView());
+    return {u_.r_obj.AsView(), {}};
   }
 
   if (taglen_ == SMALL_TAG) {
-    string tmp;
-    u_.small_str.Get(&tmp);
-    return StringOrView::FromString(std::move(tmp));
+    std::string_view arr[2];
+    u_.small_str.GetV(arr);
+    std::array<std::string_view, 2> out;  // TODO: use c++ 20 to_array
+    out[0] = arr[0];
+    out[1] = arr[1];
+    return out;
   }
 
   LOG(FATAL) << "Unsupported tag for GetRawString(): " << int(taglen_);
diff --git a/src/core/compact_object.h b/src/core/compact_object.h
@@ -429,9 +429,9 @@ class CompactObj {
     memory_resource()->deallocate(ptr, sizeof(T), alignof(T));
   }
 
-  // returns raw (non-decoded) string. Used to bypass decoding layer.
+  // Return raw (non-decoded) string as two views. First is guaranteed to be non-empty.
   // Precondition: the object is a non-inline string.
-  StringOrView GetRawString() const;
+  std::array<std::string_view, 2> GetRawString() const;
 
   StrEncoding GetStrEncoding() const {
     return StrEncoding{mask_bits_.encoding, bool(huffman_domain_)};
diff --git a/src/core/compact_object_test.cc b/src/core/compact_object_test.cc
@@ -613,7 +613,8 @@ TEST_F(CompactObjectTest, StrEncodingAndMaterialize) {
       obj.SetString(test_str, false);
 
       // Test StrEncoding helper
-      string raw_str = obj.GetRawString().Take();
+      auto strs = obj.GetRawString();
+      string raw_str = string{strs[0]} + string{strs[1]};
       CompactObj::StrEncoding enc = obj.GetStrEncoding();
       EXPECT_EQ(test_str, enc.Decode(raw_str).Take());
 
diff --git a/src/core/detail/listpack_wrap.cc b/src/core/detail/listpack_wrap.cc
@@ -118,6 +118,10 @@ ListpackWrap::Iterator ListpackWrap::end() const {
   return Iterator{lp_, nullptr, intbuf_};
 }
 
+size_t ListpackWrap::Bytes() const {
+  return lpBytes(lp_);
+}
+
 std::string_view ListpackWrap::GetView(uint8_t* lp_it, uint8_t int_buf[]) {
   int64_t ele_len = 0;
   uint8_t* elem = lpGet(lp_it, &ele_len, int_buf);
diff --git a/src/core/detail/listpack_wrap.h b/src/core/detail/listpack_wrap.h
@@ -59,6 +59,7 @@ struct ListpackWrap {
   Iterator begin() const;
   Iterator end() const;
   size_t size() const;  // number of entries
+  size_t Bytes() const;
 
   // Get view from raw listpack iterator
   static std::string_view GetView(uint8_t* lp_it, uint8_t int_buf[]);
diff --git a/src/server/tiered_storage.cc b/src/server/tiered_storage.cc
@@ -14,6 +14,7 @@
 
 #include "absl/cleanup/cleanup.h"
 #include "absl/flags/internal/flag.h"
+#include "absl/functional/bind_front.h"
 #include "base/flag_utils.h"
 #include "base/flags.h"
 #include "base/logging.h"
@@ -82,23 +83,24 @@ tiering::DiskSegment FromCoolItem(const PrimeValue::CoolItem& item) {
 }
 
 // Determine required byte size and encoding type based on value.
-// TODO(vlad): Maybe split into different accessors?
 // Do NOT enforce rules depending on dynamic runtime values as this is called
 // when scheduling stash and just before succeeeding and is expected to return the same results
 pair<size_t /*size*/, CompactObj::ExternalRep> DetermineSerializationParams(const PrimeValue& pv) {
   switch (pv.ObjType()) {
-    case OBJ_STRING:
+    case OBJ_STRING: {
       if (pv.IsInline())
         return {};
-      return std::make_pair(pv.GetRawString().view().size(), CompactObj::ExternalRep::STRING);
-    case OBJ_HASH:
+      auto strs = pv.GetRawString();
+      return std::make_pair(strs[0].size() + strs[1].size(), CompactObj::ExternalRep::STRING);
+    }
+    case OBJ_HASH: {
       if (pv.Encoding() == kEncodingListPack) {
-        auto* lp = static_cast<uint8_t*>(pv.RObjPtr());
-        size_t bytes = 4 + lpBytes(lp);  // encoded length and data bytes
-        bytes += lpLength(lp) * 2 * 4;   // 4 bytes for encoded key/value lengths
-        return std::make_pair(bytes, CompactObj::ExternalRep::SERIALIZED_MAP);
+        detail::ListpackWrap lw{static_cast<uint8_t*>(pv.RObjPtr())};
+        return std::make_pair(tiering::SerializedMap::EstimateSize(lw.Bytes(), lw.size()),
+                              CompactObj::ExternalRep::SERIALIZED_MAP);
       }
       return {};
+    }
     default:
       return {};
   };
@@ -108,18 +110,17 @@ size_t Serialize(CompactObj::ExternalRep rep, const PrimeValue& pv, io::MutableB
   DCHECK_LE(DetermineSerializationParams(pv).first, buffer.size());
   switch (rep) {
     case CompactObj::ExternalRep::STRING: {
-      auto sv = pv.GetRawString();
-      memcpy(buffer.data(), sv.view().data(), sv.view().size());
-      return sv.view().size();
+      auto strs = pv.GetRawString();
+      memcpy(buffer.data(), strs[0].data(), strs[0].size());
+      if (!strs[1].empty())
+        memcpy(buffer.data() + strs[0].size(), strs[1].data(), strs[1].size());
+      return strs[0].size() + strs[1].size();
     }
     case CompactObj::ExternalRep::SERIALIZED_MAP: {
       DCHECK_EQ(pv.Encoding(), kEncodingListPack);
-
-      // TODO(vlad): Optimize copy for serialization
       detail::ListpackWrap lw{static_cast<uint8_t*>(pv.RObjPtr())};
-      vector<pair<string, string>> entries(lw.begin(), lw.end());
       return tiering::SerializedMap::Serialize(
-          entries, {reinterpret_cast<char*>(buffer.data()), buffer.length()});
+          lw, {reinterpret_cast<char*>(buffer.data()), buffer.length()});
     }
   };
   return 0;
@@ -447,11 +448,9 @@ std::optional<util::fb2::Future<bool>> TieredStorage::TryStash(DbIndex dbid, str
   if (!ShouldStash(*value))
     return {};
 
-  // This invariant should always hold because ShouldStash tests for IoPending flag.
-  CHECK(!bins_->IsPending(dbid, key));
+  CHECK(!bins_->IsPending(dbid, key));  // Because has stash pending is false (ShouldStash checks)
 
-  // TODO: When we are low on memory we should introduce a back-pressure, to avoid OOMs
-  // with a lot of underutilized disk space.
+  // Limit write depth. TODO: Provide backpressure?
   if (op_manager_->GetStats().pending_stash_cnt >= config_.write_depth_limit) {
     ++stats_.stash_overflow_cnt;
     return {};
@@ -463,35 +462,22 @@ std::optional<util::fb2::Future<bool>> TieredStorage::TryStash(DbIndex dbid, str
   tiering::OpManager::EntryId id;
   error_code ec;
 
-  value->SetStashPending(true);
+  value->SetStashPending(true);  // Optimistically set ahead, unset in case of error
+
   if (OccupiesWholePages(est_size)) {  // large enough for own page
     id = KeyRef(dbid, key);
-    if (auto prepared = op_manager_->PrepareStash(est_size); prepared) {
-      auto [offset, buf] = *prepared;
-      size_t written = Serialize(rep, *value, buf.bytes);
-      tiering::DiskSegment segment{offset, written};
-      op_manager_->Stash(id, segment, buf);
-    } else {
-      ec = prepared.error();
-    }
+    auto serialize = absl::bind_front(Serialize, rep, cref(*value));
+    ec = op_manager_->PrepareAndStash(id, est_size, serialize);
   } else if (auto bin = bins_->Stash(dbid, key, SerializeToString(*value)); bin) {
-    id = bin->first;
-    // TODO(vlad): Write bin to prepared buffer instead of allocating one
-    if (auto prepared = op_manager_->PrepareStash(bin->second.length()); prepared) {
-      auto [offset, buf] = *prepared;
-      memcpy(buf.bytes.data(), bin->second.data(), bin->second.size());
-      tiering::DiskSegment segment{offset, bin->second.size()};
-      op_manager_->Stash(id, segment, buf);
-    } else {
-      ec = prepared.error();
-      bins_->ReportStashAborted(bin->first);
-    }
+    id = bin->id;
+    auto serialize = absl::bind_front(&tiering::SmallBins::SerializeBin, bins_.get(), &*bin);
+    ec = op_manager_->PrepareAndStash(id, 4_KB, serialize);
   } else {
-    return {};  // silently added to bin
+    return {};  // added to bin, no operations pending
   }
 
+  // Set stash pending to false on single value or whole bin
   if (ec) {
-    value->SetStashPending(false);
     LOG_IF(ERROR, ec != errc::file_too_large) << "Stash failed immediately" << ec.message();
     visit([this](auto id) { op_manager_->ClearIoPending(id); }, id);
     return {};
diff --git a/src/server/tiering/op_manager.cc b/src/server/tiering/op_manager.cc
@@ -99,6 +99,16 @@ void OpManager::Stash(EntryId id_ref, tiering::DiskSegment segment, util::fb2::U
   storage_.Stash(segment, buf, std::move(io_cb));
 }
 
+std::error_code OpManager::PrepareAndStash(EntryId id, size_t length,
+                                           const std::function<size_t(io::MutableBytes)>& writer) {
+  auto buf = PrepareStash(length);
+  if (!buf.has_value())
+    return buf.error();
+
+  size_t written = writer(buf->second.bytes);
+  Stash(id, {buf->first, written}, buf->second);
+}
+
 OpManager::ReadOp& OpManager::PrepareRead(DiskSegment aligned_segment) {
   DCHECK_EQ(aligned_segment.offset % kPageSize, 0u);
   DCHECK_EQ(aligned_segment.length % kPageSize, 0u);
diff --git a/src/server/tiering/op_manager.h b/src/server/tiering/op_manager.h
@@ -68,6 +68,11 @@ class OpManager {
   // Stash value to be offloaded. It is opaque to OpManager.
   void Stash(EntryId id, tiering::DiskSegment segment, util::fb2::UringBuf buf);
 
+  // PrepareStash + Stash via function
+  std::error_code PrepareAndStash(
+      EntryId id, size_t length,
+      const std::function<size_t /*written*/ (io::MutableBytes)>& writer);
+
   Stats GetStats() const;
 
  protected:
diff --git a/src/server/tiering/op_manager_test.cc b/src/server/tiering/op_manager_test.cc
@@ -83,15 +83,10 @@ struct OpManagerTest : PoolTestBase, OpManager {
   }
 
   std::error_code Stash(EntryId id, std::string_view value) {
-    auto prepared = storage_.PrepareStash(value.size());
-    if (!prepared.has_value())
-      return prepared.error();
-
-    auto [offset, buf] = *prepared;
-    memcpy(buf.bytes.data(), value.data(), value.size());
-    DiskSegment segment{offset, value.size()};
-    OpManager::Stash(id, segment, buf);
-    return {};
+    return PrepareAndStash(id, value.size(), [=](io::MutableBytes bytes) {
+      memcpy(bytes.data(), value.data(), value.size());
+      return value.size();
+    });
   }
 
   absl::flat_hash_map<EntryId, std::string> fetched_;
diff --git a/src/server/tiering/serialized_map.cc b/src/server/tiering/serialized_map.cc
@@ -3,11 +3,14 @@
 #include <absl/base/internal/endian.h>
 
 #include "base/logging.h"
+#include "core/detail/listpack_wrap.h"
 
 namespace dfly::tiering {
 
+constexpr size_t kLenBytes = 4;
+
 SerializedMap::Iterator& SerializedMap::Iterator::operator++() {
-  slice_.remove_prefix(8 + key_.size() + value_.size());
+  slice_.remove_prefix(2 * kLenBytes + key_.size() + value_.size());
   Read();
   return *this;
 }
@@ -37,7 +40,7 @@ SerializedMap::Iterator SerializedMap::Find(std::string_view key) const {
 }
 
 SerializedMap::Iterator SerializedMap::begin() const {
-  return Iterator{slice_.substr(4)};
+  return Iterator{slice_.substr(kLenBytes)};
 }
 
 SerializedMap::Iterator SerializedMap::end() const {
@@ -52,22 +55,18 @@ size_t SerializedMap::DataBytes() const {
   return slice_.size() - 4 - size() * 2 * 4;
 }
 
-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::EstimateSize(size_t data_bytes, size_t entries) {
+  return kLenBytes /* entry number */ + data_bytes + entries * 2 * kLenBytes /* string lenghts */;
 }
 
-size_t SerializedMap::Serialize(Input input, absl::Span<char> buffer) {
-  DCHECK_GE(buffer.size(), SerializeSize(input));
+size_t SerializedMap::Serialize(const detail::ListpackWrap& lw, absl::Span<char> buffer) {
+  DCHECK_GE(buffer.size(), EstimateSize(lw.Bytes(), lw.size()));
+
   char* ptr = buffer.data();
-  absl::little_endian::Store32(ptr, input.size());
+  absl::little_endian::Store32(ptr, lw.size());
   ptr += kLenBytes;
 
-  for (const auto& [key, value] : input) {
+  for (const auto& [key, value] : lw) {
     absl::little_endian::Store32(ptr, key.length());
     ptr += kLenBytes;
     absl::little_endian::Store32(ptr, value.length());
diff --git a/src/server/tiering/serialized_map.h b/src/server/tiering/serialized_map.h
@@ -4,6 +4,10 @@
 
 #include <string_view>
 
+namespace dfly::detail {
+struct ListpackWrap;
+}
+
 namespace dfly::tiering {
 
 // Map built over single continuous byte slice to allow easy read operations.
@@ -49,15 +53,12 @@ struct SerializedMap {
   // Number of bytes of pure keys or values
   size_t DataBytes() const;
 
-  // Input for serialization
-  using Input = const absl::Span<const std::pair<std::string, std::string>>;
-
-  // Buffer size required for serialization
-  static size_t SerializeSize(Input);
+  // Estimate upper bound for serialization size
+  static size_t EstimateSize(size_t data_bytes, size_t entries);
 
   // 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);
+  static size_t Serialize(const ::dfly::detail::ListpackWrap& lw, absl::Span<char> buffer);
 
  private:
   size_t size_;
diff --git a/src/server/tiering/serialized_map_test.cc b/src/server/tiering/serialized_map_test.cc
@@ -1,28 +1,43 @@
 #include "server/tiering/serialized_map.h"
 
+#include <mimalloc.h>
+
 #include <map>
 
 #include "base/logging.h"
+#include "core/detail/listpack_wrap.h"
 #include "gmock/gmock.h"
 
+extern "C" {
+#include "redis/zmalloc.h"
+}
+
 namespace dfly::tiering {
 
 using namespace std;
 
-class SerializedMapTest : public ::testing::Test {};
+struct SerializedMapTest : public ::testing::Test {
+  static void SetUpTestSuite() {
+    init_zmalloc_threadlocal(mi_heap_get_backing());  // to use ListpackWrap
+  }
+};
 
 TEST_F(SerializedMapTest, TestBasic) {
   const vector<std::pair<string, string>> kBase = {{"first key", "first value"},
                                                    {"second key", "second value"},
                                                    {"third key", "third value"},
                                                    {"fourth key", "fourth value"},
                                                    {"fifth key", "fifth value"}};
+  auto lw = detail::ListpackWrap::WithCapacity(100);
+  for (const auto& [k, v] : kBase)
+    lw.Insert(k, v, false);
 
   // Serialize kBase to buffer
   std::string buffer;
-  buffer.resize(SerializedMap::SerializeSize(kBase));
-  size_t written = SerializedMap::Serialize(kBase, absl::MakeSpan(buffer));
+  buffer.resize(SerializedMap::EstimateSize(lw.Bytes(), lw.size()));
+  size_t written = SerializedMap::Serialize(lw, absl::MakeSpan(buffer));
   EXPECT_GT(written, 0u);
+  buffer.resize(written);
 
   // Build map over buffer and check size
   SerializedMap map{buffer};
diff --git a/src/server/tiering/small_bins.cc b/src/server/tiering/small_bins.cc
diff --git a/src/server/tiering/small_bins.h b/src/server/tiering/small_bins.h
diff --git a/src/server/tiering/small_bins_test.cc b/src/server/tiering/small_bins_test.cc
