Add mimalloc as the default allocator for redis structures

core/compact_object_test.cc

@@ -9,6 +9,7 @@
 
 extern "C" {
 #include "redis/object.h"
+#include "redis/zmalloc.h"
 }
 
 namespace dfly {
@@ -25,6 +26,7 @@ void PrintTo(const CompactObj& cobj, std::ostream* os) {
 class CompactObjectTest : public ::testing::Test {
  protected:
   static void SetUpTestCase() {
+    init_zmalloc_threadlocal();
   }
 
   CompactObj cs_;

core/dash_test.cc

@@ -20,6 +20,7 @@
 extern "C" {
 #include "redis/dict.h"
 #include "redis/sds.h"
+#include "redis/zmalloc.h"
 }
 
 namespace dfly {
@@ -34,7 +35,7 @@ static uint64_t dictSdsHash(const void* key) {
   return dictGenHashFunction((unsigned char*)key, sdslen((char*)key));
 }
 
-static int dictSdsKeyCompare(dict* , const void* key1, const void* key2) {
+static int dictSdsKeyCompare(dict*, const void* key1, const void* key2) {
   int l1, l2;
 
   l1 = sdslen((sds)key1);
@@ -80,6 +81,10 @@ constexpr size_t foo = Segment::kBucketSz;
 
 class DashTest : public testing::Test {
  protected:
+  static void SetUpTestCase() {
+    init_zmalloc_threadlocal();
+  }
+
   DashTest() : segment_(1) {
   }
 
@@ -548,7 +553,6 @@ struct SdsDashPolicy {
   }
 };
 
-
 TEST_F(DashTest, Sds) {
   DashTable<sds, uint64_t, SdsDashPolicy> dt;
 

redis/CMakeLists.txt

@@ -1,5 +1,19 @@
+option(REDIS_ZMALLOC_MI "Implement zmalloc layer using mimalloc allocator" ON)
+
+if (REDIS_ZMALLOC_MI)
+  set(ZMALLOC_SRC "zmalloc_mi.c")
+  set(ZMALLOC_DEPS "TRDP::mimalloc")
+else()
+  set(ZMALLOC_SRC "zmalloc.c")
+  set(ZMALLOC_DEPS "")
+endif()
+
 add_library(redis_lib crc64.c crcspeed.c debug.c dict.c endianconv.c intset.c 
             listpack.c mt19937-64.c object.c lzf_c.c lzf_d.c sds.c sha256.c
-            quicklist.c redis_aux.c siphash.c t_zset.c util.c zmalloc.c)
+            quicklist.c redis_aux.c siphash.c t_zset.c util.c ${ZMALLOC_SRC})
 
-cxx_link(redis_lib)
+cxx_link(redis_lib  ${ZMALLOC_DEPS})
+
+if (REDIS_ZMALLOC_MI)
+  target_compile_definitions(redis_lib PUBLIC USE_ZMALLOC_MI)
+endif()

redis/quicklist.c

@@ -8,7 +8,7 @@
  *
  *   * Redistributions of source code must start the above copyright notice,
  *     this quicklist of conditions and the following disclaimer.
- *   * Redistributions in binary form must reproduce the above copyright
+ *   * Redistributions inquicklistBookmarksClear binary form must reproduce the above copyright
  *     notice, this quicklist of conditions and the following disclaimer in the
  *     documentation and/or other materials provided with the distribution.
  *   * Neither the name of Redis nor the names of its contributors may be used
@@ -98,9 +98,17 @@ int quicklistisSetPackedThreshold(size_t sz) {
 
 /* Bookmarks forward declarations */
 #define QL_MAX_BM ((1 << QL_BM_BITS)-1)
-quicklistBookmark *_quicklistBookmarkFindByName(quicklist *ql, const char *name);
-quicklistBookmark *_quicklistBookmarkFindByNode(quicklist *ql, quicklistNode *node);
-void _quicklistBookmarkDelete(quicklist *ql, quicklistBookmark *bm);
+
+// static quicklistBookmark *_quicklistBookmarkFindByName(quicklist *ql, const char *name);
+static quicklistBookmark *_quicklistBookmarkFindByNode(quicklist *ql, quicklistNode *node);
+static void _quicklistBookmarkDelete(quicklist *ql, quicklistBookmark *bm);
+
+static void quicklistBookmarksClear(quicklist *ql) {
+    while (ql->bookmark_count)
+        zfree(ql->bookmarks[--ql->bookmark_count].name);
+    /* NOTE: We do not shrink (realloc) the quick list. main use case for this
+     * function is just before releasing the allocation. */
+}
 
 /* Simple way to give quicklistEntry structs default values with one call. */
 #define initEntry(e)                                                           \
@@ -1634,6 +1642,7 @@ void quicklistRepr(unsigned char *ql, int full) {
     }
 }
 
+#if ROMAN_ENABLE 
 /* Create or update a bookmark in the list which will be updated to the next node
  * automatically when the one referenced gets deleted.
  * Returns 1 on success (creation of new bookmark or override of an existing one).
@@ -1679,6 +1688,7 @@ int quicklistBookmarkDelete(quicklist *ql, const char *name) {
     return 1;
 }
 
+
 quicklistBookmark *_quicklistBookmarkFindByName(quicklist *ql, const char *name) {
     unsigned i;
     for (i=0; i<ql->bookmark_count; i++) {
@@ -1689,6 +1699,8 @@ quicklistBookmark *_quicklistBookmarkFindByName(quicklist *ql, const char *name)
     return NULL;
 }
 
+#endif 
+
 quicklistBookmark *_quicklistBookmarkFindByNode(quicklist *ql, quicklistNode *node) {
     unsigned i;
     for (i=0; i<ql->bookmark_count; i++) {
@@ -1708,13 +1720,6 @@ void _quicklistBookmarkDelete(quicklist *ql, quicklistBookmark *bm) {
      * it may be re-used later (a call to realloc may NOP). */
 }
 
-void quicklistBookmarksClear(quicklist *ql) {
-    while (ql->bookmark_count)
-        zfree(ql->bookmarks[--ql->bookmark_count].name);
-    /* NOTE: We do not shrink (realloc) the quick list. main use case for this
-     * function is just before releasing the allocation. */
-}
-
 /* The rest of this file is test cases and test helpers. */
 #ifdef REDIS_TEST
 #include <stdint.h>

redis/quicklist.h

@@ -194,10 +194,6 @@ size_t quicklistGetLzf(const quicklistNode *node, void **data);
 void quicklistRepr(unsigned char *ql, int full);
 
 /* bookmarks */
-int quicklistBookmarkCreate(quicklist **ql_ref, const char *name, quicklistNode *node);
-int quicklistBookmarkDelete(quicklist *ql, const char *name);
-quicklistNode *quicklistBookmarkFind(quicklist *ql, const char *name);
-void quicklistBookmarksClear(quicklist *ql);
 int quicklistisSetPackedThreshold(size_t sz);
 
 #ifdef REDIS_TEST

redis/zmalloc.c

@@ -38,9 +38,9 @@
  * for instance to free results obtained by backtrace_symbols(). We need
  * to define this function before including zmalloc.h that may shadow the
  * free implementation if we use jemalloc or another non standard allocator. */
-void zlibc_free(void *ptr) {
+/*void zlibc_free(void *ptr) {
     free(ptr);
-}
+}*/
 
 #include <string.h>
 #include <pthread.h>
@@ -82,12 +82,8 @@ void zlibc_free(void *ptr) {
 #define update_zmalloc_stat_alloc(__n) used_memory_tl += (__n)
 #define update_zmalloc_stat_free(__n)  used_memory_tl -= (__n)
 
-static redisAtomic size_t used_memory_cached = 0;
 __thread ssize_t used_memory_tl = 0;
 
-void set_used_memory_cached(size_t val) {
-  atomicSet(used_memory_cached, val);
-}
 
 
 static void zmalloc_default_oom(size_t size) {
@@ -99,6 +95,9 @@ static void zmalloc_default_oom(size_t size) {
 
 static void (*zmalloc_oom_handler)(size_t) = zmalloc_default_oom;
 
+void init_zmalloc_threadlocal() {
+}
+
 /* Try allocating memory, and return NULL if failed.
  * '*usable' is set to the usable size if non NULL. */
 void *ztrymalloc_usable(size_t size, size_t *usable) {
@@ -320,39 +319,6 @@ void zfree(void *ptr) {
 #endif
 }
 
-/* Similar to zfree, '*usable' is set to the usable size being freed. */
-void zfree_usable(void *ptr, size_t *usable) {
-#ifndef HAVE_MALLOC_SIZE
-    void *realptr;
-    size_t oldsize;
-#endif
-
-    if (ptr == NULL) return;
-#ifdef HAVE_MALLOC_SIZE
-    update_zmalloc_stat_free(*usable = zmalloc_size(ptr));
-    free(ptr);
-#else
-    realptr = (char*)ptr-PREFIX_SIZE;
-    *usable = oldsize = *((size_t*)realptr);
-    update_zmalloc_stat_free(oldsize+PREFIX_SIZE);
-    free(realptr);
-#endif
-}
-
-char *zstrdup(const char *s) {
-    size_t l = strlen(s)+1;
-    char *p = zmalloc(l);
-
-    memcpy(p,s,l);
-    return p;
-}
-
-size_t zmalloc_used_memory(void) {
-    size_t um;
-    atomicGet(used_memory_cached,um);
-    return um;
-}
-
 void zmalloc_set_oom_handler(void (*oom_handler)(size_t)) {
     zmalloc_oom_handler = oom_handler;
 }

redis/zmalloc.h

@@ -117,9 +117,8 @@ void *zrealloc_usable(void *ptr, size_t size, size_t *usable);
 void *ztrymalloc_usable(size_t size, size_t *usable);
 void *ztrycalloc_usable(size_t size, size_t *usable);
 void *ztryrealloc_usable(void *ptr, size_t size, size_t *usable);
-void zfree_usable(void *ptr, size_t *usable);
-char *zstrdup(const char *s);
-size_t zmalloc_used_memory(void);
+
+// size_t zmalloc_used_memory(void);
 void zmalloc_set_oom_handler(void (*oom_handler)(size_t));
 size_t zmalloc_get_rss(void);
 int zmalloc_get_allocator_info(size_t *allocated, size_t *active, size_t *resident);
@@ -130,7 +129,7 @@ size_t zmalloc_get_smap_bytes_by_field(char *field, long pid);
 size_t zmalloc_get_memory_size(void);
 size_t zmalloc_usable_size(const void* p);
 
-void zlibc_free(void *ptr);
+// roman: void zlibc_free(void *ptr);
 
 #ifdef HAVE_DEFRAG
 void zfree_no_tcache(void *ptr);
@@ -149,7 +148,7 @@ int zmalloc_test(int argc, char **argv, int accurate);
 #endif
 
 extern __thread ssize_t used_memory_tl;
-void set_used_memory_cached(size_t val);
+void init_zmalloc_threadlocal();
 
 #undef __zm_str
 #undef __xstr

redis/zmalloc_mi.c

@@ -0,0 +1,90 @@
+// Copyright 2022, Roman Gershman.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#include <assert.h>
+#include <mimalloc.h>
+#include <string.h>
+
+#include "atomicvar.h"
+#include "zmalloc.h"
+
+__thread ssize_t used_memory_tl = 0;
+__thread mi_heap_t* zmalloc_heap = NULL;
+
+/* Allocate memory or panic */
+void* zmalloc(size_t size) {
+  size_t usable;
+  return zmalloc_usable(size, &usable);
+}
+
+void* ztrymalloc_usable(size_t size, size_t* usable) {
+  return zmalloc_usable(size, usable);
+}
+
+size_t zmalloc_usable_size(const void* p) {
+  return mi_usable_size(p);
+}
+
+void zfree(void* ptr) {
+  size_t usable = mi_usable_size(ptr);
+  used_memory_tl -= usable;
+
+  return mi_free_size(ptr, usable);
+}
+
+void* zrealloc(void* ptr, size_t size) {
+  size_t usable;
+  return zrealloc_usable(ptr, size, &usable);
+}
+
+void* zcalloc(size_t size) {
+  size_t usable = mi_good_size(size);
+
+  used_memory_tl += usable;
+
+  return mi_heap_calloc(zmalloc_heap, 1, usable);
+}
+
+void* zmalloc_usable(size_t size, size_t* usable) {
+  size_t g = mi_good_size(size);
+  *usable = g;
+
+  used_memory_tl += g;
+  assert(zmalloc_heap);
+  return mi_heap_malloc(zmalloc_heap, g);
+}
+
+void* zrealloc_usable(void* ptr, size_t size, size_t* usable) {
+  size_t g = mi_good_size(size);
+  size_t prev = mi_usable_size(ptr);
+  *usable = g;
+
+  used_memory_tl += (g - prev);
+  return mi_heap_realloc(zmalloc_heap, ptr, g);
+}
+
+size_t znallocx(size_t size) {
+  return mi_good_size(size);
+}
+
+void zfree_size(void* ptr, size_t size) {
+  mi_free_size(ptr, size);
+}
+
+void* ztrymalloc(size_t size) {
+  size_t usable;
+  return zmalloc_usable(size, &usable);
+}
+
+void* ztrycalloc(size_t size) {
+  size_t g = mi_good_size(size);
+  used_memory_tl += g;
+  return mi_heap_calloc(zmalloc_heap, 1, size);
+}
+
+void init_zmalloc_threadlocal() {
+  if (zmalloc_heap)
+    return;
+  zmalloc_heap = mi_heap_get_backing();
+}

server/CMakeLists.txt

@@ -20,6 +20,6 @@ cxx_test(string_family_test dfly_test_lib LABELS DFLY)
 cxx_test(generic_family_test dfly_test_lib LABELS DFLY)
 cxx_test(memcache_parser_test dfly_test_lib LABELS DFLY)
 
-add_custom_target(check_dfly DEPENDS COMMAND ctest -L DFLY)
-add_dependencies(check_dfly dragonfly_test  list_family_test
+add_custom_target(check_dfly WORKING_DIRECTORY .. COMMAND ctest -L DFLY)
+add_dependencies(check_dfly dragonfly_test list_family_test
                  generic_family_test memcache_parser_test redis_parser_test string_family_test)

server/engine_shard_set.cc

@@ -4,6 +4,10 @@
 
 #include "server/engine_shard_set.h"
 
+extern "C" {
+ #include "redis/zmalloc.h"
+}
+
 #include "base/logging.h"
 #include "server/transaction.h"
 #include "util/fiber_sched_algo.h"
@@ -79,6 +83,8 @@ EngineShard::~EngineShard() {
 
 void EngineShard::InitThreadLocal(ProactorBase* pb, bool update_db_time) {
   CHECK(shard_ == nullptr) << pb->GetIndex();
+
+  init_zmalloc_threadlocal();
   shard_ = new EngineShard(pb, update_db_time);
 }
 

server/redis_parser_test.cc

@@ -6,6 +6,7 @@
 
 extern "C" {
  #include "redis/sds.h"
+ #include "redis/zmalloc.h"
 }
 
 #include <absl/strings/str_cat.h>
@@ -37,6 +38,10 @@ MATCHER_P(ArrArg, expected, absl::StrCat(negation ? "is not" : "is", " equal to:
 
 class RedisParserTest : public testing::Test {
  protected:
+  static void SetUpTestCase() {
+    init_zmalloc_threadlocal();
+  }
+
   RedisParser::Result Parse(std::string_view str);
 
   RedisParser parser_;