fix(cluster): Use transactions for migration start

Author: dranikpg

src/server/cluster/outgoing_slot_migration.cc

@@ -82,7 +82,9 @@ OutgoingMigration::OutgoingMigration(MigrationInfo info, ClusterFamily* cf, Serv
       migration_info_(std::move(info)),
       slot_migrations_(shard_set->size()),
       server_family_(sf),
-      cf_(cf) {
+      cf_(cf),
+      tx_(new Transaction{sf->service().FindCmd("DFLYCLUSTER")}) {
+  tx_->InitByArgs(0, {});
 }
 
 OutgoingMigration::~OutgoingMigration() {
@@ -207,15 +209,22 @@ void OutgoingMigration::SyncFb() {
       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 (CheckFlowsForErrors()) {
+      LOG(WARNING) << "Preparation error detected, retrying outgoing migration";
+      continue;
+    }
+
+    // Global transactional cut for migration to register db_slice and journal listeners
+    {
+      Transaction::Guard tg{tx_.get()};
+      OnAllShards([](auto& migration) { migration->PrepareSync(); });
+    }
+
+    OnAllShards([this](auto& migration) {
+      migration->StartSync();
+      if (migration->GetError())
+        Finish(true);
     });
 
     if (CheckFlowsForErrors()) {
@@ -255,7 +264,7 @@ bool OutgoingMigration::FinalizeMigration(long attempt) {
       return false;
     }
   }
-  // TODO implement blocking on migrated slots only
+
   bool is_block_active = true;
   auto is_pause_in_progress = [&is_block_active] { return is_block_active; };
   auto pause_fb_opt = Pause(server_family_->GetNonPriviligedListeners(), nullptr,
@@ -270,14 +279,8 @@ bool OutgoingMigration::FinalizeMigration(long attempt) {
     pause_fb_opt->JoinIfNeeded();
   });
 
-  auto cb = [this, attempt](util::ProactorBase* pb) {
-    if (const auto* shard = EngineShard::tlocal(); shard) {
-      slot_migrations_[shard->shard_id()]->Finalize(attempt);
-    }
-  };
-
   VLOG(1) << "FINALIZE flows for " << cf_->MyID() << " : " << migration_info_.node_info.id;
-  shard_set->pool()->AwaitFiberOnAll(std::move(cb));
+  OnAllShards([attempt](auto& migration) { migration->Finalize(attempt); });
 
   auto cmd = absl::StrCat("DFLYMIGRATE ACK ", cf_->MyID(), " ", attempt);
   VLOG(1) << "send " << cmd;

src/server/cluster/outgoing_slot_migration.h

@@ -3,9 +3,12 @@
 //
 #pragma once
 
+#include <boost/smart_ptr/intrusive_ptr.hpp>
+
 #include "io/io.h"
 #include "server/cluster/cluster_defs.h"
 #include "server/protocol_client.h"
+#include "server/transaction.h"
 
 namespace dfly {
 class DbSlice;
@@ -87,6 +90,8 @@ class OutgoingMigration : private ProtocolClient {
   mutable util::fb2::Mutex state_mu_;
   MigrationState state_ ABSL_GUARDED_BY(state_mu_) = MigrationState::C_NO_STATE;
 
+  boost::intrusive_ptr<Transaction> tx_;
+
   // when migration is finished we need to store number of migrated keys
   // because new request can add or remove keys and we get incorrect statistic
   size_t keys_number_ = 0;

src/server/dflycmd.cc

@@ -71,32 +71,6 @@ std::string_view SyncStateName(DflyCmd::SyncState sync_state) {
   return "unsupported";
 }
 
-struct TransactionGuard {
-  static OpStatus ExitGuardCb(Transaction* t, EngineShard* shard) {
-    t->GetDbSlice(shard->shard_id()).SetExpireAllowed(true);
-    return OpStatus::OK;
-  };
-
-  explicit TransactionGuard(Transaction* t, bool disable_expirations = false) : t(t) {
-    t->Execute(
-        [disable_expirations](Transaction* t, EngineShard* shard) {
-          if (disable_expirations) {
-            t->GetDbSlice(shard->shard_id()).SetExpireAllowed(!disable_expirations);
-          }
-          return OpStatus::OK;
-        },
-        false);
-    VLOG(2) << "Transaction guard engaged";
-  }
-
-  ~TransactionGuard() {
-    VLOG(2) << "Releasing transaction guard";
-    t->Execute(ExitGuardCb, true);
-  }
-
-  Transaction* t;
-};
-
 OpStatus WaitReplicaFlowToCatchup(absl::Time end_time, shared_ptr<DflyCmd::ReplicaInfo> replica,
                                   EngineShard* shard) {
   // We don't want any writes to the journal after we send the `PING`,
@@ -299,7 +273,7 @@ void DflyCmd::Sync(CmdArgList args, ConnectionContext* cntx) {
 
   // Start full sync.
   {
-    TransactionGuard tg{cntx->transaction};
+    Transaction::Guard tg{cntx->transaction};
     AggregateStatus status;
 
     // Use explicit assignment for replica_ptr, because capturing structured bindings is C++20.
@@ -337,7 +311,7 @@ void DflyCmd::StartStable(CmdArgList args, ConnectionContext* cntx) {
     return;
 
   {
-    TransactionGuard tg{cntx->transaction};
+    Transaction::Guard tg{cntx->transaction};
     AggregateStatus status;
 
     auto cb = [this, &status, replica_ptr = replica_ptr](EngineShard* shard) {

src/server/transaction.cc

@@ -127,6 +127,16 @@ cv_status Transaction::BatonBarrier::Wait(time_point tp) {
   return cv_status::no_timeout;
 }
 
+Transaction::Guard::Guard(Transaction* tx) : tx(tx) {
+  DCHECK(tx->cid_->opt_mask() & CO::GLOBAL_TRANS);
+  tx->Execute([](auto*, auto*) { return OpStatus::OK; }, false);
+}
+
+Transaction::Guard::~Guard() {
+  tx->Conclude();
+  tx->Refurbish();
+}
+
 Transaction::Transaction(const CommandId* cid) : cid_{cid} {
   InitTxTime();
   string_view cmd_name(cid_->name());

src/server/transaction.h

@@ -170,6 +170,14 @@ class Transaction {
     RAN_IMMEDIATELY = 1 << 7,  // Whether the shard executed immediately (during schedule)
   };
 
+  struct Guard {
+    explicit Guard(Transaction* tx);
+    ~Guard();
+
+   private:
+    Transaction* tx;
+  };
+
   explicit Transaction(const CommandId* cid);
 
   // Initialize transaction for squashing placed on a specific shard with a given parent tx