chore(tiering): Optimize serialization (#6164)

Author: dranikpg

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"
@@ -1463,8 +1465,7 @@ bool CompactObj::CmpEncoded(string_view sv) const {
     //    has only 9-10 bytes in its inline prefix storage.
     DCHECK_GT(sv.size(), 16u);  // we would not be in SMALL_TAG, otherwise.
 
-    string_view slice[2];
-    u_.small_str.Get(slice);
+    auto slice = u_.small_str.Get();
     DCHECK_LT(slice[0].size(), 14u);
 
     uint8_t tmpbuf[14];
@@ -1587,19 +1588,17 @@ 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));
+    return u_.small_str.Get();
   }
 
   LOG(FATAL) << "Unsupported tag for GetRawString(): " << int(taglen_);

src/core/compact_object.h

@@ -426,9 +426,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_)};

src/core/compact_object_test.cc

@@ -612,7 +612,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());
 

src/core/detail/listpack_wrap.cc

@@ -118,6 +118,10 @@ ListpackWrap::Iterator ListpackWrap::end() const {
   return Iterator{lp_, nullptr, intbuf_};
 }
 
+size_t ListpackWrap::UsedBytes() 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);

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 UsedBytes() const;
 
   // Get view from raw listpack iterator
   static std::string_view GetView(uint8_t* lp_it, uint8_t int_buf[]);

src/core/small_string.cc

@@ -107,16 +107,11 @@ bool SmallString::Equal(const SmallString& os) const {
   if (size_ != os.size_)
     return false;
 
-  string_view me[2], other[2];
-  Get(me);
-  os.Get(other);
-
-  return me[0] == other[0] && me[1] == other[1];
+  return Get() == os.Get();
 }
 
 uint64_t SmallString::HashCode() const {
-  string_view slice[2];
-  Get(slice);
+  array<string_view, 2> slice = Get();
 
   XXH3_state_t* state = tl.xxh_state.get();
   XXH3_64bits_reset_withSeed(state, kHashSeed);
@@ -126,17 +121,18 @@ uint64_t SmallString::HashCode() const {
   return XXH3_64bits_digest(state);
 }
 
-void SmallString::Get(string_view dest[2]) const {
+array<string_view, 2> SmallString::Get() const {
   DCHECK(size_);
 
+  array<string_view, 2> dest;
   dest[0] = string_view{prefix_, kPrefLen};
   uint8_t* ptr = tl.seg_alloc->Translate(small_ptr_);
   dest[1] = string_view{reinterpret_cast<char*>(ptr), size_ - kPrefLen};
+  return dest;
 }
 
 void SmallString::Get(char* out) const {
-  string_view strs[2];
-  Get(strs);
+  auto strs = Get();
   memcpy(out, strs[0].data(), strs[0].size());
   memcpy(out + strs[0].size(), strs[1].data(), strs[1].size());
 }

src/core/small_string.h

@@ -3,6 +3,7 @@
 //
 #pragma once
 
+#include <array>
 #include <cstdint>
 #include <string_view>
 
@@ -31,7 +32,7 @@ class SmallString {
   uint64_t HashCode() const;
   uint16_t MallocUsed() const;
 
-  void Get(std::string_view dest[2]) const;
+  std::array<std::string_view, 2> Get() const;
   void Get(char* out) const;
   void Get(std::string* dest) const;
 

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.UsedBytes(), 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 {};

src/server/tiering/op_manager.cc

@@ -99,6 +99,17 @@ 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);
+  return {};
+}
+
 OpManager::ReadOp& OpManager::PrepareRead(DiskSegment aligned_segment) {
   DCHECK_EQ(aligned_segment.offset % kPageSize, 0u);
   DCHECK_EQ(aligned_segment.length % kPageSize, 0u);

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: