fix(snapshot): Refactor the big_value_mu_ mutex

fixes #4152

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

Author: BagritsevichStepan

src/server/rdb_save.h

@@ -225,6 +225,7 @@ class RdbSerializer : public SerializerBase {
   // Must be called in the thread to which `it` belongs.
   // Returns the serialized rdb_type or the error.
   // expire_ms = 0 means no expiry.
+  // This function might preempt if flush_fun_ is used.
   io::Result<uint8_t> SaveEntry(const PrimeKey& pk, const PrimeValue& pv, uint64_t expire_ms,
                                 uint32_t mc_flags, DbIndex dbid);
 

src/server/snapshot.cc

@@ -41,10 +41,10 @@ SliceSnapshot::SliceSnapshot(DbSlice* slice, CompressionMode compression_mode,
                              std::function<void(std::string)> on_push_record,
                              std::function<void()> on_snapshot_finish)
     : db_slice_(slice),
+      db_array_(slice->databases()),
       compression_mode_(compression_mode),
       on_push_(on_push_record),
       on_snapshot_finish_(on_snapshot_finish) {
-  db_array_ = slice->databases();
   tl_slice_snapshots.insert(this);
 }
 
@@ -65,7 +65,8 @@ bool SliceSnapshot::IsSnaphotInProgress() {
   return tl_slice_snapshots.size() > 0;
 }
 
-void SliceSnapshot::Start(bool stream_journal, const Cancellation* cll, SnapshotFlush allow_flush) {
+void SliceSnapshot::Start(bool stream_journal, const Cancellation* cll, SnapshotFlush allow_flush)
+    ABSL_LOCKS_EXCLUDED(big_value_mu_) {
   DCHECK(!snapshot_fb_.IsJoinable());
 
   auto db_cb = absl::bind_front(&SliceSnapshot::OnDbChange, this);
@@ -81,15 +82,16 @@ void SliceSnapshot::Start(bool stream_journal, const Cancellation* cll, Snapshot
   const auto flush_threshold = ServerState::tlocal()->serialization_max_chunk_size;
   std::function<void(size_t, RdbSerializer::FlushState)> flush_fun;
   if (flush_threshold != 0 && allow_flush == SnapshotFlush::kAllow) {
-    flush_fun = [this, flush_threshold](size_t bytes_serialized,
-                                        RdbSerializer::FlushState flush_state) {
-      if (bytes_serialized > flush_threshold) {
-        size_t serialized = FlushSerialized(flush_state);
-        VLOG(2) << "FlushSerialized " << serialized << " bytes";
-        auto& stats = ServerState::tlocal()->stats;
-        ++stats.big_value_preemptions;
-      }
-    };
+    flush_fun =
+        [this, flush_threshold](size_t bytes_serialized, RdbSerializer::FlushState flush_state)
+            ABSL_LOCKS_EXCLUDED(big_value_mu_) {
+              if (bytes_serialized > flush_threshold) {
+                size_t serialized = FlushSerialized(flush_state);
+                VLOG(2) << "FlushSerialized " << serialized << " bytes";
+                auto& stats = ServerState::tlocal()->stats;
+                ++stats.big_value_preemptions;
+              }
+            };
   }
   serializer_ = std::make_unique<RdbSerializer>(compression_mode_, flush_fun);
 
@@ -102,7 +104,8 @@ void SliceSnapshot::Start(bool stream_journal, const Cancellation* cll, Snapshot
   });
 }
 
-void SliceSnapshot::StartIncremental(Context* cntx, LSN start_lsn) {
+void SliceSnapshot::StartIncremental(Context* cntx, LSN start_lsn)
+    ABSL_LOCKS_EXCLUDED(big_value_mu_) {
   serializer_ = std::make_unique<RdbSerializer>(compression_mode_);
 
   snapshot_fb_ = fb2::Fiber("incremental_snapshot", [cntx, start_lsn, this] {
@@ -127,8 +130,9 @@ void SliceSnapshot::FinalizeJournalStream(bool cancel) {
 
   journal->UnregisterOnChange(cb_id);
   if (!cancel) {
+    util::fb2::LockGuard guard(big_value_mu_);
     serializer_->SendJournalOffset(journal->GetLsn());
-    PushSerialized(true);
+    PushSerializedUnderLock(true);
   }
 }
 
@@ -164,16 +168,15 @@ void SliceSnapshot::IterateBucketsFb(const Cancellation* cll, bool send_full_syn
 
     uint64_t last_yield = 0;
     PrimeTable* pt = &db_array_[db_indx]->prime;
-    current_db_ = db_indx;
 
     VLOG(1) << "Start traversing " << pt->size() << " items for index " << db_indx;
     do {
       if (cll->IsCancelled()) {
         return;
       }
 
-      PrimeTable::Cursor next =
-          pt->TraverseBuckets(cursor, absl::bind_front(&SliceSnapshot::BucketSaveCb, this));
+      PrimeTable::Cursor next = pt->TraverseBuckets(
+          cursor, absl::bind_front(&SliceSnapshot::BucketSaveCb, this, db_indx));
       cursor = next;
       PushSerialized(false);
 
@@ -195,8 +198,9 @@ void SliceSnapshot::IterateBucketsFb(const Cancellation* cll, bool send_full_syn
 
   CHECK(!serialize_bucket_running_);
   if (send_full_sync_cut) {
+    util::fb2::LockGuard guard(big_value_mu_);
     CHECK(!serializer_->SendFullSyncCut());
-    PushSerialized(true);
+    PushSerializedUnderLock(true);
   }
 
   // serialized + side_saved must be equal to the total saved.
@@ -213,8 +217,11 @@ void SliceSnapshot::SwitchIncrementalFb(Context* cntx, 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));
-    PushSerialized(false);
+    {
+      util::fb2::LockGuard guard(big_value_mu_);
+      serializer_->WriteJournalEntry(journal->GetEntry(lsn));
+      PushSerializedUnderLock(false);
+    }
     lsn++;
   }
 
@@ -231,7 +238,7 @@ void SliceSnapshot::SwitchIncrementalFb(Context* cntx, LSN lsn) {
   // GetLsn() is always the next lsn that we expect to create.
   if (journal->GetLsn() == lsn) {
     {
-      FiberAtomicGuard fg;
+      util::fb2::LockGuard guard(big_value_mu_);
       serializer_->SendFullSyncCut();
     }
     auto journal_cb = absl::bind_front(&SliceSnapshot::OnJournalEntry, this);
@@ -247,9 +254,7 @@ void SliceSnapshot::SwitchIncrementalFb(Context* cntx, LSN lsn) {
   }
 }
 
-bool SliceSnapshot::BucketSaveCb(PrimeTable::bucket_iterator it) {
-  std::lock_guard guard(big_value_mu_);
-
+bool SliceSnapshot::BucketSaveCb(DbIndex db_index, PrimeTable::bucket_iterator it) {
   ++stats_.savecb_calls;
 
   auto check = [&](auto v) {
@@ -266,7 +271,7 @@ bool SliceSnapshot::BucketSaveCb(PrimeTable::bucket_iterator it) {
     return false;
   }
 
-  db_slice_->FlushChangeToEarlierCallbacks(current_db_, DbSlice::Iterator::FromPrime(it),
+  db_slice_->FlushChangeToEarlierCallbacks(db_index, DbSlice::Iterator::FromPrime(it),
                                            snapshot_version_);
 
   auto* blocking_counter = db_slice_->BlockingCounter();
@@ -275,7 +280,7 @@ bool SliceSnapshot::BucketSaveCb(PrimeTable::bucket_iterator it) {
   // zero.
   std::lock_guard blocking_counter_guard(*blocking_counter);
 
-  stats_.loop_serialized += SerializeBucket(current_db_, it);
+  stats_.loop_serialized += SerializeBucket(db_index, it);
 
   return false;
 }
@@ -291,20 +296,19 @@ unsigned SliceSnapshot::SerializeBucket(DbIndex db_index, PrimeTable::bucket_ite
   while (!it.is_done()) {
     ++result;
     // might preempt due to big value serialization.
-    SerializeEntry(db_index, it->first, it->second, nullopt, serializer_.get());
+    SerializeEntry(db_index, it->first, it->second);
     ++it;
   }
   serialize_bucket_running_ = false;
   return result;
 }
 
-void SliceSnapshot::SerializeEntry(DbIndex db_indx, const PrimeKey& pk, const PrimeValue& pv,
-                                   optional<uint64_t> expire, RdbSerializer* serializer) {
+void SliceSnapshot::SerializeEntry(DbIndex db_indx, const PrimeKey& pk, const PrimeValue& pv) {
   if (pv.IsExternal() && pv.IsCool())
-    return SerializeEntry(db_indx, pk, pv.GetCool().record->value, expire, serializer);
+    return SerializeEntry(db_indx, pk, pv.GetCool().record->value);
 
-  time_t expire_time = expire.value_or(0);
-  if (!expire && pv.HasExpire()) {
+  time_t expire_time = 0;
+  if (pv.HasExpire()) {
     auto eit = db_array_[db_indx]->expire.Find(pk);
     expire_time = db_slice_->ExpireTime(eit);
   }
@@ -317,11 +321,15 @@ void SliceSnapshot::SerializeEntry(DbIndex db_indx, const PrimeKey& pk, const Pr
     EngineShard::tlocal()->tiered_storage()->Read(
         db_indx, pk.ToString(), pv,
         [future](const std::string& v) mutable { future.Resolve(PrimeValue(v)); });
-    delayed_entries_.push_back(
-        {db_indx, PrimeKey(pk.ToString()), std::move(future), expire_time, mc_flags});
+
+    auto key = PrimeKey(pk.ToString());
+
+    util::fb2::LockGuard guard(big_value_mu_);
+    delayed_entries_.push_back({db_indx, std::move(key), std::move(future), expire_time, mc_flags});
     ++type_freq_map_[RDB_TYPE_STRING];
   } else {
-    io::Result<uint8_t> res = serializer->SaveEntry(pk, pv, expire_time, mc_flags, db_indx);
+    util::fb2::LockGuard guard(big_value_mu_);
+    io::Result<uint8_t> res = serializer_->SaveEntry(pk, pv, expire_time, mc_flags, db_indx);
     CHECK(res);
     ++type_freq_map_[*res];
   }
@@ -360,6 +368,11 @@ size_t SliceSnapshot::FlushSerialized(SerializerBase::FlushState flush_state) {
 }
 
 bool SliceSnapshot::PushSerialized(bool force) {
+  util::fb2::LockGuard guard(big_value_mu_);
+  return PushSerializedUnderLock(force);
+}
+
+bool SliceSnapshot::PushSerializedUnderLock(bool force) {
   if (!force && serializer_->SerializedLen() < kMinBlobSize)
     return false;
 
@@ -383,8 +396,6 @@ bool SliceSnapshot::PushSerialized(bool force) {
 }
 
 void SliceSnapshot::OnDbChange(DbIndex db_index, const DbSlice::ChangeReq& req) {
-  std::lock_guard guard(big_value_mu_);
-
   PrimeTable* table = db_slice_->GetTables(db_index).first;
   const PrimeTable::bucket_iterator* bit = req.update();
 
@@ -409,9 +420,9 @@ void SliceSnapshot::OnJournalEntry(const journal::JournalItem& item, bool await)
   // To enable journal flushing to sync after non auto journal command is executed we call
   // TriggerJournalWriteToSink. This call uses the NOOP opcode with await=true. Since there is no
   // additional journal change to serialize, it simply invokes PushSerialized.
-  std::lock_guard guard(big_value_mu_);
   if (item.opcode != journal::Op::NOOP) {
-    serializer_->WriteJournalEntry(item.data);
+    util::fb2::LockGuard guard(big_value_mu_);
+    CHECK(!serializer_->WriteJournalEntry(item.data));
   }
 
   if (await) {
@@ -422,6 +433,8 @@ void SliceSnapshot::OnJournalEntry(const journal::JournalItem& item, bool await)
 }
 
 size_t SliceSnapshot::GetBufferCapacity() const {
+  util::fb2::LockGuard guard(big_value_mu_);
+
   if (serializer_ == nullptr) {
     return 0;
   }
@@ -430,6 +443,8 @@ size_t SliceSnapshot::GetBufferCapacity() const {
 }
 
 size_t SliceSnapshot::GetTempBuffersSize() const {
+  util::fb2::LockGuard guard(big_value_mu_);
+
   if (serializer_ == nullptr) {
     return 0;
   }

src/server/snapshot.h

@@ -61,14 +61,15 @@ class SliceSnapshot {
   enum class SnapshotFlush { kAllow, kDisallow };
 
   void Start(bool stream_journal, const Cancellation* cll,
-             SnapshotFlush allow_flush = SnapshotFlush::kDisallow);
+             SnapshotFlush allow_flush = SnapshotFlush::kDisallow)
+      ABSL_LOCKS_EXCLUDED(big_value_mu_);
 
   // Initialize a snapshot that sends only the missing journal updates
   // since start_lsn and then registers a callback switches into the
   // journal streaming mode until stopped.
   // If we're slower than the buffer and can't continue, `Cancel()` is
   // called.
-  void StartIncremental(Context* cntx, LSN start_lsn);
+  void StartIncremental(Context* cntx, LSN start_lsn) ABSL_LOCKS_EXCLUDED(big_value_mu_);
 
   // Finalizes journal streaming writes. Only called for replication.
   // Blocking. Must be called from the Snapshot thread.
@@ -89,15 +90,14 @@ class SliceSnapshot {
   void SwitchIncrementalFb(Context* cntx, LSN lsn);
 
   // Called on traversing cursor by IterateBucketsFb.
-  bool BucketSaveCb(PrimeTable::bucket_iterator it);
+  bool BucketSaveCb(DbIndex db_index, PrimeTable::bucket_iterator it);
 
   // Serialize single bucket.
   // Returns number of serialized entries, updates bucket version to snapshot version.
   unsigned SerializeBucket(DbIndex db_index, PrimeTable::bucket_iterator bucket_it);
 
   // Serialize entry into passed serializer.
-  void SerializeEntry(DbIndex db_index, const PrimeKey& pk, const PrimeValue& pv,
-                      std::optional<uint64_t> expire, RdbSerializer* serializer);
+  void SerializeEntry(DbIndex db_index, const PrimeKey& pk, const PrimeValue& pv);
 
   // DbChange listener
   void OnDbChange(DbIndex db_index, const DbSlice::ChangeReq& req);
@@ -109,11 +109,12 @@ 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);
+  size_t FlushSerialized(FlushState flush_state) ABSL_EXCLUSIVE_LOCKS_REQUIRED(big_value_mu_);
 
  public:
   uint64_t snapshot_version() const {
@@ -141,24 +142,25 @@ class SliceSnapshot {
   };
 
   DbSlice* db_slice_;
-  DbTableArray db_array_;
+  const DbTableArray db_array_;
 
-  DbIndex current_db_;
+  std::unique_ptr<RdbSerializer> serializer_ ABSL_GUARDED_BY(big_value_mu_);
 
-  std::unique_ptr<RdbSerializer> serializer_;
-  std::vector<DelayedEntry> delayed_entries_;  // collected during atomic bucket traversal
+  // collected during atomic bucket traversal
+  std::vector<DelayedEntry> delayed_entries_ ABSL_GUARDED_BY(big_value_mu_);
 
   // Used for sanity checks.
   bool serialize_bucket_running_ = false;
   util::fb2::Fiber snapshot_fb_;  // IterateEntriesFb
   util::fb2::CondVarAny seq_cond_;
-  CompressionMode compression_mode_;
-  RdbTypeFreqMap type_freq_map_;
+  const CompressionMode compression_mode_;
+  RdbTypeFreqMap type_freq_map_ ABSL_GUARDED_BY(big_value_mu_);
 
   // version upper bound for entries that should be saved (not included).
   uint64_t snapshot_version_ = 0;
   uint32_t journal_cb_id_ = 0;
 
+  // TODO: remove or use mutex
   uint64_t rec_id_ = 1, last_pushed_id_ = 0;
 
   struct Stats {
@@ -167,9 +169,9 @@ class SliceSnapshot {
     size_t side_saved = 0;
     size_t savecb_calls = 0;
     size_t keys_total = 0;
-  } stats_;
+  } stats_;  // TODO: maybe need to lock
 
-  ThreadLocalMutex big_value_mu_;
+  mutable ThreadLocalMutex big_value_mu_;
 
   std::function<void(std::string)> on_push_;
   std::function<void()> on_snapshot_finish_;