Set flag to connection context if OOM is reported

Author: mkaruza

src/server/cluster/cluster_family.cc

@@ -1039,7 +1039,7 @@ void ClusterFamily::DflyMigrateAck(CmdArgList args, SinkReplyBuilder* builder) {
                    [source_id = source_id](const auto& m) { return m.node_info.id == source_id; });
   if (m_it == in_migrations.end()) {
     LOG(WARNING) << "migration isn't in config";
-    return builder->SendError(kUnknownMigration);
+    return builder->SendSimpleString(kUnknownMigration);
   }
 
   auto migration = GetIncomingMigration(source_id);

src/server/cluster/incoming_slot_migration.cc

@@ -76,14 +76,6 @@ class ClusterShardMigration {
         break;
       }
 
-      auto used_mem = used_mem_current.load(memory_order_relaxed);
-      // If aplying transaction data will reach 90% of max_memory_limit we end migration.
-      if ((used_mem + tx_data->command.cmd_len) > (0.9 * max_memory_limit)) {
-        cntx->ReportError(std::string{kIncomingMigrationOOM});
-        in_migration_->ReportFatalError(std::string{kIncomingMigrationOOM});
-        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);
@@ -112,6 +104,12 @@ class ClusterShardMigration {
         // TODO check about ping logic
       } else {
         ExecuteTx(std::move(*tx_data), cntx);
+        // Break incoming slot migration if command reported OOM
+        if (executor_.connection_context()->IsOOM()) {
+          cntx->ReportError(std::string{kIncomingMigrationOOM});
+          in_migration_->ReportFatalError(std::string{kIncomingMigrationOOM});
+          break;
+        }
       }
     }
 

src/server/conn_context.h

@@ -295,6 +295,10 @@ class ConnectionContext : public facade::ConnectionContext {
     return conn_state.db_index;
   }
 
+  bool IsOOM() {
+    return std::exchange(is_oom_, false);
+  }
+
   void ChangeSubscription(bool to_add, bool to_reply, CmdArgList args,
                           facade::RedisReplyBuilder* rb);
 
@@ -323,6 +327,9 @@ class ConnectionContext : public facade::ConnectionContext {
   // The related connection is bound to main listener or serves the memcached protocol
   bool has_main_or_memcache_listener = false;
 
+  // OOM reported while executing
+  bool is_oom_ = false;
+
  private:
   void EnableMonitoring(bool enable) {
     subscriptions++;  // required to support the monitoring

src/server/main_service.cc

@@ -1405,6 +1405,10 @@ bool Service::InvokeCmd(const CommandId* cid, CmdArgList tail_args,
   }
 
   if (std::string reason = builder->ConsumeLastError(); !reason.empty()) {
+    // Set flag if OOM reported
+    if (reason == kOutOfMemory) {
+      cmd_cntx.conn_cntx->is_oom_ = true;
+    }
     VLOG(2) << FailedCommandToString(cid->name(), tail_args, reason);
     LOG_EVERY_T(WARNING, 1) << FailedCommandToString(cid->name(), tail_args, reason);
   }