addres review comments

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

Author: abhijat

src/server/rdb_load.cc

@@ -2187,6 +2187,12 @@ 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);
@@ -2251,8 +2257,7 @@ error_code RdbLoader::Load(io::Source* src) {
           return RdbError(errc::rdb_chunk_payload_remaining);
         }
 
-        if (!stream_states_.empty())
-          LOG(ERROR) << "eof seen while pending stream states: " << stream_states_.size();
+        LOG(ERROR) << "eof seen while pending stream states: " << stream_states_.size();
         return RdbError(errc::rdb_chunk_payload_remaining);
       }
       /* EOF: End of file, exit the main loop. */
@@ -2397,14 +2402,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());
-        if (!stop_early_.load(memory_order_relaxed) &&
-            current_chunk_state_->remaining_payload_bytes != 0) {
-          LOG(ERROR) << "chunk fully consumed but payload bytes remain "
-                     << current_chunk_state_->remaining_payload_bytes;
-          return RdbError(errc::rdb_chunk_payload_remaining);
-        }
-        // This chunk is fully consumed, clear the state
-        current_chunk_state_.reset();
+        RETURN_ON_ERR(finalize_curr_chunk());
       }
       continue;
     }
@@ -2425,16 +2423,7 @@ error_code RdbLoader::Load(io::Source* src) {
     VLOG(2) << "LoadKeyValPair key=" << last_key_loaded_ << " rdb_type=" << type
             << " db= " << cur_db_index_;
     settings.Reset();
-    if (!stop_early_.load(memory_order_relaxed) && current_chunk_state_ &&
-        current_chunk_state_->remaining_payload_bytes != 0) {
-      LOG(ERROR) << "chunk fully consumed but payload bytes remain "
-                 << current_chunk_state_->remaining_payload_bytes;
-      return RdbError(errc::rdb_chunk_payload_remaining);
-    }
-
-    // If we just read the first chunk of a key, then reset state here because LoadKeyValPair will
-    // only return when the chunk finishes
-    current_chunk_state_.reset();
+    RETURN_ON_ERR(finalize_curr_chunk());
   }  // main load loop
 
   DVLOG(1) << "RdbLoad loop finished";
@@ -2633,12 +2622,7 @@ error_code RdbLoaderBase::HandleCompressedBlob(int op_type) {
   // Stop counting payload bytes on decompressed data. At this point the entire payload size must be
   // consumed as it was the compressed blob. We switch to another buffer and must be able to read
   // everything from it without any checks
-  if (current_chunk_state_ && current_chunk_state_->remaining_payload_bytes > 0) {
-    LOG(ERROR) << "Compressed blob not fully consumed, remaining bytes "
-               << current_chunk_state_->remaining_payload_bytes;
-    return RdbError(errc::rdb_chunk_payload_remaining);
-  }
-  current_chunk_state_.reset();
+  RETURN_ON_ERR(FinishCurrentChunk());
 
   // Decompress blob and switch membuf pointer
   // Last type in the compressed blob is RDB_OPCODE_COMPRESSED_BLOB_END
@@ -2853,6 +2837,20 @@ std::error_code RdbLoaderBase::FromOpaque(const OpaqueObj& opaque, LoadConfig co
   return visitor.ec();
 }
 
+std::error_code RdbLoaderBase::FinishCurrentChunk() {
+  if (!current_chunk_state_)
+    return kOk;
+
+  if (!ChunkBudgetExhausted()) {
+    LOG(ERROR) << "chunk fully consumed but payload bytes remain "
+               << current_chunk_state_->remaining_payload_bytes;
+    return RdbError(errc::rdb_chunk_payload_remaining);
+  }
+
+  current_chunk_state_.reset();
+  return kOk;
+}
+
 void RdbLoaderBase::CopyStreamId(const StreamID& src, struct streamID* dest) {
   dest->ms = src.ms;
   dest->seq = src.seq;
@@ -3011,21 +3009,23 @@ error_code RdbLoader::LoadKeyValPair(int type, ObjSettings* settings) {
   // If the first tagged chunk did not finish the object, save enough state to resume it when
   // the next chunk with the same stream id arrives.
   if (!finalized && current_chunk_state_) {
-    if (stream_states_.contains(current_chunk_state_->stream_id)) {
-      LOG(ERROR) << "attempt to add first chunk for id " << current_chunk_state_->stream_id
-                 << " which already exists";
-      return RdbError(errc::rdb_file_corrupted);
-    }
-    stream_states_[current_chunk_state_->stream_id] = {
+    StreamState stream_state{
         .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;
+    if (!inserted) {
+      LOG(ERROR) << "attempt to add first chunk for id " << current_chunk_state_->stream_id
+                 << " which already exists";
+      return RdbError(errc::rdb_file_corrupted);
+    }
   }
 
-  int delta_ms = (absl::GetCurrentTimeNanos() - start) / 1000'000;
+  const int delta_ms = (absl::GetCurrentTimeNanos() - start) / 1000'000;
   LOG_IF(INFO, delta_ms > 1000) << "Took " << delta_ms << " ms to load rdb_type " << type;
 
   pending_read_ = {};

src/server/rdb_load.h

@@ -236,6 +236,10 @@ class RdbLoaderBase {
     return current_chunk_state_ && current_chunk_state_->remaining_payload_bytes == 0;
   }
 
+  // Called to validate that the current chunk is fully consumed, after validation resets current
+  // chunk state.
+  std::error_code FinishCurrentChunk();
+
   static void CopyStreamId(const StreamID& src, struct streamID* dest);
 
   base::IoBuf* mem_buf_ = nullptr;