fix(takeover): pick correct replica to reconcile slots (#6034)

Signed-off-by: Kostas Kyrimis <[email protected]>

Author: kostasrim

src/server/cluster/cluster_family.cc

@@ -1098,27 +1098,50 @@ size_t ClusterFamily::MigrationsErrorsCount() const {
   return error_num;
 }
 
-void ClusterFamily::ReconcileMasterFlow() {
-  auto config = cluster::ClusterConfig::Current();
+void ClusterFamily::ReconcileMasterSlots(std::string_view repl_id) {
+  util::fb2::LockGuard gu(set_config_mu);
+  util::fb2::LockGuard lk(migration_mu_);
+
+  auto config = ClusterConfig::Current();
+
+  // Sanity -- we should not reach there
+  if (!config) {
+    LOG(ERROR) << "Cluster config after takeover is empty";
+    return;
+  }
 
   for (auto& info : config->GetMutableConfig()) {
     if (info.master.id == id_) {
       if (!info.replicas.empty()) {
-        LOG_IF(ERROR, info.replicas.size() > 1)
-            << "More than one replica found, slot redirection after takeover corrupted";
+        auto target = std::find_if(info.replicas.begin(), info.replicas.end(),
+                                   [repl_id](const auto& e) { return e.id == repl_id; });
 
-        // assumes there is one replica per master node
-        // TODO figure a smart way to get the replica id_ so
-        // we can find it here
-        info.master = info.replicas.front();
+        if (target == info.replicas.end()) {
+          LOG(ERROR) << "Could not find repl_id: " << repl_id
+                     << " in cluster topology. Slot redirection after takeover corrupted";
+          return;
+        }
+
+        info.master = *target;
         info.replicas.clear();
       }
       break;
     }
   }
 }
 
-void ClusterFamily::ReconcileReplicaFlow() {
+void ClusterFamily::ReconcileReplicaSlots() {
+  util::fb2::LockGuard gu(set_config_mu);
+  util::fb2::LockGuard lk(migration_mu_);
+
+  auto config = ClusterConfig::Current();
+
+  // Sanity -- we should not reach there
+  if (!config) {
+    LOG(ERROR) << "Cluster config after takeover is empty";
+    return;
+  }
+
   auto new_config = ClusterConfig::Current()->CloneWithChanges({}, {});
   // Replace master with replica in shard config.
   bool found = false;
@@ -1142,26 +1165,6 @@ void ClusterFamily::ReconcileReplicaFlow() {
       [&new_config](util::ProactorBase*) { ClusterConfig::SetCurrent(new_config); });
 }
 
-void ClusterFamily::ReconcileMasterReplicaTakeoverSlots(bool was_master) {
-  util::fb2::LockGuard gu(set_config_mu);
-  util::fb2::LockGuard lk(migration_mu_);
-
-  auto config = ClusterConfig::Current();
-
-  // Sanity -- we should not reach there
-  if (!config) {
-    LOG(ERROR) << "Cluster config after takeover is empty";
-    return;
-  }
-
-  if (was_master) {
-    ReconcileMasterFlow();
-    return;
-  }
-
-  ReconcileReplicaFlow();
-}
-
 using EngineFunc = void (ClusterFamily::*)(CmdArgList args, const CommandContext& cmd_cntx);
 
 inline CommandId::Handler3 HandlerFunc(ClusterFamily* se, EngineFunc f) {

src/server/cluster/cluster_family.h

@@ -44,14 +44,16 @@ class ClusterFamily {
 
   size_t MigrationsErrorsCount() const ABSL_LOCKS_EXCLUDED(migration_mu_);
 
-  // Helper function to be used during takeover from both nodes (master and replica).
+  // Helper functions to be used during takeover from both nodes (master and replica).
   // It reconciles the cluster configuration for both nodes to reflect the node
   // role changes after the takeover.
   // For the taking over node it's called at the end of the ReplTakeOver flow
   // and for the taken over node it's called at the end of the dflycmd::TakeOver
-  void ReconcileMasterReplicaTakeoverSlots(bool was_master)
+  void ReconcileMasterSlots(std::string_view repl_id)
       ABSL_LOCKS_EXCLUDED(set_config_mu, migration_mu_);
 
+  void ReconcileReplicaSlots() ABSL_LOCKS_EXCLUDED(set_config_mu, migration_mu_);
+
  private:
   using SinkReplyBuilder = facade::SinkReplyBuilder;
 
@@ -125,9 +127,6 @@ class ClusterFamily {
 
   ClusterShardInfo GetEmulatedShardInfo(ConnectionContext* cntx) const;
 
-  void ReconcileMasterFlow() ABSL_EXCLUSIVE_LOCKS_REQUIRED(set_config_mu, migration_mu_);
-  void ReconcileReplicaFlow() ABSL_EXCLUSIVE_LOCKS_REQUIRED(set_config_mu, migration_mu_);
-
   // Guards set configuration, so that we won't handle 2 in parallel.
   mutable util::fb2::Mutex set_config_mu;
 

src/server/dflycmd.cc

@@ -596,7 +596,7 @@ void DflyCmd::TakeOver(CmdArgList args, RedisReplyBuilder* rb, ConnectionContext
     return;
   }
 
-  sf_->service().cluster_family().ReconcileMasterReplicaTakeoverSlots(true);
+  sf_->service().cluster_family().ReconcileMasterSlots(replica_ptr->id);
 }
 
 void DflyCmd::Expire(CmdArgList args, Transaction* tx, RedisReplyBuilder* rb) {

src/server/server_family.cc

@@ -3766,7 +3766,7 @@ void ServerFamily::ReplTakeOver(CmdArgList args, const CommandContext& cmd_cntx)
   LOG(INFO) << "Takeover successful, promoting this instance to master.";
 
   if (IsClusterEnabled()) {
-    service().cluster_family().ReconcileMasterReplicaTakeoverSlots(false);
+    service().cluster_family().ReconcileReplicaSlots();
   }
 
   last_master_data_ = replica_->Stop();