feat(server): Extend malloc-stats command functionality. (#775)

1. Remove dangerous check-fail in SegmentAllocator.
2. Allow printing heap stats for both backing and data heaps.

Signed-off-by: Roman Gershman <roman@dragonflydb.io>

src/core/segment_allocator.cc

@@ -3,20 +3,26 @@
 //
 #include "core/segment_allocator.h"
 
+#include <mimalloc-types.h>
+
 #include "base/logging.h"
 
+constexpr size_t kSegmentSize = MI_SEGMENT_SIZE;
+
+// mimalloc uses 32MiB segments and we might need change this code if it changes.
+static_assert(kSegmentSize == 1 << 25);
+
 namespace dfly {
 
 SegmentAllocator::SegmentAllocator(mi_heap_t* heap) : heap_(heap) {
 }
 
 void SegmentAllocator::ValidateMapSize() {
-  CHECK_LT(address_table_.size(), 1u << 12)
-      << "TODO: to monitor address_table_ map, it should not grow to such sizes";
-
-  // TODO: we should learn how large this maps can grow for very large databases.
-  // We should learn if mimalloc drops (deallocates) segments and we need to perform GC
-  // to protect ourselves from bloated address table.
+  if (address_table_.size() > 1u << 12) {
+    // This can happen if we restrict dragonfly to small number of threads on high-memory machine,
+    // for example.
+    LOG(WARNING) << "address_table_ map is growing too large: " << address_table_.size();
+  }
 }
 
 }  // namespace dfly

src/core/segment_allocator.h

@@ -47,7 +47,9 @@ class SegmentAllocator {
     return heap_;
   }
 
-  size_t used() const { return used_; }
+  size_t used() const {
+    return used_;
+  }
 
  private:
   static uint32_t Offset(Ptr p) {
@@ -77,7 +79,8 @@ inline auto SegmentAllocator::Allocate(uint32_t size) -> std::pair<Ptr, uint8_t*
     address_table_.push_back((uint8_t*)seg_ptr);
   }
 
-  Ptr res = (((iptr - seg_ptr) / 8) << kSegmentIdBits) | it->second;
+  uint32_t seg_offset = (iptr - seg_ptr) / 8;
+  Ptr res = (seg_offset << kSegmentIdBits) | it->second;
   used_ += mi_good_size(size);
 
   return std::make_pair(res, (uint8_t*)ptr);

src/server/memory_cmd.cc

@@ -36,32 +36,7 @@ bool MiArenaVisit(const mi_heap_t* heap, const mi_heap_area_t* area, void* block
   return true;
 };
 
-}  // namespace
-
-MemoryCmd::MemoryCmd(ServerFamily* owner, ConnectionContext* cntx) : sf_(*owner), cntx_(cntx) {
-}
-
-void MemoryCmd::Run(CmdArgList args) {
-  string_view sub_cmd = ArgS(args, 1);
-  if (sub_cmd == "USAGE") {
-    // dummy output, in practice not implemented yet.
-    return (*cntx_)->SendLong(1);
-  } else if (sub_cmd == "MALLOC-STATS") {
-    uint32_t tid = 0;
-    if (args.size() >= 3 && !absl::SimpleAtoi(ArgS(args, 2), &tid)) {
-      return (*cntx_)->SendError(kInvalidIntErr);
-    }
-    tid = tid % shard_set->pool()->size();
-    string res = shard_set->pool()->at(tid)->AwaitBrief([&] { return MallocStats(tid); });
-
-    return (*cntx_)->SendBulkString(res);
-  }
-
-  string err = UnknownSubCmd(sub_cmd, "MEMORY");
-  return (*cntx_)->SendError(err, kSyntaxErrType);
-}
-
-string MemoryCmd::MallocStats(unsigned tid) {
+std::string MallocStats(bool backing, unsigned tid) {
   string str;
 
   uint64_t start = absl::GetCurrentTimeNanos();
@@ -71,7 +46,7 @@ string MemoryCmd::MallocStats(unsigned tid) {
   absl::StrAppend(&str, "\nArena statistics from thread:", tid, "\n");
   absl::StrAppend(&str, "Count BlockSize Reserved Committed Used\n");
 
-  mi_heap_t* data_heap = ServerState::tlocal()->data_heap();
+  mi_heap_t* data_heap = backing ? mi_heap_get_backing() : ServerState::tlocal()->data_heap();
   BlockMap block_map;
 
   mi_heap_visit_blocks(data_heap, false /* visit all blocks*/, MiArenaVisit, &block_map);
@@ -92,4 +67,56 @@ string MemoryCmd::MallocStats(unsigned tid) {
   return str;
 }
 
+}  // namespace
+
+MemoryCmd::MemoryCmd(ServerFamily* owner, ConnectionContext* cntx) : sf_(*owner), cntx_(cntx) {
+}
+
+void MemoryCmd::Run(CmdArgList args) {
+  string_view sub_cmd = ArgS(args, 1);
+
+  if (sub_cmd == "HELP") {
+    string_view help_arr[] = {
+        "MEMORY <subcommand> [<arg> ...]. Subcommands are:",
+        "MALLOC-STATS [BACKING] [thread-id]",
+        "    Show malloc stats for a heap residing in specified thread-id. 0 by default.",
+        "    If BACKING is specified, show stats for the backing heap.",
+        "USAGE",
+        "    (not implemented).",
+    };
+    return (*cntx_)->SendSimpleStrArr(help_arr, ABSL_ARRAYSIZE(help_arr));
+  };
+
+  if (sub_cmd == "USAGE") {
+    // dummy output, in practice not implemented yet.
+    return (*cntx_)->SendLong(1);
+  }
+
+  if (sub_cmd == "MALLOC-STATS") {
+    uint32_t tid = 0;
+    bool backing = false;
+    if (args.size() >= 3) {
+      ToUpper(&args[2]);
+
+      unsigned tid_indx = 2;
+      if (ArgS(args, tid_indx) == "BACKING") {
+        ++tid_indx;
+        backing = true;
+      }
+
+      if (args.size() > tid_indx && !absl::SimpleAtoi(ArgS(args, tid_indx), &tid)) {
+        return (*cntx_)->SendError(kInvalidIntErr);
+      }
+    }
+
+    tid = tid % shard_set->pool()->size();
+    string res = shard_set->pool()->at(tid)->AwaitBrief([=] { return MallocStats(backing, tid); });
+
+    return (*cntx_)->SendBulkString(res);
+  }
+
+  string err = UnknownSubCmd(sub_cmd, "MEMORY");
+  return (*cntx_)->SendError(err, kSyntaxErrType);
+}
+
 }  // namespace dfly

src/server/memory_cmd.h

@@ -17,8 +17,6 @@ class MemoryCmd {
   void Run(CmdArgList args);
 
  private:
-  std::string MallocStats(unsigned tid);
-
   ServerFamily& sf_;
   ConnectionContext* cntx_;
 };