From 07b92841fe0ac0fee033cf27caf8642a28866981 Mon Sep 17 00:00:00 2001
From: Roman Gershman <roman@dragonflydb.io>
Date: Fri, 17 Jun 2022 06:38:07 +0300
Subject: [PATCH] fix(zmalloc): fix memory accounting of redis types

Signed-off-by: Roman Gershman <roman@dragonflydb.io>
---
 src/redis/zmalloc_mi.c     | 29 +++++++++++++++++++++++++----
 src/server/main_service.cc |  3 ++-
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/src/redis/zmalloc_mi.c b/src/redis/zmalloc_mi.c
index 196179b..d58c918 100644
--- a/src/redis/zmalloc_mi.c
+++ b/src/redis/zmalloc_mi.c
@@ -15,7 +15,15 @@ __thread mi_heap_t* zmalloc_heap = NULL;
 /* Allocate memory or panic */
 void* zmalloc(size_t size) {
   assert(zmalloc_heap);
-  return mi_heap_malloc(zmalloc_heap, size);
+  void* res = mi_heap_malloc(zmalloc_heap, size);
+  size_t usable = mi_usable_size(res);
+
+  // assertion does not hold. Basically mi_good_size is not a good function for
+  // doing accounting.
+  // assert(usable == mi_good_size(size));
+  zmalloc_used_memory_tl += usable;
+
+  return res;
 }
 
 void* ztrymalloc_usable(size_t size, size_t* usable) {
@@ -28,7 +36,12 @@ 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;
+
   mi_free_size(ptr, usable);
 }
 
@@ -51,7 +64,10 @@ void* zmalloc_usable(size_t size, size_t* usable) {
 
   zmalloc_used_memory_tl += g;
   assert(zmalloc_heap);
-  return mi_heap_malloc(zmalloc_heap, g);
+  void* ptr = mi_heap_malloc(zmalloc_heap, g);
+  assert(mi_usable_size(ptr) == g);
+
+  return ptr;
 }
 
 void* zrealloc_usable(void* ptr, size_t size, size_t* usable) {
@@ -60,7 +76,10 @@ void* zrealloc_usable(void* ptr, size_t size, size_t* usable) {
   *usable = g;
 
   zmalloc_used_memory_tl += (g - prev);
-  return mi_heap_realloc(zmalloc_heap, ptr, g);
+  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;
 }
 
 size_t znallocx(size_t size) {
@@ -80,7 +99,9 @@ void* ztrymalloc(size_t size) {
 void* ztrycalloc(size_t size) {
   size_t g = mi_good_size(size);
   zmalloc_used_memory_tl += g;
-  return mi_heap_calloc(zmalloc_heap, 1, size);
+  void* ptr = mi_heap_calloc(zmalloc_heap, 1, size);
+  assert(mi_usable_size(ptr) == g);
+  return ptr;
 }
 
 typedef struct Sum_s {
diff --git a/src/server/main_service.cc b/src/server/main_service.cc
index d28cb91..223d190 100644
--- a/src/server/main_service.cc
+++ b/src/server/main_service.cc
@@ -432,7 +432,8 @@ void Service::DispatchCommand(CmdArgList args, facade::ConnectionContext* cntx)
     return;
   }
 
-  if (etl.gstate() == GlobalState::LOADING || etl.gstate() == GlobalState::SHUTTING_DOWN) {
+  if ((etl.gstate() == GlobalState::LOADING && (cid->opt_mask() & CO::LOADING) == 0) ||
+      etl.gstate() == GlobalState::SHUTTING_DOWN) {
     string err = StrCat("Can not execute during ", GlobalStateName(etl.gstate()));
     (*cntx)->SendError(err);
     return;