Separate RDB snapshotting from atomic slot migration (#2533)

When we adding atomic slot migration in #1949, we reused a lot of rdb save code,
it was an easier way to implement ASM in the first time, but it comes with some
side effect. Like we are using CHILD_TYPE_RDB to do the fork, we use rdb.c/rdb.h
function to save the snapshot, these mess up the logs (we will print some logs
saying we are doing RDB stuff) and mess up the info fields (we will say we are
rdb_bgsave_in_progress but actually we are doing slot migration).

In addition, it makes the code difficult to maintain. The rdb_save method uses
a lot of rdb_* variables, but we are actually doing slot migration. If we want
to support one fork with multiple target nodes, we need to rewrite these code
for a better cleanup.

Note that the changes to rdb.c/rdb.h are reverting previous changes from when
we was reusing this code for slot migration. The slot migration snapshot logic
is similar to the previous diskless replication. We use pipe to transfer the
snapshot data from the child process to the parent process.

Interface changes:
- New slot_migration_fork_in_progress info field.
- New cow_size field in CLUSTER GETSLOTMIGRATIONS command.
- Also add slot migration fork to the cluster class trace latency.

Signed-off-by: Binbin <[email protected]>
Signed-off-by: Jacob Murphy <[email protected]>
Co-authored-by: Jacob Murphy <[email protected]>

Author: enjoy-binbin

src/aof.c

@@ -2280,6 +2280,8 @@ int rewriteObjectRio(rio *aof, robj *o, int db_num) {
     return C_OK;
 }
 
+/* This function is currently used in slot migration to rewrite the corresponding
+ * slot hashtable to rio. */
 int rewriteSlotToAppendOnlyFileRio(rio *aof, int db_num, int hashslot, size_t *key_count) {
     long long updated_time = 0;
 
@@ -2302,7 +2304,7 @@ int rewriteSlotToAppendOnlyFileRio(rio *aof, int db_num, int hashslot, size_t *k
         if (key_count && ((*key_count)++ & 1023) == 0) {
             long long now = mstime();
             if (now - updated_time >= 1000) {
-                sendChildInfo(CHILD_INFO_TYPE_CURRENT_INFO, *key_count, "AOF rewrite");
+                sendChildInfo(CHILD_INFO_TYPE_CURRENT_INFO, *key_count, "Slot migration");
                 updated_time = now;
             }
         }

src/childinfo.c

@@ -131,6 +131,8 @@ void updateChildInfo(childInfoType information_type, size_t cow, monotime cow_up
         server.stat_rdb_cow_bytes = server.stat_current_cow_peak;
     } else if (information_type == CHILD_INFO_TYPE_MODULE_COW_SIZE) {
         server.stat_module_cow_bytes = server.stat_current_cow_peak;
+    } else if (information_type == CHILD_INFO_TYPE_SLOT_MIGRATION_COW_SIZE) {
+        server.stat_slot_migration_cow_bytes = server.stat_current_cow_peak;
     } else if (information_type == CHILD_INFO_TYPE_REPL_OUTPUT_BYTES) {
         server.stat_net_repl_output_bytes += (long long)repl_output_bytes;
     }

src/cluster_migrateslots.c

@@ -9,6 +9,9 @@
 #include "module.h"
 #include "functions.h"
 
+#include <sys/wait.h>
+#include <fcntl.h>
+
 typedef enum slotMigrationJobState {
     /* Importing states */
     SLOT_IMPORT_WAIT_ACK,
@@ -76,6 +79,7 @@ typedef struct slotMigrationJob {
                                                * cleanup is done. */
     sds description;                          /* Description, used for
                                                * logging. */
+    size_t stat_cow_bytes;                    /* Copy on write bytes during slot migration fork. */
 
     /* State needed during client establishment */
     connection *conn; /* Connection to slot import source node. */
@@ -1264,9 +1268,8 @@ bool shouldRewriteHashtableIndex(int didx, hashtable *ht, void *privdata) {
 }
 
 /* Contains the logic run on the child process during the snapshot phase. */
-int childSnapshotForSyncSlot(int req, rio *rdb, void *privdata) {
-    UNUSED(req);
-    list *slot_ranges = privdata;
+int childSnapshotForSyncSlot(rio *aof, slotMigrationJob *job) {
+    list *slot_ranges = job->slot_ranges;
     size_t key_count = 0;
     for (int db_num = 0; db_num < server.dbnum; db_num++) {
         listIter li;
@@ -1276,40 +1279,127 @@ int childSnapshotForSyncSlot(int req, rio *rdb, void *privdata) {
             slotRange *r = (slotRange *)ln->value;
             for (int slot = r->start_slot; slot <= r->end_slot; slot++) {
                 if (rewriteSlotToAppendOnlyFileRio(
-                        rdb, db_num, slot, &key_count) == C_ERR) return C_ERR;
+                        aof, db_num, slot, &key_count) == C_ERR) return C_ERR;
             }
         }
     }
-    rioWrite(rdb, "*3\r\n", 4);
-    rioWriteBulkString(rdb, "CLUSTER", 7);
-    rioWriteBulkString(rdb, "SYNCSLOTS", 9);
-    rioWriteBulkString(rdb, "SNAPSHOT-EOF", 12);
+    rioWrite(aof, "*3\r\n", 4);
+    rioWriteBulkString(aof, "CLUSTER", 7);
+    rioWriteBulkString(aof, "SYNCSLOTS", 9);
+    rioWriteBulkString(aof, "SNAPSHOT-EOF", 12);
     return C_OK;
 }
 
+/* Kill the slot migration child using SIGUSR1 (so that the parent will know
+ * the child did not exit for an error, but because we wanted), and performs
+ * the cleanup needed. */
+void killSlotMigrationChild(void) {
+    /* No slot migration child? return. */
+    if (server.child_type != CHILD_TYPE_SLOT_MIGRATION) return;
+    serverLog(LL_NOTICE, "Killing running slot migration child: %ld", (long)server.child_pid);
+
+    /* Because we are not using here waitpid (like we have in killAppendOnlyChild
+     * and TerminateModuleForkChild), all the cleanup operations is done by
+     * checkChildrenDone, that later will find that the process killed. */
+    kill(server.child_pid, SIGUSR1);
+}
+
 /* Begin the snapshot for the provided job in a child process. */
-int slotExportJobBeginSnapshot(slotMigrationJob *job) {
-    connection **conns = zmalloc(sizeof(connection *));
-    *conns = job->client->conn;
-    rdbSnapshotOptions opts = {
-        .connsnum = 1,
-        .conns = conns,
-        .use_pipe = 1,
-        .req = REPLICA_REQ_NONE,
-        .skip_checksum = 1,
-        .snapshot_func = childSnapshotForSyncSlot,
-        .privdata = job->slot_ranges};
-    if (saveSnapshotToConnectionSockets(opts) != C_OK) {
+int slotExportJobBeginSnapshotToTargetSocket(slotMigrationJob *job) {
+    if (hasActiveChildProcess()) return C_ERR;
+
+    pid_t childpid;
+    int pipefds[2], slot_migration_pipe_write = -1, safe_to_exit_pipe = -1;
+    serverAssert(server.slot_migration_pipe_read == -1 && server.slot_migration_child_exit_pipe == -1);
+
+    /* Before to fork, create a pipe that is used to transfer the slot data bytes to
+     * the parent, we can't let it write directly to the sockets, since in case
+     * of TLS we must let the parent handle a continuous TLS state when the
+     * child terminates and parent takes over. */
+    if (anetPipe(pipefds, O_NONBLOCK, 0) == -1) return C_ERR;
+    server.slot_migration_pipe_read = pipefds[0]; /* read end */
+    slot_migration_pipe_write = pipefds[1];       /* write end */
+
+    /* create another pipe that is used by the parent to signal to the child
+     * that it can exit. */
+    if (anetPipe(pipefds, 0, 0) == -1) {
+        close(slot_migration_pipe_write);
+        close(server.slot_migration_pipe_read);
+        server.slot_migration_pipe_read = -1;
         return C_ERR;
     }
-    if (server.debug_pause_after_fork) debugPauseProcess();
-    return C_OK;
+    safe_to_exit_pipe = pipefds[0];                     /* read end */
+    server.slot_migration_child_exit_pipe = pipefds[1]; /* write end */
+
+    server.slot_migration_pipe_conn = job->client->conn;
+
+    if ((childpid = serverFork(CHILD_TYPE_SLOT_MIGRATION)) == 0) {
+        /* Child */
+        rio aof;
+        rioInitWithFd(&aof, slot_migration_pipe_write);
+        /* Close the reading part, so that if the parent crashes, the child will
+         * get a write error and exit. */
+        close(server.rdb_pipe_read);
+
+        serverSetProcTitle("valkey-slot-migration-to-target");
+        serverSetCpuAffinity(server.bgsave_cpulist);
+
+        int retval = childSnapshotForSyncSlot(&aof, job);
+        if (retval == C_OK && rioFlush(&aof) == 0) retval = C_ERR;
+        if (retval == C_OK) {
+            sendChildCowInfo(CHILD_INFO_TYPE_SLOT_MIGRATION_COW_SIZE, "Slot migration");
+        }
+        rioFreeFd(&aof);
+        /* wake up the reader, tell it we're done. */
+        close(slot_migration_pipe_write);
+        close(server.slot_migration_child_exit_pipe); /* close write end so that we can detect the close on the parent. */
+        ssize_t dummy = read(safe_to_exit_pipe, pipefds, 1);
+        UNUSED(dummy);
+        exitFromChild((retval == C_OK) ? 0 : 1);
+    } else {
+        /* Parent */
+        if (childpid == -1) {
+            serverLog(LL_WARNING, "Can't begin slot migration snapshot in background: fork: %s", strerror(errno));
+            close(slot_migration_pipe_write);
+            close(server.slot_migration_pipe_read);
+            close(server.slot_migration_child_exit_pipe);
+            server.slot_migration_pipe_conn = NULL;
+            return C_ERR;
+        }
+
+        serverLog(LL_NOTICE, "Started child process %ld for slot migration %s", (long)childpid, job->description);
+        close(slot_migration_pipe_write); /* close write in parent so that it can detect the close on the child. */
+        if (aeCreateFileEvent(server.el, server.slot_migration_pipe_read, AE_READABLE, slotMigrationPipeReadHandler, NULL) == AE_ERR) {
+            serverPanic("Unrecoverable error creating server.rdb_pipe_read file event.");
+        }
+        close(safe_to_exit_pipe);
+        if (server.debug_pause_after_fork) debugPauseProcess();
+        return C_OK;
+    }
+    return C_OK; /* Unreached. */
 }
 
-/* Callback triggered after snapshot is finished. We either begin sending the
- * incremental contents or fail the associated migration. */
-void clusterHandleSlotExportBackgroundSaveDone(int bgsaveerr) {
-    if (!server.cluster_enabled) return;
+/* When a background slot migration terminates, call the right handler. */
+void backgroundSlotMigrationDoneHandler(int exitcode, int bysignal) {
+    if (!bysignal && exitcode == 0) {
+        serverLog(LL_NOTICE, "Background SLOT MIGRATION transfer terminated with success");
+    } else if (!bysignal && exitcode != 0) {
+        serverLog(LL_WARNING, "Background SLOT MIGRATION transfer error");
+    } else {
+        serverLog(LL_WARNING, "Background SLOT MIGRATION transfer terminated by signal %d", bysignal);
+    }
+    if (server.slot_migration_child_exit_pipe != -1) close(server.slot_migration_child_exit_pipe);
+    if (server.slot_migration_pipe_read > 0) {
+        aeDeleteFileEvent(server.el, server.slot_migration_pipe_read, AE_READABLE);
+        close(server.slot_migration_pipe_read);
+    }
+    server.slot_migration_child_exit_pipe = -1;
+    server.slot_migration_pipe_read = -1;
+    server.slot_migration_pipe_conn = NULL;
+    zfree(server.slot_migration_pipe_buff);
+    server.slot_migration_pipe_buff = NULL;
+    server.slot_migration_pipe_bufflen = 0;
+
     listIter li;
     listNode *ln;
     listRewind(server.cluster->slot_migration_jobs, &li);
@@ -1319,8 +1409,9 @@ void clusterHandleSlotExportBackgroundSaveDone(int bgsaveerr) {
         if (job->state != SLOT_EXPORT_SNAPSHOTTING) {
             continue;
         }
-        if (bgsaveerr == C_OK) {
+        if (!bysignal && exitcode == 0) {
             slotExportBeginStreaming(job);
+            job->stat_cow_bytes = server.stat_slot_migration_cow_bytes;
         } else {
             serverLog(LL_WARNING,
                       "Child process failed to snapshot slot migration %s",
@@ -1683,7 +1774,7 @@ void proceedWithSlotMigration(slotMigrationJob *job) {
             serverLog(LL_NOTICE,
                       "Beginning snapshot of slot migration %s.",
                       job->description);
-            if (slotExportJobBeginSnapshot(job) == C_ERR) {
+            if (slotExportJobBeginSnapshotToTargetSocket(job) == C_ERR) {
                 serverLog(LL_WARNING,
                           "Slot migration %s failed to start slot snapshot",
                           job->description);
@@ -1780,10 +1871,6 @@ void resetSlotMigrationJob(slotMigrationJob *job) {
 
     sdsfree(job->response_buf);
     job->response_buf = NULL;
-
-    /* Description is not needed once migration is finished */
-    sdsfree(job->description);
-    job->description = NULL;
 }
 
 void freeSlotMigrationJob(void *o) {
@@ -1793,6 +1880,7 @@ void freeSlotMigrationJob(void *o) {
     sdsfree(job->slot_ranges_str);
     sdsfree(job->status_msg);
     sdsfree(job->response_buf);
+    sdsfree(job->description);
     zfree(o);
 }
 
@@ -1940,7 +2028,7 @@ void finishSlotMigrationJob(slotMigrationJob *job,
         slotExportTryUnpause();
         /* Fast fail the child process, which will be cleaned up fully in
          * checkChildrenDone. */
-        if (job->state == SLOT_EXPORT_SNAPSHOTTING) killRDBChild();
+        if (job->state == SLOT_EXPORT_SNAPSHOTTING) killSlotMigrationChild();
     }
     if (job->type == SLOT_MIGRATION_IMPORT &&
         state != SLOT_MIGRATION_JOB_SUCCESS) {
@@ -1997,7 +2085,7 @@ void clusterCommandGetSlotMigrations(client *c) {
     listRewind(server.cluster->slot_migration_jobs, &li);
     while ((ln = listNext(&li)) != NULL) {
         slotMigrationJob *job = ln->value;
-        addReplyMapLen(c, 10);
+        addReplyMapLen(c, 11);
         addReplyBulkCString(c, "name");
         addReplyBulkCBuffer(c, job->name, CLUSTER_NAMELEN);
         addReplyBulkCString(c, "operation");
@@ -2020,6 +2108,8 @@ void clusterCommandGetSlotMigrations(client *c) {
         addReplyBulkCString(c, slotMigrationJobStateToString(job->state));
         addReplyBulkCString(c, "message");
         addReplyBulkCString(c, job->status_msg ? job->status_msg : "");
+        addReplyBulkCString(c, "cow_size");
+        addReplyLongLong(c, (long long)job->stat_cow_bytes);
     }
 }
 

src/cluster_migrateslots.h

@@ -24,7 +24,7 @@ void clusterCommandMigrateSlots(client *c);
 void clusterCommandSyncSlots(client *c);
 void clusterCommandGetSlotMigrations(client *c);
 void clusterCommandCancelSlotMigrations(client *c);
-void clusterHandleSlotExportBackgroundSaveDone(int bgsaveerr);
+void backgroundSlotMigrationDoneHandler(int exitcode, int bysignal);
 void clusterUpdateSlotExportsOnOwnershipChange(void);
 void clusterUpdateSlotImportsOnOwnershipChange(void);
 void clusterCleanupSlotMigrationLog(void);
@@ -33,5 +33,6 @@ size_t clusterGetTotalSlotExportBufferMemory(void);
 bool clusterSlotFailoverGranted(int slot);
 void clusterFailAllSlotExportsWithMessage(char *message);
 void clusterHandleSlotMigrationErrorResponse(slotMigrationJob *job);
+void killSlotMigrationChild(void);
 
 #endif /* __CLUSTER_MIGRATESLOTS_H */

src/commands/cluster-getslotmigrations.json

@@ -39,8 +39,15 @@
                         "type": "string",
                         "pattern": "^([0-9]+-[0-9]+)( [0-9]+-[0-9]+)*$"
                     },
-                    "node": {
-                        "type": "string"
+                    "target_node": {
+                        "description": "The target node name in the migration job.",
+                        "type": "string",
+                        "pattern": "^[0-9a-fA-F]{40}$"
+                    },
+                    "source_node": {
+                        "description": "The source node name in the migration job.",
+                        "type": "string",
+                        "pattern": "^[0-9a-fA-F]{40}$"
                     },
                     "create_time": {
                         "description": "Creation time, in seconds since the unix epoch.",
@@ -55,10 +62,16 @@
                         "type": "integer"
                     },
                     "state": {
+                        "description": "Human readable string representing the migration job state.",
                         "type": "string"
                     },
                     "message": {
+                        "description": "Human readable status message with more details.",
                         "type": "string"
+                    },
+                    "cow_size": {
+                        "description": "Copy on write bytes during slot migration fork.",
+                        "type": "integer"
                     }
                 }
             }

src/db.c

@@ -802,6 +802,7 @@ int getFlushCommandFlags(client *c, int *flags) {
 void flushAllDataAndResetRDB(int flags) {
     server.dirty += emptyData(-1, flags, NULL);
     if (server.child_type == CHILD_TYPE_RDB) killRDBChild();
+    if (server.child_type == CHILD_TYPE_SLOT_MIGRATION) killSlotMigrationChild();
     if (server.saveparamslen > 0) {
         rdbSaveInfo rsi, *rsiptr;
         rsiptr = rdbPopulateSaveInfo(&rsi);

src/module.c

@@ -13355,6 +13355,10 @@ int VM_RdbLoad(ValkeyModuleCtx *ctx, ValkeyModuleRdbStream *stream, int flags) {
      * will prevent COW memory issue. */
     if (server.child_type == CHILD_TYPE_RDB) killRDBChild();
 
+    /* Kill existing slot migration fork as it is saving outdated data. Also killing it
+     * will prevent COW memory issue. */
+    if (server.child_type == CHILD_TYPE_SLOT_MIGRATION) killSlotMigrationChild();
+
     emptyData(-1, EMPTYDB_NO_FLAGS, NULL);
 
     /* rdbLoad() can go back to the networking and process network events. If