diff options
Diffstat (limited to 'examples/redis-unstable/tests/modules/defragtest.c')
| -rw-r--r-- | examples/redis-unstable/tests/modules/defragtest.c | 475 |
1 files changed, 475 insertions, 0 deletions
diff --git a/examples/redis-unstable/tests/modules/defragtest.c b/examples/redis-unstable/tests/modules/defragtest.c new file mode 100644 index 0000000..8b9e182 --- /dev/null +++ b/examples/redis-unstable/tests/modules/defragtest.c @@ -0,0 +1,475 @@ +/* A module that implements defrag callback mechanisms. + */ + +#include "redismodule.h" +#include <stdlib.h> +#include <string.h> + +#define UNUSED(V) ((void) V) + +static RedisModuleType *FragType; + +struct FragObject { + unsigned long len; + void **values; + int maxstep; +}; + +/* Make sure we get the expected cursor */ +unsigned long int last_set_cursor = 0; + +unsigned long int datatype_attempts = 0; +unsigned long int datatype_defragged = 0; +unsigned long int datatype_raw_defragged = 0; +unsigned long int datatype_resumes = 0; +unsigned long int datatype_wrong_cursor = 0; +unsigned long int defrag_started = 0; +unsigned long int defrag_ended = 0; +unsigned long int global_strings_attempts = 0; +unsigned long int global_strings_defragged = 0; +unsigned long int global_dicts_resumes = 0; /* Number of dict defragmentation resumed from a previous break */ +unsigned long int global_subdicts_resumes = 0; /* Number of subdict defragmentation resumed from a previous break */ +unsigned long int global_dicts_attempts = 0; /* Number of attempts to defragment dictionary */ +unsigned long int global_dicts_defragged = 0; /* Number of dictionaries successfully defragmented */ +unsigned long int global_dicts_items_defragged = 0; /* Number of dictionaries items successfully defragmented */ + +unsigned long global_strings_len = 0; +RedisModuleString **global_strings = NULL; + +unsigned long global_dicts_len = 0; +RedisModuleDict **global_dicts = NULL; + +static void createGlobalStrings(RedisModuleCtx *ctx, unsigned long count) +{ + global_strings_len = count; + global_strings = RedisModule_Alloc(sizeof(RedisModuleString *) * count); + + for (unsigned long i = 0; i < count; i++) { + global_strings[i] = RedisModule_CreateStringFromLongLong(ctx, i); + } +} + +static int defragGlobalStrings(RedisModuleDefragCtx *ctx) +{ + unsigned long cursor = 0; + RedisModule_DefragCursorGet(ctx, &cursor); + + if (!global_strings_len) return 0; /* strings is empty. */ + RedisModule_Assert(cursor < global_strings_len); + for (; cursor < global_strings_len; cursor++) { + RedisModuleString *str = global_strings[cursor]; + if (!str) continue; + RedisModuleString *new = RedisModule_DefragRedisModuleString(ctx, str); + global_strings_attempts++; + if (new != NULL) { + global_strings[cursor] = new; + global_strings_defragged++; + } + + if (RedisModule_DefragShouldStop(ctx)) { + RedisModule_DefragCursorSet(ctx, cursor); + return 1; + } + } + return 0; +} + +static void createFragGlobalStrings(RedisModuleCtx *ctx) { + for (unsigned long i = 0; i < global_strings_len; i++) { + if (i % 2 == 1) { + RedisModule_FreeString(ctx, global_strings[i]); + global_strings[i] = NULL; + } + } +} + +static void createGlobalDicts(RedisModuleCtx *ctx, unsigned long count) { + global_dicts_len = count; + global_dicts = RedisModule_Alloc(sizeof(RedisModuleDict *) * count); + + /* Create some nested dictionaries: + * - Each main dict contains some subdicts. + * - Each sub-dict contains some strings. */ + for (unsigned long i = 0; i < count; i++) { + RedisModuleDict *dict = RedisModule_CreateDict(ctx); + for (unsigned long j = 0; j < 10; j++) { + /* Create sub dict. */ + RedisModuleDict *subdict = RedisModule_CreateDict(ctx); + for (unsigned long k = 0; k < 10; k++) { + RedisModuleString *str = RedisModule_CreateStringFromULongLong(ctx, k); + RedisModule_DictSet(subdict, str, str); + } + + RedisModuleString *key = RedisModule_CreateStringFromULongLong(ctx, j); + RedisModule_DictSet(dict, key, subdict); + RedisModule_FreeString(ctx, key); + } + global_dicts[i] = dict; + } +} + +static void freeFragGlobalSubDict(RedisModuleCtx *ctx, RedisModuleDict *subdict) { + char *key; + size_t keylen; + RedisModuleString *str; + RedisModuleDictIter *iter = RedisModule_DictIteratorStartC(subdict, "^", NULL, 0); + while ((key = RedisModule_DictNextC(iter, &keylen, (void**)&str))) { + RedisModule_FreeString(ctx, str); + } + RedisModule_FreeDict(ctx, subdict); + RedisModule_DictIteratorStop(iter); +} + +static void createFragGlobalDicts(RedisModuleCtx *ctx) { + char *key; + size_t keylen; + RedisModuleDict *subdict; + + for (unsigned long i = 0; i < global_dicts_len; i++) { + RedisModuleDict *dict = global_dicts[i]; + if (!dict) continue; + + /* Handle dictionaries differently based on their index in global_dicts array: + * 1. For odd indices (i % 2 == 1): Remove the entire dictionary. + * 2. For even indices: Keep the dictionary but remove half of its items. */ + if (i % 2 == 1) { + RedisModuleDictIter *iter = RedisModule_DictIteratorStartC(dict, "^", NULL, 0); + while ((key = RedisModule_DictNextC(iter, &keylen, (void**)&subdict))) { + freeFragGlobalSubDict(ctx, subdict); + } + RedisModule_FreeDict(ctx, dict); + global_dicts[i] = NULL; + RedisModule_DictIteratorStop(iter); + } else { + int key_index = 0; + RedisModuleDictIter *iter = RedisModule_DictIteratorStartC(dict, "^", NULL, 0); + while ((key = RedisModule_DictNextC(iter, &keylen, (void**)&subdict))) { + if (key_index++ % 2 == 1) { + freeFragGlobalSubDict(ctx, subdict); + RedisModule_DictReplaceC(dict, key, keylen, NULL); + } + } + RedisModule_DictIteratorStop(iter); + } + } +} + +static int defragGlobalSubDictValueCB(RedisModuleDefragCtx *ctx, void *data, unsigned char *key, size_t keylen, void **newptr) { + REDISMODULE_NOT_USED(key); + REDISMODULE_NOT_USED(keylen); + if (!data) return 0; + *newptr = RedisModule_DefragAlloc(ctx, data); + return 0; +} + +static int defragGlobalDictValueCB(RedisModuleDefragCtx *ctx, void *data, unsigned char *key, size_t keylen, void **newptr) { + REDISMODULE_NOT_USED(key); + REDISMODULE_NOT_USED(keylen); + static RedisModuleString *seekTo = NULL; + RedisModuleDict *subdict = data; + if (!subdict) return 0; + if (seekTo != NULL) global_subdicts_resumes++; + + *newptr = RedisModule_DefragRedisModuleDict(ctx, subdict, defragGlobalSubDictValueCB, &seekTo); + if (*newptr) global_dicts_items_defragged++; + /* Return 1 if seekTo is not NULL, indicating this node needs more defrag work. */ + return seekTo != NULL; +} + +static int defragGlobalDicts(RedisModuleDefragCtx *ctx) { + static RedisModuleString *seekTo = NULL; + static unsigned long dict_index = 0; + unsigned long cursor = 0; + + RedisModule_DefragCursorGet(ctx, &cursor); + if (cursor == 0) { /* Start a new defrag. */ + if (seekTo) { + RedisModule_FreeString(NULL, seekTo); + seekTo = NULL; + } + dict_index = 0; + } else { + global_dicts_resumes++; + } + + if (!global_dicts_len) return 0; /* dicts is empty. */ + RedisModule_Assert(dict_index < global_dicts_len); + for (; dict_index < global_dicts_len; dict_index++) { + RedisModuleDict *dict = global_dicts[dict_index]; + if (!dict) continue; + RedisModuleDict *new = RedisModule_DefragRedisModuleDict(ctx, dict, defragGlobalDictValueCB, &seekTo); + global_dicts_attempts++; + if (new != NULL) { + global_dicts[dict_index] = new; + global_dicts_defragged++; + } + + if (seekTo != NULL) { + /* Set cursor to 1 to indicate defragmentation is not finished. */ + RedisModule_DefragCursorSet(ctx, 1); + return 1; + } + } + + /* Set cursor to 0 to indicate completion. */ + dict_index = 0; + RedisModule_DefragCursorSet(ctx, 0); + return 0; +} + +typedef enum { DEFRAG_NOT_START, DEFRAG_STRING, DEFRAG_DICT } defrag_module_stage; +static int defragGlobal(RedisModuleDefragCtx *ctx) { + static defrag_module_stage stage = DEFRAG_NOT_START; + if (stage == DEFRAG_NOT_START) { + stage = DEFRAG_STRING; /* Start a new global defrag. */ + } + + if (stage == DEFRAG_STRING) { + if (defragGlobalStrings(ctx) != 0) return 1; + stage = DEFRAG_DICT; + } + if (stage == DEFRAG_DICT) { + if (defragGlobalDicts(ctx) != 0) return 1; + stage = DEFRAG_NOT_START; + } + return 0; +} + +static void defragStart(RedisModuleDefragCtx *ctx) { + REDISMODULE_NOT_USED(ctx); + defrag_started++; +} + +static void defragEnd(RedisModuleDefragCtx *ctx) { + REDISMODULE_NOT_USED(ctx); + defrag_ended++; +} + +static void FragInfo(RedisModuleInfoCtx *ctx, int for_crash_report) { + REDISMODULE_NOT_USED(for_crash_report); + + RedisModule_InfoAddSection(ctx, "stats"); + RedisModule_InfoAddFieldLongLong(ctx, "datatype_attempts", datatype_attempts); + RedisModule_InfoAddFieldLongLong(ctx, "datatype_defragged", datatype_defragged); + RedisModule_InfoAddFieldLongLong(ctx, "datatype_raw_defragged", datatype_raw_defragged); + RedisModule_InfoAddFieldLongLong(ctx, "datatype_resumes", datatype_resumes); + RedisModule_InfoAddFieldLongLong(ctx, "datatype_wrong_cursor", datatype_wrong_cursor); + RedisModule_InfoAddFieldLongLong(ctx, "global_strings_attempts", global_strings_attempts); + RedisModule_InfoAddFieldLongLong(ctx, "global_strings_defragged", global_strings_defragged); + RedisModule_InfoAddFieldLongLong(ctx, "global_dicts_resumes", global_dicts_resumes); + RedisModule_InfoAddFieldLongLong(ctx, "global_subdicts_resumes", global_subdicts_resumes); + RedisModule_InfoAddFieldLongLong(ctx, "global_dicts_attempts", global_dicts_attempts); + RedisModule_InfoAddFieldLongLong(ctx, "global_dicts_defragged", global_dicts_defragged); + RedisModule_InfoAddFieldLongLong(ctx, "global_dicts_items_defragged", global_dicts_items_defragged); + RedisModule_InfoAddFieldLongLong(ctx, "defrag_started", defrag_started); + RedisModule_InfoAddFieldLongLong(ctx, "defrag_ended", defrag_ended); +} + +struct FragObject *createFragObject(unsigned long len, unsigned long size, int maxstep) { + struct FragObject *o = RedisModule_Alloc(sizeof(*o)); + o->len = len; + o->values = RedisModule_Alloc(sizeof(RedisModuleString*) * len); + o->maxstep = maxstep; + + for (unsigned long i = 0; i < len; i++) { + o->values[i] = RedisModule_Calloc(1, size); + } + + return o; +} + +/* FRAG.RESETSTATS */ +static int fragResetStatsCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + REDISMODULE_NOT_USED(argv); + REDISMODULE_NOT_USED(argc); + + datatype_attempts = 0; + datatype_defragged = 0; + datatype_raw_defragged = 0; + datatype_resumes = 0; + datatype_wrong_cursor = 0; + global_strings_attempts = 0; + global_strings_defragged = 0; + global_dicts_resumes = 0; + global_subdicts_resumes = 0; + global_dicts_attempts = 0; + global_dicts_defragged = 0; + global_dicts_items_defragged = 0; + defrag_started = 0; + defrag_ended = 0; + + RedisModule_ReplyWithSimpleString(ctx, "OK"); + return REDISMODULE_OK; +} + +/* FRAG.CREATE key len size maxstep */ +static int fragCreateCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + if (argc != 5) + return RedisModule_WrongArity(ctx); + + RedisModuleKey *key = RedisModule_OpenKey(ctx,argv[1], + REDISMODULE_READ|REDISMODULE_WRITE); + int type = RedisModule_KeyType(key); + if (type != REDISMODULE_KEYTYPE_EMPTY) + { + return RedisModule_ReplyWithError(ctx, "ERR key exists"); + } + + long long len; + if ((RedisModule_StringToLongLong(argv[2], &len) != REDISMODULE_OK)) { + return RedisModule_ReplyWithError(ctx, "ERR invalid len"); + } + + long long size; + if ((RedisModule_StringToLongLong(argv[3], &size) != REDISMODULE_OK)) { + return RedisModule_ReplyWithError(ctx, "ERR invalid size"); + } + + long long maxstep; + if ((RedisModule_StringToLongLong(argv[4], &maxstep) != REDISMODULE_OK)) { + return RedisModule_ReplyWithError(ctx, "ERR invalid maxstep"); + } + + struct FragObject *o = createFragObject(len, size, maxstep); + RedisModule_ModuleTypeSetValue(key, FragType, o); + RedisModule_ReplyWithSimpleString(ctx, "OK"); + RedisModule_CloseKey(key); + + return REDISMODULE_OK; +} + +/* FRAG.create_frag_global len */ +static int fragCreateGlobalCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + UNUSED(argv); + if (argc != 2) + return RedisModule_WrongArity(ctx); + + long long glen; + if ((RedisModule_StringToLongLong(argv[1], &glen) != REDISMODULE_OK)) { + return RedisModule_ReplyWithError(ctx, "ERR invalid len"); + } + + createGlobalStrings(ctx, glen); + createGlobalDicts(ctx, glen); + createFragGlobalStrings(ctx); + createFragGlobalDicts(ctx); + RedisModule_ReplyWithSimpleString(ctx, "OK"); + return REDISMODULE_OK; +} + +void FragFree(void *value) { + struct FragObject *o = value; + + for (unsigned long i = 0; i < o->len; i++) + RedisModule_Free(o->values[i]); + RedisModule_Free(o->values); + RedisModule_Free(o); +} + +size_t FragFreeEffort(RedisModuleString *key, const void *value) { + REDISMODULE_NOT_USED(key); + + const struct FragObject *o = value; + return o->len; +} + +int FragDefrag(RedisModuleDefragCtx *ctx, RedisModuleString *key, void **value) { + unsigned long i = 0; + int steps = 0; + + int dbid = RedisModule_GetDbIdFromDefragCtx(ctx); + RedisModule_Assert(dbid != -1); + + RedisModule_Log(NULL, "notice", "Defrag key: %s", RedisModule_StringPtrLen(key, NULL)); + + /* Attempt to get cursor, validate it's what we're exepcting */ + if (RedisModule_DefragCursorGet(ctx, &i) == REDISMODULE_OK) { + if (i > 0) datatype_resumes++; + + /* Validate we're expecting this cursor */ + if (i != last_set_cursor) datatype_wrong_cursor++; + } else { + if (last_set_cursor != 0) datatype_wrong_cursor++; + } + + /* Attempt to defrag the object itself */ + datatype_attempts++; + struct FragObject *o = RedisModule_DefragAlloc(ctx, *value); + if (o == NULL) { + /* Not defragged */ + o = *value; + } else { + /* Defragged */ + *value = o; + datatype_defragged++; + } + + /* Deep defrag now */ + for (; i < o->len; i++) { + datatype_attempts++; + void *new = RedisModule_DefragAlloc(ctx, o->values[i]); + if (new) { + o->values[i] = new; + datatype_defragged++; + } + + if ((o->maxstep && ++steps > o->maxstep) || + ((i % 64 == 0) && RedisModule_DefragShouldStop(ctx))) + { + RedisModule_DefragCursorSet(ctx, i); + last_set_cursor = i; + return 1; + } + } + + /* Defrag the values array itself using RedisModule_DefragAllocRaw + * and RedisModule_DefragFreeRaw for testing purposes. */ + void *new_values = RedisModule_DefragAllocRaw(ctx, o->len * sizeof(void*)); + memcpy(new_values, o->values, o->len * sizeof(void*)); + RedisModule_DefragFreeRaw(ctx, o->values); + o->values = new_values; + datatype_raw_defragged++; + + last_set_cursor = 0; + return 0; +} + +int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + REDISMODULE_NOT_USED(argv); + REDISMODULE_NOT_USED(argc); + + if (RedisModule_Init(ctx, "defragtest", 1, REDISMODULE_APIVER_1) + == REDISMODULE_ERR) return REDISMODULE_ERR; + + if (RedisModule_GetTypeMethodVersion() < REDISMODULE_TYPE_METHOD_VERSION) { + return REDISMODULE_ERR; + } + + RedisModuleTypeMethods tm = { + .version = REDISMODULE_TYPE_METHOD_VERSION, + .free = FragFree, + .free_effort = FragFreeEffort, + .defrag = FragDefrag + }; + + FragType = RedisModule_CreateDataType(ctx, "frag_type", 0, &tm); + if (FragType == NULL) return REDISMODULE_ERR; + + if (RedisModule_CreateCommand(ctx, "frag.create", + fragCreateCommand, "write deny-oom", 1, 1, 1) == REDISMODULE_ERR) + return REDISMODULE_ERR; + + if (RedisModule_CreateCommand(ctx, "frag.create_frag_global", + fragCreateGlobalCommand, "write deny-oom", 1, 1, 1) == REDISMODULE_ERR) + return REDISMODULE_ERR; + + if (RedisModule_CreateCommand(ctx, "frag.resetstats", + fragResetStatsCommand, "write deny-oom", 1, 1, 1) == REDISMODULE_ERR) + return REDISMODULE_ERR; + + RedisModule_RegisterInfoFunc(ctx, FragInfo); + RedisModule_RegisterDefragFunc2(ctx, defragGlobal); + RedisModule_RegisterDefragCallbacks(ctx, defragStart, defragEnd); + + return REDISMODULE_OK; +} |