diff options
author | guybe7 <guy.benoish@redislabs.com> | 2022-04-17 07:31:57 +0200 |
---|---|---|
committer | GitHub <noreply@github.com> | 2022-04-17 08:31:57 +0300 |
commit | fe1c096b1849b633f41f69f432c4236428035950 (patch) | |
tree | 1c698865eabc25e86045e7d19a2653ad84015e7e /tests | |
parent | effa707e9db3e9d00fff06d45e2a5a81d3d55fa9 (diff) | |
download | redis-fe1c096b1849b633f41f69f432c4236428035950.tar.gz |
Add RM_MallocSizeString, RM_MallocSizeDict (#10542)
Add APIs to allow modules to compute the memory consumption of opaque objects owned by redis.
Without these, the mem_usage callbacks of module data types are useless in many cases.
Other changes:
Fix streamRadixTreeMemoryUsage to include the size of the rax structure itself
Diffstat (limited to 'tests')
-rw-r--r-- | tests/modules/Makefile | 1 | ||||
-rw-r--r-- | tests/modules/mallocsize.c | 237 | ||||
-rw-r--r-- | tests/unit/moduleapi/mallocsize.tcl | 21 |
3 files changed, 259 insertions, 0 deletions
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" + } +} |