chore: add defrag logic for sets (#3833)

* add defrag logic for sets
* add small test


---------

Signed-off-by: kostas <kostas@dragonflydb.io>

src/core/compact_object.cc

@@ -138,39 +138,74 @@ inline void FreeObjZset(unsigned encoding, void* ptr) {
   }
 }
 
+pair<void*, bool> DefragStrMap2(StringMap* sm, float ratio) {
+  bool realloced = false;
+
+  for (auto it = sm->begin(); it != sm->end(); ++it)
+    realloced |= it.ReallocIfNeeded(ratio);
+
+  return {sm, realloced};
+}
+
+pair<void*, bool> DefragListPack(uint8_t* lp, float ratio) {
+  if (!zmalloc_page_is_underutilized(lp, ratio))
+    return {lp, false};
+
+  size_t lp_bytes = lpBytes(lp);
+  uint8_t* replacement = lpNew(lpBytes(lp));
+  memcpy(replacement, lp, lp_bytes);
+  lpFree(lp);
+
+  return {replacement, true};
+}
+
+pair<void*, bool> DefragIntSet(intset* is, float ratio) {
+  if (!zmalloc_page_is_underutilized(is, ratio))
+    return {is, false};
+
+  const size_t blob_len = intsetBlobLen(is);
+  intset* replacement = (intset*)zmalloc(blob_len);
+  memcpy(replacement, is, blob_len);
+
+  zfree(is);
+  return {replacement, true};
+}
+
 // Iterates over allocations of internal hash data structures and re-allocates
 // them if their pages are underutilized.
 // Returns pointer to new object ptr and whether any re-allocations happened.
-pair<void*, bool> DefragHash(MemoryResource* mr, unsigned encoding, void* ptr, float ratio) {
+pair<void*, bool> DefragHash(unsigned encoding, void* ptr, float ratio) {
   switch (encoding) {
     // Listpack is stored as a single contiguous array
     case kEncodingListPack: {
-      uint8_t* lp = (uint8_t*)ptr;
+      return DefragListPack((uint8_t*)ptr, ratio);
-      if (!zmalloc_page_is_underutilized(lp, ratio))
+    }
-        return {lp, false};
-
-      size_t lp_bytes = lpBytes(lp);
-      uint8_t* replacement = lpNew(lpBytes(lp));
-      memcpy(replacement, lp, lp_bytes);
-      lpFree(lp);
-
-      return {replacement, true};
-    };
 
     // StringMap supports re-allocation of it's internal nodes
     case kEncodingStrMap2: {
-      bool realloced = false;
+      return DefragStrMap2((StringMap*)ptr, ratio);
-
-      StringMap* sm = (StringMap*)ptr;
-      for (auto it = sm->begin(); it != sm->end(); ++it)
-        realloced |= it.ReallocIfNeeded(ratio);
-
-      return {sm, realloced};
     }
 
     default:
       ABSL_UNREACHABLE();
-  };
+  }
+}
+
+pair<void*, bool> DefragSet(unsigned encoding, void* ptr, float ratio) {
+  switch (encoding) {
+    // Int sets have flat storage
+    case kEncodingIntSet: {
+      return DefragIntSet((intset*)ptr, ratio);
+    }
+
+    // StringMap supports re-allocation of it's internal nodes
+    case kEncodingStrMap2: {
+      return DefragStrMap2((StringMap*)ptr, ratio);
+    }
+
+    default:
+      ABSL_UNREACHABLE();
+  }
 }
 
 inline void FreeObjStream(void* ptr) {
@@ -391,7 +426,11 @@ bool RobjWrapper::DefragIfNeeded(float ratio) {
       return true;
     }
   } else if (type() == OBJ_HASH) {
-    auto [new_ptr, realloced] = DefragHash(tl.local_mr, encoding_, inner_obj_, ratio);
+    auto [new_ptr, realloced] = DefragHash(encoding_, inner_obj_, ratio);
+    inner_obj_ = new_ptr;
+    return realloced;
+  } else if (type() == OBJ_SET) {
+    auto [new_ptr, realloced] = DefragSet(encoding_, inner_obj_, ratio);
     inner_obj_ = new_ptr;
     return realloced;
   }

src/server/set_family_test.cc

@@ -9,6 +9,10 @@
 #include "facade/facade_test.h"
 #include "server/command_registry.h"
 #include "server/test_utils.h"
+extern "C" {
+#include "redis/intset.h"
+#include "redis/zmalloc.h"
+}
 
 using namespace testing;
 using namespace std;
@@ -349,4 +353,29 @@ TEST_F(SetFamilyTest, SScan) {
   EXPECT_THAT(vec.size(), 0);
 }
 
+TEST_F(SetFamilyTest, IntSetMemcpy) {
+  // This logic is used in CompactObject::DefragIntSet
+  intset* original = intsetNew();
+  uint8_t success = 0;
+  for (int i = 0; i < 250; ++i) {
+    original = intsetAdd(original, i, &success);
+    ASSERT_THAT(success, 1);
+  }
+  const size_t blob_len = intsetBlobLen(original);
+  intset* replacement = (intset*)zmalloc(blob_len);
+  memcpy(replacement, original, blob_len);
+
+  ASSERT_THAT(original->encoding, replacement->encoding);
+  ASSERT_THAT(original->length, replacement->length);
+
+  for (int i = 0; i < 250; ++i) {
+    int64_t value;
+    ASSERT_THAT(intsetGet(replacement, i, &value), 1);
+    ASSERT_THAT(value, i);
+  }
+
+  zfree(original);
+  zfree(replacement);
+}
+
 }  // namespace dfly