diff --git a/src/core/compact_object.cc b/src/core/compact_object.cc index b3f76cb6a..1e5aa0bcd 100644 --- a/src/core/compact_object.cc +++ b/src/core/compact_object.cc @@ -138,39 +138,74 @@ inline void FreeObjZset(unsigned encoding, void* ptr) { } } +pair 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 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 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 DefragHash(MemoryResource* mr, unsigned encoding, void* ptr, float ratio) { +pair 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; - 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}; - }; + return DefragListPack((uint8_t*)ptr, ratio); + } // 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) - realloced |= it.ReallocIfNeeded(ratio); - - return {sm, realloced}; + return DefragStrMap2((StringMap*)ptr, ratio); } default: ABSL_UNREACHABLE(); - }; + } +} + +pair 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; } diff --git a/src/server/set_family_test.cc b/src/server/set_family_test.cc index d9750bf53..33381bc26 100644 --- a/src/server/set_family_test.cc +++ b/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