Add RM_DefragRedisModuleDict module API

Author: sundb

src/module.c

@@ -13913,6 +13913,42 @@ RedisModuleString *RM_DefragRedisModuleString(RedisModuleDefragCtx *ctx, RedisMo
     return activeDefragStringOb(str);
 }
 
+/* Defrag callback for radix tree iterator, called for each node,
+ * used in order to defrag the nodes allocations. */
+int moduleDefragRaxNode(raxNode **noderef) {
+    raxNode *newnode = activeDefragAlloc(*noderef);
+    if (newnode) {
+        *noderef = newnode;
+        return 1;
+    }
+    return 0;
+}
+
+/* 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 *newdict = NULL;
+    rax* newrax;
+    raxIterator ri;
+
+    if ((newdict = activeDefragAlloc(dict)))
+        dict = newdict;
+    if ((newrax = activeDefragAlloc(dict->rax)))
+        dict->rax = newrax;
+    raxStart(&ri,dict->rax);
+    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);
+    }
+    raxStop(&ri);
+    return dict;
+}
 
 /* Perform a late defrag of a module datatype key.
  *

src/redismodule.h

@@ -927,6 +927,7 @@ 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;
@@ -1310,6 +1311,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 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;