From 5350e7669e9dfe7b2b4bcf663171920c441c19e1 Mon Sep 17 00:00:00 2001 From: Yossi Gottlieb Date: Mon, 11 Jul 2016 16:47:37 +0300 Subject: [PATCH] Add ModuleDataType to/from string serialization. Add two new functions that leverage the RedisModuleDataType mechanism for RDB serialization/deserialization and make it possible to use it to/from arbitrary strings: * RM_SaveDataTypeToString() * RM_LoadDataTypeFromString() --- runtest-moduleapi | 1 + src/module.c | 55 ++++++++++ src/redismodule.h | 4 + tests/modules/Makefile | 3 +- tests/modules/datatype.c | 161 ++++++++++++++++++++++++++++++ tests/unit/moduleapi/datatype.tcl | 27 +++++ 6 files changed, 250 insertions(+), 1 deletion(-) create mode 100644 tests/modules/datatype.c create mode 100644 tests/unit/moduleapi/datatype.tcl diff --git a/runtest-moduleapi b/runtest-moduleapi index e48535126..e5a5a1200 100755 --- a/runtest-moduleapi +++ b/runtest-moduleapi @@ -22,4 +22,5 @@ $TCLSH tests/test_helper.tcl \ --single unit/moduleapi/hooks \ --single unit/moduleapi/misc \ --single unit/moduleapi/blockonkeys \ +--single unit/moduleapi/datatype "${@}" diff --git a/src/module.c b/src/module.c index ad34e7b64..ca9fb5804 100644 --- a/src/module.c +++ b/src/module.c @@ -3884,6 +3884,59 @@ void RM_DigestEndSequence(RedisModuleDigest *md) { memset(md->o,0,sizeof(md->o)); } +/* Decode a serialized representation of a module data type 'mt' from string + * 'str' and return a newly allocated value, or NULL if decoding failed. + * + * This call basically reuses the 'rdb_load' callback which module data types + * implement in order to allow a module to arbitrarily serialize/de-serialize + * keys, similar to how the Redis 'DUMP' and 'RESTORE' commands are implemented. + * + * Modules should generally use the REDISMODULE_OPTIONS_HANDLE_IO_ERRORS flag and + * make sure the de-serialization code properly checks and handles IO errors + * (freeing allocated buffers and returning a NULL). + * + * If this is NOT done, Redis will handle corrupted (or just truncated) serialized + * data by producing an error message and terminating the process. + */ + +void *RM_LoadDataTypeFromString(const RedisModuleString *str, const moduleType *mt) { + rio payload; + RedisModuleIO io; + + rioInitWithBuffer(&payload, str->ptr); + moduleInitIOContext(io,(moduleType *)mt,&payload,NULL); + + /* All RM_Save*() calls always write a version 2 compatible format, so we + * need to make sure we read the same. + */ + io.ver = 2; + return mt->rdb_load(&io,0); +} + +/* Encode a module data type 'mt' value 'data' into serialized form, and return it + * as a newly allocated RedisModuleString. + * + * This call basically reuses the 'rdb_save' callback which module data types + * implement in order to allow a module to arbitrarily serialize/de-serialize + * keys, similar to how the Redis 'DUMP' and 'RESTORE' commands are implemented. + */ + +RedisModuleString *RM_SaveDataTypeToString(RedisModuleCtx *ctx, void *data, const moduleType *mt) { + rio payload; + RedisModuleIO io; + + rioInitWithBuffer(&payload,sdsempty()); + moduleInitIOContext(io,(moduleType *)mt,&payload,NULL); + mt->rdb_save(&io,data); + if (io.error) { + return NULL; + } else { + robj *str = createObject(OBJ_STRING,payload.io.buffer.ptr); + autoMemoryAdd(ctx,REDISMODULE_AM_STRING,str); + return str; + } +} + /* -------------------------------------------------------------------------- * AOF API for modules data types * -------------------------------------------------------------------------- */ @@ -6876,6 +6929,8 @@ void moduleRegisterCoreAPI(void) { REGISTER_API(LoadFloat); REGISTER_API(SaveLongDouble); REGISTER_API(LoadLongDouble); + REGISTER_API(SaveDataTypeToString); + REGISTER_API(LoadDataTypeFromString); REGISTER_API(EmitAOF); REGISTER_API(Log); REGISTER_API(LogIOError); diff --git a/src/redismodule.h b/src/redismodule.h index 728c7f584..21f3988cf 100644 --- a/src/redismodule.h +++ b/src/redismodule.h @@ -537,6 +537,8 @@ void REDISMODULE_API_FUNC(RedisModule_SaveFloat)(RedisModuleIO *io, float value) float REDISMODULE_API_FUNC(RedisModule_LoadFloat)(RedisModuleIO *io); void REDISMODULE_API_FUNC(RedisModule_SaveLongDouble)(RedisModuleIO *io, long double value); long double REDISMODULE_API_FUNC(RedisModule_LoadLongDouble)(RedisModuleIO *io); +void *REDISMODULE_API_FUNC(RedisModule_LoadDataTypeFromString)(const RedisModuleString *str, const RedisModuleType *mt); +RedisModuleString *REDISMODULE_API_FUNC(RedisModule_SaveDataTypeToString)(RedisModuleCtx *ctx, void *data, const RedisModuleType *mt); void REDISMODULE_API_FUNC(RedisModule_Log)(RedisModuleCtx *ctx, const char *level, const char *fmt, ...); void REDISMODULE_API_FUNC(RedisModule_LogIOError)(RedisModuleIO *io, const char *levelstr, const char *fmt, ...); void REDISMODULE_API_FUNC(RedisModule__Assert)(const char *estr, const char *file, int line); @@ -745,6 +747,8 @@ static int RedisModule_Init(RedisModuleCtx *ctx, const char *name, int ver, int REDISMODULE_GET_API(LoadFloat); REDISMODULE_GET_API(SaveLongDouble); REDISMODULE_GET_API(LoadLongDouble); + REDISMODULE_GET_API(SaveDataTypeToString); + REDISMODULE_GET_API(LoadDataTypeFromString); REDISMODULE_GET_API(EmitAOF); REDISMODULE_GET_API(Log); REDISMODULE_GET_API(LogIOError); diff --git a/tests/modules/Makefile b/tests/modules/Makefile index 9e27758a2..26627c30c 100644 --- a/tests/modules/Makefile +++ b/tests/modules/Makefile @@ -19,7 +19,8 @@ TEST_MODULES = \ propagate.so \ misc.so \ hooks.so \ - blockonkeys.so + blockonkeys.so \ + datatype.so .PHONY: all diff --git a/tests/modules/datatype.c b/tests/modules/datatype.c new file mode 100644 index 000000000..7c39ab457 --- /dev/null +++ b/tests/modules/datatype.c @@ -0,0 +1,161 @@ +/* This module current tests a small subset but should be extended in the future + * for general ModuleDataType coverage. + */ + +#include "redismodule.h" + +static RedisModuleType *datatype = NULL; + +typedef struct { + long long intval; + RedisModuleString *strval; +} DataType; + +static void *datatype_load(RedisModuleIO *io, int encver) { + (void) encver; + + int intval = RedisModule_LoadSigned(io); + if (RedisModule_IsIOError(io)) return NULL; + + RedisModuleString *strval = RedisModule_LoadString(io); + if (RedisModule_IsIOError(io)) return NULL; + + DataType *dt = (DataType *) RedisModule_Alloc(sizeof(DataType)); + dt->intval = intval; + dt->strval = strval; + return dt; +} + +static void datatype_save(RedisModuleIO *io, void *value) { + DataType *dt = (DataType *) value; + RedisModule_SaveSigned(io, dt->intval); + RedisModule_SaveString(io, dt->strval); +} + +static void datatype_free(void *value) { + if (value) { + DataType *dt = (DataType *) value; + + if (dt->strval) RedisModule_FreeString(NULL, dt->strval); + RedisModule_Free(dt); + } +} + +static int datatype_set(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + if (argc != 4) { + RedisModule_WrongArity(ctx); + return REDISMODULE_OK; + } + + long long intval; + + if (RedisModule_StringToLongLong(argv[2], &intval) != REDISMODULE_OK) { + RedisModule_ReplyWithError(ctx, "Invalid integr value"); + return REDISMODULE_OK; + } + + RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_WRITE); + DataType *dt = RedisModule_Calloc(sizeof(DataType), 1); + dt->intval = intval; + dt->strval = argv[3]; + RedisModule_RetainString(ctx, dt->strval); + + RedisModule_ModuleTypeSetValue(key, datatype, dt); + RedisModule_CloseKey(key); + RedisModule_ReplyWithSimpleString(ctx, "OK"); + + return REDISMODULE_OK; +} + +static int datatype_restore(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + if (argc != 3) { + RedisModule_WrongArity(ctx); + return REDISMODULE_OK; + } + + DataType *dt = RedisModule_LoadDataTypeFromString(argv[2], datatype); + if (!dt) { + RedisModule_ReplyWithError(ctx, "Invalid data"); + return REDISMODULE_OK; + } + + RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_WRITE); + RedisModule_ModuleTypeSetValue(key, datatype, dt); + RedisModule_CloseKey(key); + RedisModule_ReplyWithSimpleString(ctx, "OK"); + + return REDISMODULE_OK; +} + +static int datatype_get(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + if (argc != 2) { + RedisModule_WrongArity(ctx); + return REDISMODULE_OK; + } + + RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ); + DataType *dt = RedisModule_ModuleTypeGetValue(key); + RedisModule_CloseKey(key); + + RedisModule_ReplyWithArray(ctx, 2); + RedisModule_ReplyWithLongLong(ctx, dt->intval); + RedisModule_ReplyWithString(ctx, dt->strval); + return REDISMODULE_OK; +} + +static int datatype_dump(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + if (argc != 2) { + RedisModule_WrongArity(ctx); + return REDISMODULE_OK; + } + + RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ); + DataType *dt = RedisModule_ModuleTypeGetValue(key); + RedisModule_CloseKey(key); + + RedisModuleString *reply = RedisModule_SaveDataTypeToString(ctx, dt, datatype); + if (!reply) { + RedisModule_ReplyWithError(ctx, "Failed to save"); + return REDISMODULE_OK; + } + + RedisModule_ReplyWithString(ctx, reply); + RedisModule_FreeString(ctx, reply); + return REDISMODULE_OK; +} + + +int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + REDISMODULE_NOT_USED(argv); + REDISMODULE_NOT_USED(argc); + + if (RedisModule_Init(ctx,"datatype",1,REDISMODULE_APIVER_1) == REDISMODULE_ERR) + return REDISMODULE_ERR; + + RedisModule_SetModuleOptions(ctx, REDISMODULE_OPTIONS_HANDLE_IO_ERRORS); + + RedisModuleTypeMethods datatype_methods = { + .version = REDISMODULE_TYPE_METHOD_VERSION, + .rdb_load = datatype_load, + .rdb_save = datatype_save, + .free = datatype_free, + }; + + datatype = RedisModule_CreateDataType(ctx, "test___dt", 1, &datatype_methods); + if (datatype == NULL) + return REDISMODULE_ERR; + + if (RedisModule_CreateCommand(ctx,"datatype.set", datatype_set,"deny-oom",1,1,1) == REDISMODULE_ERR) + return REDISMODULE_ERR; + + if (RedisModule_CreateCommand(ctx,"datatype.get", datatype_get,"",1,1,1) == REDISMODULE_ERR) + return REDISMODULE_ERR; + + if (RedisModule_CreateCommand(ctx,"datatype.restore", datatype_restore,"deny-oom",1,1,1) == REDISMODULE_ERR) + return REDISMODULE_ERR; + + if (RedisModule_CreateCommand(ctx,"datatype.dump", datatype_dump,"",1,1,1) == REDISMODULE_ERR) + return REDISMODULE_ERR; + + return REDISMODULE_OK; +} diff --git a/tests/unit/moduleapi/datatype.tcl b/tests/unit/moduleapi/datatype.tcl new file mode 100644 index 000000000..c1da696d3 --- /dev/null +++ b/tests/unit/moduleapi/datatype.tcl @@ -0,0 +1,27 @@ +set testmodule [file normalize tests/modules/datatype.so] + +start_server {tags {"modules"}} { + r module load $testmodule + + test {DataType: Test module is sane, GET/SET work.} { + r datatype.set dtkey 100 stringval + assert {[r datatype.get dtkey] eq {100 stringval}} + } + + test {DataType: RM_SaveDataTypeToString(), RM_LoadDataTypeFromString() work} { + r datatype.set dtkey -1111 MyString + set encoded [r datatype.dump dtkey] + + r datatype.restore dtkeycopy $encoded + assert {[r datatype.get dtkeycopy] eq {-1111 MyString}} + } + + test {DataType: Handle truncated RM_LoadDataTypeFromString()} { + r datatype.set dtkey -1111 MyString + set encoded [r datatype.dump dtkey] + set truncated [string range $encoded 0 end-1] + + catch {r datatype.restore dtkeycopy $truncated} e + set e + } {*Invalid*} +} -- GitLab