feat: add KEY dimenstion for memory usage (#6043)

Before: key memory and string memory were reported as one.
Now: metrics and info memory will have a separate dimension for keys
using heap memory.

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

Author: romange

src/core/compact_object.cc

@@ -1660,10 +1660,12 @@ bool CompactObj::JsonWrapper::DefragIfNeeded(PageUsage* page_usage) {
   return flat.DefragIfNeeded(page_usage);
 }
 
-constexpr std::pair<CompactObjType, std::string_view> kObjTypeToString[8] = {
-    {OBJ_STRING, "string"sv},  {OBJ_LIST, "list"sv},    {OBJ_SET, "set"sv},
-    {OBJ_ZSET, "zset"sv},      {OBJ_HASH, "hash"sv},    {OBJ_STREAM, "stream"sv},
-    {OBJ_JSON, "ReJSON-RL"sv}, {OBJ_SBF, "MBbloom--"sv}};
+constexpr std::pair<CompactObjType, std::string_view> kObjTypeToString[] =
+    {
+        {OBJ_STRING, "string"sv},  {OBJ_LIST, "list"sv},    {OBJ_SET, "set"sv},
+        {OBJ_ZSET, "zset"sv},      {OBJ_HASH, "hash"sv},    {OBJ_STREAM, "stream"sv},
+        {OBJ_KEY, "key"sv},  // pseudo-type used for memory tracking
+        {OBJ_JSON, "ReJSON-RL"sv}, {OBJ_SBF, "MBbloom--"sv}};
 
 std::string_view ObjTypeToString(CompactObjType type) {
   for (auto& p : kObjTypeToString) {

src/redis/redis_aux.h

@@ -9,6 +9,9 @@
 #define OBJ_JSON 15U
 #define OBJ_SBF  16U
 
+// A pseudo type for keys stored in the db, same as OBJ_MODULE which is not used in Dragonfly.
+#define OBJ_KEY  5U
+
 /* How many types of objects exist */
 #define OBJ_TYPE_MAX 17U
 

src/server/db_slice.cc

@@ -483,9 +483,9 @@ DbSlice::AutoUpdater::~AutoUpdater() {
 }
 
 void DbSlice::AutoUpdater::ReduceHeapUsage() {
-  AccountObjectMemory(fields_.key, fields_.it->second.ObjType(), -fields_.orig_heap_size,
+  AccountObjectMemory(fields_.key, fields_.it->second.ObjType(), -fields_.orig_value_heap_size,
                       fields_.db_slice->GetDBTable(fields_.db_ind));
-  fields_.orig_heap_size = 0;  // Reset to avoid double accounting.
+  fields_.orig_value_heap_size = 0;  // Reset to avoid double accounting.
 }
 
 void DbSlice::AutoUpdater::Run() {
@@ -501,7 +501,7 @@ void DbSlice::AutoUpdater::Run() {
   CHECK_NE(fields_.db_slice, nullptr);
 
   ssize_t delta = static_cast<int64_t>(fields_.it->second.MallocUsed()) -
-                  static_cast<int64_t>(fields_.orig_heap_size);
+                  static_cast<int64_t>(fields_.orig_value_heap_size);
   AccountObjectMemory(fields_.key, fields_.it->second.ObjType(), delta,
                       fields_.db_slice->GetDBTable(fields_.db_ind));
   fields_.db_slice->PostUpdate(fields_.db_ind, fields_.key);
@@ -518,7 +518,7 @@ DbSlice::AutoUpdater::AutoUpdater(DbIndex db_ind, std::string_view key, const It
               .db_ind = db_ind,
               .it = it,
               .key = key,
-              .orig_heap_size = it->second.MallocUsed()} {
+              .orig_value_heap_size = it->second.MallocUsed()} {
   DCHECK(IsValid(it));
 }
 
@@ -770,7 +770,7 @@ OpResult<DbSlice::ItAndUpdater> DbSlice::AddOrFindInternal(const Context& cntx,
   if (it->first.IsInline()) {
     ++db.stats.inline_keys;
   } else {
-    AccountObjectMemory(key, it->first.ObjType(), it->first.MallocUsed(), &db);  // Account for key
+    AccountObjectMemory(key, OBJ_KEY, it->first.MallocUsed(), &db);  // Account for key
   }
 
   DCHECK_EQ(it->second.MallocUsed(), 0UL);  // Make sure accounting is no-op
@@ -1553,7 +1553,7 @@ void DbSlice::InvalidateSlotWatches(const cluster::SlotSet& slot_ids) {
 }
 
 void DbSlice::RemoveOffloadedEntriesFromTieredStorage(absl::Span<const DbIndex> indices,
-                                                      const DbTableArray& db_arr) {
+                                                      const DbTableArray& db_arr) const {
   // Currently being used only for tiered storage.
   TieredStorage* tiered_storage = shard_owner()->tiered_storage();
   string scratch;
@@ -1787,8 +1787,7 @@ void DbSlice::PerformDeletionAtomic(const Iterator& del_it, const ExpIterator& e
   if (del_it->first.IsInline()) {
     --stats.inline_keys;
   } else {
-    AccountObjectMemory(del_it.key(), del_it->first.ObjType(), -key_size_used,
-                        table);  // Key
+    AccountObjectMemory(del_it.key(), OBJ_KEY, -key_size_used, table);  // Key
   }
   AccountObjectMemory(del_it.key(), pv.ObjType(), -value_heap_size, table);  // Value
 

src/server/db_slice.h

@@ -174,7 +174,7 @@ class DbSlice {
       std::string_view key;
 
       // The following fields are calculated at init time
-      size_t orig_heap_size = 0;
+      size_t orig_value_heap_size = 0;
     };
 
     AutoUpdater(DbIndex db_ind, std::string_view key, const Iterator& it, DbSlice* db_slice);
@@ -570,7 +570,7 @@ class DbSlice {
 
   // Clear tiered storage entries for the specified indices. Called during flushing some indices.
   void RemoveOffloadedEntriesFromTieredStorage(absl::Span<const DbIndex> indices,
-                                               const DbTableArray& db_arr);
+                                               const DbTableArray& db_arr) const;
 
   void PerformDeletionAtomic(const Iterator& del_it, const ExpIterator& exp_it, DbTable* table);
 

src/server/dragonfly_test.cc

@@ -953,6 +953,12 @@ TEST_F(DflyEngineTest, Huffman) {
   EXPECT_LT(metrics.heap_used_bytes, 14'000'000);         // less than 15mb
 }
 
+TEST_F(DflyEngineTest, MemoryKeys) {
+  Run({"debug", "populate", "10000", "abcd_efgh_ijkl_mnop", "10"});
+  auto metrics = GetMetrics();
+  EXPECT_GT(metrics.db_stats[0].memory_usage_by_type[OBJ_KEY], 100000);
+}
+
 class DflyCommandAliasTest : public DflyEngineTest {
  protected:
   DflyCommandAliasTest() {

src/server/engine_shard.cc

@@ -346,7 +346,7 @@ void EngineShard::Shutdown() {
 }
 
 void EngineShard::StopPeriodicFiber() {
-  ProactorBase::me()->RemoveOnIdleTask(defrag_task_);
+  ProactorBase::me()->RemoveOnIdleTask(defrag_task_id_);
   fiber_heartbeat_periodic_done_.Notify();
   if (fiber_heartbeat_periodic_.IsJoinable()) {
     fiber_heartbeat_periodic_.Join();
@@ -393,7 +393,7 @@ void EngineShard::StartPeriodicHeartbeatFiber(util::ProactorBase* pb) {
   fiber_heartbeat_periodic_ = fb2::Fiber(fb_opts, [this, period_ms, heartbeat]() mutable {
     RunFPeriodically(heartbeat, period_ms, "heartbeat", &fiber_heartbeat_periodic_done_);
   });
-  defrag_task_ = pb->AddOnIdleTask([this]() { return DefragTask(); });
+  defrag_task_id_ = pb->AddOnIdleTask([this]() { return DefragTask(); });
 }
 
 void EngineShard::StartPeriodicShardHandlerFiber(util::ProactorBase* pb,

src/server/engine_shard.h

@@ -294,7 +294,7 @@ class EngineShard {
   journal::Journal* journal_ = nullptr;
   IntentLock shard_lock_;
 
-  uint32_t defrag_task_ = 0;
+  uint32_t defrag_task_id_ = 0;
   EvictionTaskState eviction_state_;  // Used on eviction fiber
   util::fb2::Fiber fiber_heartbeat_periodic_;
   util::fb2::Done fiber_heartbeat_periodic_done_;