core: Record change in json object size on defragmentation

Signed-off-by: Abhijat Malviya <[email protected]>

Author: abhijat

src/core/compact_object.cc

@@ -1623,21 +1623,31 @@ MemoryResource* CompactObj::memory_resource() {
 }
 
 bool CompactObj::JsonConsT::DefragIfNeeded(PageUsage* page_usage) {
-  if (JsonType* old = json_ptr; ShouldDefragment(page_usage)) {
+  JsonType* old = json_ptr;
+  if (auto [should_defragment, old_size] = ShouldDefragment(page_usage); should_defragment) {
+    const MiMemoryResource* mr = static_cast<MiMemoryResource*>(CompactObj::memory_resource());
+    const ssize_t before = mr->used();
     json_ptr = AllocateMR<JsonType>(DeepCopyJSON(old, memory_resource()));
     DeleteMR<JsonType>(old);
+    const ssize_t after = mr->used();
+    if (const ssize_t delta = after - before; delta != 0) {
+      bytes_used += delta;
+      page_usage->RecordChangeInSize(OBJ_JSON, delta);
+    }
     return true;
   }
+
   return false;
 }
 
-bool CompactObj::JsonConsT::ShouldDefragment(PageUsage* page_usage) const {
+std::pair<bool, ssize_t> CompactObj::JsonConsT::ShouldDefragment(PageUsage* page_usage) const {
   bool should_defragment = false;
-  json_ptr->compute_memory_size([&page_usage, &should_defragment](const void* p) {
-    should_defragment |= page_usage->IsPageForObjectUnderUtilized(const_cast<void*>(p));
-    return 0;
-  });
-  return should_defragment;
+  const size_t old_size =
+      json_ptr->compute_memory_size([&page_usage, &should_defragment](const void* p) {
+        should_defragment |= page_usage->IsPageForObjectUnderUtilized(const_cast<void*>(p));
+        return p ? mi_usable_size(p) : 0;
+      });
+  return {should_defragment, old_size};
 }
 
 bool CompactObj::FlatJsonT::DefragIfNeeded(PageUsage* page_usage) {

src/core/compact_object.h

@@ -497,7 +497,7 @@ class CompactObj {
 
     // Computes if the contained object should be defragmented, by examining pointers within it and
     // returning true if any of them reside in an underutilized page.
-    bool ShouldDefragment(PageUsage* page_usage) const;
+    std::pair<bool, ssize_t> ShouldDefragment(PageUsage* page_usage) const;
   };
 
   struct FlatJsonT {

src/server/engine_shard.cc

@@ -268,6 +268,7 @@ std::optional<CollectedPageStats> EngineShard::DoDefrag(CollectPageStats collect
   uint64_t attempts = 0;
 
   PageUsage page_usage{collect_page_stats, threshold};
+  DbTableStats& table_stats = slice.GetDBTable(defrag_state_.dbid)->stats;
   do {
     cur = prime_table->Traverse(cur, [&](PrimeIterator it) {
       // for each value check whether we should move it because it
@@ -276,6 +277,9 @@ std::optional<CollectedPageStats> EngineShard::DoDefrag(CollectPageStats collect
       attempts++;
       if (did) {
         reallocations++;
+        if (const auto [type, delta] = page_usage.FetchAndResetDelta(); delta != 0) {
+          table_stats.AddTypeMemoryUsage(type, delta);
+        }
       }
     });
     traverses_count++;