fix: fix defrag stats (#1740)

* fix: fix defrag stats

Signed-off-by: Vladislav <[email protected]>

Author: dranikpg

src/core/compact_object.cc

@@ -141,31 +141,34 @@ inline void FreeObjZset(unsigned encoding, void* ptr) {
   }
 }
 
-// Iterates over allocations of internal hash data structed and re-allocates
+// Iterates over allocations of internal hash data structures and re-allocates
 // them if their pages are underutilized.
-void* DefragHash(MemoryResource* mr, unsigned encoding, void* ptr, float ratio) {
+// Returns pointer to new object ptr and whether any re-allocations happened.
+pair<void*, bool> DefragHash(MemoryResource* mr, unsigned encoding, void* ptr, float ratio) {
   switch (encoding) {
     // Listpack is stored as a single contiguous array
     case kEncodingListPack: {
       uint8_t* lp = (uint8_t*)ptr;
       if (!zmalloc_page_is_underutilized(lp, ratio))
-        return lp;
+        return {lp, false};
 
       size_t lp_bytes = lpBytes(lp);
       uint8_t* replacement = lpNew(lpBytes(lp));
       memcpy(replacement, lp, lp_bytes);
       lpFree(lp);
 
-      return replacement;
+      return {replacement, true};
     };
 
     // StringMap supports re-allocation of it's internal nodes
     case kEncodingStrMap2: {
+      bool realloced = false;
+
       StringMap* sm = (StringMap*)ptr;
       for (auto it = sm->begin(); it != sm->end(); ++it)
-        it.ReallocIfNeeded(ratio);
+        realloced |= it.ReallocIfNeeded(ratio);
 
-      return sm;
+      return {sm, realloced};
     }
 
     default:
@@ -394,10 +397,13 @@ void RobjWrapper::SetString(string_view s, MemoryResource* mr) {
 bool RobjWrapper::DefragIfNeeded(float ratio) {
   if (type() == OBJ_STRING) {
     if (zmalloc_page_is_underutilized(inner_obj(), ratio)) {
-      return Reallocate(tl.local_mr);
+      ReallocateString(tl.local_mr);
+      return true;
     }
   } else if (type() == OBJ_HASH) {
-    inner_obj_ = DefragHash(tl.local_mr, encoding_, inner_obj_, ratio);
+    auto [new_ptr, realloced] = DefragHash(tl.local_mr, encoding_, inner_obj_, ratio);
+    inner_obj_ = new_ptr;
+    return realloced;
   }
   return false;
 }
@@ -487,12 +493,12 @@ int RobjWrapper::ZsetAdd(double score, sds ele, int in_flags, int* out_flags, do
   return ss->Add(score, ele, in_flags, out_flags, newscore);
 }
 
-bool RobjWrapper::Reallocate(MemoryResource* mr) {
+void RobjWrapper::ReallocateString(MemoryResource* mr) {
+  DCHECK_EQ(type(), OBJ_STRING);
   void* old_ptr = inner_obj_;
   inner_obj_ = mr->allocate(sz_, kAlignSize);
   memcpy(inner_obj_, old_ptr, sz_);
   mr->deallocate(old_ptr, 0, kAlignSize);
-  return true;
 }
 
 void RobjWrapper::Init(unsigned type, unsigned encoding, void* inner) {

src/core/compact_object.h

@@ -59,13 +59,16 @@ class RobjWrapper {
     return std::string_view{reinterpret_cast<char*>(inner_obj_), sz_};
   }
 
+  // Try reducing memory fragmentation by re-allocating values from underutilized pages.
+  // Returns true if re-allocated.
   bool DefragIfNeeded(float ratio);
 
   // as defined in zset.h
   int ZsetAdd(double score, char* ele, int in_flags, int* out_flags, double* newscore);
 
  private:
-  bool Reallocate(MemoryResource* mr);
+  void ReallocateString(MemoryResource* mr);
+
   size_t InnerObjMallocUsed() const;
   void MakeInnerRoom(size_t current_cap, size_t desired, MemoryResource* mr);
 

src/core/compact_object_test.cc

@@ -634,11 +634,11 @@ TEST_F(CompactObjectTest, DefragHash) {
 
   // Trigger re-allocation
   cobj_.InitRobj(OBJ_HASH, kEncodingListPack, target_lp);
-  cobj_.DefragIfNeeded(0.8);
+  ASSERT_TRUE(cobj_.DefragIfNeeded(0.8));
 
   // Check the pointer changes as the listpack needed defragmentation
   auto lp = (uint8_t*)cobj_.RObjPtr();
-  CHECK_NE(lp, target_lp);
+  EXPECT_NE(lp, target_lp) << "must have changed due to realloc";
 
   uint8_t* fptr = lpFirst(lp);
   for (size_t i = 0; i < 100; i++) {

src/core/string_map.cc

@@ -101,14 +101,16 @@ sds StringMap::Find(std::string_view key) {
   return GetValue(str);
 }
 
-sds StringMap::ReallocIfNeeded(void* obj, float ratio) {
+pair<sds, bool> StringMap::ReallocIfNeeded(void* obj, float ratio) {
   sds key = (sds)obj;
   size_t key_len = sdslen(key);
 
   auto* value_ptr = key + key_len + 1;
   uint64_t value_tag = absl::little_endian::Load64(value_ptr);
   sds value = (sds)(uint64_t(value_tag) & kValMask);
 
+  bool realloced_value = false;
+
   // If the allocated value is underutilized, re-allocate it and update the pointer inside the key
   if (zmalloc_page_is_underutilized(value, ratio)) {
     size_t value_len = sdslen(value);
@@ -117,18 +119,19 @@ sds StringMap::ReallocIfNeeded(void* obj, float ratio) {
     uint64_t new_value_tag = (uint64_t(new_value) & kValMask) | (value_tag & ~kValMask);
     absl::little_endian::Store64(value_ptr, new_value_tag);
     sdsfree(value);
+    realloced_value = true;
   }
 
   if (!zmalloc_page_is_underutilized(key, ratio))
-    return key;
+    return {key, realloced_value};
 
   size_t space_size = 8 /* value ptr */ + ((value_tag & kValTtlBit) ? 4 : 0) /* optional expiry */;
 
   sds new_key = AllocSdsWithSpace(key_len, space_size);
   memcpy(new_key, key, key_len + 1 /* \0 */ + space_size);
   sdsfree(key);
 
-  return new_key;
+  return {new_key, true};
 }
 
 uint64_t StringMap::Hash(const void* obj, uint32_t cookie) const {

src/core/string_map.h

@@ -62,14 +62,18 @@ class StringMap : public DenseSet {
       return BreakToPair(ptr);
     }
 
-    void ReallocIfNeeded(float ratio) {
+    // Try reducing memory fragmentation of the value by re-allocating. Returns true if
+    // re-allocation happened.
+    bool ReallocIfNeeded(float ratio) {
       // Unwrap all links to correctly call SetObject()
       auto* ptr = curr_entry_;
       while (ptr->IsLink())
         ptr = ptr->AsLink();
 
       auto* obj = ptr->GetObject();
-      ptr->SetObject(static_cast<StringMap*>(owner_)->ReallocIfNeeded(obj, ratio));
+      auto [new_obj, realloced] = static_cast<StringMap*>(owner_)->ReallocIfNeeded(obj, ratio);
+      ptr->SetObject(new_obj);
+      return realloced;
     }
 
     iterator& operator++() {
@@ -115,8 +119,8 @@ class StringMap : public DenseSet {
 
  private:
   // Reallocate key and/or value if their pages are underutilized.
-  // Returns new object pointer (stays the same if utilization is enough).
-  sds ReallocIfNeeded(void* obj, float ratio);
+  // Returns new pointer (stays same if key utilization is enough) and if reallocation happened.
+  std::pair<sds, bool> ReallocIfNeeded(void* obj, float ratio);
 
   uint64_t Hash(const void* obj, uint32_t cookie) const final;
   bool ObjEqual(const void* left, const void* right, uint32_t right_cookie) const final;