clean up

Author: kostasrim

src/server/server_family.cc

@@ -1234,10 +1234,7 @@ void ServerFamily::Shutdown() {
     DebugCmd::Shutdown();
   });
 
-  service_.proactor_pool().AwaitFiberOnAll([](auto index, auto* cntx) {
-    tl_replica = nullptr;
-    tl_cluster_replicas.clear();
-  });
+  UpdateReplicationThreadLocals(nullptr);
 }
 
 bool ServerFamily::HasPrivilegedInterface() {
@@ -3457,10 +3454,9 @@ void ServerFamily::ReplicaOfNoOne(SinkReplyBuilder* builder) {
     SetMasterFlagOnAllThreads(true);
     // TODO we should not allow partial sync after NO-ONE. Only after Takeover.
     last_master_data_ = replica_->Stop();
-    // TODO set thread locals to nullptr
     replica_.reset();
     StopAllClusterReplicas();
-    service_.proactor_pool().AwaitFiberOnAll([](auto index, auto* cntx) { tl_replica = nullptr; });
+    UpdateReplicationThreadLocals(nullptr);
   }
 
   // May not switch to ACTIVE if the process is, for example, shutting down at the same time.
@@ -3523,6 +3519,11 @@ void ServerFamily::ReplicaOfInternal(CmdArgList args, Transaction* tx, SinkReply
     return builder->SendError(ec ? ec.Format() : "replication cancelled");
   }
 
+  // Critical section.
+  // 1. Stop the old replica_ if it exists
+  // 2. Update all the pointers to the new replica and update master flag
+  // 3. Start the main replication fiber
+  // 4. Send OK
   util::fb2::LockGuard lk(replicaof_mu_);
   if (replica_)
     last_master_data = replica_->Stop();
@@ -3534,10 +3535,7 @@ void ServerFamily::ReplicaOfInternal(CmdArgList args, Transaction* tx, SinkReply
 
   // Update thread locals. That way INFO never blocks
   replica_ = new_replica;
-  service_.proactor_pool().AwaitFiberOnAll([new_replica](auto index, auto* context) {
-    tl_replica = new_replica;
-    tl_cluster_replicas.clear();
-  });
+  UpdateReplicationThreadLocals(new_replica);
   SetMasterFlagOnAllThreads(false);
 
   if (on_err == ActionOnConnectionFail::kReturnOnError) {
@@ -3622,7 +3620,7 @@ void ServerFamily::ReplTakeOver(CmdArgList args, const CommandContext& cmd_cntx)
   last_master_data_ = replica_->Stop();
   replica_.reset();
 
-  service_.proactor_pool().AwaitFiberOnAll([](auto index, auto* context) { tl_replica = nullptr; });
+  UpdateReplicationThreadLocals(nullptr);
 
   return builder->SendOk();
 }
@@ -3916,6 +3914,13 @@ void ServerFamily::ClientPauseCmd(CmdArgList args, SinkReplyBuilder* builder,
   }
 }
 
+void ServerFamily::UpdateReplicationThreadLocals(std::shared_ptr<Replica> repl) {
+  service_.proactor_pool().AwaitFiberOnAll([repl](auto index, auto* context) {
+    tl_replica = repl;
+    tl_cluster_replicas.clear();
+  });
+}
+
 #define HFUNC(x) SetHandler(HandlerFunc(this, &ServerFamily::x))
 
 namespace acl {

src/server/server_family.h

@@ -414,6 +414,7 @@ class ServerFamily {
 
   void ReplicaOfNoOne(SinkReplyBuilder* builder);
   bool IsDragonflyLoadingAtomic();
+  void UpdateReplicationThreadLocals(std::shared_ptr<Replica> repl);
 
   util::fb2::Fiber snapshot_schedule_fb_;
   util::fb2::Fiber load_fiber_;

tests/dragonfly/replication_test.py

@@ -466,12 +466,8 @@ async def ping_status():
             await asyncio.sleep(0.05)
 
     async def replicate():
-        try:
-            await c_replica.execute_command(f"REPLICAOF localhost {master.port}")
-            return True
-        except redis.exceptions.ResponseError as e:
-            assert e.args[0] == "replication cancelled"
-            return False
+        await c_replica.execute_command(f"REPLICAOF localhost {master.port}")
+        return True
 
     ping_job = asyncio.create_task(ping_status())
     replication_commands = [asyncio.create_task(replicate()) for _ in range(COMMANDS_TO_ISSUE)]
@@ -492,6 +488,9 @@ async def replicate():
 
     replica.stop()
     lines = replica.find_in_logs("Stopping replication")
+    # Cancelled 99 times by REPLICAOF command and once by Shutdown() because
+    # we stopped the instance
+    assert len(lines) == COMMANDS_TO_ISSUE
 
 
 """