chore: to move some of the command-specific state into CommandContext (#6199)

Main change is around command latency measurement, where we concentrate the logic in one central
place: CommandContext::RecordLatency.

Related to #6196

Signed-off-by: Roman Gershman <[email protected]>

Author: romange

src/server/command_registry.cc

@@ -41,6 +41,7 @@ using namespace facade;
 
 using absl::AsciiStrToUpper;
 using absl::GetFlag;
+using absl::StrCat;
 using absl::StrSplit;
 
 namespace {
@@ -128,16 +129,56 @@ constexpr int32_t kLatencyHistogramPrecision = 2;
 
 }  // namespace
 
+void CommandContext::RecordLatency(facade::ArgSlice tail_args) const {
+  DCHECK_GT(start_time_ns, 0u);
+  int64_t after = absl::GetCurrentTimeNanos();
+
+  ServerState* ss = ServerState::tlocal();  // Might have migrated thread, read after invocation
+  int64_t execution_time_usec = (after - start_time_ns) / 1000;
+
+  cid->RecordLatency(ss->thread_index(), execution_time_usec);
+
+  DCHECK(conn_cntx != nullptr);
+
+  // TODO: we should probably discard more commands here,
+  // not just the blocking ones
+  const auto* conn = conn_cntx->conn();
+  if (conn_cntx->conn_state.squashing_info) {
+    // We run from the squashed transaction.
+    conn = conn_cntx->conn_state.squashing_info->owner->conn();
+  }
+
+  if (!(cid->opt_mask() & CO::BLOCKING) && conn != nullptr &&
+      // Use SafeTLocal() to avoid accessing the wrong thread local instance
+      ServerState::SafeTLocal()->ShouldLogSlowCmd(execution_time_usec)) {
+    vector<string> aux_params;
+    CmdArgVec aux_slices;
+
+    if (tail_args.empty() && cid->name() == "EXEC") {
+      // abuse tail_args to pass more information about the slow EXEC.
+      aux_params.emplace_back(StrCat("CMDCOUNT/", exec_body_len));
+      aux_slices.emplace_back(aux_params.back());
+      tail_args = absl::MakeSpan(aux_slices);
+    }
+    ServerState::SafeTLocal()->GetSlowLog().Add(cid->name(), tail_args, conn->GetName(),
+                                                conn->RemoteEndpointStr(), execution_time_usec,
+                                                absl::GetCurrentTimeNanos() / 1000);
+  }
+}
+
 CommandId::CommandId(const char* name, uint32_t mask, int8_t arity, int8_t first_key,
                      int8_t last_key, std::optional<uint32_t> acl_categories)
     : facade::CommandId(name, ImplicitCategories(mask), arity, first_key, last_key,
                         acl_categories.value_or(ImplicitAclCategories(mask))) {
   implicit_acl_ = !acl_categories.has_value();
-  hdr_histogram* hist = nullptr;
-  const int init_result = hdr_init(kLatencyHistogramMinValue, kLatencyHistogramMaxValue,
-                                   kLatencyHistogramPrecision, &hist);
-  CHECK_EQ(init_result, 0) << "failed to initialize histogram for command " << name;
-  latency_histogram_ = hist;
+  bool is_latency_tracked = GetFlag(FLAGS_latency_tracking);
+  if (is_latency_tracked) {
+    hdr_histogram* hist = nullptr;
+    const int init_result = hdr_init(kLatencyHistogramMinValue, kLatencyHistogramMaxValue,
+                                     kLatencyHistogramPrecision, &hist);
+    CHECK_EQ(init_result, 0) << "failed to initialize histogram for command " << name;
+    latency_histogram_ = hist;
+  }
 
   if (name_.rfind("EVAL", 0) == 0)
     kind_multi_ctr_ = CO::MultiControlKind::EVAL;
@@ -149,6 +190,7 @@ CommandId::CommandId(const char* name, uint32_t mask, int8_t arity, int8_t first
     kind_pubsub_ = CO::PubSubKind::PATTERN;
   else if (base::_in(name_, {"SPUBLISH", "SSUBSCRIBE", "SUNSUBSCRIBE"}))
     kind_pubsub_ = CO::PubSubKind::SHARDED;
+  can_be_monitored_ = (opt_mask_ & CO::ADMIN) == 0 && name_ != "EXEC";
 }
 
 CommandId::~CommandId() {
@@ -169,8 +211,10 @@ CommandId CommandId::Clone(const std::string_view name) const {
 
   // explicit sharing of the object since it's an alias we can do that.
   // I am assuming that the source object lifetime is at least as of the cloned object.
-  hdr_close(cloned.latency_histogram_);  // Free the histogram in the cloned object.
-  cloned.latency_histogram_ = static_cast<hdr_histogram*>(latency_histogram_);
+  if (cloned.latency_histogram_) {
+    hdr_close(cloned.latency_histogram_);  // Free the histogram in the cloned object.
+    cloned.latency_histogram_ = static_cast<hdr_histogram*>(latency_histogram_);
+  }
   return cloned;
 }
 
@@ -189,28 +233,6 @@ bool CommandId::IsMultiTransactional() const {
   return kind_multi_ctr_.has_value();
 }
 
-uint64_t CommandId::Invoke(CmdArgList args, const CommandContext& cmd_cntx) const {
-  uint64_t before = cmd_cntx.conn_cntx->conn_state.cmd_start_time_ns;
-  DCHECK_GT(before, 0u);
-  handler_(args, cmd_cntx);
-  int64_t after = absl::GetCurrentTimeNanos();
-
-  ServerState* ss = ServerState::tlocal();  // Might have migrated thread, read after invocation
-  int64_t execution_time_usec = (after - before) / 1000;
-
-  auto& ent = command_stats_[ss->thread_index()];
-
-  ++ent.first;
-  ent.second += execution_time_usec;
-  static const bool is_latency_tracked = GetFlag(FLAGS_latency_tracking);
-  if (is_latency_tracked) {
-    if (hdr_histogram* cmd_histogram = latency_histogram_; cmd_histogram != nullptr) {
-      hdr_record_value(cmd_histogram, execution_time_usec);
-    }
-  }
-  return execution_time_usec;
-}
-
 optional<facade::ErrorReply> CommandId::Validate(CmdArgList tail_args) const {
   if ((arity() > 0 && tail_args.size() + 1 != size_t(arity())) ||
       (arity() < 0 && tail_args.size() + 1 < size_t(-arity()))) {
@@ -238,8 +260,15 @@ void CommandId::ResetStats(unsigned thread_index) {
   }
 }
 
-hdr_histogram* CommandId::LatencyHist() const {
-  return latency_histogram_;
+void CommandId::RecordLatency(unsigned tid, uint64_t latency_usec) const {
+  auto& ent = command_stats_[tid];
+
+  ++ent.first;
+  ent.second += latency_usec;
+
+  if (latency_histogram_) {
+    hdr_record_value(latency_histogram_, latency_usec);
+  }
 }
 
 CommandRegistry::CommandRegistry() {
@@ -280,7 +309,7 @@ CommandRegistry& CommandRegistry::operator<<(CommandId cmd) {
     if (it->second.empty()) {
       return *this;  // Incase of empty string we want to remove the command from registry.
     }
-    k = is_sub_command ? absl::StrCat(it->second, " ", maybe_subcommand[1]) : it->second;
+    k = is_sub_command ? StrCat(it->second, " ", maybe_subcommand[1]) : it->second;
   }
 
   if (restricted_cmds_.find(k) != restricted_cmds_.end()) {
@@ -333,7 +362,7 @@ std::pair<const CommandId*, ParsedArgs> CommandRegistry::FindExtended(string_vie
     }
 
     auto second_cmd = absl::AsciiStrToUpper(tail_args.Front());
-    string full_cmd = absl::StrCat(cmd, " ", second_cmd);
+    string full_cmd = StrCat(cmd, " ", second_cmd);
 
     return {Find(full_cmd), tail_args.Tail()};
   }
@@ -355,7 +384,7 @@ absl::flat_hash_map<std::string, hdr_histogram*> CommandRegistry::LatencyMap() c
   absl::flat_hash_map<std::string, hdr_histogram*> cmd_latencies;
   cmd_latencies.reserve(cmd_map_.size());
   for (const auto& [cmd_name, cmd] : cmd_map_) {
-    cmd_latencies.insert({absl::AsciiStrToLower(cmd_name), cmd.LatencyHist()});
+    cmd_latencies.insert({absl::AsciiStrToLower(cmd_name), cmd.GetLatencyHist()});
   }
   return cmd_latencies;
 }

src/server/command_registry.h

@@ -70,14 +70,25 @@ enum class MultiControlKind : uint8_t {
 // Per thread vector of command stats. Each entry is {cmd_calls, cmd_latency_agg in usec}.
 using CmdCallStats = std::pair<uint64_t, uint64_t>;
 
+class CommandId;
+
 struct CommandContext {
-  CommandContext(Transaction* _tx, facade::SinkReplyBuilder* _rb, ConnectionContext* cntx)
-      : tx(_tx), rb(_rb), conn_cntx(cntx) {
+  CommandContext(const CommandId* _cid, Transaction* _tx, facade::SinkReplyBuilder* _rb,
+                 ConnectionContext* cntx)
+      : cid(_cid), tx(_tx), rb(_rb), conn_cntx(cntx) {
   }
 
+  const CommandId* cid;
   Transaction* tx;
   facade::SinkReplyBuilder* rb;
   ConnectionContext* conn_cntx;
+
+  uint64_t start_time_ns = 0;
+
+  // number of commands in the last exec body.
+  unsigned exec_body_len = 0;
+
+  void RecordLatency(facade::ArgSlice tail_args) const;
 };
 
 // TODO: move it to helio
@@ -103,7 +114,7 @@ template <typename T> class MoveOnly {
   }
 
  private:
-  T value_;
+  T value_{};
 };
 
 class CommandId : public facade::CommandId {
@@ -131,7 +142,9 @@ class CommandId : public facade::CommandId {
                                           std::optional<facade::ErrorReply>(CmdArgList) const>;
 
   // Returns the invoke time in usec.
-  uint64_t Invoke(CmdArgList args, const CommandContext& cmd_cntx) const;
+  void Invoke(CmdArgList args, const CommandContext& cmd_cntx) const {
+    handler_(args, cmd_cntx);
+  }
 
   // Returns error if validation failed, otherwise nullopt
   std::optional<facade::ErrorReply> Validate(CmdArgList tail_args) const;
@@ -152,6 +165,10 @@ class CommandId : public facade::CommandId {
     return opt_mask_ & CO::BLOCKING;
   }
 
+  bool CanBeMonitored() const {
+    return can_be_monitored_;
+  }
+
   CommandId&& SetHandler(Handler3 f) && {
     handler_ = std::move(f);
     return std::move(*this);
@@ -181,7 +198,9 @@ class CommandId : public facade::CommandId {
     return is_alias_;
   }
 
-  hdr_histogram* LatencyHist() const;
+  hdr_histogram* GetLatencyHist() const {
+    return latency_histogram_;
+  }
 
   std::optional<CO::PubSubKind> PubSubKind() const {
     return kind_pubsub_;
@@ -192,13 +211,17 @@ class CommandId : public facade::CommandId {
     return kind_multi_ctr_;
   }
 
+  void RecordLatency(unsigned tid, uint64_t latency_usec) const;
+
  private:
   std::optional<CO::PubSubKind> kind_pubsub_;
   std::optional<CO::MultiControlKind> kind_multi_ctr_;
 
   // The following fields must copy manually in the move constructor.
   bool implicit_acl_;
   bool is_alias_{false};
+  bool can_be_monitored_{true};
+
   std::unique_ptr<CmdCallStats[]> command_stats_;
   Handler3 handler_;
   ArgValidator validator_;

src/server/conn_context.h

@@ -284,7 +284,6 @@ struct ConnectionState {
   std::unique_ptr<ScriptInfo> script_info;
   std::unique_ptr<SubscribeInfo> subscribe_info;
   ClientTracking tracking_info_;
-  uint64_t cmd_start_time_ns = 0;  // time when the last command started executing
 };
 
 class ConnectionContext : public facade::ConnectionContext {
@@ -295,17 +294,14 @@ class ConnectionContext : public facade::ConnectionContext {
   struct DebugInfo {
     uint32_t shards_count = 0;
     TxClock clock = 0;
-
-    // number of commands in the last exec body.
-    unsigned exec_body_len = 0;
   };
 
   DebugInfo last_command_debug;
 
   // TODO: to introduce proper accessors.
   Namespace* ns = nullptr;
   Transaction* transaction = nullptr;
-  const CommandId* cid = nullptr;
+  // const CommandId* cid = nullptr;
 
   ConnectionState conn_state;
 

src/server/debugcmd.cc

@@ -1632,8 +1632,8 @@ void DebugCmd::DoPopulateBatch(const PopulateOptions& options, const PopulateBat
       crb.SetReplyMode(ReplyMode::NONE);
       stub_tx->InitByArgs(cntx_->ns, local_cntx.conn_state.db_index, args_span);
 
-      sf_.service().InvokeCmd(cid, args_span,
-                              CommandContext{local_cntx.transaction, &crb, &local_cntx});
+      CommandContext cmd_cntx{cid, local_cntx.transaction, &crb, &local_cntx};
+      sf_.service().InvokeCmd(args_span, &cmd_cntx);
     }
 
     if (options.expire_ttl_range.has_value()) {
@@ -1653,8 +1653,8 @@ void DebugCmd::DoPopulateBatch(const PopulateOptions& options, const PopulateBat
       crb.SetReplyMode(ReplyMode::NONE);
       stub_tx->MultiSwitchCmd(cid);
       stub_tx->InitByArgs(cntx_->ns, local_cntx.conn_state.db_index, args_span);
-      sf_.service().InvokeCmd(cid, args_span,
-                              CommandContext{local_cntx.transaction, &crb, &local_cntx});
+      CommandContext cmd_cntx{cid, local_cntx.transaction, &crb, &local_cntx};
+      sf_.service().InvokeCmd(args_span, &cmd_cntx);
     }
   }
 

src/server/dflycmd.cc

@@ -592,7 +592,7 @@ void DflyCmd::TakeOver(CmdArgList args, RedisReplyBuilder* rb, ConnectionContext
     VLOG(1) << "Takeover accepted, shutting down.";
     std::string save_arg = "NOSAVE";
     MutableSlice sargs(save_arg);
-    sf_->ShutdownCmd(CmdArgList(&sargs, 1), CommandContext{nullptr, rb, nullptr});
+    sf_->ShutdownCmd(CmdArgList(&sargs, 1), CommandContext{nullptr, nullptr, rb, nullptr});
     return;
   }
 

src/server/hset_family.cc

@@ -547,7 +547,7 @@ void HSetEx(CmdArgList args, const CommandContext& cmd_cntx) {
   OpSetParams op_sp;
 
   const auto option_already_set = [&cmd_cntx] {
-    return cmd_cntx.rb->SendError(WrongNumArgsError(cmd_cntx.conn_cntx->cid->name()));
+    return cmd_cntx.rb->SendError(WrongNumArgsError(cmd_cntx.cid->name()));
   };
 
   while (true) {
@@ -580,8 +580,7 @@ void HSetEx(CmdArgList args, const CommandContext& cmd_cntx) {
   CmdArgList fields = parser.Tail();
 
   if (fields.size() % 2 != 0) {
-    return cmd_cntx.rb->SendError(facade::WrongNumArgsError(cmd_cntx.conn_cntx->cid->name()),
-                                  kSyntaxErrType);
+    return cmd_cntx.rb->SendError(facade::WrongNumArgsError(cmd_cntx.cid->name()), kSyntaxErrType);
   }
 
   auto cb = [&](Transaction* t, EngineShard* shard) {
@@ -851,7 +850,7 @@ void HSetFamily::HScan(CmdArgList args, const CommandContext& cmd_cntx) {
 void HSetFamily::HSet(CmdArgList args, const CommandContext& cmd_cntx) {
   string_view key = ArgS(args, 0);
 
-  string_view cmd{cmd_cntx.conn_cntx->cid->name()};
+  string_view cmd{cmd_cntx.cid->name()};
 
   if (args.size() % 2 != 1) {
     return cmd_cntx.rb->SendError(facade::WrongNumArgsError(cmd), kSyntaxErrType);