summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
Diffstat
-rwxr-xr-xruntest-moduleapi1
-rw-r--r--src/module.c23
-rw-r--r--src/object.c2
-rw-r--r--src/redismodule.h4
-rw-r--r--tests/modules/Makefile1
-rw-r--r--tests/modules/mallocsize.c237
-rw-r--r--tests/unit/moduleapi/mallocsize.tcl21
7 files changed, 287 insertions, 2 deletions
diff --git a/runtest-moduleapi b/runtest-moduleapi
index 38c0d5434..81af306e5 100755
--- a/runtest-moduleapi
+++ b/runtest-moduleapi
@@ -40,6 +40,7 @@ $TCLSH tests/test_helper.tcl \
--single unit/moduleapi/zset \
--single unit/moduleapi/list \
--single unit/moduleapi/stream \
+--single unit/moduleapi/mallocsize \
--single unit/moduleapi/datatype2 \
--single unit/moduleapi/cluster \
--single unit/moduleapi/aclcheck \
diff --git a/src/module.c b/src/module.c
index 38f85ec04..163104d8a 100644
--- a/src/module.c
+++ b/src/module.c
@@ -9764,10 +9764,29 @@ int RM_CommandFilterArgDelete(RedisModuleCommandFilterCtx *fctx, int pos)
* with the allocation calls, since sometimes the underlying allocator
* will allocate more memory.
*/
-size_t RM_MallocSize(void* ptr){
+size_t RM_MallocSize(void* ptr) {
return zmalloc_size(ptr);
}
+/* Same as RM_MallocSize, except it works on RedisModuleString pointers.
+ */
+size_t RM_MallocSizeString(RedisModuleString* str) {
+ serverAssert(str->type == OBJ_STRING);
+ return sizeof(*str) + getStringObjectSdsUsedMemory(str);
+}
+
+/* Same as RM_MallocSize, except it works on RedisModuleDict pointers.
+ * Note that the returned value is only the overhead of the underlying structures,
+ * it does not include the allocation size of the keys and values.
+ */
+size_t RM_MallocSizeDict(RedisModuleDict* dict) {
+ size_t size = sizeof(RedisModuleDict) + sizeof(rax);
+ size += dict->rax->numnodes * sizeof(raxNode);
+ /* For more info about this weird line, see streamRadixTreeMemoryUsage */
+ size += dict->rax->numnodes * sizeof(long)*30;
+ return size;
+}
+
/* Return the a number between 0 to 1 indicating the amount of memory
* currently used, relative to the Redis "maxmemory" configuration.
*
@@ -12536,6 +12555,8 @@ void moduleRegisterCoreAPI(void) {
REGISTER_API(GetBlockedClientReadyKey);
REGISTER_API(GetUsedMemoryRatio);
REGISTER_API(MallocSize);
+ REGISTER_API(MallocSizeString);
+ REGISTER_API(MallocSizeDict);
REGISTER_API(ScanCursorCreate);
REGISTER_API(ScanCursorDestroy);
REGISTER_API(ScanCursorRestart);
diff --git a/src/object.c b/src/object.c
index a60a27e90..093e2619e 100644
--- a/src/object.c
+++ b/src/object.c
@@ -958,7 +958,7 @@ char *strEncoding(int encoding) {
* on the insertion speed and thus the ability of the radix tree
* to compress prefixes. */
size_t streamRadixTreeMemoryUsage(rax *rax) {
- size_t size;
+ size_t size = sizeof(*rax);
size = rax->numele * sizeof(streamID);
size += rax->numnodes * sizeof(raxNode);
/* Add a fixed overhead due to the aux data pointer, children, ... */
diff --git a/src/redismodule.h b/src/redismodule.h
index 954a4c7c6..846967a62 100644
--- a/src/redismodule.h
+++ b/src/redismodule.h
@@ -1158,6 +1158,8 @@ REDISMODULE_API int (*RedisModule_ExitFromChild)(int retcode) REDISMODULE_ATTR;
REDISMODULE_API int (*RedisModule_KillForkChild)(int child_pid) REDISMODULE_ATTR;
REDISMODULE_API float (*RedisModule_GetUsedMemoryRatio)() REDISMODULE_ATTR;
REDISMODULE_API size_t (*RedisModule_MallocSize)(void* ptr) REDISMODULE_ATTR;
+REDISMODULE_API size_t (*RedisModule_MallocSizeString)(RedisModuleString* str) REDISMODULE_ATTR;
+REDISMODULE_API size_t (*RedisModule_MallocSizeDict)(RedisModuleDict* dict) REDISMODULE_ATTR;
REDISMODULE_API RedisModuleUser * (*RedisModule_CreateModuleUser)(const char *name) REDISMODULE_ATTR;
REDISMODULE_API void (*RedisModule_FreeModuleUser)(RedisModuleUser *user) REDISMODULE_ATTR;
REDISMODULE_API int (*RedisModule_SetModuleUserACL)(RedisModuleUser *user, const char* acl) REDISMODULE_ATTR;
@@ -1488,6 +1490,8 @@ static int RedisModule_Init(RedisModuleCtx *ctx, const char *name, int ver, int
REDISMODULE_GET_API(KillForkChild);
REDISMODULE_GET_API(GetUsedMemoryRatio);
REDISMODULE_GET_API(MallocSize);
+ REDISMODULE_GET_API(MallocSizeString);
+ REDISMODULE_GET_API(MallocSizeDict);
REDISMODULE_GET_API(CreateModuleUser);
REDISMODULE_GET_API(FreeModuleUser);
REDISMODULE_GET_API(SetModuleUserACL);
diff --git a/tests/modules/Makefile b/tests/modules/Makefile
index a43e1e9a5..16b5570aa 100644
--- a/tests/modules/Makefile
+++ b/tests/modules/Makefile
@@ -49,6 +49,7 @@ TEST_MODULES = \
hash.so \
zset.so \
stream.so \
+ mallocsize.so \
aclcheck.so \
list.so \
subcommands.so \
diff --git a/tests/modules/mallocsize.c b/tests/modules/mallocsize.c
new file mode 100644
index 000000000..a1d31c136
--- /dev/null
+++ b/tests/modules/mallocsize.c
@@ -0,0 +1,237 @@
+#include "redismodule.h"
+#include <string.h>
+#include <assert.h>
+#include <unistd.h>
+
+#define UNUSED(V) ((void) V)
+
+/* Registered type */
+RedisModuleType *mallocsize_type = NULL;
+
+typedef enum {
+ UDT_RAW,
+ UDT_STRING,
+ UDT_DICT
+} udt_type_t;
+
+typedef struct {
+ void *ptr;
+ size_t len;
+} raw_t;
+
+typedef struct {
+ udt_type_t type;
+ union {
+ raw_t raw;
+ RedisModuleString *str;
+ RedisModuleDict *dict;
+ } data;
+} udt_t;
+
+void udt_free(void *value) {
+ udt_t *udt = value;
+ switch (udt->type) {
+ case (UDT_RAW): {
+ RedisModule_Free(udt->data.raw.ptr);
+ break;
+ }
+ case (UDT_STRING): {
+ RedisModule_FreeString(NULL, udt->data.str);
+ break;
+ }
+ case (UDT_DICT): {
+ RedisModuleString *dk, *dv;
+ RedisModuleDictIter *iter = RedisModule_DictIteratorStartC(udt->data.dict, "^", NULL, 0);
+ while((dk = RedisModule_DictNext(NULL, iter, (void **)&dv)) != NULL) {
+ RedisModule_FreeString(NULL, dk);
+ RedisModule_FreeString(NULL, dv);
+ }
+ RedisModule_DictIteratorStop(iter);
+ RedisModule_FreeDict(NULL, udt->data.dict);
+ break;
+ }
+ }
+ RedisModule_Free(udt);
+}
+
+void udt_rdb_save(RedisModuleIO *rdb, void *value) {
+ udt_t *udt = value;
+ RedisModule_SaveUnsigned(rdb, udt->type);
+ switch (udt->type) {
+ case (UDT_RAW): {
+ RedisModule_SaveStringBuffer(rdb, udt->data.raw.ptr, udt->data.raw.len);
+ break;
+ }
+ case (UDT_STRING): {
+ RedisModule_SaveString(rdb, udt->data.str);
+ break;
+ }
+ case (UDT_DICT): {
+ RedisModule_SaveUnsigned(rdb, RedisModule_DictSize(udt->data.dict));
+ RedisModuleString *dk, *dv;
+ RedisModuleDictIter *iter = RedisModule_DictIteratorStartC(udt->data.dict, "^", NULL, 0);
+ while((dk = RedisModule_DictNext(NULL, iter, (void **)&dv)) != NULL) {
+ RedisModule_SaveString(rdb, dk);
+ RedisModule_SaveString(rdb, dv);
+ RedisModule_FreeString(NULL, dk); /* Allocated by RedisModule_DictNext */
+ }
+ RedisModule_DictIteratorStop(iter);
+ break;
+ }
+ }
+}
+
+void *udt_rdb_load(RedisModuleIO *rdb, int encver) {
+ if (encver != 0)
+ return NULL;
+ udt_t *udt = RedisModule_Alloc(sizeof(*udt));
+ udt->type = RedisModule_LoadUnsigned(rdb);
+ switch (udt->type) {
+ case (UDT_RAW): {
+ udt->data.raw.ptr = RedisModule_LoadStringBuffer(rdb, &udt->data.raw.len);
+ break;
+ }
+ case (UDT_STRING): {
+ udt->data.str = RedisModule_LoadString(rdb);
+ break;
+ }
+ case (UDT_DICT): {
+ long long dict_len = RedisModule_LoadUnsigned(rdb);
+ udt->data.dict = RedisModule_CreateDict(NULL);
+ for (int i = 0; i < dict_len; i += 2) {
+ RedisModuleString *key = RedisModule_LoadString(rdb);
+ RedisModuleString *val = RedisModule_LoadString(rdb);
+ RedisModule_DictSet(udt->data.dict, key, val);
+ }
+ break;
+ }
+ }
+
+ return udt;
+}
+
+size_t udt_mem_usage(RedisModuleKeyOptCtx *ctx, const void *value, size_t sample_size) {
+ UNUSED(ctx);
+ UNUSED(sample_size);
+
+ const udt_t *udt = value;
+ size_t size = sizeof(*udt);
+
+ switch (udt->type) {
+ case (UDT_RAW): {
+ size += RedisModule_MallocSize(udt->data.raw.ptr);
+ break;
+ }
+ case (UDT_STRING): {
+ size += RedisModule_MallocSizeString(udt->data.str);
+ break;
+ }
+ case (UDT_DICT): {
+ void *dk;
+ size_t keylen;
+ RedisModuleString *dv;
+ RedisModuleDictIter *iter = RedisModule_DictIteratorStartC(udt->data.dict, "^", NULL, 0);
+ while((dk = RedisModule_DictNextC(iter, &keylen, (void **)&dv)) != NULL) {
+ size += keylen;
+ size += RedisModule_MallocSizeString(dv);
+ }
+ RedisModule_DictIteratorStop(iter);
+ break;
+ }
+ }
+
+ return size;
+}
+
+/* MALLOCSIZE.SETRAW key len */
+int cmd_setraw(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+ if (argc != 3)
+ return RedisModule_WrongArity(ctx);
+
+ RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_WRITE);
+
+ udt_t *udt = RedisModule_Alloc(sizeof(*udt));
+ udt->type = UDT_RAW;
+
+ long long raw_len;
+ RedisModule_StringToLongLong(argv[2], &raw_len);
+ udt->data.raw.ptr = RedisModule_Alloc(raw_len);
+ udt->data.raw.len = raw_len;
+
+ RedisModule_ModuleTypeSetValue(key, mallocsize_type, udt);
+ RedisModule_CloseKey(key);
+
+ return RedisModule_ReplyWithSimpleString(ctx, "OK");
+}
+
+/* MALLOCSIZE.SETSTR key string */
+int cmd_setstr(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+ if (argc != 3)
+ return RedisModule_WrongArity(ctx);
+
+ RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_WRITE);
+
+ udt_t *udt = RedisModule_Alloc(sizeof(*udt));
+ udt->type = UDT_STRING;
+
+ udt->data.str = argv[2];
+ RedisModule_RetainString(ctx, argv[2]);
+
+ RedisModule_ModuleTypeSetValue(key, mallocsize_type, udt);
+ RedisModule_CloseKey(key);
+
+ return RedisModule_ReplyWithSimpleString(ctx, "OK");
+}
+
+/* MALLOCSIZE.SETDICT key field value [field value ...] */
+int cmd_setdict(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+ if (argc < 4 || argc % 2)
+ return RedisModule_WrongArity(ctx);
+
+ RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_WRITE);
+
+ udt_t *udt = RedisModule_Alloc(sizeof(*udt));
+ udt->type = UDT_DICT;
+
+ udt->data.dict = RedisModule_CreateDict(ctx);
+ for (int i = 2; i < argc; i += 2) {
+ RedisModule_DictSet(udt->data.dict, argv[i], argv[i+1]);
+ /* No need to retain argv[i], it is copied as the rax key */
+ RedisModule_RetainString(ctx, argv[i+1]);
+ }
+
+ RedisModule_ModuleTypeSetValue(key, mallocsize_type, udt);
+ RedisModule_CloseKey(key);
+
+ return RedisModule_ReplyWithSimpleString(ctx, "OK");
+}
+
+int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+ UNUSED(argv);
+ UNUSED(argc);
+ if (RedisModule_Init(ctx,"mallocsize",1,REDISMODULE_APIVER_1)== REDISMODULE_ERR)
+ return REDISMODULE_ERR;
+
+ RedisModuleTypeMethods tm = {
+ .version = REDISMODULE_TYPE_METHOD_VERSION,
+ .rdb_load = udt_rdb_load,
+ .rdb_save = udt_rdb_save,
+ .free = udt_free,
+ .mem_usage2 = udt_mem_usage,
+ };
+
+ mallocsize_type = RedisModule_CreateDataType(ctx, "allocsize", 0, &tm);
+ if (mallocsize_type == NULL)
+ return REDISMODULE_ERR;
+
+ if (RedisModule_CreateCommand(ctx, "mallocsize.setraw", cmd_setraw, "", 1, 1, 1) == REDISMODULE_ERR)
+ return REDISMODULE_ERR;
+
+ if (RedisModule_CreateCommand(ctx, "mallocsize.setstr", cmd_setstr, "", 1, 1, 1) == REDISMODULE_ERR)
+ return REDISMODULE_ERR;
+
+ if (RedisModule_CreateCommand(ctx, "mallocsize.setdict", cmd_setdict, "", 1, 1, 1) == REDISMODULE_ERR)
+ return REDISMODULE_ERR;
+
+ return REDISMODULE_OK;
+}
diff --git a/tests/unit/moduleapi/mallocsize.tcl b/tests/unit/moduleapi/mallocsize.tcl
new file mode 100644
index 000000000..359a7ae44
--- /dev/null
+++ b/tests/unit/moduleapi/mallocsize.tcl
@@ -0,0 +1,21 @@
+set testmodule [file normalize tests/modules/mallocsize.so]
+
+
+start_server {tags {"modules"}} {
+ r module load $testmodule
+
+ test {MallocSize of raw bytes} {
+ assert_equal [r mallocsize.setraw key 40] {OK}
+ assert_morethan [r memory usage key] 40
+ }
+
+ test {MallocSize of string} {
+ assert_equal [r mallocsize.setstr key abcdefg] {OK}
+ assert_morethan [r memory usage key] 7 ;# Length of "abcdefg"
+ }
+
+ test {MallocSize of dict} {
+ assert_equal [r mallocsize.setdict key f1 v1 f2 v2] {OK}
+ assert_morethan [r memory usage key] 8 ;# Length of "f1v1f2v2"
+ }
+}
View Lorry Controller Status