feat(server): Serializer bucket states

Signed-off-by: Vladislav Oleshko <vlad@dragonflydb.io>

Author: dranikpg

src/server/CMakeLists.txt

@@ -160,7 +160,6 @@ helio_cxx_test(cluster/cluster_config_test dfly_test_lib LABELS DFLY)
 helio_cxx_test(cluster/cluster_family_test dfly_test_lib LABELS DFLY)
 helio_cxx_test(acl/acl_family_test dfly_test_lib LABELS DFLY)
 helio_cxx_test(engine_shard_set_test dfly_test_lib LABELS DFLY)
-helio_cxx_test(serializer_base_test dfly_test_lib LABELS DFLY)
 
 add_dependencies(check_dfly dragonfly_test json_family_test list_family_test
                  generic_family_test memcache_parser_test rdb_test journal_test

src/server/serializer_base.cc

@@ -15,6 +15,35 @@
 
 namespace dfly {
 
+void BucketDependencies::Increment(BucketIdentity bucket) {
+  auto& counter = deps_[bucket];
+  if (!counter)
+    counter = std::make_shared<LocalLatch>();
+  counter->lock();
+}
+
+void BucketDependencies::Decrement(BucketIdentity bucket) {
+  auto it = deps_.find(bucket);
+  DCHECK(it != deps_.end());
+  it->second->unlock();
+
+  if (!it->second->IsBlocked())
+    deps_.erase(it);
+}
+
+void BucketDependencies::Wait(BucketIdentity bucket) const {
+  auto it = deps_.find(bucket);
+  if (it == deps_.end())
+    return;
+
+  auto counter = it->second;  // copy value for address stability
+  counter->Wait();
+}
+
+bool BucketDependencies::DEBUG_IsBusy(BucketIdentity bucket) const {
+  return deps_.contains(bucket);
+}
+
 void DelayedEntryHandler::EnqueueOffloaded(BucketIdentity bucket, DbIndex db_index, PrimeKey pk,
                                            const PrimeValue& pv, time_t expire_time,
                                            uint32_t mc_flags) {
@@ -26,6 +55,8 @@ void DelayedEntryHandler::EnqueueOffloaded(BucketIdentity bucket, DbIndex db_ind
   auto future = ReadTieredString(db_index, key, pv, EngineShard::tlocal()->tiered_storage());
   auto entry = std::make_unique<TieredDelayedEntry>(db_index, std::move(pk), std::move(future),
                                                     expire_time, mc_flags);
+
+  deps_.Increment(bucket);
   delayed_entries_.emplace(bucket, std::move(entry));
 }
 
@@ -35,7 +66,7 @@ void DelayedEntryHandler::ProcessDelayedEntries(bool force, BucketIdentity flush
   if (delayed_entries_.size() > kMaxDelayedEntries)
     force |= true;
 
-  auto serialize_entry = [&](auto it) {
+  auto serialize_entry = [&](decltype(delayed_entries_)::iterator it) {
     auto& entry = it->second;
     auto value = entry->value.Get();
 
@@ -47,6 +78,8 @@ void DelayedEntryHandler::ProcessDelayedEntries(bool force, BucketIdentity flush
 
     PrimeValue pv{*value};
     SerializeFetchedEntry(*entry, pv);
+
+    deps_.Decrement(it->first);
     delayed_entries_.erase(it++);
   };
 
@@ -68,8 +101,9 @@ void DelayedEntryHandler::ProcessDelayedEntries(bool force, BucketIdentity flush
 }
 
 SerializerBase::SerializerBase(DbSlice* slice, ExecutionState* cntx)
-    : db_slice_(slice), base_cntx_(cntx) {
-  DCHECK(db_slice_);
+    : DelayedEntryHandler(static_cast<BucketDependencies&>(*this)),
+      db_slice_(slice),
+      base_cntx_(cntx) {
   DCHECK(base_cntx_);
 }
 
@@ -97,45 +131,19 @@ void SerializerBase::UnregisterChangeListener() {
     db_slice_->UnregisterOnChange(version);
 }
 
-void SerializerBase::MarkBucketSerializing(BucketIdentity bid) {
-  DCHECK(!bucket_states_.contains(bid)) << "Bucket already in transient state";
-  bucket_states_[bid] = BucketPhase::kSerializing;
-}
-
-void SerializerBase::FinishBucketIteration(BucketIdentity bid) {
-  auto it = bucket_states_.find(bid);
-  DCHECK(it != bucket_states_.end());
-  DCHECK(it->second == BucketPhase::kSerializing);
-
-  bucket_states_.erase(it);
-  ++stats_.buckets_serialized;
-}
-
-bool SerializerBase::ShouldProcessBucket(PrimeTable::bucket_iterator it) {
-  // Check if bucket is invalid or was already serialized
-  if (it.is_done() || it.GetVersion() >= snapshot_version_) {
-    ++stats_.buckets_skipped;
-    return false;
-  }
-
-  // Check if this bucket is currently being serialized
-  if (bucket_states_.contains(it.bucket_address())) {
-    ++stats_.change_during_serialization;
-    return false;
-  }
-
-  return true;
-}
-
 bool SerializerBase::ProcessBucket(DbIndex db_index, PrimeTable::bucket_iterator it,
                                    bool on_update) {
   std::lock_guard guard(big_value_mu_);
 
-  // Check if this bucket should be serialized
-  if (!ShouldProcessBucket(it)) {
+  // Check if this bucket is stale
+  if (it.is_done() || it.GetVersion() >= snapshot_version_) {
     // Force flush all delayed entries in the touched bucket
     if (EngineShard::tlocal()->tiered_storage() != nullptr && on_update && !it.is_done())
       ProcessDelayedEntries(false, it.bucket_address(), base_cntx_);
+
+    // Expected to be fully serialized due to big_value_mu_ guarding all paths
+    // Otherwise, this needs to be changed to a wait
+    DCHECK(!BucketDependencies::DEBUG_IsBusy(it.bucket_address()));
     return false;
   }
 
@@ -148,9 +156,11 @@ bool SerializerBase::ProcessBucket(DbIndex db_index, PrimeTable::bucket_iterator
   }
 
   it.SetVersion(snapshot_version_);
-  MarkBucketSerializing(it.bucket_address());
+  BucketDependencies::Increment(it.bucket_address());
+
   stats_.keys_serialized += SerializeBucket(db_index, it, on_update);
-  FinishBucketIteration(it.bucket_address());
+  stats_.buckets_serialized++;
+  BucketDependencies::Decrement(it.bucket_address());
 
   if (EngineShard::tlocal()->tiered_storage() != nullptr)
     ProcessDelayedEntries(false, on_update ? it.bucket_address() : 0, base_cntx_);

src/server/serializer_base.h

@@ -6,13 +6,15 @@
 
 #include <absl/container/flat_hash_map.h>
 
+#include <memory>
 #include <vector>
 
 #include "server/db_slice.h"
 #include "server/journal/types.h"
 #include "server/synchronization.h"
 #include "server/table.h"
 #include "server/tiered_storage.h"
+#include "util/fibers/synchronization.h"
 
 namespace dfly {
 
@@ -22,6 +24,21 @@ class ExecutionState;
 // Unique across all databases/segments for the lifetime of a serialization.
 using BucketIdentity = uintptr_t;
 
+// Track dependencies for buckets
+struct BucketDependencies {
+  // Increase number of dependencies for bucket
+  void Increment(BucketIdentity bucket);
+  void Decrement(BucketIdentity bucket);
+
+  // Wait for all bucket dependencies to resolve
+  void Wait(BucketIdentity bucket) const;
+  bool DEBUG_IsBusy(BucketIdentity) const;
+
+ private:
+  using SharedLatch = std::shared_ptr<LocalLatch>;
+  absl::flat_hash_map<BucketIdentity, SharedLatch> deps_;
+};
+
 // Tracks serialization progress of offloaded (delayed) entries.
 struct DelayedEntryHandler {
   void EnqueueOffloaded(BucketIdentity bucket, DbIndex db_index, PrimeKey pk, const PrimeValue& pv,
@@ -35,7 +52,13 @@ struct DelayedEntryHandler {
   // Serialize delayed entry that was fetched with serializer specific implementation
   virtual void SerializeFetchedEntry(const TieredDelayedEntry& tde, const PrimeValue& pv) = 0;
 
+ protected:
+  explicit DelayedEntryHandler(BucketDependencies& deps) : deps_{deps} {
+  }
+
  private:
+  BucketDependencies& deps_;
+
   // Entries that are waiting for tiered storage reads to complete before they can be serialized.
   std::multimap<BucketIdentity, std::unique_ptr<TieredDelayedEntry>> delayed_entries_;
 };
@@ -50,7 +73,7 @@ struct DelayedEntryHandler {
 //
 // State tracking is purely observational in early PRs: it drives DCHECKs and
 // stats but does not alter the serialization control flow.
-class SerializerBase : public DelayedEntryHandler {
+class SerializerBase : public BucketDependencies, public DelayedEntryHandler {
  public:
   struct Stats {
     uint64_t keys_serialized = 0;              // total number of keys serialized
@@ -108,21 +131,7 @@ class SerializerBase : public DelayedEntryHandler {
 
  private:
   friend class SerializerBaseTest;
-  SerializerBase() : db_slice_(nullptr), base_cntx_(nullptr) {
-  }
-
-  // Return identity if bucket should be processed.
-  // Checks bucket validity, version and state
-  bool ShouldProcessBucket(PrimeTable::bucket_iterator);
-
-  // Transition bucket from NotVisited -> Serializing.
-  void MarkBucketSerializing(BucketIdentity bid);
-
-  // Transition bucket from Serializing -> Covered (empty delayed) or
-  // Serializing -> DelayedPending (non-empty delayed).
-  void FinishBucketIteration(BucketIdentity bid);
 
-  absl::flat_hash_map<BucketIdentity, BucketPhase> bucket_states_;
   uint64_t change_cb_id_ = 0;
 };