fix(migration): Use transactions!

Signed-off-by: Vladislav Oleshko <[email protected]>

Author: dranikpg

src/server/cluster/outgoing_slot_migration.cc

@@ -16,7 +16,9 @@
 #include "server/engine_shard_set.h"
 #include "server/error.h"
 #include "server/journal/streamer.h"
+#include "server/main_service.h"
 #include "server/server_family.h"
+#include "server/transaction.h"
 
 ABSL_FLAG(int, slot_migration_connection_timeout_ms, 2000, "Timeout for network operations");
 
@@ -33,7 +35,9 @@ class OutgoingMigration::SliceSlotMigration : private ProtocolClient {
       : ProtocolClient(server_context), streamer_(slice, std::move(slots), journal, &cntx_) {
   }
 
-  void Sync(const std::string& node_id, uint32_t shard_id) {
+  void PrepareFlow(const std::string& node_id) {
+    uint32_t shard_id = EngineShard::tlocal()->shard_id();
+
     VLOG(1) << "Connecting to source node_id " << node_id << " shard_id " << shard_id;
     auto timeout = absl::GetFlag(FLAGS_slot_migration_connection_timeout_ms) * 1ms;
     if (auto ec = ConnectAndAuth(timeout, &cntx_); ec) {
@@ -57,10 +61,16 @@ class OutgoingMigration::SliceSlotMigration : private ProtocolClient {
       cntx_.ReportError("Incorrect response for FLOW cmd");
       return;
     }
+  }
 
+  void PrepareSync() {
     streamer_.Start(Sock());
   }
 
+  void StartSync() {
+    streamer_.Run();
+  }
+
   void Cancel() {
     streamer_.Cancel();
   }
@@ -89,11 +99,7 @@ OutgoingMigration::~OutgoingMigration() {
   main_sync_fb_.JoinIfNeeded();
 
   // Destroy each flow in its dedicated thread, because we could be the last owner of the db tables
-  shard_set->pool()->AwaitFiberOnAll([this](util::ProactorBase* pb) {
-    if (const auto* shard = EngineShard::tlocal(); shard) {
-      slot_migrations_[shard->shard_id()].reset();
-    }
-  });
+  OnAllShards([](auto& flow) { flow.reset(); });
 }
 
 bool OutgoingMigration::ChangeState(MigrationState new_state) {
@@ -132,13 +138,7 @@ void OutgoingMigration::Finish(bool is_error) {
   }
 
   if (should_cancel_flows) {
-    shard_set->pool()->AwaitFiberOnAll([this](util::ProactorBase* pb) {
-      if (const auto* shard = EngineShard::tlocal(); shard) {
-        auto& flow = slot_migrations_[shard->shard_id()];
-        CHECK(flow != nullptr);
-        flow->Cancel();
-      }
-    });
+    OnAllShards([](auto& flow) { flow->Cancel(); });
   }
 }
 
@@ -195,29 +195,42 @@ void OutgoingMigration::SyncFb() {
       continue;
     }
 
-    shard_set->pool()->AwaitFiberOnAll([this](util::ProactorBase* pb) {
-      if (auto* shard = EngineShard::tlocal(); shard) {
-        server_family_->journal()->StartInThread();
-        slot_migrations_[shard->shard_id()] = std::make_unique<SliceSlotMigration>(
-            &shard->db_slice(), server(), migration_info_.slot_ranges, server_family_->journal());
-      }
+    OnAllShards([this](auto& flow) {
+      server_family_->journal()->StartInThread();
+      auto& slice = EngineShard::tlocal()->db_slice();
+      flow = std::make_unique<SliceSlotMigration>(&slice, server(), migration_info_.slot_ranges,
+                                                  server_family_->journal());
     });
 
     if (!ChangeState(MigrationState::C_SYNC)) {
       break;
     }
 
-    shard_set->pool()->AwaitFiberOnAll([this](util::ProactorBase* pb) {
-      if (auto* shard = EngineShard::tlocal(); shard) {
-        auto& migration = slot_migrations_[shard->shard_id()];
-        CHECK(migration != nullptr);
-        migration->Sync(cf_->MyID(), shard->shard_id());
-        if (migration->GetError()) {
-          Finish(true);
-        }
-      }
+    OnAllShards([this](auto& migration) {
+      migration->PrepareFlow(cf_->MyID());
+      if (migration->GetError())
+        Finish(true);
     });
 
+    // Global transactional cut for migration to register db_slice and journal listeners
+    {
+      const CommandId* cid = server_family_->service().FindCmd("DFLYCLUSTER");
+      DCHECK(cid->opt_mask() & CO::GLOBAL_TRANS);
+      boost::intrusive_ptr<Transaction> lock_tx{new Transaction(cid)};
+      lock_tx->InitByArgs(0, {});
+      lock_tx->Execute(
+          [this](auto* tx, EngineShard* shard) {
+            slot_migrations_[shard->shard_id()]->PrepareSync();
+            return OpStatus::OK;
+          },
+          true);
+    }
+
+    if (GetState() != MigrationState::C_SYNC)
+      break;
+
+    OnAllShards([](auto& migration) { migration->StartSync(); });
+
     if (CheckFlowsForErrors()) {
       LOG(WARNING) << "Errors detected, retrying outgoing migration";
       continue;
@@ -270,14 +283,8 @@ bool OutgoingMigration::FinalizeMigration(long attempt) {
     pause_fb_opt->JoinIfNeeded();
   });
 
-  auto cb = [this](util::ProactorBase* pb) {
-    if (const auto* shard = EngineShard::tlocal(); shard) {
-      slot_migrations_[shard->shard_id()]->Finalize();
-    }
-  };
-
   VLOG(1) << "FINALIZE flows for " << cf_->MyID() << " : " << migration_info_.node_info.id;
-  shard_set->pool()->AwaitFiberOnAll(std::move(cb));
+  OnAllShards([](auto& migration) { migration->Finalize(); });
 
   auto cmd = absl::StrCat("DFLYMIGRATE ACK ", cf_->MyID(), " ", attempt);
   VLOG(1) << "send " << cmd;
@@ -326,6 +333,7 @@ void OutgoingMigration::Start() {
   main_sync_fb_ = fb2::Fiber("outgoing_migration", &OutgoingMigration::SyncFb, this);
 }
 
+// TODO: better parent context binding?
 bool OutgoingMigration::CheckFlowsForErrors() {
   for (const auto& flow : slot_migrations_) {
     if (flow->GetError()) {
@@ -336,6 +344,13 @@ bool OutgoingMigration::CheckFlowsForErrors() {
   return false;
 }
 
+void OutgoingMigration::OnAllShards(std::function<void(std::unique_ptr<SliceSlotMigration>&)> f) {
+  shard_set->pool()->AwaitFiberOnAll([this, &f](auto* pb) {
+    if (auto* shard = EngineShard::tlocal(); shard)
+      f(slot_migrations_[shard->shard_id()]);
+  });
+}
+
 size_t OutgoingMigration::GetKeyCount() const {
   if (state_ == MigrationState::C_FINISHED) {
     return keys_number_;

src/server/cluster/outgoing_slot_migration.h

@@ -75,6 +75,8 @@ class OutgoingMigration : private ProtocolClient {
 
   bool ChangeState(MigrationState new_state) ABSL_LOCKS_EXCLUDED(state_mu_);
 
+  void OnAllShards(std::function<void(std::unique_ptr<SliceSlotMigration>&)>);
+
  private:
   MigrationInfo migration_info_;
   std::vector<std::unique_ptr<SliceSlotMigration>> slot_migrations_;

src/server/db_slice.cc

@@ -464,12 +464,12 @@ OpResult<DbSlice::PrimeItAndExp> DbSlice::FindInternal(const Context& cntx, std:
   }
 
   if (caching_mode_ && IsValid(res.it)) {
-    if (!change_cb_.empty()) {
+    if (!change_cbs_.empty()) {
       auto bump_cb = [&](PrimeTable::bucket_iterator bit) {
         DVLOG(2) << "Running callbacks for key " << key << " in dbid " << cntx.db_index;
-        CallChangeCallbacks(cntx.db_index, bit);
+        CallOnChange(cntx.db_index, bit);
       };
-      db.prime.CVCUponBump(change_cb_.back().first, res.it, bump_cb);
+      db.prime.CVCUponBump(change_cbs_.back().first, res.it, bump_cb);
     }
     auto bump_it = db.prime.BumpUp(res.it, PrimeBumpPolicy{fetched_items_});
     if (bump_it != res.it) {  // the item was bumped
@@ -531,7 +531,7 @@ OpResult<DbSlice::AddOrFindResult> DbSlice::AddOrFindInternal(const Context& cnt
 
   // It's a new entry.
   DVLOG(2) << "Running callbacks for key " << key << " in dbid " << cntx.db_index;
-  CallChangeCallbacks(cntx.db_index, key);
+  CallOnChange(cntx.db_index, key);
 
   // In case we are loading from rdb file or replicating we want to disable conservative memory
   // checks (inside PrimeEvictionPolicy::CanGrow) and reject insertions only after we pass max
@@ -980,7 +980,7 @@ void DbSlice::PreUpdate(DbIndex db_ind, Iterator it, std::string_view key) {
   FiberAtomicGuard fg;
 
   DVLOG(2) << "Running callbacks in dbid " << db_ind;
-  CallChangeCallbacks(db_ind, ChangeReq{it.GetInnerIt()});
+  CallOnChange(db_ind, ChangeReq{it.GetInnerIt()});
 
   // If the value has a pending stash, cancel it before any modification are applied.
   // Note: we don't delete offloaded values before updates, because a read-modify operation (like
@@ -1095,19 +1095,14 @@ void DbSlice::ExpireAllIfNeeded() {
 }
 
 uint64_t DbSlice::RegisterOnChange(ChangeCallback cb) {
-  // TODO rewrite this logic to be more clear
-  // this mutex lock is needed to check that this method is not called simultaneously with
-  // change_cb_ calls and journal_slice::change_cb_arr_ calls.
-  // It can be unlocked anytime because DbSlice::RegisterOnChange
-  // and journal_slice::RegisterOnChange calls without preemption
-  std::lock_guard lk(cb_mu_);
-
-  uint64_t ver = NextVersion();
-  change_cb_.emplace_back(ver, std::move(cb));
-  DCHECK(std::is_sorted(change_cb_.begin(), change_cb_.end(),
-                        [](auto& a, auto& b) { return a.first < b.first; }));
+  return change_cbs_.emplace_back(NextVersion(), std::move(cb)).first;
+}
 
-  return ver;
+void DbSlice::UnregisterOnChange(uint64_t version) {
+  auto it = find_if(change_cbs_.begin(), change_cbs_.end(),
+                    [version](const auto& cb) { return cb.first == version; });
+  DCHECK(it != change_cbs_.end());
+  change_cbs_.erase(it);
 }
 
 void DbSlice::FlushChangeToEarlierCallbacks(DbIndex db_ind, Iterator it, uint64_t upper_bound) {
@@ -1117,7 +1112,7 @@ void DbSlice::FlushChangeToEarlierCallbacks(DbIndex db_ind, Iterator it, uint64_
   DVLOG(2) << "Running callbacks in dbid " << db_ind << " with bucket_version=" << bucket_version
            << ", upper_bound=" << upper_bound;
 
-  for (const auto& ccb : change_cb_) {
+  for (const auto& ccb : change_cbs_) {
     uint64_t cb_version = ccb.first;
     DCHECK_LE(cb_version, upper_bound);
     if (cb_version == upper_bound) {
@@ -1129,16 +1124,14 @@ void DbSlice::FlushChangeToEarlierCallbacks(DbIndex db_ind, Iterator it, uint64_
   }
 }
 
-//! Unregisters the callback.
-void DbSlice::UnregisterOnChange(uint64_t id) {
-  lock_guard lk(cb_mu_);  // we need to wait until callback is finished before remove it
-  for (auto it = change_cb_.begin(); it != change_cb_.end(); ++it) {
-    if (it->first == id) {
-      change_cb_.erase(it);
-      return;
-    }
-  }
-  LOG(DFATAL) << "Could not find " << id << " to unregister";
+void DbSlice::CallOnChange(DbIndex id, const ChangeReq& cr) const {
+  FiberAtomicGuard fg;  // Callbacks don't preemept
+  // Callbacks are always sorted by version
+  DCHECK(is_sorted(change_cbs_.begin(), change_cbs_.end(),
+                   [](const auto& l, const auto& r) { return l.first < r.first; }));
+
+  for (const auto& ccb : change_cbs_)
+    ccb.second(id, cr);
 }
 
 auto DbSlice::DeleteExpiredStep(const Context& cntx, unsigned count) -> DeleteExpiredStats {
@@ -1533,10 +1526,4 @@ void DbSlice::OnCbFinish() {
   fetched_items_.clear();
 }
 
-void DbSlice::CallChangeCallbacks(DbIndex id, const ChangeReq& cr) const {
-  for (const auto& ccb : change_cb_) {
-    ccb.second(id, cr);
-  }
-}
-
 }  // namespace dfly

src/server/db_slice.h

@@ -405,17 +405,15 @@ class DbSlice {
 
   using ChangeCallback = std::function<void(DbIndex, const ChangeReq&)>;
 
-  //! Registers the callback to be called for each change.
-  //! Returns the registration id which is also the unique version of the dbslice
-  //! at a time of the call.
+  // Called before every access to an entry with a FindMutable call. Returns version
   uint64_t RegisterOnChange(ChangeCallback cb);
 
+  // Unregister by version. Can't be called from a callback!
+  void UnregisterOnChange(uint64_t version);
+
   // Call registered callbacks with version less than upper_bound.
   void FlushChangeToEarlierCallbacks(DbIndex db_ind, Iterator it, uint64_t upper_bound);
 
-  //! Unregisters the callback.
-  void UnregisterOnChange(uint64_t id);
-
   struct DeleteExpiredStats {
     uint32_t deleted = 0;         // number of deleted items due to expiry (less than traversed).
     uint32_t traversed = 0;       // number of traversed items that have ttl bit
@@ -469,14 +467,6 @@ class DbSlice {
   void PerformDeletion(Iterator del_it, DbTable* table);
   void PerformDeletion(PrimeIterator del_it, DbTable* table);
 
-  void LockChangeCb() const {
-    return cb_mu_.lock_shared();
-  }
-
-  void UnlockChangeCb() const {
-    return cb_mu_.unlock_shared();
-  }
-
  private:
   void PreUpdate(DbIndex db_ind, Iterator it, std::string_view key);
   void PostUpdate(DbIndex db_ind, Iterator it, std::string_view key, size_t orig_size);
@@ -531,7 +521,7 @@ class DbSlice {
     return version_++;
   }
 
-  void CallChangeCallbacks(DbIndex id, const ChangeReq& cr) const;
+  void CallOnChange(DbIndex id, const ChangeReq& cr) const;
 
  private:
   ShardId shard_id_;
@@ -554,14 +544,8 @@ class DbSlice {
   // Used in temporary computations in Acquire/Release.
   mutable absl::flat_hash_set<uint64_t> uniq_fps_;
 
-  // To ensure correct data replication, we must serialize the buckets that each running command
-  // will modify, followed by serializing the command to the journal. We use a mutex to prevent
-  // interleaving between bucket and journal registrations, and the command execution with its
-  // journaling. LockChangeCb is called before the callback, and UnlockChangeCb is called after
-  // journaling is completed. Register to bucket and journal changes is also does without preemption
-  mutable util::fb2::SharedMutex cb_mu_;
   // ordered from the smallest to largest version.
-  std::vector<std::pair<uint64_t, ChangeCallback>> change_cb_;
+  std::vector<std::pair<uint64_t, ChangeCallback>> change_cbs_;
 
   // Used in temporary computations in Find item and CbFinish
   mutable absl::flat_hash_set<CompactObjectView> fetched_items_;

src/server/detail/save_stages_controller.cc

@@ -254,10 +254,7 @@ void SaveStagesController::SaveDfs() {
 
   // Save shard files.
   auto cb = [this](Transaction* t, EngineShard* shard) {
-    // a hack to avoid deadlock in Transaction::RunCallback(...)
-    shard->db_slice().UnlockChangeCb();
     SaveDfsSingle(shard);
-    shard->db_slice().LockChangeCb();
     return OpStatus::OK;
   };
   trans_->ScheduleSingleHop(std::move(cb));
@@ -297,10 +294,7 @@ void SaveStagesController::SaveRdb() {
   }
 
   auto cb = [snapshot = snapshot.get()](Transaction* t, EngineShard* shard) {
-    // a hack to avoid deadlock in Transaction::RunCallback(...)
-    shard->db_slice().UnlockChangeCb();
     snapshot->StartInShard(shard);
-    shard->db_slice().LockChangeCb();
     return OpStatus::OK;
   };
   trans_->ScheduleSingleHop(std::move(cb));

src/server/journal/journal_slice.cc

@@ -190,28 +190,25 @@ void JournalSlice::AddLogRecord(const Entry& entry, bool await) {
   {
     std::shared_lock lk(cb_mu_);
     DVLOG(2) << "AddLogRecord: run callbacks for " << entry.ToString()
-             << " num callbacks: " << change_cb_arr_.size();
+             << " num callbacks: " << change_cbs_.size();
 
-    for (const auto& k_v : change_cb_arr_) {
+    for (const auto& k_v : change_cbs_) {
       k_v.second(*item, await);
     }
   }
 }
 
 uint32_t JournalSlice::RegisterOnChange(ChangeCallback cb) {
-  // mutex lock isn't needed due to iterators are not invalidated
-  uint32_t id = next_cb_id_++;
-  change_cb_arr_.emplace_back(id, std::move(cb));
-  return id;
+  lock_guard lk(cb_mu_);
+  return change_cbs_.emplace_back(next_cb_id_++, std::move(cb)).first;
 }
 
 void JournalSlice::UnregisterOnChange(uint32_t id) {
-  // we need to wait until callback is finished before remove it
   lock_guard lk(cb_mu_);
-  auto it = find_if(change_cb_arr_.begin(), change_cb_arr_.end(),
+  auto it = find_if(change_cbs_.begin(), change_cbs_.end(),
                     [id](const auto& e) { return e.first == id; });
-  CHECK(it != change_cb_arr_.end());
-  change_cb_arr_.erase(it);
+  CHECK(it != change_cbs_.end());
+  change_cbs_.erase(it);
 }
 
 }  // namespace journal