refactor: remove is_oom flag from ConnectionContext

src/facade/ok_main.cc

@@ -22,8 +22,10 @@ namespace {
 
 class OkService : public ServiceInterface {
  public:
-  void DispatchCommand(ArgSlice args, SinkReplyBuilder* builder, ConnectionContext* cntx) final {
+  DispatchResult DispatchCommand(ArgSlice args, SinkReplyBuilder* builder,
+                                 ConnectionContext* cntx) final {
     builder->SendOk();
+    return DispatchResult::OK;
   }
 
   size_t DispatchManyCommands(absl::Span<ArgSlice> args_lists, SinkReplyBuilder* builder,
@@ -47,9 +49,7 @@ void RunEngine(ProactorPool* pool, AcceptServer* acceptor) {
   OkService service;
 
   Connection::Init(pool->size());
-  pool->Await([](auto*) {
-    tl_facade_stats = new FacadeStats;
-  });
+  pool->Await([](auto*) { tl_facade_stats = new FacadeStats; });
 
   acceptor->AddListener(GetFlag(FLAGS_port), new Listener{Protocol::REDIS, &service});
 

src/facade/service_interface.h

@@ -21,13 +21,15 @@ class Connection;
 class SinkReplyBuilder;
 class MCReplyBuilder;
 
+enum class DispatchResult { OK, OOM, ERROR };
+
 class ServiceInterface {
  public:
   virtual ~ServiceInterface() {
   }
 
-  virtual void DispatchCommand(ArgSlice args, SinkReplyBuilder* builder,
-                               ConnectionContext* cntx) = 0;
+  virtual DispatchResult DispatchCommand(ArgSlice args, SinkReplyBuilder* builder,
+                                         ConnectionContext* cntx) = 0;
 
   // Returns number of processed commands
   virtual size_t DispatchManyCommands(absl::Span<ArgSlice> args_list, SinkReplyBuilder* builder,

src/server/cluster/incoming_slot_migration.cc

@@ -103,9 +103,9 @@ class ClusterShardMigration {
       if (tx_data->opcode == journal::Op::PING) {
         // TODO check about ping logic
       } else {
-        ExecuteTx(std::move(*tx_data), cntx);
+        auto err = ExecuteTx(std::move(*tx_data), cntx);
         // Break incoming slot migration if command reported OOM
-        if (executor_.connection_context()->IsOOM()) {
+        if (err == std::errc::not_enough_memory) {
           cntx->ReportError(std::string{kIncomingMigrationOOM});
           in_migration_->ReportFatalError(std::string{kIncomingMigrationOOM});
           break;
@@ -143,12 +143,13 @@ class ClusterShardMigration {
   }
 
  private:
-  void ExecuteTx(TransactionData&& tx_data, ExecutionState* cntx) {
+  std::error_code ExecuteTx(TransactionData&& tx_data, ExecutionState* cntx) {
     if (!cntx->IsRunning()) {
-      return;
+      return {};
     }
+
     if (!tx_data.IsGlobalCmd()) {
-      executor_.Execute(tx_data.dbid, tx_data.command);
+      return executor_.Execute(tx_data.dbid, tx_data.command);
     } else {
       // TODO check which global commands should be supported
       std::string error =
@@ -158,6 +159,8 @@ class ClusterShardMigration {
       cntx->ReportError(error);
       in_migration_->ReportError(error);
     }
+
+    return {};
   }
 
   uint32_t source_shard_id_;

src/server/conn_context.h

@@ -298,10 +298,6 @@ 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);
 
@@ -329,9 +325,6 @@ 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/journal/executor.cc

@@ -55,9 +55,9 @@ void JournalExecutor::Execute(DbIndex dbid, absl::Span<journal::ParsedEntry::Cmd
   }
 }
 
-void JournalExecutor::Execute(DbIndex dbid, journal::ParsedEntry::CmdData& cmd) {
+std::error_code JournalExecutor::Execute(DbIndex dbid, journal::ParsedEntry::CmdData& cmd) {
   SelectDb(dbid);
-  Execute(cmd);
+  return Execute(cmd);
 }
 
 void JournalExecutor::FlushAll() {
@@ -70,9 +70,11 @@ void JournalExecutor::FlushSlots(const cluster::SlotRange& slot_range) {
   Execute(cmd);
 }
 
-void JournalExecutor::Execute(journal::ParsedEntry::CmdData& cmd) {
+std::error_code JournalExecutor::Execute(journal::ParsedEntry::CmdData& cmd) {
   auto span = CmdArgList{cmd.cmd_args.data(), cmd.cmd_args.size()};
-  service_->DispatchCommand(span, &reply_builder_, &conn_context_);
+  auto res = service_->DispatchCommand(span, &reply_builder_, &conn_context_);
+  return res == facade::DispatchResult::OOM ? make_error_code(errc::not_enough_memory)
+                                            : error_code();
 }
 
 void JournalExecutor::SelectDb(DbIndex dbid) {

src/server/journal/executor.h

@@ -23,7 +23,7 @@ class JournalExecutor {
   JournalExecutor(JournalExecutor&&) = delete;
 
   void Execute(DbIndex dbid, absl::Span<journal::ParsedEntry::CmdData> cmds);
-  void Execute(DbIndex dbid, journal::ParsedEntry::CmdData& cmd);
+  std::error_code Execute(DbIndex dbid, journal::ParsedEntry::CmdData& cmd);
 
   void FlushAll();  // Execute FLUSHALL.
   void FlushSlots(const cluster::SlotRange& slot_range);
@@ -33,7 +33,7 @@ class JournalExecutor {
   }
 
  private:
-  void Execute(journal::ParsedEntry::CmdData& cmd);
+  std::error_code Execute(journal::ParsedEntry::CmdData& cmd);
 
   // Select database. Ensure it exists if accessed for first time.
   void SelectDb(DbIndex dbid);

src/server/main_service.cc

@@ -1066,7 +1066,7 @@ optional<ErrorReply> Service::VerifyCommandExecution(const CommandId* cid,
                                                      const ConnectionContext* cntx,
                                                      CmdArgList tail_args) {
   if (ShouldDenyOnOOM(cid)) {
-    return facade::ErrorReply{kOutOfMemory};
+    return facade::ErrorReply{OpStatus::OUT_OF_MEMORY};
   }
 
   return VerifyConnectionAclStatus(cid, cntx, "ACL rules changed between the MULTI and EXEC",
@@ -1184,8 +1184,8 @@ std::optional<ErrorReply> Service::VerifyCommandState(const CommandId* cid, CmdA
   return VerifyConnectionAclStatus(cid, &dfly_cntx, "has no ACL permissions", tail_args);
 }
 
-void Service::DispatchCommand(ArgSlice args, SinkReplyBuilder* builder,
-                              facade::ConnectionContext* cntx) {
+DispatchResult Service::DispatchCommand(ArgSlice args, SinkReplyBuilder* builder,
+                                        facade::ConnectionContext* cntx) {
   DCHECK(!args.empty());
   DCHECK_NE(0u, shard_set->size()) << "Init was not called";
 
@@ -1196,7 +1196,8 @@ void Service::DispatchCommand(ArgSlice args, SinkReplyBuilder* builder,
   const auto [cid, args_no_cmd] = registry_.FindExtended(cmd, args.subspan(1));
 
   if (cid == nullptr) {
-    return builder->SendError(ReportUnknownCmd(cmd));
+    builder->SendError(ReportUnknownCmd(cmd));
+    return DispatchResult::ERROR;
   }
 
   ConnectionContext* dfly_cntx = static_cast<ConnectionContext*>(cntx);
@@ -1229,10 +1230,10 @@ void Service::DispatchCommand(ArgSlice args, SinkReplyBuilder* builder,
     // their access revoked and reinstated
     if (cid->name() == "REPLCONF" && absl::EqualsIgnoreCase(ArgS(args_no_cmd, 0), "ACK")) {
       server_family_.GetDflyCmd()->OnClose(dfly_cntx->conn_state.replication_info.repl_session_id);
-      return;
+      return DispatchResult::ERROR;
     }
     builder->SendError(std::move(*err));
-    return;
+    return DispatchResult::ERROR;
   }
 
   VLOG_IF(1, cid->opt_mask() & CO::CommandOpt::DANGEROUS)
@@ -1246,7 +1247,8 @@ void Service::DispatchCommand(ArgSlice args, SinkReplyBuilder* builder,
     if (cid->IsWriteOnly()) {
       dfly_cntx->conn_state.exec_info.is_write = true;
     }
-    return builder->SendSimpleString("QUEUED");
+    builder->SendSimpleString("QUEUED");
+    return DispatchResult::OK;
   }
 
   // Create command transaction
@@ -1259,8 +1261,10 @@ void Service::DispatchCommand(ArgSlice args, SinkReplyBuilder* builder,
       OpStatus status = dfly_cntx->transaction->InitByArgs(
           dfly_cntx->ns, dfly_cntx->conn_state.db_index, args_no_cmd);
 
-      if (status != OpStatus::OK)
-        return builder->SendError(status);
+      if (status != OpStatus::OK) {
+        builder->SendError(status);
+        return DispatchResult::ERROR;
+      }
     }
   } else {
     DCHECK(dfly_cntx->transaction == nullptr);
@@ -1272,8 +1276,10 @@ void Service::DispatchCommand(ArgSlice args, SinkReplyBuilder* builder,
         CHECK(dfly_cntx->ns != nullptr);
         if (auto st =
                 dist_trans->InitByArgs(dfly_cntx->ns, dfly_cntx->conn_state.db_index, args_no_cmd);
-            st != OpStatus::OK)
-          return builder->SendError(st);
+            st != OpStatus::OK) {
+          builder->SendError(st);
+          return DispatchResult::ERROR;
+        }
       }
 
       dfly_cntx->transaction = dist_trans.get();
@@ -1285,14 +1291,17 @@ void Service::DispatchCommand(ArgSlice args, SinkReplyBuilder* builder,
 
   dfly_cntx->cid = cid;
 
-  if (!InvokeCmd(cid, args_no_cmd, CommandContext{dfly_cntx->transaction, builder, dfly_cntx})) {
+  auto res =
+      InvokeCmd(cid, args_no_cmd, CommandContext{dfly_cntx->transaction, builder, dfly_cntx});
+  if ((res != DispatchResult::OK) && (res != DispatchResult::OOM)) {
     builder->SendError("Internal Error");
     builder->CloseConnection();
   }
 
   if (!dispatching_in_multi) {
     dfly_cntx->transaction = nullptr;
   }
+  return res;
 }
 
 class ReplyGuard {
@@ -1345,8 +1354,8 @@ OpResult<void> OpTrackKeys(const OpArgs slice_args, const facade::Connection::We
   return OpStatus::OK;
 }
 
-bool Service::InvokeCmd(const CommandId* cid, CmdArgList tail_args,
-                        const CommandContext& cmd_cntx) {
+DispatchResult Service::InvokeCmd(const CommandId* cid, CmdArgList tail_args,
+                                  const CommandContext& cmd_cntx) {
   DCHECK(cid);
   DCHECK(!cid->Validate(tail_args));
 
@@ -1361,11 +1370,16 @@ bool Service::InvokeCmd(const CommandId* cid, CmdArgList tail_args,
     // Bonus points because this allows to continue replication with ACL users who got
     // their access revoked and reinstated
     if (cid->name() == "REPLCONF" && absl::EqualsIgnoreCase(ArgS(tail_args, 0), "ACK")) {
-      return true;
+      return DispatchResult::OK;
     }
+    auto res = err->status;
     builder->SendError(std::move(*err));
     builder->ConsumeLastError();
-    return true;  // return false only for internal error aborts
+    if (res == OpStatus::OUT_OF_MEMORY) {
+      return DispatchResult::OOM;
+    }
+
+    return DispatchResult::OK;  // return ERROR only for internal error aborts
   }
 
   // We are not sending any admin command in the monitor, and we do not want to
@@ -1404,13 +1418,14 @@ bool Service::InvokeCmd(const CommandId* cid, CmdArgList tail_args,
     invoke_time_usec = cid->Invoke(tail_args, cmd_cntx);
   } catch (std::exception& e) {
     LOG(ERROR) << "Internal error, system probably unstable " << e.what();
-    return false;
+    return DispatchResult::ERROR;
   }
 
+  DispatchResult res = DispatchResult::OK;
   if (std::string reason = builder->ConsumeLastError(); !reason.empty()) {
     // Set flag if OOM reported
     if (reason == kOutOfMemory) {
-      cmd_cntx.conn_cntx->is_oom_ = true;
+      res = DispatchResult::OOM;
     }
     VLOG(2) << FailedCommandToString(cid->name(), tail_args, reason);
     LOG_EVERY_T(WARNING, 1) << FailedCommandToString(cid->name(), tail_args, reason);
@@ -1454,7 +1469,7 @@ bool Service::InvokeCmd(const CommandId* cid, CmdArgList tail_args,
     cntx->last_command_debug.clock = tx->txid();
   }
 
-  return true;
+  return res;
 }
 
 size_t Service::DispatchManyCommands(absl::Span<CmdArgList> args_list, SinkReplyBuilder* builder,
@@ -2295,8 +2310,9 @@ void Service::Exec(CmdArgList args, const CommandContext& cmd_cntx) {
           }
         }
 
-        bool ok = InvokeCmd(scmd.Cid(), args, cmd_cntx);
-        if (!ok || rb->GetError())  // checks for i/o error, not logical error.
+        auto invoke_res = InvokeCmd(scmd.Cid(), args, cmd_cntx);
+        if ((invoke_res != DispatchResult::OK) ||
+            rb->GetError())  // checks for i/o error, not logical error.
           break;
       }
     }

src/server/main_service.h

@@ -36,15 +36,16 @@ class Service : public facade::ServiceInterface {
   void Shutdown();
 
   // Prepare command execution, verify and execute, reply to context
-  void DispatchCommand(ArgSlice args, facade::SinkReplyBuilder* builder,
-                       facade::ConnectionContext* cntx) final;
+  facade::DispatchResult DispatchCommand(ArgSlice args, facade::SinkReplyBuilder* builder,
+                                         facade::ConnectionContext* cntx) final;
 
   // Execute multiple consecutive commands, possibly in parallel by squashing
   size_t DispatchManyCommands(absl::Span<ArgSlice> args_list, facade::SinkReplyBuilder* builder,
                               facade::ConnectionContext* cntx) final;
 
   // Check VerifyCommandExecution and invoke command with args
-  bool InvokeCmd(const CommandId* cid, CmdArgList tail_args, const CommandContext& cmd_cntx);
+  facade::DispatchResult InvokeCmd(const CommandId* cid, CmdArgList tail_args,
+                                   const CommandContext& cmd_cntx);
 
   // Verify command can be executed now (check out of memory), always called immediately before
   // execution