Unify mimalloc memory management

src/core/compact_object.cc

@@ -481,7 +481,6 @@ auto CompactObj::GetStats() -> Stats {
 void CompactObj::InitThreadLocal(pmr::memory_resource* mr) {
   tl.local_mr = mr;
   tl.tmp_buf = base::PODArray<uint8_t>{mr};
-  SmallString::InitThreadLocal();
 }
 
 CompactObj::~CompactObj() {

src/core/compact_object_test.cc

@@ -38,7 +38,9 @@ class CompactObjectTest : public ::testing::Test {
   static void SetUpTestSuite() {
     InitRedisTables();  // to initialize server struct.
 
-    init_zmalloc_threadlocal();
+    auto* tlh = mi_heap_get_backing();
+    init_zmalloc_threadlocal(tlh);
+    SmallString::InitThreadLocal(tlh);
     CompactObj::InitThreadLocal(pmr::get_default_resource());
   }
 

src/core/mi_memory_resource.h

@@ -15,6 +15,8 @@ class MiMemoryResource final : public std::pmr::memory_resource {
   explicit MiMemoryResource(mi_heap_t* heap) : heap_(heap) {
   }
 
+  mi_heap_t* heap() { return heap_;}
+
  private:
   void* do_allocate(std::size_t size, std::size_t align) {
     return mi_heap_malloc_aligned(heap_, size, align);

src/core/small_string.cc

@@ -34,8 +34,8 @@ constexpr XXH64_hash_t kHashSeed = 24061983;  // same as in compact_object.cc
 
 }  // namespace
 
-void SmallString::InitThreadLocal() {
-  SegmentAllocator* ns = new SegmentAllocator(mi_heap_get_backing());
+void SmallString::InitThreadLocal(void * heap) {
+  SegmentAllocator* ns = new SegmentAllocator((mi_heap_t*)heap);
 
   tl.seg_alloc.reset(ns);
   tl.xxh_state.reset(XXH3_createState());

src/core/small_string.h

@@ -18,7 +18,7 @@ class SmallString {
 
  public:
 
-  static void InitThreadLocal();
+  static void InitThreadLocal(void * heap);
 
   void Reset() {
     size_ = 0;

src/redis/zmalloc.h

@@ -113,7 +113,7 @@ size_t zmalloc_usable_size(const void* p);
 
 // roman: void zlibc_free(void *ptr);
 
-void init_zmalloc_threadlocal();
+void init_zmalloc_threadlocal(void* heap);
 
 #undef __zm_str
 #undef __xstr

src/redis/zmalloc_mi.c

@@ -112,8 +112,8 @@ int zmalloc_get_allocator_info(size_t* allocated, size_t* active, size_t* reside
   return 1;
 }
 
-void init_zmalloc_threadlocal() {
+void init_zmalloc_threadlocal(void* heap) {
   if (zmalloc_heap)
     return;
-  zmalloc_heap = mi_heap_get_backing();
+  zmalloc_heap = heap;
 }

src/server/dfly_main.cc

@@ -2,6 +2,8 @@
 // See LICENSE for licensing terms.
 //
 
+#include <mimalloc.h>
+
 #include "base/init.h"
 #include "facade/dragonfly_listener.h"
 #include "server/main_service.h"
@@ -23,6 +25,7 @@ void RunEngine(ProactorPool* pool, AcceptServer* acceptor) {
 
   service.Init(acceptor);
   acceptor->AddListener(FLAGS_port, new Listener{Protocol::REDIS, &service});
+
   if (FLAGS_memcache_port > 0) {
     acceptor->AddListener(FLAGS_memcache_port, new Listener{Protocol::MEMCACHE, &service});
   }
@@ -35,6 +38,8 @@ void RunEngine(ProactorPool* pool, AcceptServer* acceptor) {
 
 }  // namespace dfly
 
+extern "C" void _mi_options_init();
+
 int main(int argc, char* argv[]) {
   gflags::SetUsageMessage("dragonfly [FLAGS]");
   gflags::SetVersionString("v0.2");
@@ -43,6 +48,12 @@ int main(int argc, char* argv[]) {
 
   CHECK_GT(FLAGS_port, 0u);
 
+  mi_option_enable(mi_option_large_os_pages);
+  mi_option_enable(mi_option_show_errors);
+  mi_option_set(mi_option_max_warnings, 0);
+
+  _mi_options_init();
+
   uring::UringPool pp{1024};
   pp.Run();
 

src/server/dragonfly_test.cc

@@ -48,11 +48,6 @@ class DflyEngineTest : public BaseFamilyTest {
   DflyEngineTest() : BaseFamilyTest() {
     num_threads_ = kPoolThreadCount;
   }
-
-  static void SetUpTestSuite() {
-    init_zmalloc_threadlocal();
-  }
-
 };
 
 // TODO: to implement equivalent parsing in redis parser.

src/server/engine_shard_set.cc

@@ -90,14 +90,16 @@ EngineShard::~EngineShard() {
 
 void EngineShard::InitThreadLocal(ProactorBase* pb, bool update_db_time) {
   CHECK(shard_ == nullptr) << pb->GetIndex();
+  CHECK(mi_heap_get_backing() == mi_heap_get_default());
 
-  init_zmalloc_threadlocal();
+  mi_heap_t* tlh = mi_heap_new();
+  init_zmalloc_threadlocal(tlh);
 
-  mi_heap_t* tlh = mi_heap_get_backing();
   void* ptr = mi_heap_malloc_aligned(tlh, sizeof(EngineShard), alignof(EngineShard));
   shard_ = new (ptr) EngineShard(pb, update_db_time, tlh);
 
   CompactObj::InitThreadLocal(shard_->memory_resource());
+  SmallString::InitThreadLocal(tlh);
 }
 
 void EngineShard::DestroyThreadLocal() {
@@ -474,7 +476,7 @@ bool EngineShard::HasResultConverged(TxId notifyid) const {
 }
 
 void EngineShard::CacheStats() {
-  mi_heap_t* tlh = mi_heap_get_backing();
+  mi_heap_t* tlh = mi_resource_.heap();
   struct Sum {
     size_t used = 0;
     size_t comitted = 0;

src/server/hset_family_test.cc

@@ -8,7 +8,6 @@ extern "C" {
 #include "redis/listpack.h"
 #include "redis/object.h"
 #include "redis/sds.h"
-#include "redis/zmalloc.h"
 }
 
 #include "base/gtest.h"
@@ -26,9 +25,6 @@ namespace dfly {
 
 class HSetFamilyTest : public BaseFamilyTest {
  protected:
-  static void SetUpTestSuite() {
-    init_zmalloc_threadlocal();
-  }
 };
 
 TEST_F(HSetFamilyTest, Hash) {

src/server/rdb_test.cc

@@ -8,6 +8,8 @@ extern "C" {
 #include "redis/zmalloc.h"
 }
 
+#include <mimalloc.h>
+
 #include "base/gtest.h"
 #include "base/logging.h"
 #include "io/file.h"
@@ -33,7 +35,7 @@ class RdbTest : public testing::Test {
 
   static void SetUpTestSuite() {
     crc64_init();
-    init_zmalloc_threadlocal();
+    init_zmalloc_threadlocal(mi_heap_get_backing());
   }
 
  protected:

src/server/server_family.cc

@@ -447,6 +447,7 @@ tcp_port:)";
       Replica::Info rinfo = replica_ptr->GetInfo();
       absl::StrAppend(&info, "master_host:", rinfo.host, "\n");
       absl::StrAppend(&info, "master_port:", rinfo.port, "\n");
+
       const char* link = rinfo.master_link_established ? "up" : "down";
       absl::StrAppend(&info, "master_link_status:", link, "\n");
       absl::StrAppend(&info, "master_last_io_seconds_ago:", rinfo.master_last_io_sec, "\n");

src/server/test_utils.cc

@@ -4,8 +4,14 @@
 
 #include "server/test_utils.h"
 
+extern "C" {
+#include "redis/zmalloc.h"
+}
+
+
 #include <absl/strings/match.h>
 #include <absl/strings/str_split.h>
+#include <mimalloc.h>
 
 #include "base/logging.h"
 #include "base/stl_util.h"
@@ -54,6 +60,10 @@ BaseFamilyTest::BaseFamilyTest() {
 BaseFamilyTest::~BaseFamilyTest() {
 }
 
+void BaseFamilyTest::SetUpTestSuite() {
+  init_zmalloc_threadlocal(mi_heap_get_backing());
+}
+
 void BaseFamilyTest::SetUp() {
   pp_.reset(new uring::UringPool(16, num_threads_));
   pp_->Run();

src/server/test_utils.h

@@ -23,6 +23,8 @@ class BaseFamilyTest : public ::testing::Test {
   BaseFamilyTest();
   ~BaseFamilyTest();
 
+  static void SetUpTestSuite();
+
   void SetUp() override;
   void TearDown() override;