Get OOM status from SinkReplyBuilder in incoming slot migration

Modify SinkReplyBuilder to contain also OpStatus when error is set.
Don't reset error status after command is done, it will be done on next
invoicecmd/dispatchcmd call.

Author: mkaruza

src/facade/op_status.h

@@ -5,6 +5,7 @@
 #pragma once
 
 #include <cstdint>
+#include <limits>
 #include <ostream>
 
 namespace facade {
@@ -32,6 +33,7 @@ enum class OpStatus : uint16_t {
   MEMBER_NOTFOUND,
   INVALID_JSON_PATH,
   INVALID_JSON,
+  CUSTOM_ERROR = std::numeric_limits<uint16_t>::max()  // Custom status error
 };
 
 class OpResultBase {

src/facade/reply_builder.cc

@@ -90,6 +90,7 @@ void SinkReplyBuilder::SendError(ErrorReply error) {
 void SinkReplyBuilder::SendError(OpStatus status) {
   if (status == OpStatus::OK)
     return SendSimpleString("OK");
+  last_error_status_ = status;
   SendError(StatusToMsg(status));
 }
 
@@ -258,6 +259,7 @@ void MCReplyBuilder::SendLong(long val) {
 
 void MCReplyBuilder::SendError(string_view str, std::string_view type) {
   last_error_ = str;
+  last_error_status_ = OpStatus::CUSTOM_ERROR;
   SendSimpleString(absl::StrCat("SERVER_ERROR ", str));
 }
 
@@ -368,6 +370,7 @@ void RedisReplyBuilderBase::SendError(std::string_view str, std::string_view typ
   }
   tl_facade_stats->reply_stats.err_count[type]++;
   last_error_ = str;
+  last_error_status_ = OpStatus::CUSTOM_ERROR;
 
   if (str[0] != '-') {
     WritePieces("-ERR ");

src/facade/reply_builder.h

@@ -120,8 +120,17 @@ class SinkReplyBuilder {
   void SendError(ErrorReply error);
   virtual void SendProtocolError(std::string_view str) = 0;
 
-  std::string ConsumeLastError() {
-    return std::exchange(last_error_, {});
+  void ResetLastError() {
+    last_error_status_ = OpStatus::OK;
+    std::exchange(last_error_, {});
+  }
+
+  std::string_view GetLastError() {
+    return last_error_;
+  }
+
+  OpStatus GetLastErrorStatus() const {
+    return last_error_status_;
   }
 
   uint64_t GetLastSendTimeNs() const;
@@ -137,6 +146,7 @@ class SinkReplyBuilder {
  protected:
   size_t replies_recorded_ = 0;
   std::string last_error_;
+  OpStatus last_error_status_;
 
  private:
   io::Sink* sink_;

src/facade/reply_capture.cc

@@ -65,7 +65,7 @@ CapturingReplyBuilder::Payload CapturingReplyBuilder::Take() {
   CHECK(stack_.empty());
   Payload pl = std::move(current_);
   current_ = monostate{};
-  ConsumeLastError();
+  ResetLastError();
   return pl;
 }
 
@@ -165,7 +165,7 @@ void CapturingReplyBuilder::Apply(Payload&& pl, RedisReplyBuilder* rb) {
   CaptureVisitor cv{rb};
   visit(cv, std::move(pl));
   // Consumed and printed by InvokeCmd. We just send the actual error here
-  rb->ConsumeLastError();
+  rb->ResetLastError();
 }
 
 void CapturingReplyBuilder::SetReplyMode(ReplyMode mode) {

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);
+        // Check for OOM error
+        if (executor_.ExecutionLastErrorStatus() == facade::OpStatus::OUT_OF_MEMORY) {
+          cntx->ReportError(std::string{kIncomingMigrationOOM});
+          in_migration_->ReportFatalError(std::string{kIncomingMigrationOOM});
+          break;
+        }
       }
     }
 

src/server/journal/executor.h

@@ -32,6 +32,10 @@ class JournalExecutor {
     return &conn_context_;
   }
 
+  facade::OpStatus ExecutionLastErrorStatus() {
+    return reply_builder_.GetLastErrorStatus();
+  }
+
  private:
   void Execute(journal::ParsedEntry::CmdData& cmd);
 

src/server/main_service.cc

@@ -1146,9 +1146,10 @@ void Service::DispatchCommand(ArgSlice args, SinkReplyBuilder* builder,
   DCHECK(!args.empty());
   DCHECK_NE(0u, shard_set->size()) << "Init was not called";
 
-  absl::Cleanup clear_last_error([builder]() { builder->ConsumeLastError(); });
   ServerState& etl = *ServerState::tlocal();
 
+  builder->ResetLastError();
+
   string cmd = absl::AsciiStrToUpper(args[0]);
   const auto [cid, args_no_cmd] = registry_.FindExtended(cmd, args.subspan(1));
 
@@ -1312,6 +1313,9 @@ bool Service::InvokeCmd(const CommandId* cid, CmdArgList tail_args,
   DCHECK(builder);
   DCHECK(cntx);
 
+  // Reset
+  builder->ResetLastError();
+
   if (auto err = VerifyCommandExecution(cid, cntx, tail_args); err) {
     // We need to skip this because ACK's should not be replied to
     // Bonus points because this allows to continue replication with ACL users who got
@@ -1320,7 +1324,6 @@ bool Service::InvokeCmd(const CommandId* cid, CmdArgList tail_args,
       return true;
     }
     builder->SendError(std::move(*err));
-    builder->ConsumeLastError();
     return true;  // return false only for internal error aborts
   }
 
@@ -1354,16 +1357,15 @@ bool Service::InvokeCmd(const CommandId* cid, CmdArgList tail_args,
   ReplyGuard reply_guard(cid->name(), builder, cntx);
 #endif
   uint64_t invoke_time_usec = 0;
-  auto last_error = builder->ConsumeLastError();
-  DCHECK(last_error.empty());
   try {
     invoke_time_usec = cid->Invoke(tail_args, cmd_cntx);
   } catch (std::exception& e) {
     LOG(ERROR) << "Internal error, system probably unstable " << e.what();
     return false;
   }
 
-  if (std::string reason = builder->ConsumeLastError(); !reason.empty()) {
+  if (builder->GetLastErrorStatus() != facade::OpStatus::OK) {
+    std::string_view reason = builder->GetLastError();
     VLOG(2) << FailedCommandToString(cid->name(), tail_args, reason);
     LOG_EVERY_T(WARNING, 1) << FailedCommandToString(cid->name(), tail_args, reason);
   }

src/server/multi_command_squasher.cc

@@ -145,7 +145,7 @@ bool MultiCommandSquasher::ExecuteStandalone(RedisReplyBuilder* rb, const Stored
   if (opts_.verify_commands) {
     if (auto err = service_->VerifyCommandState(cmd->Cid(), args, *cntx_); err) {
       rb->SendError(std::move(*err));
-      rb->ConsumeLastError();
+      rb->ResetLastError();
       return !opts_.error_abort;
     }
   }