fixes

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

Author: abhijat

src/server/rdb_load.cc

@@ -2100,11 +2100,11 @@ struct RdbLoader::ObjSettings {
 // The first chunk carries the object type and key. Continuation chunks carry only payload,
 // so we keep the object type, key, db, settings, and pending read.
 struct RdbLoader::StreamState {
+  std::string key;
   DbIndex db_index;
   int type;
   PendingRead pending_read;
   ObjSettings settings;
-  std::string key;
 };
 
 RdbLoader::RdbLoader(Service* service, RdbLoadContext* load_context, std::string snapshot_id)
@@ -2187,12 +2187,6 @@ error_code RdbLoader::Load(io::Source* src) {
     GetCurrentDbSlice().IncrLoadInProgress();
   }
 
-  auto finalize_curr_chunk = [&] {
-    if (!stop_early_.load(memory_order_relaxed))
-      return FinishCurrentChunk();
-    return kOk;
-  };
-
   while (!stop_early_.load(memory_order_relaxed)) {
     if (pause_) {
       ThisFiber::SleepFor(100ms);
@@ -2248,15 +2242,15 @@ error_code RdbLoader::Load(io::Source* src) {
     }
 
     if (type == RDB_OPCODE_EOF) {
-      if (current_chunk_state_.has_value() || !stream_states_.empty()) {
-        if (current_chunk_state_.has_value()) {
-          LOG(ERROR) << "eof seen while a previous chunk is not yet finished, stream id "
-                     << current_chunk_state_->stream_id << ", remaining bytes "
-                     << current_chunk_state_->remaining_payload_bytes
-                     << ", pending stream states: " << stream_states_.size();
-          return RdbError(errc::rdb_chunk_payload_remaining);
-        }
+      if (current_chunk_state_) {
+        LOG(ERROR) << "eof seen while a previous chunk is not yet finished, stream id "
+                   << current_chunk_state_->stream_id << ", remaining bytes "
+                   << current_chunk_state_->remaining_payload_bytes
+                   << ", pending stream states: " << stream_states_.size();
+        return RdbError(errc::rdb_chunk_payload_remaining);
+      }
 
+      if (!stream_states_.empty()) {
         LOG(ERROR) << "eof seen while pending stream states: " << stream_states_.size();
         return RdbError(errc::rdb_chunk_payload_remaining);
       }
@@ -2402,7 +2396,7 @@ error_code RdbLoader::Load(io::Source* src) {
       // path below will read its type and key.
       if (stream_states_.contains(current_chunk_state_->stream_id)) {
         RETURN_ON_ERR(LoadValueChunk());
-        RETURN_ON_ERR(finalize_curr_chunk());
+        RETURN_ON_ERR(FinalizeCurrentChunkIfNeeded());
       }
       continue;
     }
@@ -2423,7 +2417,7 @@ error_code RdbLoader::Load(io::Source* src) {
     VLOG(2) << "LoadKeyValPair key=" << last_key_loaded_ << " rdb_type=" << type
             << " db= " << cur_db_index_;
     settings.Reset();
-    RETURN_ON_ERR(finalize_curr_chunk());
+    RETURN_ON_ERR(FinalizeCurrentChunkIfNeeded());
   }  // main load loop
 
   DVLOG(1) << "RdbLoad loop finished";
@@ -3010,11 +3004,11 @@ error_code RdbLoader::LoadKeyValPair(int type, ObjSettings* settings) {
   // the next chunk with the same stream id arrives.
   if (!finalized && current_chunk_state_) {
     StreamState stream_state{
+        .key = std::move(key),
         .db_index = cur_db_index_,
         .type = type,
         .pending_read = std::move(pending_read_),
         .settings = *settings,
-        .key = std::move(key),
     };
     const bool inserted =
         stream_states_.try_emplace(current_chunk_state_->stream_id, std::move(stream_state)).second;
@@ -3286,6 +3280,12 @@ error_code RdbLoader::HandleShardDocIndex() {
   return kOk;
 }
 
+std::error_code RdbLoader::FinalizeCurrentChunkIfNeeded() {
+  if (stop_early_.load(memory_order_relaxed))
+    return kOk;
+  return FinishCurrentChunk();
+}
+
 error_code RdbLoader::LoadVectorIndexNodes(uint64_t elements_number,
                                            std::vector<search::HnswNodeData>* nodes) {
   nodes->reserve(elements_number);

src/server/rdb_load.h

@@ -427,6 +427,10 @@ class RdbLoader : protected RdbLoaderBase {
   std::error_code HandleVectorIndex();
   std::error_code HandleShardDocIndex();
 
+  // validates if the current chunk is fully read, resets the state. returns early if stop_early_ is
+  // requested.
+  std::error_code FinalizeCurrentChunkIfNeeded();
+
   Service* service_;
   RdbLoadContext* load_context_;
 

src/server/rdb_test.cc

@@ -1544,8 +1544,8 @@ TEST_F(RdbTest, SplitSBF) {
 
   // split the blob of the second filter into three chunks. this exercises the loader path where we
   // first try to load the incomplete filter, and return early before that finishes
-  constexpr size_t first_split = 17;
-  constexpr size_t second_split = 13;
+  constexpr size_t kFirstSplit = 17;
+  constexpr size_t kSecondSplit = 13;
 
   pp_->at(0)->Await([&] {
     const DbContext ctx{&namespaces->GetDefaultNamespace(), 0, GetCurrentTimeMs()};
@@ -1559,7 +1559,7 @@ TEST_F(RdbTest, SplitSBF) {
     const std::string blob0{sbf->data(0)};
 
     blob1 = std::string{sbf->data(1)};
-    ASSERT_GT(blob1.size(), first_split + second_split);
+    ASSERT_GT(blob1.size(), kFirstSplit + kSecondSplit);
 
     first.push_back(RDB_TYPE_SBF2);
     // brand new key whose shape is copied off bf_src
@@ -1583,8 +1583,8 @@ TEST_F(RdbTest, SplitSBF) {
     // total size of blob1
     AppendLen(&first, blob1.size());
     // only 17 bytes from blob1 in this chunk
-    AppendLen(&first, first_split);
-    first.append(blob1.data(), first_split);
+    AppendLen(&first, kFirstSplit);
+    first.append(blob1.data(), kFirstSplit);
   });
 
   // add this plain string between chunks of blob1 filter
@@ -1595,14 +1595,14 @@ TEST_F(RdbTest, SplitSBF) {
 
   // p2 of blob1
   std::string second;
-  AppendLen(&second, second_split);
-  second.append(blob1.data() + first_split, second_split);
+  AppendLen(&second, kSecondSplit);
+  second.append(blob1.data() + kFirstSplit, kSecondSplit);
 
-  // p2 of blob1
+  // p3 of blob1
   std::string third;
-  constexpr auto sum = first_split + second_split;
-  AppendLen(&third, blob1.size() - sum);
-  third.append(blob1.data() + sum, blob1.size() - sum);
+  constexpr auto kPrefixConsumed = kFirstSplit + kSecondSplit;
+  AppendLen(&third, blob1.size() - kPrefixConsumed);
+  third.append(blob1.data() + kPrefixConsumed, blob1.size() - kPrefixConsumed);
 
   std::string body;
   body += MakeTaggedChunk(1, first);