fix(cluster_family): Cancel slot migration in case of OOM errors on target node

We should cancel slot migration in case of OOM reported back.

Signed-off-by: mkaruza <[email protected]>

Author: mkaruza

src/server/cluster/cluster_config.cc

@@ -20,6 +20,7 @@ namespace dfly::cluster {
 namespace {
 
 thread_local shared_ptr<ClusterConfig> tl_cluster_config;
+thread_local shared_ptr<ClusterConfig> tl_previous_cluster_config;
 
 bool HasValidNodeIds(const ClusterShardInfos& new_config) {
   absl::flat_hash_set<string_view> nodes;
@@ -424,7 +425,12 @@ std::shared_ptr<ClusterConfig> ClusterConfig::Current() {
 }
 
 void ClusterConfig::SetCurrent(std::shared_ptr<ClusterConfig> config) {
+  tl_previous_cluster_config = std::move(tl_cluster_config);
   tl_cluster_config = std::move(config);
 }
 
+void ClusterConfig::RestorePrevious() {
+  tl_cluster_config = std::move(tl_previous_cluster_config);
+}
+
 }  // namespace dfly::cluster

src/server/cluster/cluster_config.h

@@ -58,6 +58,9 @@ class ClusterConfig {
   // Set a thread-local pointer.
   static void SetCurrent(std::shared_ptr<ClusterConfig> config);
 
+  // Restore previous configuration in case of error
+  static void RestorePrevious();
+
  private:
   struct SlotEntry {
     const ClusterShardInfo* shard = nullptr;

src/server/cluster/cluster_family.cc

@@ -16,6 +16,8 @@
 #include "facade/error.h"
 #include "server/acl/acl_commands_def.h"
 #include "server/channel_store.h"
+#include "server/cluster/cluster_defs.h"
+#include "server/cluster/outgoing_slot_migration.h"
 #include "server/command_registry.h"
 #include "server/conn_context.h"
 #include "server/dflycmd.h"
@@ -769,7 +771,6 @@ void ClusterFamily::DflySlotMigrationStatus(CmdArgList args, SinkReplyBuilder* b
   };
 
   for (const auto& m : incoming_migrations_jobs_) {
-    // TODO add error status
     append_answer("in", m->GetSourceID(), node_id, m->GetState(), m->GetKeyCount(),
                   m->GetErrorStr());
   }
@@ -1045,7 +1046,23 @@ void ClusterFamily::DflyMigrateAck(CmdArgList args, SinkReplyBuilder* builder) {
     return builder->SendError(kIdNotFound);
 
   if (!migration->Join(attempt)) {
-    return builder->SendError("Join timeout happened");
+    // Cleanup migraton that resulted in error
+    if (migration->GetState() == MigrationState::C_ERROR) {
+      // Restore previous cluster configuration
+      server_family_->service().proactor_pool().AwaitFiberOnAll(
+          [](util::ProactorBase*) { ClusterConfig::RestorePrevious(); });
+      // Cleanup slots if they were added before error
+      const auto& migration_slots = migration->GetSlots();
+      DeleteSlots(migration_slots);
+      LOG_IF(INFO, !migration->GetSlots().Empty())
+          << "Flushing incoming slots on error: " << migration_slots.ToString();
+      WriteFlushSlotsToJournal(migration_slots);
+      // Remove migration
+      RemoveIncomingMigrations(in_migrations);
+      return builder->SendError(migration->GetErrorStr());
+    } else {
+      return builder->SendError("Join timeout happened");
+    }
   }
 
   ApplyMigrationSlotRangeToConfig(migration->GetSourceID(), migration->GetSlots(), true);

src/server/cluster/incoming_slot_migration.cc

@@ -7,9 +7,13 @@
 #include <absl/cleanup/cleanup.h>
 #include <absl/strings/str_cat.h>
 
+#include <utility>
+
 #include "base/flags.h"
 #include "base/logging.h"
 #include "cluster_utility.h"
+#include "server/cluster/cluster_defs.h"
+#include "server/common.h"
 #include "server/error.h"
 #include "server/journal/executor.h"
 #include "server/journal/tx_executor.h"
@@ -70,6 +74,22 @@ class ClusterShardMigration {
         break;
       }
 
+      auto memory_limit_check = [&]() -> bool {
+        auto used_mem = used_mem_current.load(memory_order_relaxed);
+        if ((used_mem + tx_data->command.cmd_len) > max_memory_limit) {
+          std::string error =
+              absl::StrCat("Applying incoming slot data is overflowing max memory limit. Closing.");
+          cntx->ReportError(error);
+          in_migration_->ChangeToErrorState(error);
+          return true;
+        }
+        return false;
+      };
+
+      if (memory_limit_check()) {
+        break;
+      }
+
       while (tx_data->opcode == journal::Op::LSN) {
         VLOG(2) << "Attempt to finalize flow " << source_shard_id_ << " attempt " << tx_data->lsn;
         last_attempt_.store(tx_data->lsn);
@@ -79,6 +99,11 @@ class ClusterShardMigration {
           VLOG(1) << "Finalized flow " << source_shard_id_;
           return;
         }
+        if (memory_limit_check()) {
+          VLOG(2) << "Flow finalization " << source_shard_id_
+                  << " canceled due memory limit reached";
+          return;
+        }
         if (!tx_data->command.cmd_args.empty()) {
           VLOG(1) << "Flow finalization failed " << source_shard_id_ << " by "
                   << tx_data->command.cmd_args[0];
@@ -181,6 +206,13 @@ bool IncomingSlotMigration::Join(long attempt) {
       return false;
     }
 
+    // If any of migration shards reported ERROR (OOM) we can return error
+    if (GetState() == MigrationState::C_ERROR) {
+      LOG(WARNING) << "Error in incoming slot migration. Can't join migration for " << source_id_;
+      ReportError(GenericError("Error in incoming slot migration."));
+      return false;
+    }
+
     // if data was sent after LSN, WaitFor() always returns false so to reduce wait time
     // we check current state and if WaitFor false but GetLastAttempt() == attempt
     // the Join is failed and we can return false
@@ -251,7 +283,9 @@ void IncomingSlotMigration::Init(uint32_t shards_num) {
 
 void IncomingSlotMigration::StartFlow(uint32_t shard, util::FiberSocketBase* source) {
   shard_flows_[shard]->Start(&cntx_, source);
-  VLOG(1) << "Incoming flow " << shard << " finished for " << source_id_;
+  VLOG(1) << "Incoming flow " << shard
+          << (GetState() == MigrationState::C_ERROR ? " cancelled " : " finished ") << "for "
+          << source_id_;
 }
 
 size_t IncomingSlotMigration::GetKeyCount() const {

src/server/cluster/incoming_slot_migration.h

@@ -53,6 +53,16 @@ class IncomingSlotMigration {
   void ReportError(dfly::GenericError err) ABSL_LOCKS_EXCLUDED(error_mu_) {
     errors_count_.fetch_add(1, std::memory_order_relaxed);
     util::fb2::LockGuard lk(error_mu_);
+    // We don't want to update error message if we are in ERROR state
+    if (GetState() != MigrationState::C_ERROR)
+      last_error_ = std::move(err);
+  }
+
+  void ChangeToErrorState(dfly::GenericError err) ABSL_LOCKS_EXCLUDED(state_mu_, error_mu_) {
+    errors_count_.fetch_add(1, std::memory_order_relaxed);
+    util::fb2::LockGuard lk_error(error_mu_);
+    util::fb2::LockGuard lk_state(state_mu_);
+    state_ = MigrationState::C_ERROR;
     last_error_ = std::move(err);
   }
 

src/server/cluster/outgoing_slot_migration.cc

@@ -12,6 +12,8 @@
 #include "base/logging.h"
 #include "cluster_family.h"
 #include "cluster_utility.h"
+#include "facade/resp_expr.h"
+#include "server/cluster/cluster_defs.h"
 #include "server/db_slice.h"
 #include "server/engine_shard_set.h"
 #include "server/error.h"
@@ -276,6 +278,13 @@ void OutgoingMigration::SyncFb() {
       VLOG(1) << "Waiting for migration to finalize...";
       ThisFiber::SleepFor(500ms);
     }
+    // Break if outgoing migration resulted in ERROR
+    if (GetState() == MigrationState::C_ERROR) {
+      // Restore previous configuration
+      server_family_->service().proactor_pool().AwaitFiberOnAll(
+          [](util::ProactorBase*) { ClusterConfig::RestorePrevious(); });
+      break;
+    }
     if (!exec_st_.IsRunning()) {
       continue;
     }
@@ -354,6 +363,14 @@ bool OutgoingMigration::FinalizeMigration(long attempt) {
       return false;
     }
 
+    if (CheckRespFirstTypes({RespExpr::ERROR})) {
+      auto error = facade::ToSV(LastResponseArgs().front().GetBuf());
+      LOG(WARNING) << "Error response for " << cf_->MyID() << " : " << migration_info_.node_info.id
+                   << " attempt " << attempt << " msg: " << error;
+      Finish(std::string(error));
+      return false;
+    }
+
     if (!CheckRespFirstTypes({RespExpr::INT64})) {
       LOG(WARNING) << "Incorrect response type for " << cf_->MyID() << " : "
                    << migration_info_.node_info.id << " attempt " << attempt

src/server/journal/executor.cc

@@ -33,7 +33,7 @@ template <typename... Ts> journal::ParsedEntry::CmdData BuildFromParts(Ts... par
     start += part.size();
   }
 
-  return {std::move(buf), std::move(slice_parts)};
+  return {std::move(buf), std::move(slice_parts), cmd_str.size()};
 }
 }  // namespace
 

src/server/journal/serializer.cc

@@ -162,6 +162,7 @@ std::error_code JournalReader::ReadCommand(journal::ParsedEntry::CmdData* data)
 
   size_t cmd_size = 0;
   SET_OR_RETURN(ReadUInt<uint64_t>(), cmd_size);
+  data->cmd_len = cmd_size;
 
   // Read all strings consecutively.
   data->command_buf = make_unique<uint8_t[]>(cmd_size);
@@ -174,6 +175,9 @@ std::error_code JournalReader::ReadCommand(journal::ParsedEntry::CmdData* data)
     ptr += size;
     cmd_size -= size;
   }
+
+  data->cmd_len -= cmd_size;
+
   return {};
 }
 

src/server/journal/tx_executor.cc

@@ -53,8 +53,10 @@ void TransactionData::AddEntry(journal::ParsedEntry&& entry) {
   switch (entry.opcode) {
     case journal::Op::LSN:
       lsn = entry.lsn;
+      command.cmd_len = 0;
       return;
     case journal::Op::PING:
+      command.cmd_len = 0;
       return;
     case journal::Op::EXPIRED:
     case journal::Op::COMMAND:

src/server/journal/types.h

@@ -73,6 +73,7 @@ struct ParsedEntry : public EntryBase {
   struct CmdData {
     std::unique_ptr<uint8_t[]> command_buf;
     CmdArgVec cmd_args;  // represents the parsed command.
+    size_t cmd_len;
   };
   CmdData cmd;