Simplify RedisModule DefragRedisModuleDict API

Author: sundb

src/module.c

@@ -13930,8 +13930,8 @@ int moduleDefragRaxNode(raxNode **noderef) {
 }
 
 /* Defrag a RedisModuleDict previously allocated by RM_CreateDict. */
-typedef void *(*RedisModuleDefragDictFunc)(RedisModuleDefragCtx *ctx, RedisModuleDict *dict, void *data, void *privdata);
-RedisModuleDict *RM_DefragRedisModuleDict(RedisModuleDefragCtx *ctx, RedisModuleDict *dict, RedisModuleDefragDictFunc callback, void *privdata) {
+RedisModuleDict *RM_DefragRedisModuleDict(RedisModuleDefragCtx *ctx, RedisModuleDict *dict) {
+    UNUSED(ctx);
     RedisModuleDict *newdict = NULL;
     rax* newrax;
     raxIterator ri;
@@ -13944,13 +13944,7 @@ RedisModuleDict *RM_DefragRedisModuleDict(RedisModuleDefragCtx *ctx, RedisModule
     ri.node_cb = moduleDefragRaxNode;
     moduleDefragRaxNode(&dict->rax->head);
     raxSeek(&ri,"^",NULL,0);
-    while (raxNext(&ri)) {
-        void *newdata = NULL;
-        if (callback)
-            newdata = callback(ctx, dict, ri.data, privdata);
-        if (newdata)
-            raxSetData(ri.node, ri.data=newdata);
-    }
+    while (raxNext(&ri)) {}
     raxStop(&ri);
     return dict;
 }

src/redismodule.h

@@ -927,7 +927,6 @@ typedef int (*RedisModuleConfigSetEnumFunc)(const char *name, int val, void *pri
 typedef int (*RedisModuleConfigApplyFunc)(RedisModuleCtx *ctx, void *privdata, RedisModuleString **err);
 typedef void (*RedisModuleOnUnblocked)(RedisModuleCtx *ctx, RedisModuleCallReply *reply, void *private_data);
 typedef int (*RedisModuleAuthCallback)(RedisModuleCtx *ctx, RedisModuleString *username, RedisModuleString *password, RedisModuleString **err);
-typedef void (*RedisModuleDefragDictFunc)(RedisModuleDefragCtx *ctx, RedisModuleDict *dict, void *data, void *privdata);
 
 typedef struct RedisModuleTypeMethods {
     uint64_t version;
@@ -1311,7 +1310,7 @@ REDISMODULE_API void *(*RedisModule_DefragAlloc)(RedisModuleDefragCtx *ctx, void
 REDISMODULE_API void *(*RedisModule_DefragAllocRaw)(RedisModuleDefragCtx *ctx, size_t size) REDISMODULE_ATTR;
 REDISMODULE_API void (*RedisModule_DefragFreeRaw)(RedisModuleDefragCtx *ctx, void *ptr) REDISMODULE_ATTR;
 REDISMODULE_API RedisModuleString *(*RedisModule_DefragRedisModuleString)(RedisModuleDefragCtx *ctx, RedisModuleString *str) REDISMODULE_ATTR;
-REDISMODULE_API RedisModuleDict *(*RedisModule_DefragRedisModuleDict)(RedisModuleDefragCtx *ctx, RedisModuleDict *dict, RedisModuleDefragDictFunc *callback, void *privdata) REDISMODULE_ATTR;
+REDISMODULE_API RedisModuleDict *(*RedisModule_DefragRedisModuleDict)(RedisModuleDefragCtx *ctx, RedisModuleDict *dict) REDISMODULE_ATTR;
 REDISMODULE_API int (*RedisModule_DefragShouldStop)(RedisModuleDefragCtx *ctx) REDISMODULE_ATTR;
 REDISMODULE_API int (*RedisModule_DefragCursorSet)(RedisModuleDefragCtx *ctx, unsigned long cursor) REDISMODULE_ATTR;
 REDISMODULE_API int (*RedisModule_DefragCursorGet)(RedisModuleDefragCtx *ctx, unsigned long *cursor) REDISMODULE_ATTR;

tests/modules/defragtest.c

@@ -29,6 +29,9 @@ unsigned long int global_defragged = 0;
 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;
@@ -39,7 +42,7 @@ static void createGlobalStrings(RedisModuleCtx *ctx, unsigned long count)
     }
 }
 
-static void defragGlobalStrings(RedisModuleDefragCtx *ctx)
+static int defragGlobalStrings(RedisModuleDefragCtx *ctx)
 {
     unsigned long cursor = 0;
     RedisModule_DefragCursorGet(ctx, &cursor);
@@ -54,6 +57,56 @@ static void defragGlobalStrings(RedisModuleDefragCtx *ctx)
 
     cursor++;
     RedisModule_DefragCursorSet(ctx, cursor == global_strings_len ? 0 : cursor);
+    return cursor;
+}
+
+static void createGlobalDicts(RedisModuleCtx *ctx, unsigned long count) {
+    global_dicts_len = count;
+    global_dicts = RedisModule_Alloc(sizeof(RedisModuleDict *) * count);
+
+    for (unsigned long i = 0; i < count; i++) {
+        RedisModuleDict *dict = RedisModule_CreateDict(ctx);
+        for (unsigned long j = 0; j < 100; j ++) {
+            RedisModuleString *str = RedisModule_CreateStringFromULongLong(ctx, j);
+            RedisModule_DictSet(dict, str, str);
+        }
+        global_dicts[i] = dict;
+    }
+}
+
+static int defragGlobalDicts(RedisModuleDefragCtx *ctx) {
+    unsigned long cursor = 0;
+    RedisModule_DefragCursorGet(ctx, &cursor);
+
+    RedisModule_Assert(cursor < global_dicts_len);
+    RedisModuleDict *new = RedisModule_DefragRedisModuleDict(ctx, global_dicts[cursor]);
+    global_attempts++;
+    if (new != NULL) {
+        global_dicts[cursor] = new;
+        global_defragged++;
+    }
+
+    cursor++;
+    RedisModule_DefragCursorSet(ctx, cursor == global_dicts_len ? 0 : cursor);
+    return cursor;
+}
+
+typedef enum { DEFRAG_NOT_START, DEFRAG_STRING, DEFRAG_DICT } defrag_module_stage;
+static void 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) {
+            stage = DEFRAG_DICT;
+        }
+    } else {
+        if (defragGlobalDicts(ctx) == 0) {
+            stage = DEFRAG_NOT_START;
+        }
+    }
 }
 
 static void defragStart(RedisModuleDefragCtx *ctx) {
@@ -243,6 +296,7 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
     }
 
     createGlobalStrings(ctx, glen);
+    createGlobalDicts(ctx, glen);
 
     RedisModuleTypeMethods tm = {
             .version = REDISMODULE_TYPE_METHOD_VERSION,
@@ -263,7 +317,7 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
         return REDISMODULE_ERR;
 
     RedisModule_RegisterInfoFunc(ctx, FragInfo);
-    RedisModule_RegisterDefragFunc(ctx, defragGlobalStrings);
+    RedisModule_RegisterDefragFunc(ctx, defragGlobal);
     RedisModule_RegisterDefragCallbacks(ctx, defragStart, defragEnd);
 
     return REDISMODULE_OK;