Track quicklist memory incrementally via zmalloc_usable capture

src/listpack.c

@@ -41,6 +41,14 @@
 
 #include "listpack.h"
 #include "listpack_malloc.h"
+
+/* lp_last_alloc_size is defined as a file-local static in listpack_malloc.h
+ * (included above).  This getter exposes it to other modules (e.g. quicklist)
+ * without leaking the variable as an extern global. */
+size_t lpLastAllocSize(void) {
+    return lp_last_alloc_size;
+}
+
 #include "serverassert.h"
 #include "util.h"
 #include "config.h"
@@ -175,9 +183,13 @@ void lpFreeVoid(void *lp) {
 /* Shrink the memory to fit. */
 unsigned char *lpShrinkToFit(unsigned char *lp) {
     size_t size = lpGetTotalBytes(lp);
-    if (size < lp_malloc_size(lp)) {
+    size_t alloc_size = lp_malloc_size(lp);
+    if (size < alloc_size) {
         return lp_realloc(lp, size);
     } else {
+        /* Already at minimum allocation — no realloc needed.
+         * Update lp_last_alloc_size so callers don't see a stale value. */
+        lp_last_alloc_size = alloc_size;
         return lp;
     }
 }
@@ -777,9 +789,20 @@ unsigned char *lpInsert(unsigned char *lp,
     unsigned char *dst = lp + poff; /* May be updated after reallocation. */
 
     /* Realloc before: we need more room. */
-    if (new_listpack_bytes > old_listpack_bytes && new_listpack_bytes > lp_malloc_size(lp)) {
-        if ((lp = lp_realloc(lp, new_listpack_bytes)) == NULL) return NULL;
-        dst = lp + poff;
+    if (new_listpack_bytes > old_listpack_bytes) {
+        size_t alloc_size = lp_malloc_size(lp);
+        if (new_listpack_bytes > alloc_size) {
+            if ((lp = lp_realloc(lp, new_listpack_bytes)) == NULL) return NULL;
+            dst = lp + poff;
+        } else {
+            /* Growth fits within jemalloc slack — no realloc needed.
+             * Update lp_last_alloc_size so callers see the current value. */
+            lp_last_alloc_size = alloc_size;
+        }
+    } else if (new_listpack_bytes == old_listpack_bytes) {
+        /* Same size (e.g. replace with same-length element) — no realloc.
+         * Update lp_last_alloc_size so callers don't see a stale value. */
+        lp_last_alloc_size = lp_malloc_size(lp);
     }
 
     /* Setup the listpack relocating the elements to make the exact room
@@ -927,7 +950,10 @@ unsigned char *lpDeleteRangeWithEntry(unsigned char *lp, unsigned char **p, unsi
     unsigned char *first, *tail;
     first = tail = *p;
 
-    if (num == 0) return lp; /* Nothing to delete, return ASAP. */
+    if (num == 0) { /* Nothing to delete, return ASAP. */
+        lp_last_alloc_size = lp_malloc_size(lp);
+        return lp;
+    }
 
     /* Find the next entry to the last entry that needs to be deleted.
      * lpLength may be unreliable due to corrupt data, so we cannot
@@ -966,8 +992,14 @@ unsigned char *lpDeleteRange(unsigned char *lp, long index, unsigned long num) {
     unsigned char *p;
     uint32_t numele = lpGetNumElements(lp);
 
-    if (num == 0) return lp; /* Nothing to delete, return ASAP. */
-    if ((p = lpSeek(lp, index)) == NULL) return lp;
+    if (num == 0) { /* Nothing to delete, return ASAP. */
+        lp_last_alloc_size = lp_malloc_size(lp);
+        return lp;
+    }
+    if ((p = lpSeek(lp, index)) == NULL) {
+        lp_last_alloc_size = lp_malloc_size(lp);
+        return lp;
+    }
 
     /* If we know we're gonna delete beyond the end of the listpack, we can just move
      * the EOF marker, and there's no need to iterate through the entries,

src/listpack.h

@@ -97,5 +97,6 @@ unsigned char *
 lpNextRandom(unsigned char *lp, unsigned char *p, unsigned int *index, unsigned int remaining, int even_only);
 int lpSafeToAdd(unsigned char *lp, size_t add);
 void lpRepr(unsigned char *lp);
+size_t lpLastAllocSize(void);
 
 #endif

src/listpack_malloc.h

@@ -39,11 +39,18 @@
 #ifndef LISTPACK_ALLOC_H
 #define LISTPACK_ALLOC_H
 #include "zmalloc.h"
+
+/* File-local variable defined in listpack.c that captures the usable
+ * allocation size from the last lp_malloc / lp_realloc call.  zmalloc_usable
+ * and zrealloc_usable already compute this internally; we just stop
+ * discarding it.  Exposed to callers via lpLastAllocSize(). */
+static size_t lp_last_alloc_size = 0;
+
 /* We use zmalloc_usable/zrealloc_usable instead of zmalloc/zrealloc
  * to ensure the safe invocation of 'zmalloc_usable_size().
  * See comment in zmalloc_usable_size(). */
-#define lp_malloc(sz) zmalloc_usable(sz, NULL)
-#define lp_realloc(ptr, sz) zrealloc_usable(ptr, sz, NULL)
+#define lp_malloc(sz) zmalloc_usable(sz, &lp_last_alloc_size)
+#define lp_realloc(ptr, sz) zrealloc_usable(ptr, sz, &lp_last_alloc_size)
 #define lp_free zfree
 #define lp_malloc_size zmalloc_usable_size
 #endif

src/object.c

@@ -1356,6 +1356,163 @@ size_t objectComputeSize(robj *key, robj *o, size_t sample_size, int dbid) {
     return asize;
 }
 
+/* Same as objectComputeSize but uses tracked_size for quicklists instead of
+ * iterating through nodes. Used for testing tracked_size accuracy. */
+size_t objectComputeSizeWithTrackedSize(robj *key, robj *o, size_t sample_size, int dbid) {
+    size_t elesize = 0, samples = 0;
+    size_t asize = zmalloc_size((void *)o);
+
+    if (o->type == OBJ_STRING) {
+        if (o->encoding == OBJ_ENCODING_RAW) {
+            asize += sdsAllocSize(objectGetVal(o));
+        } else if (o->encoding != OBJ_ENCODING_INT && o->encoding != OBJ_ENCODING_EMBSTR) {
+            serverPanic("Unknown string encoding");
+        }
+    } else if (o->type == OBJ_LIST) {
+        if (o->encoding == OBJ_ENCODING_QUICKLIST) {
+            quicklist *ql = objectGetVal(o);
+            asize += ql->tracked_size;
+        } else if (o->encoding == OBJ_ENCODING_LISTPACK) {
+            asize += zmalloc_size(objectGetVal(o));
+        } else {
+            serverPanic("Unknown list encoding");
+        }
+    } else if (o->type == OBJ_SET) {
+        if (o->encoding == OBJ_ENCODING_HASHTABLE) {
+            hashtable *ht = objectGetVal(o);
+            asize += hashtableMemUsage(ht);
+
+            hashtableIterator iter;
+            hashtableInitIterator(&iter, ht, 0);
+            void *next;
+            while (hashtableNext(&iter, &next) && samples < sample_size) {
+                sds element = next;
+                elesize += sdsAllocSize(element);
+                samples++;
+            }
+            hashtableCleanupIterator(&iter);
+            if (samples) asize += (double)elesize / samples * hashtableSize(ht);
+        } else if (o->encoding == OBJ_ENCODING_INTSET) {
+            asize += zmalloc_size(objectGetVal(o));
+        } else if (o->encoding == OBJ_ENCODING_LISTPACK) {
+            asize += zmalloc_size(objectGetVal(o));
+        } else {
+            serverPanic("Unknown set encoding");
+        }
+    } else if (o->type == OBJ_ZSET) {
+        if (o->encoding == OBJ_ENCODING_LISTPACK) {
+            asize += zmalloc_size(objectGetVal(o));
+        } else if (o->encoding == OBJ_ENCODING_SKIPLIST) {
+            hashtable *ht = ((zset *)objectGetVal(o))->ht;
+            zskiplist *zsl = ((zset *)objectGetVal(o))->zsl;
+            zskiplistNode *zheader = zslGetHeader(zsl);
+            zskiplistNode *znode = zheader->level[0].forward;
+            asize += sizeof(zset) + zslGetAllocSize() + hashtableMemUsage(ht);
+            while (znode != NULL && samples < sample_size) {
+                elesize += zmalloc_size(znode);
+                samples++;
+                znode = znode->level[0].forward;
+            }
+            if (samples) asize += (double)elesize / samples * hashtableSize(ht);
+        } else {
+            serverPanic("Unknown sorted set encoding");
+        }
+    } else if (o->type == OBJ_HASH) {
+        if (o->encoding == OBJ_ENCODING_LISTPACK) {
+            asize += zmalloc_size(objectGetVal(o));
+        } else if (o->encoding == OBJ_ENCODING_HASHTABLE) {
+            hashtable *ht = objectGetVal(o);
+            hashtableIterator iter;
+            vset *volatile_fields = hashtableMetadata(ht);
+            hashtableInitIterator(&iter, ht, 0);
+            void *next;
+
+            asize += hashtableMemUsage(ht);
+            while (hashtableNext(&iter, &next) && samples < sample_size) {
+                elesize += entryMemUsage(next);
+                samples++;
+            }
+            hashtableCleanupIterator(&iter);
+            if (samples) asize += (double)elesize / samples * hashtableSize(ht);
+            if (vsetIsValid(volatile_fields)) asize += vsetMemUsage(volatile_fields);
+        } else {
+            serverPanic("Unknown hash encoding");
+        }
+    } else if (o->type == OBJ_STREAM) {
+        stream *s = objectGetVal(o);
+        asize += sizeof(*s) + raxAllocSize(s->rax);
+
+        /* Now we have to add the listpacks. The last listpack is often non
+         * complete, so we estimate the size of the first N listpacks, and
+         * use the average to compute the size of the first N-1 listpacks, and
+         * finally add the real size of the last node. */
+        raxIterator ri;
+        raxStart(&ri, s->rax);
+        raxSeek(&ri, "^", NULL, 0);
+        size_t lpsize = 0;
+        samples = 0;
+        while (samples < sample_size && raxNext(&ri)) {
+            unsigned char *lp = ri.data;
+            /* Use the allocated size, since we overprovision the node initially. */
+            lpsize += zmalloc_size(lp);
+            samples++;
+        }
+        if (s->rax->numele <= samples) {
+            asize += lpsize;
+        } else {
+            if (samples) lpsize /= samples; /* Compute the average. */
+            asize += lpsize * (s->rax->numele - 1);
+            /* No need to seek, we are already at the last element. */
+            asize += zmalloc_size(ri.data);
+        }
+        raxStop(&ri);
+
+        /* Consumer groups also have a non-trivial memory overhead if there
+         * are many consumers and many groups, let's count at least the
+         * overhead of the pending entries in the groups and consumers
+         * PELs. */
+        if (s->cgroups) {
+            raxStart(&ri, s->cgroups);
+            raxSeek(&ri, "^", NULL, 0);
+            samples = 0;
+            elesize = 0;
+            while (samples < sample_size && raxNext(&ri)) {
+                streamCG *cg = ri.data;
+                elesize += sizeof(*cg);
+                elesize += raxAllocSize(cg->pel);
+                elesize += sizeof(streamNACK) * raxSize(cg->pel);
+
+                /* For each consumer we also need to add the basic data
+                 * structures and the PEL memory usage. */
+                raxIterator cri;
+                raxStart(&cri, cg->consumers);
+                raxSeek(&cri, "^", NULL, 0);
+                size_t inner_samples = 0;
+                size_t inner_elesize = 0;
+                while (inner_samples < sample_size && raxNext(&cri)) {
+                    streamConsumer *consumer = cri.data;
+                    inner_elesize += sizeof(*consumer);
+                    inner_elesize += sdslen(consumer->name);
+                    inner_elesize += raxAllocSize(consumer->pel);
+                    /* Don't count NACKs again, they are shared with the
+                     * consumer group PEL. */
+                    inner_samples++;
+                }
+                raxStop(&cri);
+                if (inner_samples) elesize += (double)inner_elesize / inner_samples * raxSize(cg->consumers);
+                samples++;
+            }
+            raxStop(&ri);
+            if (samples) asize += (double)elesize / samples * raxSize(s->cgroups);
+        }
+    } else if (o->type == OBJ_MODULE) {
+        asize += moduleGetMemUsage(key, o, sample_size, dbid);
+    } else {
+        serverPanic("Unknown object type");
+    }
+    return asize;
+}
+
 /* Release data obtained with getMemoryOverheadData(). */
 void freeMemoryOverheadData(struct serverMemOverhead *mh) {
     zfree(mh->db);