Optimize zset memory usage by embedding element in skiplist (#2508)

By default, when the number of elements in a zset exceeds 128, the
underlying data structure adopts a skiplist. We can reduce memory usage
by embedding elements into the skiplist nodes. Change the `zskiplistNode`
memory layout as follows:

```
Before
                 +-------------+
         +-----> | element-sds |
         |       +-------------+
         |
 +------------------+-------+------------------+---------+-----+---------+
 | element--pointer | score | backward-pointer | level-0 | ... | level-N |
 +------------------+-------+------------------+---------+-----+---------+



After
 +-------+------------------+---------+-----+---------+-------------+
 + score | backward-pointer | level-0 | ... | level-N | element-sds |
 +-------+------------------+---------+-----+---------+-------------+
```

Before the embedded SDS representation, we include one byte representing
the size of the SDS header, i.e. the offset into the SDS representation
where that actual string starts.

The memory saving is therefore one pointer minus one byte = 7 bytes per
element, regardless of other factors such as element size or number of
elements.

### Benchmark step

I generated the test data using the following lua script && cli command.
And check memory usage using the `info` command.

**lua script**
```
local start_idx = tonumber(ARGV[1])
local end_idx = tonumber(ARGV[2])
local elem_count = tonumber(ARGV[3])

for i = start_idx, end_idx do
    local key = "zset:" .. string.format("%012d", i)
    local members = {}

    for j = 0, elem_count - 1 do
        table.insert(members, j)
        table.insert(members, "member:" .. j)
    end

    redis.call("ZADD", key, unpack(members))
end

return "OK: Created " .. (end_idx - start_idx + 1) .. " zsets"
```

**valkey-cli command**
`valkey-cli EVAL "$(catcreate_zsets.lua)" 0 0 100000
${ZSET_ELEMENT_NUM}`

### Benchmark result
|number of elements in a zset | memory usage before optimization |
memory usage after optimization | change |
|-------|-------|-------|-------|
| 129 | 1047MB | 943MB | -9.9% |
| 256 |  2010MB|  1803MB| -10.3%|
| 512 |  3904MB|3483MB| -10.8%|

---------

Signed-off-by: chzhoo <[email protected]>
Co-authored-by: Viktor Söderqvist <[email protected]>

Author: chzhoo

src/aof.c

@@ -1922,7 +1922,7 @@ int rewriteSortedSetObject(rio *r, robj *key, robj *o) {
                     return 0;
                 }
             }
