From 7ed808d5d093ec472de4868e02182cb1432c6967 Mon Sep 17 00:00:00 2001
From: Roman Gershman <roman@dragonflydb.io>
Date: Sun, 18 Sep 2022 13:15:39 +0300
Subject: [PATCH] fix(rdb): Fix RDB load bug when loading hset

---
 src/redis/zmalloc_mi.c | 39 +++++++++++++++++++--------------------
 src/server/rdb_load.cc |  4 ++--
 src/server/rdb_test.cc | 12 ++++++++++++
 3 files changed, 33 insertions(+), 22 deletions(-)

diff --git a/src/redis/zmalloc_mi.c b/src/redis/zmalloc_mi.c
index d58c918a73e8..66aca4529adc 100644
--- a/src/redis/zmalloc_mi.c
+++ b/src/redis/zmalloc_mi.c
@@ -37,8 +37,6 @@ size_t zmalloc_usable_size(const void* p) {
 void zfree(void* ptr) {
   size_t usable = mi_usable_size(ptr);
 
-  // I wish we can keep this assert but rdb_load creates objects in one thread and
-  // uses them in another.
   // assert(zmalloc_used_memory_tl >= (ssize_t)usable);
   zmalloc_used_memory_tl -= usable;
 
@@ -51,34 +49,34 @@ void* zrealloc(void* ptr, size_t size) {
 }
 
 void* zcalloc(size_t size) {
-  size_t usable = mi_good_size(size);
+  // mi_good_size(size) is not working. try for example, size=690557.
 
+  void* res = mi_heap_calloc(zmalloc_heap, 1, size);
+  size_t usable = mi_usable_size(res);
   zmalloc_used_memory_tl += usable;
 
-  return mi_heap_calloc(zmalloc_heap, 1, size);
+  return res;
 }
 
 void* zmalloc_usable(size_t size, size_t* usable) {
-  size_t g = mi_good_size(size);
-  *usable = g;
-
-  zmalloc_used_memory_tl += g;
   assert(zmalloc_heap);
-  void* ptr = mi_heap_malloc(zmalloc_heap, g);
-  assert(mi_usable_size(ptr) == g);
+  void* res = mi_heap_malloc(zmalloc_heap, size);
+  size_t uss = mi_usable_size(res);
+  *usable = uss;
 
-  return ptr;
+  zmalloc_used_memory_tl += uss;
+
+  return res;
 }
 
 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;
+  ssize_t prev = mi_usable_size(ptr);
+
+  void* res = mi_heap_realloc(zmalloc_heap, ptr, size);
+  ssize_t uss = mi_usable_size(res);
+  *usable = uss;
+  zmalloc_used_memory_tl += (uss - prev);
 
-  zmalloc_used_memory_tl += (g - prev);
-  void* res = mi_heap_realloc(zmalloc_heap, ptr, g);
-  // does not hold, say when prev = 16 and size = 6. mi_malloc does not shrink in this case.
-  // assert(mi_usable_size(res) == g);
   return res;
 }
 
@@ -87,8 +85,9 @@ size_t znallocx(size_t size) {
 }
 
 void zfree_size(void* ptr, size_t size) {
-  zmalloc_used_memory_tl -= size;
-  mi_free_size(ptr, size);
+  ssize_t uss = mi_usable_size(ptr);
+  zmalloc_used_memory_tl -= uss;
+  mi_free_size(ptr, uss);
 }
 
 void* ztrymalloc(size_t size) {
diff --git a/src/server/rdb_load.cc b/src/server/rdb_load.cc
index 3192df864119..d44615331566 100644
--- a/src/server/rdb_load.cc
+++ b/src/server/rdb_load.cc
@@ -395,8 +395,8 @@ void RdbLoader::OpaqueObjLoader::CreateHMap(const LoadTrace* ltrace) {
     }
 
     lp = lpShrinkToFit(lp);
-    robj* o = createObject(OBJ_HASH, lp);
-    o->encoding = OBJ_ENCODING_LISTPACK;
+    res = createObject(OBJ_HASH, lp);
+    res->encoding = OBJ_ENCODING_LISTPACK;
   } else {
     dict* hmap = dictCreate(&hashDictType);
 
diff --git a/src/server/rdb_test.cc b/src/server/rdb_test.cc
index 1466d04fe2fb..14e88daa7071 100644
--- a/src/server/rdb_test.cc
+++ b/src/server/rdb_test.cc
@@ -5,6 +5,7 @@
 
 extern "C" {
 #include "redis/crc64.h"
+#include "redis/redis_aux.h"
 #include "redis/zmalloc.h"
 }
 
@@ -257,4 +258,15 @@ TEST_F(RdbTest, SaveManyDbs) {
   }
 }
 
+TEST_F(RdbTest, HMapBugs) {
+  // Force OBJ_ENCODING_HT encoding.
+  server.hash_max_listpack_value = 0;
+  Run({"hset", "hmap1", "key1", "val", "key2", "val2"});
+  Run({"hset", "hmap2", "key1", string(690557, 'a')});
+
+  server.hash_max_listpack_value = 32;
+  Run({"debug", "reload"});
+  EXPECT_EQ(2, CheckedInt({"hlen", "hmap1"}));
+}
+
 }  // namespace dfly