fix(core): Account for size change on json object defragment (#6053)

core: Record change in size after defrag

Signed-off-by: Abhijat Malviya <abhijat@dragonflydb.io>

Author: abhijat

src/core/compact_object.cc

@@ -1610,11 +1610,25 @@ MemoryResource* CompactObj::memory_resource() {
 }
 
 bool CompactObj::JsonConsT::DefragIfNeeded(PageUsage* page_usage) {
-  if (JsonType* old = json_ptr; ShouldDefragment(page_usage)) {
+  if (ShouldDefragment(page_usage)) {
+    const MiMemoryResource* mr = static_cast<MiMemoryResource*>(memory_resource());
+
+    const int64_t before = static_cast<int64_t>(mr->used());
+    DCHECK_GE(before, 0) << "Memory usage is more than int64_t max value";
+
+    JsonType* old = json_ptr;
     json_ptr = AllocateMR<JsonType>(DeepCopyJSON(old));
     DeleteMR<JsonType>(old);
+
+    const int64_t after = static_cast<int64_t>(mr->used());
+    DCHECK_GE(after, 0) << "Memory usage is more than int64_t max value";
+
+    if (const int64_t delta = after - before; delta != 0) {
+      bytes_used = UpdateSize(bytes_used, delta);
+    }
     return true;
   }
+
   return false;
 }
 

src/core/page_usage_stats_test.cc

@@ -194,17 +194,31 @@ TEST_F(PageUsageStatsTest, JSONCons) {
   // still fail. This is because freeing the compact object code path takes the wrong branch based
   // on encoding. The flat encoding was tested manually adjusting this same test with changed
   // encoding.
-  std::string_view data{R"#({"data": "some", "count": 1, "checked": false})#"};
+  std::string data = R"({"contents":[)";
+  for (size_t i = 0; i < 1000; ++i) {
+    const auto si = std::to_string(i);
+    data += R"({"id":)" + si + R"(,"class":"v___)" + si + R"("})";
+    if (i < 999) {
+      data += ",";
+    }
+  }
+  data += R"(], "data": "some", "count": 1, "checked": false})";
+
+  auto* mr = static_cast<MiMemoryResource*>(CompactObj::memory_resource());
+  size_t before = mr->used();
 
   auto parsed = ParseJsonUsingShardHeap(data);
   EXPECT_TRUE(parsed.has_value());
 
   c_obj_.SetJson(std::move(parsed.value()));
+  c_obj_.SetJsonSize(mr->used() - before);
+  EXPECT_GT(c_obj_.MallocUsed(), 0);
 
   PageUsage p{CollectPageStats::YES, 0.1};
   p.SetForceReallocate(true);
 
   c_obj_.DefragIfNeeded(&p);
+  EXPECT_GT(c_obj_.MallocUsed(), 0);
 
   const auto stats = p.CollectedStats();
   EXPECT_GT(stats.pages_scanned, 0);

src/server/engine_shard.cc

@@ -351,14 +351,20 @@ std::optional<CollectedPageStats> EngineShard::DoDefrag(CollectPageStats collect
   uint64_t attempts = 0;
 
   PageUsage page_usage{collect_page_stats, threshold};
+  DbTable* db_table = slice.GetDBTable(defrag_state_.dbid);
   do {
     cur = prime_table->Traverse(cur, [&](PrimeIterator it) {
       // for each value check whether we should move it because it
       // seats on underutilized page of memory, and if so, do it.
-      bool did = it->second.DefragIfNeeded(&page_usage);
+      const ssize_t original_size = it->second.MallocUsed();
+      const bool did = it->second.DefragIfNeeded(&page_usage);
       attempts++;
       if (did) {
         reallocations++;
+        if (const ssize_t delta = it->second.MallocUsed() - original_size;
+            delta != 0 && db_table != nullptr) {
+          db_table->stats.AddTypeMemoryUsage(it->second.ObjType(), delta);
+        }
       }
     });
     traverses_count++;