Unify RedisModuleDefragFunc and RedisModuleTypeDefragFunc

Author: sundb

src/defrag.c

@@ -1225,8 +1225,7 @@ static doneStatus defragModuleGlobals(monotime endtime, void *ctx) {
     defrag_module_ctx->module_ctx.cursor = &defrag_module_ctx->cursor;
 
     /* Call the module's defrag callback function and check if more work remains. */
-    defrag_module_ctx->module->defrag_cb(&defrag_module_ctx->module_ctx);
-    if (defrag_module_ctx->cursor != 0)
+    if (defrag_module_ctx->module->defrag_cb(&defrag_module_ctx->module_ctx) != 0)
         return DEFRAG_NOT_DONE;
 
     return DEFRAG_DONE;

src/redismodule.h

@@ -839,7 +839,8 @@ typedef struct RedisModuleDefragCtx RedisModuleDefragCtx;
 /* Function pointers needed by both the core and modules, these needs to be
  * exposed since you can't cast a function pointer to (void *). */
 typedef void (*RedisModuleInfoFunc)(RedisModuleInfoCtx *ctx, int for_crash_report);
-typedef void (*RedisModuleDefragFunc)(RedisModuleDefragCtx *ctx);
+typedef int (*RedisModuleDefragFunc)(RedisModuleDefragCtx *ctx);
+typedef void (*RedisModuleDefragEventFunc)(RedisModuleDefragCtx *ctx);
 typedef void (*RedisModuleUserChangedFunc) (uint64_t client_id, void *privdata);
 
 /* ------------------------- End of common defines ------------------------ */
@@ -1305,7 +1306,7 @@ REDISMODULE_API int *(*RedisModule_GetCommandKeys)(RedisModuleCtx *ctx, RedisMod
 REDISMODULE_API int *(*RedisModule_GetCommandKeysWithFlags)(RedisModuleCtx *ctx, RedisModuleString **argv, int argc, int *num_keys, int **out_flags) REDISMODULE_ATTR;
 REDISMODULE_API const char *(*RedisModule_GetCurrentCommandName)(RedisModuleCtx *ctx) REDISMODULE_ATTR;
 REDISMODULE_API int (*RedisModule_RegisterDefragFunc)(RedisModuleCtx *ctx, RedisModuleDefragFunc func) REDISMODULE_ATTR;
-REDISMODULE_API int (*RedisModule_RegisterDefragCallbacks)(RedisModuleCtx *ctx, RedisModuleDefragFunc start, RedisModuleDefragFunc end) REDISMODULE_ATTR;
+REDISMODULE_API int (*RedisModule_RegisterDefragCallbacks)(RedisModuleCtx *ctx, RedisModuleDefragEventFunc start, RedisModuleDefragEventFunc end) REDISMODULE_ATTR;
 REDISMODULE_API void *(*RedisModule_DefragAlloc)(RedisModuleDefragCtx *ctx, void *ptr) REDISMODULE_ATTR;
 REDISMODULE_API void *(*RedisModule_DefragAllocRaw)(RedisModuleDefragCtx *ctx, size_t size) REDISMODULE_ATTR;
 REDISMODULE_API void (*RedisModule_DefragFreeRaw)(RedisModuleDefragCtx *ctx, void *ptr) REDISMODULE_ATTR;

tests/modules/defragtest.c

@@ -39,21 +39,26 @@ 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);
 
     RedisModule_Assert(cursor < global_strings_len);
-    RedisModuleString *new = RedisModule_DefragRedisModuleString(ctx, global_strings[cursor]);
-    global_attempts++;
-    if (new != NULL) {
-        global_strings[cursor] = new;
-        global_defragged++;
-    }
+    for (; cursor < global_strings_len; cursor++) {
+        RedisModuleString *new = RedisModule_DefragRedisModuleString(ctx, global_strings[cursor]);
+        global_attempts++;
+        if (new != NULL) {
+            global_strings[cursor] = new;
+            global_defragged++;
+        }
 
-    cursor++;
-    RedisModule_DefragCursorSet(ctx, cursor == global_strings_len ? 0 : cursor);
+        if (cursor % 16 == 0 && RedisModule_DefragShouldStop(ctx)) {
+            RedisModule_DefragCursorSet(ctx, cursor);
+            return 1;
+        }
+    }
+    return 0;
 }
 
 static void defragStart(RedisModuleDefragCtx *ctx) {