chore: add defrag logic for zsets (#3836)

* add defrag logic for zsets
* add tests

---------

Signed-off-by: kostas <kostas@dragonflydb.io>
This commit is contained in:
Kostas Kyrimis 2024-10-23 10:56:56 +03:00 committed by GitHub
parent ea9dc9c454
commit c5a8008348
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 210 additions and 12 deletions

View file

@ -171,6 +171,11 @@ pair<void*, bool> DefragIntSet(intset* is, float ratio) {
return {replacement, true};
}
pair<void*, bool> DefragSortedMap(detail::SortedMap* sm, float ratio) {
const bool reallocated = sm->DefragIfNeeded(ratio);
return {sm, reallocated};
}
// 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.
@ -208,6 +213,23 @@ pair<void*, bool> DefragSet(unsigned encoding, void* ptr, float ratio) {
}
}
pair<void*, bool> DefragZSet(unsigned encoding, void* ptr, float ratio) {
switch (encoding) {
// Listpack is stored as a single contiguous array
case OBJ_ENCODING_LISTPACK: {
return DefragListPack((uint8_t*)ptr, ratio);
}
// SKIPLIST really means ScoreMap
case OBJ_ENCODING_SKIPLIST: {
return DefragSortedMap((detail::SortedMap*)ptr, ratio);
}
default:
ABSL_UNREACHABLE();
}
}
inline void FreeObjStream(void* ptr) {
freeStream((stream*)ptr);
}
@ -420,19 +442,23 @@ void RobjWrapper::SetString(string_view s, MemoryResource* mr) {
}
bool RobjWrapper::DefragIfNeeded(float ratio) {
auto do_defrag = [this, ratio](auto defrag_fun) mutable {
auto [new_ptr, realloced] = defrag_fun(encoding_, inner_obj_, ratio);
inner_obj_ = new_ptr;
return realloced;
};
if (type() == OBJ_STRING) {
if (zmalloc_page_is_underutilized(inner_obj(), ratio)) {
ReallocateString(tl.local_mr);
return true;
}
} else if (type() == OBJ_HASH) {
auto [new_ptr, realloced] = DefragHash(encoding_, inner_obj_, ratio);
inner_obj_ = new_ptr;
return realloced;
return do_defrag(DefragHash);
} else if (type() == OBJ_SET) {
auto [new_ptr, realloced] = DefragSet(encoding_, inner_obj_, ratio);
inner_obj_ = new_ptr;
return realloced;
return do_defrag(DefragSet);
} else if (type() == OBJ_ZSET) {
return do_defrag(DefragZSet);
}
return false;
}
@ -440,11 +466,11 @@ bool RobjWrapper::DefragIfNeeded(float ratio) {
int RobjWrapper::ZsetAdd(double score, sds ele, int in_flags, int* out_flags, double* newscore) {
// copied from zsetAdd for listpack only.
/* Turn options into simple to check vars. */
int incr = (in_flags & ZADD_IN_INCR) != 0;
int nx = (in_flags & ZADD_IN_NX) != 0;
int xx = (in_flags & ZADD_IN_XX) != 0;
int gt = (in_flags & ZADD_IN_GT) != 0;
int lt = (in_flags & ZADD_IN_LT) != 0;
bool incr = (in_flags & ZADD_IN_INCR) != 0;
bool nx = (in_flags & ZADD_IN_NX) != 0;
bool xx = (in_flags & ZADD_IN_XX) != 0;
bool gt = (in_flags & ZADD_IN_GT) != 0;
bool lt = (in_flags & ZADD_IN_LT) != 0;
*out_flags = 0; /* We'll return our response flags. */
double curscore;