From 033a49e23dd1f2a53f7b88c40f2ce4db437cd2c5 Mon Sep 17 00:00:00 2001
From: antirez <antirez@gmail.com>
Date: Thu, 15 Sep 2016 15:25:05 +0200
Subject: [PATCH] MEMORY USAGE: SAMPLES option added + fixes to size
 computation.

The new SAMPLES option is added, defaulting to 5, and with 0 being a
special value to scan the whole set of elements.

Fixes to the object size computation were made since the original PR
assumed data structures still contaning robj structures, while now after
the lazyfree changes, are all SDS strings.
---
 src/object.c | 45 +++++++++++++++++++++++++++------------------
 1 file changed, 27 insertions(+), 18 deletions(-)

diff --git a/src/object.c b/src/object.c
index ddc49dfd0..e47d663e0 100644
--- a/src/object.c
+++ b/src/object.c
@@ -700,7 +700,7 @@ char *strEncoding(int encoding) {
  * are checked and averaged to estimate the total size. */
 #define OBJ_COMPUTE_SIZE_DEF_SAMPLES 5 /* Default sample size. */
 size_t objectComputeSize(robj *o, size_t sample_size) {
-    robj *ele;
+    sds ele, ele2;
     dict *d;
     dictIterator *di;
     struct dictEntry *de;
@@ -709,8 +709,7 @@ size_t objectComputeSize(robj *o, size_t sample_size) {
     if (o->type == OBJ_STRING) {
         if(o->encoding == OBJ_ENCODING_INT) {
             asize = sizeof(*o);
-        }
-        else if(o->encoding == OBJ_ENCODING_RAW) {
+        } else if(o->encoding == OBJ_ENCODING_RAW) {
             asize = sdsAllocSize(o->ptr)+sizeof(*o);
         } else if(o->encoding == OBJ_ENCODING_EMBSTR) {
             asize = sdslen(o->ptr)+2+sizeof(*o);
@@ -739,9 +738,7 @@ size_t objectComputeSize(robj *o, size_t sample_size) {
             asize = sizeof(*o)+sizeof(dict)+(sizeof(struct dictEntry*)*dictSlots(d));
             while((de = dictNext(di)) != NULL && samples < sample_size) {
                 ele = dictGetKey(de);
-                elesize += (ele->encoding == OBJ_ENCODING_RAW) ?
-                                (sizeof(*o)+sdsAllocSize(ele->ptr)) : sizeof(*o);
-                elesize += sizeof(struct dictEntry);
+                elesize += sizeof(struct dictEntry) + sdsAllocSize(ele);
                 samples++;
             }
             dictReleaseIterator(di);
@@ -761,10 +758,8 @@ size_t objectComputeSize(robj *o, size_t sample_size) {
             asize = sizeof(*o)+sizeof(zset)+(sizeof(struct dictEntry*)*dictSlots(d));
             while((de = dictNext(di)) != NULL && samples < sample_size) {
                 ele = dictGetKey(de);
-                elesize += (ele->encoding == OBJ_ENCODING_RAW) ?
-                                (sizeof(*o)+sdsAllocSize(ele->ptr)) : sizeof(*o);
-                elesize += sizeof(struct dictEntry);
-                elesize += sizeof(zskiplistNode)*dictSize(d);
+                elesize += sdsAllocSize(ele);
+                elesize += sizeof(struct dictEntry) + sizeof(zskiplistNode);
                 samples++;
             }
             dictReleaseIterator(di);
@@ -781,14 +776,10 @@ size_t objectComputeSize(robj *o, size_t sample_size) {
             asize = sizeof(*o)+sizeof(dict)+(sizeof(struct dictEntry*)*dictSlots(d));
             while((de = dictNext(di)) != NULL && samples < sample_size) {
                 ele = dictGetKey(de);
-                elesize += (ele->encoding == OBJ_ENCODING_RAW) ?
-                                (sizeof(*o)+sdsAllocSize(ele->ptr)) : sizeof(*o);
-                ele = dictGetVal(de);
-                elesize += (ele->encoding == OBJ_ENCODING_RAW) ?
-                                (sizeof(*o)+sdsAllocSize(ele->ptr)) : sizeof(*o);
+                ele2 = dictGetVal(de);
+                elesize += sdsAllocSize(ele) + sdsAllocSize(ele2);
                 elesize += sizeof(struct dictEntry);
                 samples++;
-                printf("%zu samples: %zu usage\n", samples, elesize);
             }
             dictReleaseIterator(di);
             if (samples) asize += (double)elesize/samples*dictSize(d);
@@ -955,10 +946,28 @@ struct redisMemOverhead *getMemoryOverheadData(void) {
 void memoryCommand(client *c) {
     robj *o;
 
-    if (!strcasecmp(c->argv[1]->ptr,"usage") && c->argc == 3) {
+    if (!strcasecmp(c->argv[1]->ptr,"usage") && c->argc >= 3) {
+        long long samples = OBJ_COMPUTE_SIZE_DEF_SAMPLES;
+        for (int j = 3; j < c->argc; j++) {
+            if (!strcasecmp(c->argv[j]->ptr,"samples") &&
+                j+1 < c->argc)
+            {
+                if (getLongLongFromObjectOrReply(c,c->argv[j+1],&samples,NULL)
+                     == C_ERR) return;
+                if (samples < 0) {
+                    addReply(c,shared.syntaxerr);
+                    return;
+                }
+                if (samples == 0) samples = LLONG_MAX;;
+                j++; /* skip option argument. */
+            } else {
+                addReply(c,shared.syntaxerr);
+                return;
+            }
+        }
         if ((o = objectCommandLookupOrReply(c,c->argv[2],shared.nullbulk))
                 == NULL) return;
-        size_t usage = objectComputeSize(o,OBJ_COMPUTE_SIZE_DEF_SAMPLES);
+        size_t usage = objectComputeSize(o,samples);
         usage += sdsAllocSize(c->argv[1]->ptr);
         usage += sizeof(dictEntry);
         addReplyLongLong(c,usage);