From baa98988215f005b5b57e156c68a5bdc8454a3a5 Mon Sep 17 00:00:00 2001
From: antirez <antirez@gmail.com>
Date: Thu, 12 Jan 2017 09:08:43 +0100
Subject: [PATCH] MEMORY USAGE: support for modules data types.

As a side effect of supporting it, we no longer crash when MEMORY USAGE
is called against a module data type.

Close #3637.
---
 src/module.c            |  2 +-
 src/modules/hellotype.c | 11 +++++++----
 src/object.c            |  8 ++++++++
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/src/module.c b/src/module.c
index a5b3d52ae..1fbc5094f 100644
--- a/src/module.c
+++ b/src/module.c
@@ -2743,8 +2743,8 @@ moduleType *RM_CreateDataType(RedisModuleCtx *ctx, const char *name, int encver,
         moduleTypeLoadFunc rdb_load;
         moduleTypeSaveFunc rdb_save;
         moduleTypeRewriteFunc aof_rewrite;
-        moduleTypeDigestFunc digest;
         moduleTypeMemUsageFunc mem_usage;
+        moduleTypeDigestFunc digest;
         moduleTypeFreeFunc free;
     } *tms = (struct typemethods*) typemethods_ptr;
 
diff --git a/src/modules/hellotype.c b/src/modules/hellotype.c
index 02a5bb477..1aa6c5e63 100644
--- a/src/modules/hellotype.c
+++ b/src/modules/hellotype.c
@@ -226,10 +226,12 @@ void HelloTypeAofRewrite(RedisModuleIO *aof, RedisModuleString *key, void *value
     }
 }
 
-void HelloTypeDigest(RedisModuleDigest *digest, void *value) {
-    REDISMODULE_NOT_USED(digest);
-    REDISMODULE_NOT_USED(value);
-    /* TODO: The DIGEST module interface is yet not implemented. */
+/* The goal of this function is to return the amount of memory used by
+ * the HelloType value. */
+size_t HelloTypeMemUsage(void *value) {
+    struct HelloTypeObject *hto = value;
+    struct HelloTypeNode *node = hto->head;
+    return sizeof(*hto) + sizeof(*node)*hto->len;
 }
 
 void HelloTypeFree(void *value) {
@@ -250,6 +252,7 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
         .rdb_load = HelloTypeRdbLoad,
         .rdb_save = HelloTypeRdbSave,
         .aof_rewrite = HelloTypeAofRewrite,
+        .mem_usage = HelloTypeMemUsage,
         .free = HelloTypeFree
     };
 
diff --git a/src/object.c b/src/object.c
index 1ae37c9d1..125c1a595 100644
--- a/src/object.c
+++ b/src/object.c
@@ -786,6 +786,14 @@ size_t objectComputeSize(robj *o, size_t sample_size) {
         } else {
             serverPanic("Unknown hash encoding");
         }
+    } else if (o->type == OBJ_MODULE) {
+        moduleValue *mv = o->ptr;
+        moduleType *mt = mv->type;
+        if (mt->mem_usage != NULL) {
+            asize = mt->mem_usage(mv->value);
+        } else {
+            asize = 0;
+        }
     } else {
         serverPanic("Unknown object type");
     }