-            sds ele = node->ele;
+            sds ele = zslGetNodeElement(node);
             if (!rioWriteBulkDouble(r, node->score) || !rioWriteBulkString(r, ele, sdslen(ele))) {
                 hashtableResetIterator(&iter);
                 return 0;

src/db.c

@@ -1054,7 +1054,7 @@ void hashtableScanCallback(void *privdata, void *entry) {
         key = (sds)entry;
     } else if (o->type == OBJ_ZSET) {
         zskiplistNode *node = (zskiplistNode *)entry;
-        key = node->ele;
+        key = zslGetNodeElement(node);
         /* zset data is copied after filtering by key */
     } else if (o->type == OBJ_HASH) {
         key = entryGetField(entry);
@@ -1077,7 +1077,7 @@ void hashtableScanCallback(void *privdata, void *entry) {
     if (o->type == OBJ_ZSET) {
         /* zset data is copied */
         zskiplistNode *node = (zskiplistNode *)entry;
-        key = sdsdup(node->ele);
+        key = sdsdup(zslGetNodeElement(node));
         if (!data->only_keys) {
             char buf[MAX_LONG_DOUBLE_CHARS];
             int len = ld2string(buf, sizeof(buf), node->score, LD_STR_AUTO);

src/debug.c

@@ -215,7 +215,8 @@ void xorObjectDigest(serverDb *db, robj *keyobj, unsigned char *digest, robj *o)
                 const int len = fpconv_dtoa(node->score, buf);
                 buf[len] = '\0';
                 memset(eledigest, 0, 20);
-                mixDigest(eledigest, node->ele, sdslen(node->ele));
+                sds ele = zslGetNodeElement(node);
+                mixDigest(eledigest, ele, sdslen(ele));
                 mixDigest(eledigest, buf, strlen(buf));
                 xorDigest(digest, eledigest, 20);
             }

src/defrag.c

@@ -266,12 +266,11 @@ static void activeDefragZsetNode(void *privdata, void *entry_ref) {
     zskiplistNode **node_ref = (zskiplistNode **)entry_ref;
     zskiplistNode *node = *node_ref;
 
-    /* defragment node internals */
-    sds newsds = activeDefragSds(node->ele);
-    if (newsds) node->ele = newsds;
+    size_t allocation_size;
+    zskiplistNode *newnode = activeDefragAllocWithoutFree(node, &allocation_size);
+    if (newnode == NULL) return;
 
     const double score = node->score;
-    const sds ele = node->ele;
 
     /* find skiplist pointers that need to be updated if we end up moving the
      * skiplist node. */
@@ -283,7 +282,7 @@ static void activeDefragZsetNode(void *privdata, void *entry_ref) {
         zskiplistNode *next = x->level[i].forward;
         while (next &&
                (next->score < score ||
-                (next->score == score && sdscmp(next->ele, ele) < 0))) {
+                (next->score == score && next != node))) {
             x = next;
             next = x->level[i].forward;
         }
@@ -292,12 +291,9 @@ static void activeDefragZsetNode(void *privdata, void *entry_ref) {
     /* should have arrived at intended node */
     serverAssert(x->level[0].forward == node);
 
-    /* try to defrag the skiplist record itself */
-    zskiplistNode *newnode = activeDefragAlloc(node);
-    if (newnode) {
-        zslUpdateNode(zsl, node, newnode, update);
-        *node_ref = newnode; /* update hashtable pointer */
-    }
+    zslUpdateNode(zsl, node, newnode, update);
+    *node_ref = newnode; /* update hashtable pointer */
+    allocatorDefragFree(node, allocation_size);
 }
 
 #define DEFRAG_SDS_DICT_NO_VAL 0

src/geo.c

@@ -322,7 +322,8 @@ int geoGetPointsInRange(robj *zobj, double min, double max, GeoShape *shape, geo
             if (!zslValueLteMax(ln->score, &range)) break;
             if (geoWithinShape(shape, ln->score, xy, &distance) == C_OK) {
                 /* Append the new element. */
-                geoArrayAppend(ga, xy, distance, ln->score, sdsdup(ln->ele));
+                sds ele = zslGetNodeElement(ln);
+                geoArrayAppend(ga, xy, distance, ln->score, sdsdup(ele));
             }
             if (ga->used && limit && ga->used >= limit) break;
             ln = ln->level[0].forward;
@@ -825,6 +826,7 @@ void georadiusGeneric(client *c, int srcKeyIndex, int flags) {
             totelelen += elelen;
             znode = zslInsert(zs->zsl, score, gp->member);
             serverAssert(hashtableAdd(zs->ht, znode));
+            sdsfree(gp->member);
             gp->member = NULL;
         }
 

src/module.c

@@ -5165,7 +5165,8 @@ ValkeyModuleString *VM_ZsetRangeCurrentElement(ValkeyModuleKey *key, double *sco
     } else if (key->value->encoding == OBJ_ENCODING_SKIPLIST) {
         zskiplistNode *ln = key->u.zset.current;
         if (score) *score = ln->score;
-        str = createStringObject(ln->ele, sdslen(ln->ele));
+        sds ele = zslGetNodeElement(ln);
+        str = createStringObject(ele, sdslen(ele));
     } else {
         serverPanic("Unsupported zset encoding");
     }
@@ -5222,7 +5223,7 @@ int VM_ZsetRangeNext(ValkeyModuleKey *key) {
                 key->u.zset.er = 1;
                 return 0;
             } else if (key->u.zset.type == VALKEYMODULE_ZSET_RANGE_LEX) {
-                if (!zslLexValueLteMax(next->ele, &key->u.zset.lrs)) {
+                if (!zslLexValueLteMax(zslGetNodeElement(next), &key->u.zset.lrs)) {
                     key->u.zset.er = 1;
                     return 0;
                 }
@@ -5284,7 +5285,7 @@ int VM_ZsetRangePrev(ValkeyModuleKey *key) {
                 key->u.zset.er = 1;
                 return 0;
             } else if (key->u.zset.type == VALKEYMODULE_ZSET_RANGE_LEX) {
-                if (!zslLexValueGteMin(prev->ele, &key->u.zset.lrs)) {
+                if (!zslLexValueGteMin(zslGetNodeElement(prev), &key->u.zset.lrs)) {
                     key->u.zset.er = 1;
                     return 0;
                 }
@@ -11418,7 +11419,7 @@ static void moduleScanKeyHashtableCallback(void *privdata, void *entry) {
         /* no value */
     } else if (o->type == OBJ_ZSET) {
         zskiplistNode *node = (zskiplistNode *)entry;
-        key = node->ele;
+        key = zslGetNodeElement(node);
         value = createStringObjectFromLongDouble(node->score, 0);
     } else if (o->type == OBJ_HASH) {
         key = entryGetField(entry);

src/object.c

@@ -667,8 +667,9 @@ void dismissZsetObject(robj *o, size_t size_hint) {
         if (size_hint / zsl->length >= server.page_size) {
             zskiplistNode *zn = zsl->tail;
             while (zn != NULL) {
-                dismissSds(zn->ele);
-                zn = zn->backward;
+                zskiplistNode *next = zn->backward;
+                dismissMemory(zn, 0);
+                zn = next;
             }
         }
 
@@ -1190,7 +1191,6 @@ size_t objectComputeSize(robj *key, robj *o, size_t sample_size, int dbid) {
             asize += sizeof(zset) + sizeof(zskiplist) +
                      hashtableMemUsage(ht) + zmalloc_size(zsl->header);
             while (znode != NULL && samples < sample_size) {
-                elesize += sdsAllocSize(znode->ele);
                 elesize += zmalloc_size(znode);
                 samples++;
                 znode = znode->level[0].forward;

src/rdb.c

@@ -959,7 +959,8 @@ ssize_t rdbSaveObject(rio *rdb, robj *o, robj *key, int dbid) {
              * O(1) instead of O(log(N)). */
             zskiplistNode *zn = zsl->tail;
             while (zn != NULL) {
-                if ((n = rdbSaveRawString(rdb, (unsigned char *)zn->ele, sdslen(zn->ele))) == -1) {
+                sds ele = zslGetNodeElement(zn);
+                if ((n = rdbSaveRawString(rdb, (unsigned char *)ele, sdslen(ele))) == -1) {
                     return -1;
                 }
                 nwritten += n;
@@ -2095,6 +2096,7 @@ robj *rdbLoadObject(int rdbtype, rio *rdb, sds key, int dbid, int *error) {
             totelelen += sdslen(sdsele);
 
             znode = zslInsert(zs->zsl, score, sdsele);
+            sdsfree(sdsele);
             if (!hashtableAdd(zs->ht, znode)) {
                 rdbReportCorruptRDB("Duplicate zset fields detected");
                 decrRefCount(o);

src/server.c

@@ -561,7 +561,7 @@ hashtableType setHashtableType = {
 
 const void *zsetHashtableGetKey(const void *element) {
     const zskiplistNode *node = element;
-    return node->ele;
+    return zslGetNodeElement(node);
 }
 
 /* Sorted sets hash (note: a skiplist is used in addition to the hash table) */

src/server.h

@@ -1423,7 +1423,6 @@ struct sharedObjectsStruct {
 
 /* ZSETs use a specialized version of Skiplists */
 typedef struct zskiplistNode {
-    sds ele;
     double score;
     struct zskiplistNode *backward;
     struct zskiplistLevel {
@@ -1434,6 +1433,7 @@ typedef struct zskiplistNode {
          * So we use it in order to hold the height of the node, which is the number of levels. */
         unsigned long span;
     } level[];
+    /* After the level[], sds header length (1 byte) and an embedded sds element are stored. */
 } zskiplistNode;
 
 typedef struct zskiplist {
@@ -3275,8 +3275,9 @@ typedef struct {
 
 zskiplist *zslCreate(void);
 void zslFree(zskiplist *zsl);
-zskiplistNode *zslInsert(zskiplist *zsl, double score, sds ele);
+zskiplistNode *zslInsert(zskiplist *zsl, double score, const_sds ele);
 zskiplistNode *zslNthInRange(zskiplist *zsl, zrangespec *range, long n, long *rank);
+sds zslGetNodeElement(const zskiplistNode *x);
 double zzlGetScore(unsigned char *sptr);
 void zzlNext(unsigned char *zl, unsigned char **eptr, unsigned char **sptr);
 void zzlPrev(unsigned char *zl, unsigned char **eptr, unsigned char **sptr);