From 0f2c6b6a18a39443a5260632db84f33957eafff8 Mon Sep 17 00:00:00 2001
From: antirez <antirez@gmail.com>
Date: Tue, 20 Mar 2018 17:50:37 +0100
Subject: [PATCH] Streams: improve MEMORY USAGE computation, include CGs.

---
 src/object.c | 71 +++++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 56 insertions(+), 15 deletions(-)

diff --git a/src/object.c b/src/object.c
index 8614d6e9..682185be 100644
--- a/src/object.c
+++ b/src/object.c
@@ -708,7 +708,31 @@ char *strEncoding(int encoding) {
     }
 }
 
-/* =========================== Memory introspection ========================== */
+/* =========================== Memory introspection ========================= */
+
+
+/* This is an helper function with the goal of estimating the memory
+ * size of a radix tree that is used to store Stream IDs.
+ *
+ * Note: to guess the size of the radix tree is not trivial, so we
+ * approximate it considering 128 bytes of data overhead for each
+ * key (the ID), and then adding the number of bare nodes, plus some
+ * overhead due by the data and child pointers. This secret recipe
+ * was obtained by checking the average radix tree created by real
+ * workloads, and then adjusting the constants to get numbers that
+ * more or less match the real memory usage.
+ *
+ * Actually the number of nodes and keys may be different depending
+ * on the insertion speed and thus the ability of the radix tree
+ * to compress prefixes. */
+size_t streamRadixTreeMemoryUsage(rax *rax) {
+    size_t size;
+    size = rax->numele * sizeof(streamID);
+    size += rax->numnodes * sizeof(raxNode);
+    /* Add a fixed overhead due to the aux data pointer, children, ... */
+    size += rax->numnodes * sizeof(long)*30;
+    return size;
+}
 
 /* Returns the size in bytes consumed by the key's value in RAM.
  * Note that the returned value is just an approximation, especially in the
@@ -804,21 +828,8 @@ size_t objectComputeSize(robj *o, size_t sample_size) {
         }
     } else if (o->type == OBJ_STREAM) {
         stream *s = o->ptr;
-        /* Note: to guess the size of the radix tree is not trivial, so we
-         * approximate it considering 64 bytes of data overhead for each
-         * key (the ID), and then adding the number of bare nodes, plus some
-         * overhead due by the data and child pointers. This secret recipe
-         * was obtained by checking the average radix tree created by real
-         * workloads, and then adjusting the constants to get numbers that
-         * more or less match the real memory usage.
-         *
-         * Actually the number of nodes and keys may be different depending
-         * on the insertion speed and thus the ability of the radix tree
-         * to compress prefixes. */
         asize = sizeof(*o);
-        asize += s->rax->numele * 64;
-        asize += s->rax->numnodes * sizeof(raxNode);
-        asize += s->rax->numnodes * 32*7; /* Add a few child pointers... */
+        asize += streamRadixTreeMemoryUsage(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
@@ -845,6 +856,36 @@ size_t objectComputeSize(robj *o, size_t sample_size) {
             asize += lpBytes(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);
+            while(raxNext(&ri)) {
+                streamCG *cg = ri.data;
+                asize += sizeof(*cg);
+                asize += streamRadixTreeMemoryUsage(cg->pel);
+                asize += 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);
+                while(raxNext(&cri)) {
+                    streamConsumer *consumer = cri.data;
+                    asize += sizeof(*consumer);
+                    asize += sdslen(consumer->name);
+                    asize += streamRadixTreeMemoryUsage(consumer->pel);
+                    /* Don't count NACKs again, they are shared with the
+                     * consumer group PEL. */
+                }
+                raxStop(&cri);
+            }
+            raxStop(&ri);
+        }
     } else if (o->type == OBJ_MODULE) {
         moduleValue *mv = o->ptr;
         moduleType *mt = mv->type;