diff --git a/src/module.c b/src/module.c index c54d164b7..7e2528b96 100644 --- a/src/module.c +++ b/src/module.c @@ -545,11 +545,23 @@ void *RM_Calloc(size_t nmemb, size_t size) { return zcalloc_usable(nmemb*size,NULL); } +/* Similar to RM_Calloc, but returns NULL in case of allocation failure, instead + * of panicking. */ +void *RM_TryCalloc(size_t nmemb, size_t size) { + return ztrycalloc_usable(nmemb*size,NULL); +} + /* Use like realloc() for memory obtained with RedisModule_Alloc(). */ void* RM_Realloc(void *ptr, size_t bytes) { return zrealloc_usable(ptr,bytes,NULL); } +/* Similar to RM_Realloc, but returns NULL in case of allocation failure, + * instead of panicking. */ +void *RM_TryRealloc(void *ptr, size_t bytes) { + return ztryrealloc_usable(ptr,bytes,NULL); +} + /* Use like free() for memory obtained by RedisModule_Alloc() and * RedisModule_Realloc(). However you should never try to free with * RedisModule_Free() memory allocated with malloc() inside your module. */ @@ -13590,7 +13602,9 @@ void moduleRegisterCoreAPI(void) { REGISTER_API(Alloc); REGISTER_API(TryAlloc); REGISTER_API(Calloc); + REGISTER_API(TryCalloc); REGISTER_API(Realloc); + REGISTER_API(TryRealloc); REGISTER_API(Free); REGISTER_API(Strdup); REGISTER_API(CreateCommand); diff --git a/src/redismodule.h b/src/redismodule.h index de533d049..65827e1b8 100644 --- a/src/redismodule.h +++ b/src/redismodule.h @@ -959,8 +959,10 @@ typedef struct RedisModuleTypeMethods { REDISMODULE_API void * (*RedisModule_Alloc)(size_t bytes) REDISMODULE_ATTR; REDISMODULE_API void * (*RedisModule_TryAlloc)(size_t bytes) REDISMODULE_ATTR; REDISMODULE_API void * (*RedisModule_Realloc)(void *ptr, size_t bytes) REDISMODULE_ATTR; +REDISMODULE_API void * (*RedisModule_TryRealloc)(void *ptr, size_t bytes) REDISMODULE_ATTR; REDISMODULE_API void (*RedisModule_Free)(void *ptr) REDISMODULE_ATTR; REDISMODULE_API void * (*RedisModule_Calloc)(size_t nmemb, size_t size) REDISMODULE_ATTR; +REDISMODULE_API void * (*RedisModule_TryCalloc)(size_t nmemb, size_t size) REDISMODULE_ATTR; REDISMODULE_API char * (*RedisModule_Strdup)(const char *str) REDISMODULE_ATTR; REDISMODULE_API int (*RedisModule_GetApi)(const char *, void *) REDISMODULE_ATTR; REDISMODULE_API int (*RedisModule_CreateCommand)(RedisModuleCtx *ctx, const char *name, RedisModuleCmdFunc cmdfunc, const char *strflags, int firstkey, int lastkey, int keystep) REDISMODULE_ATTR; @@ -1322,8 +1324,10 @@ static int RedisModule_Init(RedisModuleCtx *ctx, const char *name, int ver, int REDISMODULE_GET_API(Alloc); REDISMODULE_GET_API(TryAlloc); REDISMODULE_GET_API(Calloc); + REDISMODULE_GET_API(TryCalloc); REDISMODULE_GET_API(Free); REDISMODULE_GET_API(Realloc); + REDISMODULE_GET_API(TryRealloc); REDISMODULE_GET_API(Strdup); REDISMODULE_GET_API(CreateCommand); REDISMODULE_GET_API(GetCommand); diff --git a/tests/modules/misc.c b/tests/modules/misc.c index 46bfcb117..feed9ada0 100644 --- a/tests/modules/misc.c +++ b/tests/modules/misc.c @@ -503,6 +503,27 @@ final: return REDISMODULE_OK; } +int test_malloc_api(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + UNUSED(argv); + UNUSED(argc); + + void *p; + + p = RedisModule_TryAlloc(1024); + memset(p, 0, 1024); + RedisModule_Free(p); + + p = RedisModule_TryCalloc(1, 1024); + memset(p, 1, 1024); + + p = RedisModule_TryRealloc(p, 5 * 1024); + memset(p, 1, 5 * 1024); + RedisModule_Free(p); + + RedisModule_ReplyWithSimpleString(ctx, "OK"); + return REDISMODULE_OK; +} + int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { REDISMODULE_NOT_USED(argv); REDISMODULE_NOT_USED(argc); @@ -566,6 +587,8 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) return REDISMODULE_ERR; if (RedisModule_CreateCommand(ctx, "test.clear_n_events", test_clear_n_events,"", 0, 0, 0) == REDISMODULE_ERR) return REDISMODULE_ERR; + if (RedisModule_CreateCommand(ctx, "test.malloc_api", test_malloc_api,"", 0, 0, 0) == REDISMODULE_ERR) + return REDISMODULE_ERR; return REDISMODULE_OK; } diff --git a/tests/unit/moduleapi/misc.tcl b/tests/unit/moduleapi/misc.tcl index cf205462e..9f10caf39 100644 --- a/tests/unit/moduleapi/misc.tcl +++ b/tests/unit/moduleapi/misc.tcl @@ -493,6 +493,10 @@ start_server {overrides {save {900 1}} tags {"modules"}} { # server is writable again r set x y } {OK} + + test "malloc API" { + assert_equal {OK} [r test.malloc_api 0] + } } start_server {tags {"modules"}} {