server: Use simpler buffer switch design

Signed-off-by: Abhijat Malviya <abhijat@dragonflydb.io>

Author: abhijat

src/server/rdb_save.cc

@@ -214,7 +214,10 @@ uint8_t RdbObjectType(const CompactObj& pv) {
 }
 
 RdbSerializerBase::RdbSerializerBase(CompressionMode compression_mode)
-    : compression_mode_(compression_mode), mem_buf_{4_KB}, tmp_buf_(nullptr) {
+    : compression_mode_(compression_mode),
+      default_buf_{4_KB},
+      mem_buf_{&default_buf_},
+      tmp_buf_(nullptr) {
 }
 
 RdbSerializer::RdbSerializer(CompressionMode compression_mode, ConsumeFun consume_fun,
@@ -277,10 +280,7 @@ io::Result<uint8_t> RdbSerializer::SaveEntry(const PrimeKey& pk, const PrimeValu
     return 0;
   }
 
-  // If mem buf contains data, stash it before we begin this entry.
-  StashCurrentBuffer();
-
-  active_entry_.Reset(ActiveEntry::Kind::Baseline);
+  active_entry_id_ = next_stream_id_++;
 
   DVLOG(3) << "Selecting " << dbid << " previous: " << last_entry_db_index_;
   auto ec = SelectDb(dbid);
@@ -858,7 +858,6 @@ error_code RdbSerializer::SendEofAndChecksum() {
 }
 
 error_code RdbSerializer::SendJournalOffset(uint64_t journal_offset) {
-  SetRawMode();
   VLOG(2) << "SendJournalOffset";
   RETURN_ON_ERR(WriteOpcode(RDB_OPCODE_JOURNAL_OFFSET));
   uint8_t buf[sizeof(uint64_t)];
@@ -914,18 +913,18 @@ std::error_code RdbSerializerBase::WriteOpcode(uint8_t opcode) {
 }
 
 size_t RdbSerializerBase::GetBufferCapacity() const {
-  return mem_buf_.Capacity();
+  return mem_buf_->Capacity();
 }
 
 size_t RdbSerializerBase::GetTempBufferSize() const {
   return tmp_buf_.size();
 }
 
 error_code RdbSerializerBase::WriteRaw(const io::Bytes& buf) {
-  mem_buf_.Reserve(mem_buf_.InputLen() + buf.size());
-  IoBuf::Bytes dest = mem_buf_.AppendBuffer();
+  mem_buf_->Reserve(mem_buf_->InputLen() + buf.size());
+  IoBuf::Bytes dest = mem_buf_->AppendBuffer();
   memcpy(dest.data(), buf.data(), buf.size());
-  mem_buf_.CommitWrite(buf.size());
+  mem_buf_->CommitWrite(buf.size());
   return error_code{};
 }
 
@@ -942,33 +941,35 @@ string RdbSerializerBase::Flush(FlushState flush_state) {
 
   string result(io::View(bytes));
 
-  mem_buf_.ConsumeInput(bytes.size());
+  mem_buf_->ConsumeInput(bytes.size());
 
   return result;
 }
 
 string RdbSerializer::Flush(FlushState flush_state) {
-  string blob = FlushImpl(flush_state);
+  const auto bytes = PrepareFlush(flush_state);
+  auto result = PrefixDefaultBufferAndTag(bytes);
+  mem_buf_->ConsumeInput(bytes.size());
+  if (result.empty())
+    return {};
+
   if (send_tagged_entries_) {
-    if (auto res = CompressBlob(blob); res)
-      blob = std::move(*res);
+    if (auto res = CompressBlob(result); res)
+      result = std::move(*res);
   }
-  return blob;
-}
 
-string RdbSerializer::FlushImpl(FlushState flush_state) {
-  DrainMemBufIntoPendingRecords(flush_state);
+  if (result.size() > serialization_peak_bytes_) {
+    serialization_peak_bytes_ = result.size();
+  }
 
-  std::string out = absl::StrJoin(pending_records_, "");
-  pending_records_.clear();
-  pending_record_bytes_ = 0;
+  DVLOG(2) << "FlushToSink " << result.size() << " bytes";
 
   // After every flush we should write the DB index again because the blobs in the channel are
   // interleaved and multiple savers can correspond to a single writer (in case of single file rdb
   // snapshot)
   last_entry_db_index_ = kInvalidDbId;
 
-  return out;
+  return result;
 }
 
 namespace {
@@ -1038,11 +1039,14 @@ string RdbSerializerBase::DumpValue(const PrimeValue& obj, bool ignore_crc) {
 }
 
 size_t RdbSerializerBase::SerializedLen() const {
-  return mem_buf_.InputLen();
+  auto len = mem_buf_->InputLen();
+  if (mem_buf_ != &default_buf_)
+    len += default_buf_.InputLen();
+  return len;
 }
 
 io::Bytes RdbSerializerBase::PrepareFlush(FlushState flush_state) {
-  size_t sz = mem_buf_.InputLen();
+  size_t sz = mem_buf_->InputLen();
   if (sz == 0)
     return {};
 
@@ -1061,7 +1065,7 @@ io::Bytes RdbSerializerBase::PrepareFlush(FlushState flush_state) {
     number_of_chunks_ = is_last_chunk ? 0 : number_of_chunks_ + 1;
   }
 
-  return mem_buf_.InputBuffer();
+  return mem_buf_->InputBuffer();
 }
 
 error_code RdbSerializerBase::WriteJournalEntry(std::string_view serialized_entry) {
@@ -1889,59 +1893,105 @@ std::optional<std::string> RdbSerializerBase::CompressBlob(std::string_view inpu
 }
 
 void RdbSerializerBase::CompressBlob() {
-  Bytes blob_to_compress = mem_buf_.InputBuffer();
+  Bytes blob_to_compress = mem_buf_->InputBuffer();
   std::string_view input{reinterpret_cast<const char*>(blob_to_compress.data()),
                          blob_to_compress.size()};
   auto compressed = CompressBlob(input);
   if (!compressed)
     return;
 
-  mem_buf_.ConsumeInput(blob_to_compress.size());
-  mem_buf_.Reserve(compressed->size());
-  auto destination = mem_buf_.AppendBuffer();
+  mem_buf_->ConsumeInput(blob_to_compress.size());
+  mem_buf_->Reserve(compressed->size());
+  auto destination = mem_buf_->AppendBuffer();
   memcpy(destination.data(), compressed->data(), compressed->size());
-  mem_buf_.CommitWrite(compressed->size());
+  mem_buf_->CommitWrite(compressed->size());
 }
 
 size_t RdbSerializer::GetTempBufferSize() const {
   return RdbSerializerBase::GetTempBufferSize() + tmp_str_.size();
 }
 
-size_t RdbSerializer::SerializedLen() const {
-  return RdbSerializerBase::SerializedLen() + pending_record_bytes_;
-}
-
 std::error_code RdbSerializer::SendFullSyncCut() {
-  SetRawMode();
   return RdbSerializerBase::SendFullSyncCut();
 }
 
 error_code RdbSerializer::WriteJournalEntry(string_view serialized_entry) {
-  SetRawMode();
   return RdbSerializerBase::WriteJournalEntry(serialized_entry);
 }
 
+std::string RdbSerializer::PrefixDefaultBufferAndTag(Bytes bytes) {
+  const bool should_tag = send_tagged_entries_ && current_continuation_id_ != 0;
+
+  const auto current_bytes = io::View(bytes);
+  if (mem_buf_ == &default_buf_ || default_buf_.InputLen() == 0) {
+    if (!should_tag)
+      return std::string{current_bytes};
+
+    auto header = MakeTagHeader(current_bytes.size());
+    std::string out;
+    out.reserve(header.size() + current_bytes.size());
+    out.append(std::string_view(reinterpret_cast<const char*>(header.data()), header.size()));
+    out.append(current_bytes);
+    return out;
+  }
+
+  std::string out;
+  const auto prefix = default_buf_.InputBuffer();
+
+  out.reserve(prefix.size() + 9 + current_bytes.size());
+  out.append(io::View(prefix));
+  default_buf_.ConsumeInput(prefix.size());
+
+  if (should_tag)
+    out.append(io::View(MakeTagHeader(current_bytes.size())));
+  out.append(current_bytes);
+  return out;
+}
+
 void RdbSerializer::PushToConsumerIfNeeded(FlushState flush_state) {
-  if (consume_fun_ && SerializedLen() > flush_threshold_) {
+  if (!consume_fun_)
+    return;
+
+  const auto id = active_entry_id_;
+  const bool is_mid = flush_state == FlushState::kFlushMidEntry;
+  const bool is_end = flush_state == FlushState::kFlushEndEntry;
+  const bool has_continuation = continuations_.contains(id);
+  // Force tail entries which were split earlier to be pushed even if we do not have enough data, to
+  // avoid those entries lying around in continuations_
+  const bool must_push_tail = send_tagged_entries_ && is_end && has_continuation;
+
+  if (SerializedLen() <= flush_threshold_ && !must_push_tail)
+    return;
+
+  if (!send_tagged_entries_) {
     string blob = Flush(flush_state);
     DCHECK(!blob.empty());  // SerializedLen() > 0.
-    const auto saved_entry = active_entry_;
     consume_fun_(std::move(blob));
-    StashCurrentBuffer();
-    // Restore saved entry if we yielded during consume_fun_, restores the tagging information
-    // (ActiveEntry::chunked, ActiveEntry::stream_id)
-    active_entry_ = saved_entry;
+    return;
   }
-}
 
-void RdbSerializer::StashCurrentBuffer() {
-  if (!send_tagged_entries_)
-    return;
+  if (is_mid) {
+    continuations_.try_emplace(id, std::make_unique<IoBuf>());
+    current_continuation_id_ = id;
+  }
 
-  if (mem_buf_.InputLen() == 0)
-    return;
+  string blob = Flush(flush_state);
+  DCHECK(!blob.empty());  // SerializedLen() > 0.
+
+  mem_buf_ = &default_buf_;
+  current_continuation_id_ = 0;
+  consume_fun_(std::move(blob));
 
-  DrainMemBufIntoPendingRecords(FlushState::kFlushEndEntry);
+  active_entry_id_ = id;
+
+  if (is_mid) {
+    mem_buf_ = continuations_[id].get();
+    current_continuation_id_ = id;
+  } else {
+    continuations_.erase(id);
+    mem_buf_ = &default_buf_;
+    current_continuation_id_ = 0;
+  }
 }
 
 std::string RdbSerializer::TagChunk(std::string blob, uint32_t stream_id) {
@@ -1959,57 +2009,12 @@ std::string RdbSerializer::TagChunk(std::string blob, uint32_t stream_id) {
   return result;
 }
 
-void RdbSerializer::SetRawMode() {
-  if (active_entry_.kind != ActiveEntry::Kind::Raw)
-    StashCurrentBuffer();
-  active_entry_.Reset(ActiveEntry::Kind::Raw);
-}
-
-uint32_t RdbSerializer::AllocateStreamId() {
-  return next_stream_id_++;
-}
-
-std::string RdbSerializer::FinalizeCurrentRecord(FlushState flush_state) {
-  if (mem_buf_.InputLen() == 0) {
-    return {};
-  }
-
-  auto blob = RdbSerializerBase::Flush(flush_state);
-  if (!send_tagged_entries_) {
-    if (flush_state == FlushState::kFlushEndEntry)
-      active_entry_.Reset(ActiveEntry::Kind::Raw);
-    // not tagged entry always returns blob as it is
-    return blob;
-  }
-
-  if (active_entry_.kind == ActiveEntry::Kind::Raw)
-    return blob;
-
-  if (flush_state == FlushState::kFlushMidEntry) {
-    // Mark the current entry as chunked. Any future chunks going out for this entry will be tagged
-    // too
-    if (!active_entry_.chunked) {
-      active_entry_.chunked = true;
-      active_entry_.stream_id = AllocateStreamId();
-    }
-    return TagChunk(blob, *active_entry_.stream_id);
-  }
-
-  // end of entry
-  if (active_entry_.chunked) {
-    blob = TagChunk(blob, *active_entry_.stream_id);
-  }
-
-  // Reset active entry, current baseline entry is finished
-  active_entry_.Reset(ActiveEntry::Kind::Raw);
-  return blob;
-}
-
-void RdbSerializer::DrainMemBufIntoPendingRecords(FlushState flush_state) {
-  if (auto record = FinalizeCurrentRecord(flush_state); !record.empty()) {
-    pending_record_bytes_ += record.size();
-    pending_records_.push_back(std::move(record));
-  }
+std::array<uint8_t, 9> RdbSerializer::MakeTagHeader(size_t size) const {
+  std::array<uint8_t, 9> header;
+  header[0] = RDB_OPCODE_TAGGED_CHUNK;
+  absl::little_endian::Store32(header.data() + 1, current_continuation_id_);
+  absl::little_endian::Store32(header.data() + 5, size);
+  return header;
 }
 
 }  // namespace dfly