server: metrics to track stored command sizes (#5764)

* server: Track size of stored commands in exec info body

The stored command size is tracked in a field which will be used to
update metrics.

A new field is introduced instead of just calling UsedMemory on exec
info so we do not have to iterate over the entire stored commands vector
each time we read the used memory.

Signed-off-by: Abhijat Malviya <[email protected]>

* server: Add metric for stored commands used memory

Signed-off-by: Abhijat Malviya <[email protected]>

Author: abhijat

src/server/conn_context.cc

@@ -229,6 +229,18 @@ size_t ConnectionState::ExecInfo::UsedMemory() const {
   return dfly::HeapSize(body) + dfly::HeapSize(watched_keys);
 }
 
+void ConnectionState::ExecInfo::AddStoredCmd(const CommandId* cid, bool own_args, CmdArgList args) {
+  body.emplace_back(cid, own_args, args);
+  stored_cmd_bytes += body.back().UsedMemory();
+}
+
+size_t ConnectionState::ExecInfo::ClearStoredCmds() {
+  const size_t used = GetStoredCmdBytes();
+  vector<StoredCmd>{}.swap(body);
+  stored_cmd_bytes = 0;
+  return used;
+}
+
 size_t ConnectionState::ScriptInfo::UsedMemory() const {
   return dfly::HeapSize(lock_tags) + async_cmds_heap_mem;
 }
@@ -307,7 +319,8 @@ vector<unsigned> ConnectionContext::ChangeSubscriptions(CmdArgList channels, boo
 void ConnectionState::ExecInfo::Clear() {
   DCHECK(!preborrowed_interpreter);  // Must have been released properly
   state = EXEC_INACTIVE;
-  vector<StoredCmd>{}.swap(body);
+  const size_t cleared_size = ClearStoredCmds();
+  ServerState::tlocal()->stats.stored_cmd_bytes -= cleared_size;
   is_write = false;
   ClearWatched();
 }

src/server/conn_context.h

@@ -86,6 +86,18 @@ struct ConnectionState {
 
     size_t UsedMemory() const;
 
+    // Adds a StoredCmd and updates the stored_cmds_size
+    void AddStoredCmd(const CommandId* cid, bool own_args, CmdArgList args);
+
+    // Empties the body vector and resets stored_cmds_size to 0. Returns the size before data was
+    // cleared.
+    size_t ClearStoredCmds();
+
+    // Returns memory used by the body field without iterating over each stored command
+    size_t GetStoredCmdBytes() const {
+      return stored_cmd_bytes + body.capacity() * sizeof(StoredCmd);
+    }
+
     ExecState state = EXEC_INACTIVE;
     std::vector<StoredCmd> body;
     bool is_write = false;
@@ -99,6 +111,10 @@ struct ConnectionState {
     // executing the multi transaction, which can create deadlocks by blocking other transactions
     // that already borrowed all available interpreters but wait for keys to be unlocked.
     Interpreter* preborrowed_interpreter = nullptr;
+
+    // The total size of all stored commands kept in "body". Does not include memory allocated by
+    // the "body" vector.
+    size_t stored_cmd_bytes = 0;
   };
 
   // Lua-script related data.

src/server/main_service.cc

@@ -1346,9 +1346,12 @@ DispatchResult Service::DispatchCommand(ArgSlice args, SinkReplyBuilder* builder
   bool is_trans_cmd = CO::IsTransKind(cid->name());
   if (dfly_cntx->conn_state.exec_info.IsCollecting() && !is_trans_cmd) {
     // TODO: protect against aggregating huge transactions.
-    dfly_cntx->conn_state.exec_info.body.emplace_back(cid, true, args_no_cmd);
+    auto& exec_info = dfly_cntx->conn_state.exec_info;
+    const size_t old_size = exec_info.GetStoredCmdBytes();
+    exec_info.AddStoredCmd(cid, true, args_no_cmd);
+    etl.stats.stored_cmd_bytes += exec_info.GetStoredCmdBytes() - old_size;
     if (cid->IsWriteOnly()) {
-      dfly_cntx->conn_state.exec_info.is_write = true;
+      exec_info.is_write = true;
     }
     builder->SendSimpleString("QUEUED");
     return DispatchResult::OK;

src/server/multi_test.cc

@@ -1223,4 +1223,22 @@ TEST_F(MultiTest, ForceAtomicityFlag) {
   EXPECT_THAT(Run({"eval", kScript, "0"}), ErrArg("undeclared"));
 }
 
+TEST_F(MultiTest, StoredCmdBytesMetric) {
+  ASSERT_EQ(GetMetrics().coordinator_stats.stored_cmd_bytes, 0);
+
+  RespExpr resp = Run({"multi"});
+  ASSERT_EQ(resp, "OK");
+
+  for (auto i = 0; i < 100; ++i) {
+    ASSERT_EQ(Run({"get", kKey1}), "QUEUED");
+  }
+
+  ASSERT_GT(GetMetrics().coordinator_stats.stored_cmd_bytes, 0);
+
+  resp = Run({"exec"});
+  ASSERT_THAT(resp, ArrLen(100));
+  ASSERT_THAT(resp.GetVec(), Contains(ArgType(RespExpr::NIL)).Times(100));
+  ASSERT_EQ(GetMetrics().coordinator_stats.stored_cmd_bytes, 0);
+}
+
 }  // namespace dfly

src/server/server_family.cc

@@ -1778,6 +1778,9 @@ void PrintPrometheusMetrics(uint64_t uptime, const Metrics& m, DflyCmd* dfly_cmd
   AppendMetricValue("memory_by_class_bytes", total.obj_memory_usage, {"class"}, {"object_used"},
                     &memory_by_class_bytes);
 
+  AppendMetricValue("memory_by_class_bytes", m.coordinator_stats.stored_cmd_bytes, {"class"},
+                    {"conn_stored_commands"}, &memory_by_class_bytes);
+
   // Command stats
   if (!m.cmd_stats_map.empty()) {
     string command_metrics;

src/server/server_state.cc

@@ -53,7 +53,7 @@ ServerState::Stats::Stats(unsigned num_shards)
 }
 
 ServerState::Stats& ServerState::Stats::Add(const ServerState::Stats& other) {
-  static_assert(sizeof(Stats) == 24 * 8, "Stats size mismatch");
+  static_assert(sizeof(Stats) == 25 * 8, "Stats size mismatch");
 
 #define ADD(x) this->x += (other.x)
 
@@ -95,6 +95,8 @@ ServerState::Stats& ServerState::Stats::Add(const ServerState::Stats& other) {
   } else {
     this->squash_width_freq_arr = other.squash_width_freq_arr;
   }
+
+  ADD(stored_cmd_bytes);
   return *this;
 #undef ADD
 }

src/server/server_state.h

@@ -136,6 +136,9 @@ class ServerState {  // public struct - to allow initialization.
     uint32_t conn_timeout_events = 0;
     uint64_t psync_requests_total = 0;
     std::valarray<uint64_t> tx_width_freq_arr, squash_width_freq_arr;
+
+    // Memory size of stored commands during multi-exec in connections
+    size_t stored_cmd_bytes = 0;
   };
 
   // Unsafe version.