refactor(snapshot): Add CallSerializerUnderLock method

Signed-off-by: Stepan Bagritsevich <stefan@dragonflydb.io>

Author: BagritsevichStepan

src/server/snapshot.cc

@@ -38,7 +38,11 @@ constexpr size_t kMinBlobSize = 32_KB;
 
 SliceSnapshot::SliceSnapshot(CompressionMode compression_mode, DbSlice* slice,
                              SnapshotDataConsumerInterface* consumer, Context* cntx)
-    : db_slice_(slice), db_array_(slice->databases()) compression_mode_(compression_mode), consumer_(consumer), cntx_(cntx) {
+    : db_slice_(slice),
+      db_array_(slice->databases()),
+      compression_mode_(compression_mode),
+      consumer_(consumer),
+      cntx_(cntx) {
   tl_slice_snapshots.insert(this);
 }
 
@@ -125,9 +129,10 @@ void SliceSnapshot::FinalizeJournalStream(bool cancel) {
 
   journal->UnregisterOnChange(cb_id);
   if (!cancel) {
-    util::fb2::LockGuard guard(big_value_mu_);
-    serializer_->SendJournalOffset(journal->GetLsn());
-    PushSerializedUnderLock(true);
+    CallSerializerUnderLock([&journal](RdbSerializer* serializer) {
+      serializer->SendJournalOffset(journal->GetLsn());
+    });
+    PushSerialized(true);
   }
 }
 
@@ -170,8 +175,8 @@ void SliceSnapshot::IterateBucketsFb(bool send_full_sync_cut) {
         return;
       }
 
-      PrimeTable::Cursor next =
-          pt->TraverseBuckets(cursor, absl::bind_front(&SliceSnapshot::BucketSaveCb, this));
+      PrimeTable::Cursor next = pt->TraverseBuckets(
+          cursor, [this, db_indx](auto it) { return BucketSaveCb(db_indx, it); });
       cursor = next;
       PushSerialized(false);
 
@@ -193,9 +198,9 @@ void SliceSnapshot::IterateBucketsFb(bool send_full_sync_cut) {
 
   CHECK(!serialize_bucket_running_);
   if (send_full_sync_cut) {
-    util::fb2::LockGuard guard(big_value_mu_);
-    CHECK(!serializer_->SendFullSyncCut());
-    PushSerializedUnderLock(true);
+    CallSerializerUnderLock(
+        [](RdbSerializer* serializer) { CHECK(!serializer->SendFullSyncCut()); });
+    PushSerialized(true);
   }
 
   // serialized + side_saved must be equal to the total saved.
@@ -211,8 +216,10 @@ void SliceSnapshot::SwitchIncrementalFb(LSN lsn) {
   VLOG(1) << "Starting incremental snapshot from lsn=" << lsn;
 
   // The replica sends the LSN of the next entry is wants to receive.
-  while (!cntx->IsCancelled() && journal->IsLSNInBuffer(lsn)) {
-    serializer_->WriteJournalEntry(journal->GetEntry(lsn));
+  while (!cntx_->IsCancelled() && journal->IsLSNInBuffer(lsn)) {
+    CallSerializerUnderLock([entry = journal->GetEntry(lsn)](RdbSerializer* serializer) {
+      serializer->WriteJournalEntry(entry);
+    });
     PushSerialized(false);
     lsn++;
   }
@@ -229,10 +236,7 @@ void SliceSnapshot::SwitchIncrementalFb(LSN lsn) {
 
   // GetLsn() is always the next lsn that we expect to create.
   if (journal->GetLsn() == lsn) {
-    {
-      util::fb2::LockGuard guard(big_value_mu_);
-      serializer_->SendFullSyncCut();
-    }
+    CallSerializerUnderLock([](RdbSerializer* serializer) { serializer->SendFullSyncCut(); });
     auto journal_cb = [this](const journal::JournalItem& item, bool await) {
       OnJournalEntry(item, await);
     };
@@ -322,15 +326,20 @@ void SliceSnapshot::SerializeEntry(DbIndex db_indx, const PrimeKey& pk, const Pr
     delayed_entries_.push_back({db_indx, std::move(key), std::move(future), expire_time, mc_flags});
     ++type_freq_map_[RDB_TYPE_STRING];
   } else {
-    util::fb2::LockGuard guard(big_value_mu_);
-    io::Result<uint8_t> res = serializer_->SaveEntry(pk, pv, expire_time, mc_flags, db_indx);
+    io::Result<uint8_t> res;
+    CallSerializerUnderLock([&](RdbSerializer* serializer) {
+      res = serializer->SaveEntry(pk, pv, expire_time, mc_flags, db_indx);
+    });
     CHECK(res);
     ++type_freq_map_[*res];
   }
 }
 
 size_t SliceSnapshot::FlushSerialized(SerializerBase::FlushState flush_state) {
   io::StringFile sfile;
+
+  // FlushSerialized is already under lock
+  // So no locking is needed
   serializer_->FlushToSink(&sfile, flush_state);
 
   size_t serialized = sfile.val.size();
@@ -361,29 +370,45 @@ size_t SliceSnapshot::FlushSerialized(SerializerBase::FlushState flush_state) {
   return serialized;
 }
 
-bool SliceSnapshot::PushSerialized(bool force) {
+template <typename Callback> void SliceSnapshot::CallSerializerUnderLock(Callback cb) {
   util::fb2::LockGuard guard(big_value_mu_);
-  return PushSerializedUnderLock(force);
+  cb(serializer_.get());
 }
 
-bool SliceSnapshot::PushSerializedUnderLock(bool force) {
-  if (!force && serializer_->SerializedLen() < kMinBlobSize)
-    return false;
+bool SliceSnapshot::PushSerialized(bool force) {
+  if (!force) {
+    size_t len = 0;
+    CallSerializerUnderLock(
+        [&len](RdbSerializer* serializer) { len = serializer->SerializedLen(); });
+    if (len < kMinBlobSize) {
+      return false;
+    }
+  }
 
   // Flush any of the leftovers to avoid interleavings
-  size_t serialized = FlushSerialized(FlushState::kFlushEndEntry);
+  size_t serialized = 0;
+  {
+    util::fb2::LockGuard guard(big_value_mu_);
+    FlushSerialized(FlushState::kFlushEndEntry);
+  }
 
   if (!delayed_entries_.empty()) {
     // Async bucket serialization might have accumulated some delayed values.
     // Because we can finally block in this function, we'll await and serialize them
     do {
+      /* We can preempt on SaveEntry so first we need to change all data before
+       * calling it */
       auto& entry = delayed_entries_.back();
-      serializer_->SaveEntry(entry.key, entry.value.Get(), entry.expire, entry.dbid,
-                             entry.mc_flags);
       delayed_entries_.pop_back();
+
+      CallSerializerUnderLock([&entry](RdbSerializer* serializer) {
+        serializer->SaveEntry(entry.key, entry.value.Get(), entry.expire, entry.dbid,
+                              entry.mc_flags);
+      });
     } while (!delayed_entries_.empty());
 
     // blocking point.
+    util::fb2::LockGuard guard(big_value_mu_);
     serialized += FlushSerialized(FlushState::kFlushEndEntry);
   }
   return serialized > 0;
@@ -415,8 +440,8 @@ void SliceSnapshot::OnJournalEntry(const journal::JournalItem& item, bool await)
   // TriggerJournalWriteToSink. This call uses the NOOP opcode with await=true. Since there is no
   // additional journal change to serialize, it simply invokes PushSerialized.
   if (item.opcode != journal::Op::NOOP) {
-    util::fb2::LockGuard guard(big_value_mu_);
-    CHECK(!serializer_->WriteJournalEntry(item.data));
+    CallSerializerUnderLock(
+        [&item](RdbSerializer* serializer) { CHECK(!serializer->WriteJournalEntry(item.data)); });
   }
 
   if (await) {

src/server/snapshot.h

@@ -117,13 +117,15 @@ class SliceSnapshot {
   // Push regardless of buffer size if force is true.
   // Return true if pushed. Can block. Is called from the snapshot thread.
   bool PushSerialized(bool force);
-  bool PushSerializedUnderLock(bool force) ABSL_EXCLUSIVE_LOCKS_REQUIRED(big_value_mu_);
 
   // Helper function that flushes the serialized items into the RecordStream.
   // Can block.
   using FlushState = SerializerBase::FlushState;
   size_t FlushSerialized(FlushState flush_state);  // ABSL_EXCLUSIVE_LOCKS_REQUIRED(big_value_mu_)
 
+  // Calls the provided callback with the serializer under a lock.
+  template <typename Callback> void CallSerializerUnderLock(Callback cb);
+
  public:
   uint64_t snapshot_version() const {
     return snapshot_version_;
@@ -152,17 +154,18 @@ class SliceSnapshot {
   DbSlice* db_slice_;
   const DbTableArray db_array_;
 
-  std::unique_ptr<RdbSerializer> serializer_;  // ABSL_GUARDED_BY(big_value_mu_);
+  // Guarded by big_value_mu_
+  std::unique_ptr<RdbSerializer> serializer_;
 
   // collected during atomic bucket traversal
-  std::vector<DelayedEntry> delayed_entries_ ABSL_GUARDED_BY(big_value_mu_);
+  std::vector<DelayedEntry> delayed_entries_;
 
   // Used for sanity checks.
   bool serialize_bucket_running_ = false;
   util::fb2::Fiber snapshot_fb_;  // IterateEntriesFb
   util::fb2::CondVarAny seq_cond_;
   const CompressionMode compression_mode_;
-  RdbTypeFreqMap type_freq_map_ ABSL_GUARDED_BY(big_value_mu_);
+  RdbTypeFreqMap type_freq_map_;
 
   // version upper bound for entries that should be saved (not included).
   uint64_t snapshot_version_ = 0;