feat(cluster): migrate thread on same shard config (#5554)

---------

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

Author: dranikpg

.gitignore

@@ -20,3 +20,4 @@ releases
 .hypothesis
 .secrets
 cmake-build-debug
+.venv/

src/server/cluster/cluster_family.cc

@@ -977,6 +977,16 @@ void ClusterFamily::DflyMigrateFlow(CmdArgList args, SinkReplyBuilder* builder,
 
   builder->SendOk();
 
+  // Try migrating the connection if we have the same shard configuration
+  if (migration->ShardNum() == shard_set->size() &&
+      int32_t(shard_id) != fb2::ProactorBase::me()->GetPoolIndex()) {
+    DCHECK_LT(shard_id, shard_set->size());
+    if (bool success = cntx->conn()->Migrate(shard_set->pool()->at(shard_id)); !success) {
+      builder->SendError("invalid state");
+      return;
+    }
+  }
+
   migration->StartFlow(shard_id, cntx->conn()->socket());
 }
 

src/server/cluster/incoming_slot_migration.h

@@ -50,6 +50,10 @@ class IncomingSlotMigration {
     return source_id_;
   }
 
+  size_t ShardNum() const {
+    return shard_flows_.size();
+  }
+
   // Switch to  FATAL state and store error message
   void ReportFatalError(dfly::GenericError err) ABSL_LOCKS_EXCLUDED(state_mu_, error_mu_) {
     errors_count_.fetch_add(1, std::memory_order_relaxed);

src/server/db_slice.cc

@@ -802,14 +802,17 @@ void DbSlice::FlushSlotsFb(const cluster::SlotSet& slot_ids) {
   uint64_t next_version = 0;
   uint64_t del_count = 0;
 
+  // Explicitly copy table smart pointer to keep reference count up (flushall drops it)
+  boost::intrusive_ptr<DbTable> table = db_arr_.front();
+
   std::string tmp;
   auto iterate_bucket = [&](PrimeTable::bucket_iterator it) {
     it.AdvanceIfNotOccupied();
     while (!it.is_done()) {
       std::string_view key = it->first.GetSlice(&tmp);
       SlotId sid = KeySlot(key);
       if (slot_ids.Contains(sid) && it.GetVersion() < next_version) {
-        PerformDeletion(Iterator::FromPrime(it), db_arr_[0].get());
+        PerformDeletion(Iterator::FromPrime(it), table.get());
         ++del_count;
       }
       ++it;
@@ -836,7 +839,7 @@ void DbSlice::FlushSlotsFb(const cluster::SlotSet& slot_ids) {
   next_version = RegisterOnChange(std::move(on_change));
 
   ServerState& etl = *ServerState::tlocal();
-  PrimeTable* pt = &db_arr_[0]->prime;
+  PrimeTable* pt = &table->prime;
   PrimeTable::Cursor cursor;
 
   do {

tests/dragonfly/cluster_test.py

@@ -1330,6 +1330,14 @@ async def test_cluster_flushall_during_migration(
 
     assert await nodes[0].client.dbsize() == 0
 
+    # Push config that causes mass async slot deletion on nodes[1]
+    nodes[0].slots = [(0, 16383)]
+    nodes[1].slots = []
+    await push_config(json.dumps(generate_config(nodes)), [node.admin_client for node in nodes])
+
+    # Issue flushall right after pushing new config so it runs at the same time as disowned slots are flushed
+    await nodes[1].client.execute_command("flushall")
+
 
 @pytest.mark.parametrize("interrupt", [False, True])
 @dfly_args({"proactor_threads": 4, "cluster_mode": "yes"})