Prevent exposure of importing keys on replicas during atomic slot migration (#2635)

murphyjacob4 · murphyjacob4 · commit 68a666c69852 · 2025-10-08T13:19:23.000-07:00
# Problem In the current slot migration design, replicas are completely unaware of the slot migration. Because of this, they do not know to hide importing keys, which results in exposure of these keys to commands like KEYS, SCAN, RANDOMKEY, and DBSIZE. # Design The main part of the design is that we will now listen for and process the `SYNCSLOTS ESTABLISH` command on the replica. When a `SYNCSLOTS ESTABLISH` command is received from the primary, we begin tracking a new slot import in a special `SLOT_IMPORT_OCCURRING_ON_PRIMARY` state. Replicas use this state to track the import, and await for a future `SYNCSLOTS FINISH` message that tells them the import is successful/failed. ## Success Case ``` Source Target Target Replica | | | |------------ SYNCSLOTS ESTABLISH -------------->| | | |----- SYNCSLOTS ESTABLISH ------>| |<-------------------- +OK ----------------------| | | | | |~~~~~~~~~~~~~~ snapshot as AOF ~~~~~~~~~~~~~~~~>| | | |~~~~~~ forward snapshot ~~~~~~~~>| |----------- SYNCSLOTS SNAPSHOT-EOF ------------>| | | | | |<----------- SYNCSLOTS REQUEST-PAUSE -----------| | | | | |~~~~~~~~~~~~ incremental changes ~~~~~~~~~~~~~~>| | | |~~~~~~ forward changes ~~~~~~~~~>| |--------------- SYNCSLOTS PAUSED -------------->| | | | | |<---------- SYNCSLOTS REQUEST-FAILOVER ---------| | | | | |---------- SYNCSLOTS FAILOVER-GRANTED --------->| | | | | | (performs takeover & | | propagates topology) | | | | | |------- SYNCSLOTS FINISH ------->| (finds out about topology | | change & marks migration done) | | | | | ``` ## Failure Case ``` Source Target Target Replica | | | |------------ SYNCSLOTS ESTABLISH -------------->| | | |----- SYNCSLOTS ESTABLISH ------>| |<-------------------- +OK ----------------------| | ... ... ... | | | | <FAILURE> | | | | | (performs cleanup) | | | ~~~~~~ UNLINK <key> ... ~~~~~~~>| | | | | | ------ SYNCSLOTS FINISH ------->| | | | ``` ## Full Sync, Partial Sync, and RDB In order to ensure replicas that resync during the import are still aware of the import, the slot import is serialized to a new `cluster-slot-imports` aux field. The encoding includes the job name, the source node name, and the slot ranges being imported. Upon loading an RDB with the `cluster-slot-imports` aux field, replicas will add a new migration in the `SLOT_IMPORT_OCCURRING_ON_PRIMARY` state. It's important to note that a previously saved RDB file can be used as the basis for partial sync with a primary. Because of this, whenever we load an RDB file with the `cluster-slot-imports` aux field, even from disk, we will still add a new migration to track the import. If after loading the RDB, the Valkey node is a primary, it will cancel the slot migration. Having this tracking state loaded on primaries will ensure that replicas partial syncing to a restarted primary still get their `SYNCSLOTS FINISH` message in the replication stream. ## AOF Since AOF cannot be used as the basis for a partial sync, we don't necessarily need to persist the `SYNCSLOTS ESTABLISH` and `FINISH` commands to the AOF. However, considering there is work to change this (#59 #1901) this design doesn't make any assumptions about this. We will propagate the `ESTABLISH` and `FINISH` commands to the AOF, and ensure that they can be properly replayed on AOF load to get to the right state. Similar to RDB, if there are any pending "ESTABLISH" commands that don't have a "FINISH" afterwards upon becoming primary, we will make sure to fail those in `verifyClusterConfigWithData`. Additionally, there was a bug in the existing slot migration where slot import clients were not having their commands persisted to AOF. This has been fixed by ensuring we still propagate to AOF even for slot import clients. ## Promotion & Demotion Since the primary is solely responsible for cleaning up unowned slots, primaries that are demoted will not clean up previously active slot imports. The promoted replica will be responsible for both cleaning up the slot (`verifyClusterConifgWithData`) and sending a `SYNCSLOTS FINISH`. # Other Options Considered I also considered tracking "dirty" slots rather than using the slot import state machine. In this setup, primaries and replicas would simply mark each slot's hashtable in the kvstore as dirty when something is written to it and we do not currently own that slot. This approach is simpler, but has a problem in that modules loaded on the replica would still not get slot migration start/end notifications. If the modules on the replica do not get such notifications, they will not be able to properly contain these dirty keys during slot migration events. --------- Signed-off-by: Jacob Murphy <jkmurphy@google.com>
diff --git a/src/cluster_legacy.c b/src/cluster_legacy.c
@@ -6385,6 +6385,9 @@ int verifyClusterConfigWithData(void) {
      * completely depend on the replication stream. */
     if (nodeIsReplica(myself)) return C_OK;
 
+    /* Allow slot migrations to clean up after reloading */
+    clusterCleanSlotImportsAfterLoad();
+
     /* Check that all the slots we see populated memory have a corresponding
      * entry in the cluster table. Otherwise fix the table. */
     for (j = 0; j < CLUSTER_SLOTS; j++) {
@@ -6464,7 +6467,7 @@ static void clusterSetPrimary(clusterNode *n, int closeSlots, int full_sync_requ
     removeAllNotOwnedShardChannelSubscriptions();
     resetManualFailover();
 
-    /* Becoming a replica cancels all in progress imports and exports */
+    /* Perform needed slot migration state transitions */
     clusterUpdateSlotExportsOnOwnershipChange();
     clusterUpdateSlotImportsOnOwnershipChange();
 
diff --git a/src/cluster_migrateslots.c b/src/cluster_migrateslots.c
diff --git a/src/cluster_migrateslots.h b/src/cluster_migrateslots.h
@@ -27,12 +27,17 @@ void clusterCommandCancelSlotMigrations(client *c);
 void backgroundSlotMigrationDoneHandler(int exitcode, int bysignal);
 void clusterUpdateSlotExportsOnOwnershipChange(void);
 void clusterUpdateSlotImportsOnOwnershipChange(void);
-void clusterCleanupSlotMigrationLog(void);
 void clusterHandleFlushDuringSlotMigration(void);
 size_t clusterGetTotalSlotExportBufferMemory(void);
 bool clusterSlotFailoverGranted(int slot);
 void clusterFailAllSlotExportsWithMessage(char *message);
 void clusterHandleSlotMigrationErrorResponse(slotMigrationJob *job);
 void killSlotMigrationChild(void);
+void clusterCleanSlotImportsOnFullSync(void);
+void clusterCleanSlotImportsOnPromotion(void);
+void clusterCleanSlotImportsBeforeLoad(void);
+void clusterCleanSlotImportsAfterLoad(void);
+int clusterRDBSaveSlotImports(rio *rdb);
+int clusterRDBLoadSlotImport(rio *rdb);
 
 #endif /* __CLUSTER_MIGRATESLOTS_H */
diff --git a/src/commands.def b/src/commands.def
@@ -1169,7 +1169,7 @@ struct COMMAND_STRUCT CLUSTER_Subcommands[] = {
 {MAKE_CMD("slaves","Lists the replica nodes of a primary node.","O(N) where N is the number of replicas.","3.0.0",CMD_DOC_DEPRECATED,"`CLUSTER REPLICAS`","5.0.0","cluster",COMMAND_GROUP_CLUSTER,CLUSTER_SLAVES_History,0,CLUSTER_SLAVES_Tips,1,clusterCommand,3,CMD_ADMIN|CMD_STALE,0,CLUSTER_SLAVES_Keyspecs,0,NULL,1),.args=CLUSTER_SLAVES_Args},
 {MAKE_CMD("slot-stats","Return an array of slot usage statistics for slots assigned to the current node.","O(N) where N is the total number of slots based on arguments. O(N*log(N)) with ORDERBY subcommand.","8.0.0",CMD_DOC_NONE,NULL,NULL,"cluster",COMMAND_GROUP_CLUSTER,CLUSTER_SLOT_STATS_History,0,CLUSTER_SLOT_STATS_Tips,2,clusterSlotStatsCommand,-4,CMD_STALE|CMD_LOADING,0,CLUSTER_SLOT_STATS_Keyspecs,0,NULL,1),.args=CLUSTER_SLOT_STATS_Args},
 {MAKE_CMD("slots","Returns the mapping of cluster slots to nodes.","O(N) where N is the total number of Cluster nodes","3.0.0",CMD_DOC_NONE,NULL,NULL,"cluster",COMMAND_GROUP_CLUSTER,CLUSTER_SLOTS_History,2,CLUSTER_SLOTS_Tips,1,clusterCommand,2,CMD_LOADING|CMD_STALE,0,CLUSTER_SLOTS_Keyspecs,0,NULL,0)},
-{MAKE_CMD("syncslots","A container for internal slot migration commands.","Depends on subcommand.","9.0.0",CMD_DOC_SYSCMD,NULL,NULL,"cluster",COMMAND_GROUP_CLUSTER,CLUSTER_SYNCSLOTS_History,0,CLUSTER_SYNCSLOTS_Tips,0,clusterCommand,-3,CMD_NO_ASYNC_LOADING|CMD_ADMIN|CMD_STALE,0,CLUSTER_SYNCSLOTS_Keyspecs,0,NULL,0)},
+{MAKE_CMD("syncslots","A container for internal slot migration commands.","Depends on subcommand.","9.0.0",CMD_DOC_SYSCMD,NULL,NULL,"cluster",COMMAND_GROUP_CLUSTER,CLUSTER_SYNCSLOTS_History,0,CLUSTER_SYNCSLOTS_Tips,0,clusterCommand,-3,CMD_NO_ASYNC_LOADING|CMD_ADMIN|CMD_STALE|CMD_MAY_REPLICATE,0,CLUSTER_SYNCSLOTS_Keyspecs,0,NULL,0)},
 {0}
 };
 
diff --git a/src/commands/cluster-syncslots.json b/src/commands/cluster-syncslots.json
@@ -13,7 +13,8 @@
         "command_flags": [
             "NO_ASYNC_LOADING",
             "ADMIN",
-            "STALE"
+            "STALE",
+            "MAY_REPLICATE"
         ]
     }
 }
diff --git a/src/db.c b/src/db.c
@@ -671,6 +671,14 @@ long long emptyData(int dbnum, int flags, void(callback)(hashtable *)) {
      * there. */
     signalFlushedDb(dbnum, async);
 
+    if (clusterIsAnySlotImporting() || clusterIsAnySlotExporting()) {
+        /* On flush, in progress migrations will be cancelled, and should be
+         * retried by operators. We also may emptyData when reloading an RDB, in
+         * which case we will remove active slot imports. Replicas will get a
+         * new set of slot imports from their primary. */
+        clusterHandleFlushDuringSlotMigration();
+    }
+
     /* Empty the database structure. */
     removed = emptyDbStructure(server.db, dbnum, async, callback);
 
@@ -825,12 +833,6 @@ void flushdbCommand(client *c) {
 
     if (getFlushCommandFlags(c, &flags) == C_ERR) return;
 
-    if (clusterIsAnySlotImporting() || clusterIsAnySlotExporting()) {
-        /* In progress migrations will be cancelled, and should be retried by
-         * operators. */
-        clusterHandleFlushDuringSlotMigration();
-    }
-
     /* flushdb should not flush the functions */
     server.dirty += emptyData(c->db->id, flags | EMPTYDB_NOFUNCTIONS, NULL);
 
@@ -855,12 +857,6 @@ void flushallCommand(client *c) {
     int flags;
     if (getFlushCommandFlags(c, &flags) == C_ERR) return;
 
-    if (clusterIsAnySlotImporting() || clusterIsAnySlotExporting()) {
-        /* In progress migrations will be cancelled, and should be retried by
-         * operators. */
-        clusterHandleFlushDuringSlotMigration();
-    }
-
     /* flushall should not flush the functions */
     flushAllDataAndResetRDB(flags | EMPTYDB_NOFUNCTIONS);
 
diff --git a/src/rdb.c b/src/rdb.c
@@ -1472,6 +1472,9 @@ int rdbSaveRio(int req, rio *rdb, int *error, int rdbflags, rdbSaveInfo *rsi) {
 
     /* save all databases, skip this if we're in functions-only mode */
     if (!(req & REPLICA_REQ_RDB_EXCLUDE_DATA)) {
+        /* RDB slot import info is encoded in a required opcode since exposing
+         * importing slots is a consistency problem. */
+        if (clusterRDBSaveSlotImports(rdb) == C_ERR) goto werr;
         for (j = 0; j < server.dbnum; j++) {
             if (rdbSaveDb(rdb, j, rdbflags, &key_counter) == -1) goto werr;
         }
@@ -2923,6 +2926,9 @@ void startLoading(size_t size, int rdbflags, int async) {
     server.rdb_last_load_keys_loaded = 0;
     blockingOperationStarts();
 
+    /* Cleanup slot migrations (we need a clean state for the incoming load) */
+    clusterCleanSlotImportsBeforeLoad();
+
     /* Fire the loading modules start event. */
     int subevent;
     if (rdbflags & RDBFLAGS_AOF_PREAMBLE)
@@ -3194,6 +3200,9 @@ int rdbLoadRioWithLoadingCtx(rio *rdb, int rdbflags, rdbSaveInfo *rsi, rdbLoadin
                 should_expand_db = 0;
             }
             continue; /* Read next opcode. */
+        } else if (type == RDB_OPCODE_SLOT_IMPORT) {
+            if (clusterRDBLoadSlotImport(rdb) == C_ERR) goto eoferr;
+            continue; /* Read next opcode. */
         } else if (type == RDB_OPCODE_AUX) {
             /* AUX: generic string-string fields. Use to add state to RDB
              * which is backward compatible. Implementations of RDB loading
diff --git a/src/rdb.h b/src/rdb.h
@@ -131,6 +131,7 @@ enum RdbType {
 
 /* Special RDB opcodes (saved/loaded with rdbSaveType/rdbLoadType).
  * These are special RDB types, but they start from 255 and grow down. */
+#define RDB_OPCODE_SLOT_IMPORT 243     /* Slot import state. */
 #define RDB_OPCODE_SLOT_INFO 244       /* Foreign slot info, safe to ignore. */
 #define RDB_OPCODE_FUNCTION2 245       /* function library data */
 #define RDB_OPCODE_FUNCTION_PRE_GA 246 /* old function library data for 7.0 rc1 and rc2 */
diff --git a/src/replication.c b/src/replication.c
@@ -2196,6 +2196,9 @@ void replicationAttachToNewPrimary(void) {
     serverAssert(server.primary == NULL);
     replicationDiscardCachedPrimary();
 
+    /* Cancel any in progress imports (we will now use the primary's) */
+    clusterCleanSlotImportsOnFullSync();
+
     disconnectReplicas();     /* Force our replicas to resync with us as well. */
     freeReplicationBacklog(); /* Don't allow our chained replicas to PSYNC. */
 }
@@ -4422,6 +4425,9 @@ void replicationUnsetPrimary(void) {
     /* Restart the AOF subsystem in case we shut it down during a sync when
      * we were still a replica. */
     if (server.aof_enabled && server.aof_state == AOF_OFF) restartAOFAfterSYNC();
+
+    /* Cancel any ongoing atomic slot migrations */
+    clusterCleanSlotImportsOnPromotion();
 }
 
 /* This function is called when the replica lose the connection with the
diff --git a/src/server.c b/src/server.c
@@ -2151,6 +2151,8 @@ void createSharedObjects(void) {
     shared.persist = createSharedString("PERSIST");
     shared.set = createSharedString("SET");
     shared.eval = createSharedString("EVAL");
+    shared.cluster = createSharedString("CLUSTER");
+    shared.syncslots = createSharedString("SYNCSLOTS");
 
     /* Shared command argument */
     shared.left = createSharedString("left");
@@ -2174,6 +2176,12 @@ void createSharedObjects(void) {
     shared.special_equals = createSharedString("=");
     shared.redacted = createSharedString("(redacted)");
     shared.fields = createSharedString("FIELDS");
+    shared.finish = createSharedString("FINISH");
+    shared.state = createSharedString("STATE");
+    shared.success = createSharedString("SUCCESS");
+    shared.failed = createSharedString("FAILED");
+    shared.name = createSharedString("NAME");
+    shared.message = createSharedString("MESSAGE");
 
     for (j = 0; j < OBJ_SHARED_INTEGERS; j++) {
         shared.integers[j] = makeObjectShared(createObject(OBJ_STRING, (void *)(long)j));
@@ -3582,8 +3590,15 @@ void alsoPropagate(int dbid, robj **argv, int argc, int target, int slot) {
     if (!shouldPropagate(target)) return;
 
     /* Don't propagate commands on slot migration clients, these will be proxied
-     * in replicationFeedStreamFromPrimaryStream() */
-    if (server.current_client != NULL && server.current_client->slot_migration_job) return;
+     * in replicationFeedStreamFromPrimaryStream().
+     *
+     * However, if we need to propagate to AOF, we should still do that. */
+    bool propagate_aof = (target & PROPAGATE_AOF) && server.aof_state != AOF_OFF;
+    if (server.current_client != NULL && server.current_client->slot_migration_job) {
+        if (!propagate_aof) return;
+        /* Disable propagation to replication (just do the AOF) */
+        target &= ~PROPAGATE_REPL;
+    }
 
     argvcopy = zmalloc(sizeof(robj *) * argc);
     for (j = 0; j < argc; j++) {
@@ -4484,10 +4499,10 @@ int processCommand(client *c) {
         return C_OK;
     }
 
-    /* If the server is paused, block the client until
-     * the pause has ended. Replicas are never paused. */
-    if (!c->flag.replica && ((isPausedActions(PAUSE_ACTION_CLIENT_ALL)) ||
-                             ((isPausedActions(PAUSE_ACTION_CLIENT_WRITE)) && is_may_replicate_command))) {
+    /* If the server is paused, block the client until the pause has ended. Replicas and slot
+     * export clients are never paused to allow failover/slot migration to succeed. */
+    if (!c->flag.replica && (!c->slot_migration_job || isImportSlotMigrationJob(c->slot_migration_job)) &&
+        ((isPausedActions(PAUSE_ACTION_CLIENT_ALL)) || ((isPausedActions(PAUSE_ACTION_CLIENT_WRITE)) && is_may_replicate_command))) {
         blockPostponeClient(c);
         return C_OK;
     }
diff --git a/src/server.h b/src/server.h
@@ -1409,9 +1409,10 @@ struct sharedObjectsStruct {
         *execaborterr, *noautherr, *noreplicaserr, *busykeyerr, *oomerr, *plus, *messagebulk, *pmessagebulk,
         *subscribebulk, *unsubscribebulk, *psubscribebulk, *punsubscribebulk, *del, *unlink, *rpop, *lpop, *lpush,
         *rpoplpush, *lmove, *blmove, *zpopmin, *zpopmax, *emptyscan, *multi, *exec, *left, *right, *hset, *hdel, *hpexpireat, *hpersist, *srem,
-        *xgroup, *xclaim, *script, *replconf, *eval, *persist, *set, *pexpireat, *pexpire, *time, *pxat, *absttl,
+        *xgroup, *xclaim, *script, *replconf, *eval, *cluster, *syncslots, *persist, *set, *pexpireat, *pexpire, *time, *pxat, *absttl,
         *retrycount, *force, *justid, *entriesread, *lastid, *ping, *setid, *keepttl, *load, *createconsumer, *getack,
         *special_asterisk, *special_equals, *default_username, *redacted, *ssubscribebulk, *sunsubscribebulk, *fields,
+        *finish, *state, *success, *failed, *name, *message,
         *smessagebulk, *select[PROTO_SHARED_SELECT_CMDS], *integers[OBJ_SHARED_INTEGERS],
         *mbulkhdr[OBJ_SHARED_BULKHDR_LEN], /* "*<value>\r\n" */
         *bulkhdr[OBJ_SHARED_BULKHDR_LEN],  /* "$<value>\r\n" */
diff --git a/src/valkey-check-rdb.c b/src/valkey-check-rdb.c
@@ -700,6 +700,17 @@ int redis_check_rdb(char *rdbfilename, FILE *fp) {
             if (rdbLoadLen(&rdb, NULL) == RDB_LENERR) goto eoferr;
             if (rdbLoadLen(&rdb, NULL) == RDB_LENERR) goto eoferr;
             continue; /* Read type again. */
+        } else if (type == RDB_OPCODE_SLOT_IMPORT) {
+            robj *job_name;
+            if ((job_name = rdbLoadStringObject(&rdb)) == NULL) goto eoferr;
+            decrRefCount(job_name);
+            uint64_t num_slot_ranges;
+            if ((num_slot_ranges = rdbLoadLen(&rdb, NULL)) == RDB_LENERR) goto eoferr;
+            for (uint64_t i = 0; i < num_slot_ranges; i++) {
+                if (rdbLoadLen(&rdb, NULL) == RDB_LENERR) goto eoferr;
+                if (rdbLoadLen(&rdb, NULL) == RDB_LENERR) goto eoferr;
+            }
+            continue; /* Read type again. */
         } else if (type == RDB_OPCODE_AUX) {
             /* AUX: generic string-string fields. Use to add state to RDB
              * which is backward compatible. Implementations of RDB loading
diff --git a/src/valkeymodule.h b/src/valkeymodule.h
@@ -829,12 +829,10 @@ typedef struct ValkeyModuleSlotRange {
 } ValkeyModuleSlotRange;
 
 typedef struct ValkeyModuleAtomicSlotMigrationInfo {
-    uint64_t version;                                  /* Version of this structure for ABI compat. */
-    char job_name[VALKEYMODULE_NODE_ID_LEN + 1];       /* Unique ID for the migration operation. */
-    ValkeyModuleSlotRange *slot_ranges;                /* Array of slot ranges involved in the migration. */
-    uint32_t num_slot_ranges;                          /* Number of slot ranges in the array. */
-    char source_node_id[VALKEYMODULE_NODE_ID_LEN + 1]; /* Node ID of the source node. */
-    char target_node_id[VALKEYMODULE_NODE_ID_LEN + 1]; /* Node ID of the target node. */
+    uint64_t version;                            /* Version of this structure for ABI compat. */
+    char job_name[VALKEYMODULE_NODE_ID_LEN + 1]; /* Unique ID for the migration operation. */
+    ValkeyModuleSlotRange *slot_ranges;          /* Array of slot ranges involved in the migration. */
+    uint32_t num_slot_ranges;                    /* Number of slot ranges in the array. */
 } ValkeyModuleAtomicSlotMigrationInfoV1;
 
 #define ValkeyModuleAtomicSlotMigrationInfo ValkeyModuleAtomicSlotMigrationInfoV1
diff --git a/tests/modules/hooks.c b/tests/modules/hooks.c
@@ -379,14 +379,6 @@ void authAttemptCallback(ValkeyModuleCtx *ctx, ValkeyModuleEvent e, uint64_t sub
 }
 
 void logAtomicSlotMigrationInfo(ValkeyModuleCtx *ctx, const char *prefix, ValkeyModuleAtomicSlotMigrationInfo *asmi) {
-    ValkeyModuleString *target_keyname = ValkeyModule_CreateStringPrintf(ctx, "%s-target", prefix);
-    LogStringEvent(ctx, ValkeyModule_StringPtrLen(target_keyname, NULL), asmi->target_node_id);
-    ValkeyModule_FreeString(ctx, target_keyname);
-
-    ValkeyModuleString *source_keyname = ValkeyModule_CreateStringPrintf(ctx, "%s-source", prefix);
-    LogStringEvent(ctx, ValkeyModule_StringPtrLen(source_keyname, NULL), asmi->source_node_id);
-    ValkeyModule_FreeString(ctx, source_keyname);
-
     ValkeyModuleString *job_keyname = ValkeyModule_CreateStringPrintf(ctx, "%s-jobname", prefix);
     LogStringEvent(ctx, ValkeyModule_StringPtrLen(job_keyname, NULL), asmi->job_name);
     ValkeyModule_FreeString(ctx, job_keyname);
diff --git a/tests/unit/cluster/cluster-migrateslots.tcl b/tests/unit/cluster/cluster-migrateslots.tcl
diff --git a/tests/unit/moduleapi/hooks.tcl b/tests/unit/moduleapi/hooks.tcl

Original file line number	Diff line number	Diff line change
`@@ -13,7 +13,8 @@`
`13`	`13`	`"command_flags": [`
`14`	`14`	`"NO_ASYNC_LOADING",`
`15`	`15`	`"ADMIN",`
`16`		`- "STALE"`
	`16`	`+ "STALE",`
	`17`	`+ "MAY_REPLICATE"`
`17`	`18`	`]`
`18`	`19`	`}`
`19`	`20`	`}`