Making the module global-defrag callback usable

src/defrag.c

@@ -967,8 +967,16 @@ static doneStatus defragLuaScripts(monotime endtime, void *target, void *privdat
 static doneStatus defragModuleGlobals(monotime endtime, void *target, void *privdata) {
     UNUSED(target);
     UNUSED(privdata);
-    if (endtime == 0) return DEFRAG_NOT_DONE; // required initialization
-    moduleDefragGlobals();
+    if (endtime == 0) {
+        // Init: clear each module's done flag so the stage visits every module again.
+        moduleDefragGlobalsStart();
+        return DEFRAG_NOT_DONE;
+    }
+    /* Reschedule the stage if a module still has work (its cursor is non-zero, e.g. work
+     * queued on its own threads) or we ran out of time.  Modules already done this cycle are
+     * skipped on the next call, so we resume with the remaining ones. */
+    int more_work = moduleDefragGlobals(endtime);
+    if (more_work || getMonotonicUs() >= endtime) return DEFRAG_NOT_DONE;
     return DEFRAG_DONE;
 }
 

src/module.c

@@ -2449,6 +2449,8 @@ void VM_SetModuleAttribs(ValkeyModuleCtx *ctx, const char *name, int ver, int ap
     module->options = 0;
     module->info_cb = 0;
     module->defrag_cb = 0;
+    module->defrag_cursor = 0;
+    module->defrag_done_this_cycle = 0;
     module->loadmod = NULL;
     module->num_commands_with_acl_categories = 0;
     module->onload = 1;
@@ -14986,22 +14988,32 @@ struct ValkeyModuleDefragCtx {
 
 /* Register a defrag callback for global data, i.e. anything that the module
  * may allocate that is not tied to a specific data type.
+ *
+ * Unlike the per-key data type defrag callback, the global callback is invoked
+ * with a time limit: it should call VM_DefragShouldStop() periodically and
+ * return once that returns non-zero. To resume where it left off on the next
+ * call, it uses VM_DefragCursorSet()/VM_DefragCursorGet(). A stored cursor of 0
+ * means "done"; a non-zero cursor tells the defrag process there is more work
+ * and the callback will be invoked again. The callback MUST store a cursor of 0
+ * once it has finished, otherwise it will keep being invoked.
  */
 int VM_RegisterDefragFunc(ValkeyModuleCtx *ctx, ValkeyModuleDefragFunc cb) {
     ctx->module->defrag_cb = cb;
     return VALKEYMODULE_OK;
 }
 
-/* When the data type defrag callback iterates complex structures, this
- * function should be called periodically. A zero (false) return
- * indicates the callback may continue its work. A non-zero value (true)
- * indicates it should stop.
+/* When a defrag callback iterates complex structures, this function should be
+ * called periodically. A zero (false) return indicates the callback may
+ * continue its work. A non-zero value (true) indicates it should stop.
  *
- * When stopped, the callback may use VM_DefragCursorSet() to store its
+ * When stopped, the callback should use VM_DefragCursorSet() to store its
  * position so it can later use VM_DefragCursorGet() to resume defragging.
  *
- * When stopped and more work is left to be done, the callback should
- * return 1. Otherwise, it should return 0.
+ * How "more work remains" is signalled depends on the callback type:
+ *  - the per-key data type defrag callback returns 1 if stopped with more work
+ *    left, or 0 when done;
+ *  - the global defrag callback returns nothing; instead a stored cursor of 0
+ *    means done and a non-zero cursor means more work remains.
  *
  * NOTE: Modules should consider the frequency in which this function is called,
  * so it generally makes sense to do small batches of work in between calls.
@@ -15012,25 +15024,25 @@ int VM_DefragShouldStop(ValkeyModuleDefragCtx *ctx) {
 
 /* Store an arbitrary cursor value for future re-use.
  *
- * This should only be called if VM_DefragShouldStop() has returned a non-zero
- * value and the defrag callback is about to exit without fully iterating its
- * data type.
- *
- * This behavior is reserved to cases where late defrag is performed. Late
- * defrag is selected for keys that implement the `free_effort` callback and
- * return a `free_effort` value that is larger than the defrag
- * 'active-defrag-max-scan-fields' configuration directive.
- *
- * Smaller keys, keys that do not implement `free_effort` or the global
- * defrag callback are not called in late-defrag mode. In those cases, a
- * call to this function will return VALKEYMODULE_ERR.
- *
- * The cursor may be used by the module to represent some progress into the
- * module's data type. Modules may also store additional cursor-related
- * information locally and use the cursor as a flag that indicates when
- * traversal of a new key begins. This is possible because the API makes
- * a guarantee that concurrent defragmentation of multiple keys will
- * not be performed.
+ * This is used to resume defragmentation across callback invocations, and is
+ * available in two cases:
+ *  - "late defrag" of a data type key. Late defrag is selected for keys that
+ *    implement the `free_effort` callback and return a value larger than the
+ *    'active-defrag-max-scan-fields' configuration directive. Smaller keys, and
+ *    keys that do not implement `free_effort`, are not defragged in late mode,
+ *    and a call to this function for them returns VALKEYMODULE_ERR.
+ *  - the global defrag callback (registered via VM_RegisterDefragFunc), which
+ *    is always given a cursor. There, a stored cursor of 0 means the callback
+ *    is done and a non-zero value means it should be invoked again to continue.
+ *
+ * The cursor may be used by the module to represent some progress into its
+ * data. Modules may also store additional cursor-related information locally
+ * and use the cursor as a flag that indicates when traversal of a new key
+ * begins. This is possible because the API guarantees that concurrent
+ * defragmentation of multiple keys will not be performed.
+ *
+ * Returns VALKEYMODULE_ERR if no cursor is available for this callback (see
+ * above), VALKEYMODULE_OK otherwise.
  */
 int VM_DefragCursorSet(ValkeyModuleDefragCtx *ctx, unsigned long cursor) {
     if (!ctx->cursor) return VALKEYMODULE_ERR;
@@ -15041,9 +15053,9 @@ int VM_DefragCursorSet(ValkeyModuleDefragCtx *ctx, unsigned long cursor) {
 
 /* Fetch a cursor value that has been previously stored using VM_DefragCursorSet().
  *
- * If not called for a late defrag operation, VALKEYMODULE_ERR will be returned and
- * the cursor should be ignored. See VM_DefragCursorSet() for more details on
- * defrag cursors.
+ * Returns VALKEYMODULE_ERR if no cursor is available for this callback (see
+ * VM_DefragCursorSet() for when that is the case), in which case the cursor
+ * should be ignored. On the first invocation the stored cursor is 0.
  */
 int VM_DefragCursorGet(ValkeyModuleDefragCtx *ctx, unsigned long *cursor) {
     if (!ctx->cursor) return VALKEYMODULE_ERR;
@@ -15145,20 +15157,64 @@ int moduleDefragValue(robj *key, robj *value, int dbid) {
     return 1;
 }
 
-/* Call registered module API defrag functions */
-void moduleDefragGlobals(void) {
-    if (listLength(modules) == 0) return;
+/* Index of the module to start from on the next moduleDefragGlobals() call. Advanced past the
+ * module we stopped on so a module with ongoing work doesn't starve the others. */
+static unsigned long defrag_module_start_idx = 0;
+
+/* Called at stage init (endtime==0) to start a new global defrag pass.  Clears each module's
+ * done flag so every module is visited again, and resets the round-robin start position.  Cursors
+ * are not touched here: a module owns its cursor and may carry progress across cycles. */
+void moduleDefragGlobalsStart(void) {
+    defrag_module_start_idx = 0;
 
     listIter li;
     listNode *ln;
 
     listRewind(modules, &li);
     while ((ln = listNext(&li)) != NULL) {
         struct ValkeyModule *module = listNodeValue(ln);
+        module->defrag_done_this_cycle = 0;
+    }
+}
+
+/* Invoke each module's global defrag callback, forwarding 'endtime' so the callback can bound its
+ * own latency via VM_DefragShouldStop().  Each module is given a persistent cursor
+ * (module->defrag_cursor) to save progress with VM_DefragCursorSet() and resume on a later call.
+ *
+ * The cursor is also the module's "more work" signal, following the convention used elsewhere in
+ * defrag: a non-zero cursor means the module wants to be called again (scan not finished, or work
+ * still draining on its own threads); a zero cursor means it is done for this cycle.  A module done
+ * this cycle sets defrag_done_this_cycle and is skipped until the next cycle clears it.
+ *
+ * When the deadline is hit mid-iteration we resume on the next call from the module after the one
+ * we stopped on (defrag_module_start_idx), so a module that keeps consuming the deadline can't
+ * starve the modules after it.  The done flags and start index live outside the module structs, so
+ * they are unaffected if a module is unloaded between calls.
+ *
+ * Returns 1 if any module still has work to do, 0 otherwise. */
+int moduleDefragGlobals(monotime endtime) {
+    int more_work = 0;
+    unsigned long count = listLength(modules);
+    if (count == 0) return more_work;
+
+    for (unsigned long n = 0; n < count; n++) {
+        unsigned long idx = (defrag_module_start_idx + n) % count;
+        struct ValkeyModule *module = listNodeValue(listIndex(modules, idx));
         if (!module->defrag_cb) continue;
-        ValkeyModuleDefragCtx defrag_ctx = {0, NULL, NULL, -1};
+        if (module->defrag_done_this_cycle) continue;
+        ValkeyModuleDefragCtx defrag_ctx = {endtime, &module->defrag_cursor, NULL, -1};
         module->defrag_cb(&defrag_ctx);
+        if (module->defrag_cursor != 0) {
+            more_work = 1;
+        } else {
+            module->defrag_done_this_cycle = 1;
+        }
+        if (endtime != 0 && getMonotonicUs() >= endtime) {
+            defrag_module_start_idx = (idx + 1) % count;
+            break;
+        }
     }
+    return more_work;
 }
 
 /* Returns the name of the key currently being processed.

src/module.h

@@ -117,6 +117,8 @@ typedef struct ValkeyModule {
     int blocked_clients;                  /* Count of ValkeyModuleBlockedClient in this module. */
     ValkeyModuleInfoFunc info_cb;         /* Callback for module to add INFO fields. */
     ValkeyModuleDefragFunc defrag_cb;     /* Callback for global data defrag. */
+    unsigned long defrag_cursor;          /* Global defrag cursor, owned by the module, persists across cycles. */
+    int defrag_done_this_cycle;           /* Global defrag: module is done this cycle, skip until next cycle. */
     struct moduleLoadQueueEntry *loadmod; /* Module load arguments for config rewrite. */
     int num_commands_with_acl_categories; /* Number of commands in this module included in acl categories */
     int onload;                           /* Flag to identify if the call is being made from Onload (0 or 1) */
@@ -242,7 +244,8 @@ size_t moduleGetMemUsage(robj *key, robj *val, size_t sample_size, int dbid);
 robj *moduleTypeDupOrReply(client *c, robj *fromkey, robj *tokey, int todb, robj *value);
 int moduleDefragValue(robj *key, robj *obj, int dbid);
 int moduleLateDefrag(robj *key, robj *value, unsigned long *cursor, monotime endtime, int dbid);
-void moduleDefragGlobals(void);
+void moduleDefragGlobalsStart(void);
+int moduleDefragGlobals(monotime endtime);
 void *moduleGetHandleByName(char *modulename);
 int moduleIsModuleCommand(void *module_handle, struct serverCommand *cmd);
 void freeClientModuleData(client *c);

tests/modules/Makefile

@@ -47,6 +47,7 @@ TEST_MODULES = \
     test_lazyfree.so \
     timer.so \
     defragtest.so \
+    defragglobalbusy.so \
     keyspecs.so \
     hash.so \
     hash_stringref.so \

tests/modules/defragglobalbusy.c

@@ -0,0 +1,41 @@
+/* A module whose global defrag callback never finishes: it consumes the whole
+ * deadline on every invocation and always leaves a non-zero cursor. Used
+ * together with defragtest to check that such a module does not starve the
+ * global defrag callbacks of other modules (see defrag.tcl).
+ */
+
+#include "valkeymodule.h"
+
+/* Number of times our global defrag callback was invoked. Exposed via INFO so
+ * the test can confirm the busy module actually ran. */
+unsigned long long busy_calls = 0;
+
+static void defragBusyGlobal(ValkeyModuleDefragCtx *ctx) {
+    busy_calls++;
+    /* Burn the rest of the deadline, then report we still have work by leaving
+     * a non-zero cursor. This models a module that never drains within a
+     * single defrag cycle. */
+    while (!ValkeyModule_DefragShouldStop(ctx)) {
+        /* spin until the deadline is reached */
+    }
+    ValkeyModule_DefragCursorSet(ctx, 1);
+}
+
+static void BusyInfo(ValkeyModuleInfoCtx *ctx, int for_crash_report) {
+    VALKEYMODULE_NOT_USED(for_crash_report);
+    ValkeyModule_InfoAddSection(ctx, "stats");
+    ValkeyModule_InfoAddFieldULongLong(ctx, "busy_calls", busy_calls);
+}
+
+int ValkeyModule_OnLoad(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc) {
+    VALKEYMODULE_NOT_USED(argv);
+    VALKEYMODULE_NOT_USED(argc);
+
+    if (ValkeyModule_Init(ctx, "defragglobalbusy", 1, VALKEYMODULE_APIVER_1) == VALKEYMODULE_ERR)
+        return VALKEYMODULE_ERR;
+
+    ValkeyModule_RegisterInfoFunc(ctx, BusyInfo);
+    ValkeyModule_RegisterDefragFunc(ctx, defragBusyGlobal);
+
+    return VALKEYMODULE_OK;
+}

tests/modules/defragtest.c

@@ -14,16 +14,22 @@ struct FragObject {
 
 /* Make sure we get the expected cursor */
 unsigned long int last_set_cursor = 0;
+unsigned long int last_set_global_cursor = 0;
 
 unsigned long int datatype_attempts = 0;
 unsigned long int datatype_defragged = 0;
 unsigned long int datatype_resumes = 0;
 unsigned long int datatype_wrong_cursor = 0;
 unsigned long int global_attempts = 0;
 unsigned long int global_defragged = 0;
+unsigned long int global_resumes = 0;
+unsigned long int global_wrong_cursor = 0;
 
 int global_strings_len = 0;
 ValkeyModuleString **global_strings = NULL;
+/* If non-zero, the global defrag callback stops after this many strings per
+ * invocation, forcing it to resume via the cursor on later calls. */
+int global_maxstep = 0;
 
 static void createGlobalStrings(ValkeyModuleCtx *ctx, int count)
 {
@@ -37,14 +43,39 @@ static void createGlobalStrings(ValkeyModuleCtx *ctx, int count)
 
 static void defragGlobalStrings(ValkeyModuleDefragCtx *ctx)
 {
-    for (int i = 0; i < global_strings_len; i++) {
+    unsigned long i = 0;
+    int steps = 0;
+
+    /* Resume from the saved cursor, validating it's what we set last time. */
+    if (ValkeyModule_DefragCursorGet(ctx, &i) == VALKEYMODULE_OK) {
+        if (i > 0) global_resumes++;
+        if (i != last_set_global_cursor) global_wrong_cursor++;
+    } else {
+        if (last_set_global_cursor != 0) global_wrong_cursor++;
+    }
+
+    for (; i < (unsigned long)global_strings_len; i++) {
         ValkeyModuleString *new = ValkeyModule_DefragValkeyModuleString(ctx, global_strings[i]);
         global_attempts++;
         if (new != NULL) {
             global_strings[i] = new;
             global_defragged++;
         }
+
+        /* Stop after maxstep strings, or when out of time, saving progress in
+         * the cursor so the next invocation resumes here. */
+        if ((global_maxstep && ++steps >= global_maxstep) ||
+            ((i % 64 == 0) && ValkeyModule_DefragShouldStop(ctx)))
+        {
+            ValkeyModule_DefragCursorSet(ctx, i + 1);
+            last_set_global_cursor = i + 1;
+            return;
+        }
     }
+
+    /* Finished: reset the cursor to 0 so core sees this module as done. */
+    ValkeyModule_DefragCursorSet(ctx, 0);
+    last_set_global_cursor = 0;
 }
 
 static void FragInfo(ValkeyModuleInfoCtx *ctx, int for_crash_report) {
@@ -57,6 +88,8 @@ static void FragInfo(ValkeyModuleInfoCtx *ctx, int for_crash_report) {
     ValkeyModule_InfoAddFieldLongLong(ctx, "datatype_wrong_cursor", datatype_wrong_cursor);
     ValkeyModule_InfoAddFieldLongLong(ctx, "global_attempts", global_attempts);
     ValkeyModule_InfoAddFieldLongLong(ctx, "global_defragged", global_defragged);
+    ValkeyModule_InfoAddFieldLongLong(ctx, "global_resumes", global_resumes);
+    ValkeyModule_InfoAddFieldLongLong(ctx, "global_wrong_cursor", global_wrong_cursor);
 }
 
 struct FragObject *createFragObject(unsigned long len, unsigned long size, int maxstep) {
@@ -83,6 +116,8 @@ static int fragResetStatsCommand(ValkeyModuleCtx *ctx, ValkeyModuleString **argv
     datatype_wrong_cursor = 0;
     global_attempts = 0;
     global_defragged = 0;
+    global_resumes = 0;
+    global_wrong_cursor = 0;
 
     ValkeyModule_ReplyWithSimpleString(ctx, "OK");
     return VALKEYMODULE_OK;
@@ -204,10 +239,19 @@ int ValkeyModule_OnLoad(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int arg
     }
 
     long long glen;
-    if (argc != 1 || ValkeyModule_StringToLongLong(argv[0], &glen) == VALKEYMODULE_ERR) {
+    if (argc < 1 || argc > 2 || ValkeyModule_StringToLongLong(argv[0], &glen) == VALKEYMODULE_ERR) {
         return VALKEYMODULE_ERR;
     }
 
+    /* Optional 2nd arg: global defrag step limit per callback invocation. */
+    if (argc == 2) {
+        long long gmaxstep;
+        if (ValkeyModule_StringToLongLong(argv[1], &gmaxstep) == VALKEYMODULE_ERR) {
+            return VALKEYMODULE_ERR;
+        }
+        global_maxstep = gmaxstep;
+    }
+
     createGlobalStrings(ctx, glen);
 
     ValkeyModuleTypeMethods tm = {