chore: cancel slot migrations on shutdown

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

Author: dranikpg

src/server/cluster/cluster_family.cc

@@ -79,6 +79,29 @@ ClusterConfig* ClusterFamily::cluster_config() {
   return tl_cluster_config.get();
 }
 
+void ClusterFamily::Shutdown() {
+  vector<shared_ptr<IncomingSlotMigration>> incoming;
+  vector<shared_ptr<OutgoingMigration>> outgoing;
+  {
+    std::lock_guard lk(migration_mu_);
+    incoming = std::move(incoming_migrations_jobs_);
+    outgoing = std::move(outgoing_migration_jobs_);
+  }
+
+  for (auto& migration : incoming)
+    migration->Stop(true);
+  for (auto& migration : outgoing)
+    migration->Finish(true, true);
+
+  {
+    std::lock_guard lk(migration_mu_);
+    DCHECK(incoming_migrations_jobs_.empty());  // can't initiate migrations in shutdown state
+    DCHECK(outgoing_migration_jobs_.empty());
+    incoming_migrations_jobs_.clear();
+    outgoing_migration_jobs_.clear();
+  }
+}
+
 ClusterShardInfo ClusterFamily::GetEmulatedShardInfo(ConnectionContext* cntx) const {
   ClusterShardInfo info{.slot_ranges = SlotRanges({{.start = 0, .end = kMaxSlotNum}}),
                         .master = {},
@@ -784,7 +807,7 @@ bool RemoveIncomingMigrationImpl(std::vector<std::shared_ptr<IncomingSlotMigrati
   SlotSet migration_slots(migration->GetSlots());
   SlotSet removed = migration_slots.GetRemovedSlots(tl_cluster_config->GetOwnedSlots());
 
-  migration->Stop();
+  migration->Stop(false);
   // all fibers has migration shared_ptr so we don't need to join it and can erase
   jobs.erase(it);
 

src/server/cluster/cluster_family.h

@@ -28,6 +28,8 @@ class ClusterFamily {
 
   void Register(CommandRegistry* registry);
 
+  void Shutdown();
+
   // Returns a thread-local pointer.
   static ClusterConfig* cluster_config();
 

src/server/cluster/incoming_slot_migration.cc

@@ -165,17 +165,20 @@ bool IncomingSlotMigration::Join(long attempt) {
   }
 }
 
-void IncomingSlotMigration::Stop() {
+void IncomingSlotMigration::Stop(bool join) {
   string_view log_state = state_.load() == MigrationState::C_FINISHED ? "Finishing" : "Cancelling";
   LOG(INFO) << log_state << " incoming migration of slots " << slots_.ToString();
   cntx_.Cancel();
 
   for (auto& flow : shard_flows_) {
-    if (auto err = flow->Cancel(); err) {
-      VLOG(1) << "Error during flow Stop: " << err;
-    }
+    auto err = flow->Cancel();
+    VLOG_IF(1, err) << "Error during flow Stop: " << err;
   }
-  bc_->Cancel();
+
+  if (join)
+    bc_->Wait();
+  else
+    bc_->Cancel();
 }
 
 void IncomingSlotMigration::StartFlow(uint32_t shard, util::FiberSocketBase* source) {

src/server/cluster/incoming_slot_migration.h

@@ -32,8 +32,9 @@ class IncomingSlotMigration {
   // After Join we still can get data due to error situation
   [[nodiscard]] bool Join(long attempt);
 
-  // Stop migrations, can be called even after migration is finished
-  void Stop();
+  // Stop migrations, can be called even after migration is finished.
+  // If join is set, wait for flows to finish
+  void Stop(bool join);
 
   MigrationState GetState() const {
     return state_.load();

src/server/cluster/outgoing_slot_migration.cc

@@ -130,7 +130,7 @@ void OutgoingMigration::OnAllShards(
   });
 }
 
-void OutgoingMigration::Finish(bool is_error) {
+void OutgoingMigration::Finish(bool is_error, bool force_finish) {
   VLOG(1) << "Finish outgoing migration for " << cf_->MyID() << " : "
           << migration_info_.node_info.id;
   bool should_cancel_flows = false;
@@ -155,10 +155,12 @@ void OutgoingMigration::Finish(bool is_error) {
     state_ = is_error ? MigrationState::C_ERROR : MigrationState::C_FINISHED;
   }
 
-  if (should_cancel_flows) {
+  if (force_finish || should_cancel_flows) {
     OnAllShards([](auto& migration) {
-      CHECK(migration != nullptr);
-      migration->Cancel();
+      if (migration) {
+        CHECK(migration != nullptr);
+        migration->Cancel();
+      }
     });
   }
 }

src/server/cluster/outgoing_slot_migration.h

@@ -31,8 +31,9 @@ class OutgoingMigration : private ProtocolClient {
   void Start();
 
   // mark migration as FINISHED and cancel migration if it's not finished yet
-  // can be called from any thread, but only after Start()
-  void Finish(bool is_error = false);
+  // can be called from any thread, but only after Start().
+  // force_finish forecfully cancels shard flows
+  void Finish(bool is_error = false, bool force_finish = false);
 
   MigrationState GetState() const;
 

src/server/main_service.cc

@@ -910,7 +910,9 @@ void Service::Shutdown() {
 
   config_registry.Reset();
 
-  // to shutdown all the runtime components that depend on EngineShard.
+  // to shutdown all the runtime components that depend on EngineShard
+  cluster_family_.Shutdown();
+
   server_family_.Shutdown();
   StringFamily::Shutdown();
   GenericFamily::Shutdown();

tests/dragonfly/cluster_test.py

@@ -1117,8 +1117,9 @@ async def test_cluster_flushall_during_migration(
     await close_clients(*[node.client for node in nodes], *[node.admin_client for node in nodes])
 
 
+@pytest.mark.parametrize("interrupt", [False, True])
 @dfly_args({"proactor_threads": 4, "cluster_mode": "yes"})
-async def test_cluster_data_migration(df_factory: DflyInstanceFactory):
+async def test_cluster_data_migration(df_factory: DflyInstanceFactory, interrupt: bool):
     # Check data migration from one node to another
     instances = [
         df_factory.create(port=BASE_PORT + i, admin_port=BASE_PORT + i + 1000) for i in range(2)
@@ -1145,6 +1146,13 @@ async def test_cluster_data_migration(df_factory: DflyInstanceFactory):
     logging.debug("Start migration")
     await push_config(json.dumps(generate_config(nodes)), [node.admin_client for node in nodes])
 
+    if interrupt:  # Test nodes properly shut down with pending migration
+        await asyncio.sleep(random.random())
+        await close_clients(
+            *[node.client for node in nodes], *[node.admin_client for node in nodes]
+        )
+        return
+
     await wait_for_status(nodes[1].admin_client, nodes[0].id, "FINISHED")
 
     for i in range(20, 22):