fix: make slot migration work end to end

CHANGELOG.md

@@ -11,6 +11,11 @@
 - raft and the `openraft` dependency. raft only ever ran on the bootstrap node, and gossip with config epochs already decided slot ownership and failover, as in Redis Cluster. `raft-port-offset` and `--cluster-raft-port-offset` are still accepted and ignored, so existing configs keep working (#443)
 
 ### fixes
+- slot migration (#451):
+  - MIGRATE sends ASKING before RESTORE, so the target accepts the key for a slot it is importing. every MIGRATE during a reshard failed with MOVED before
+  - MIGRATE takes `AUTH password` and `AUTH2 username password` for a target with a password, honors `COPY`, and replies `NOKEY` for a missing key, as in Redis
+  - while a slot migrates, the source sends ASK for any key it no longer has, not only for keys MIGRATE moved, so keys written during the move land on the target instead of being left behind. a command with keys on both sides gets TRYAGAIN
+  - `ember-cli cluster reshard` passes its password to MIGRATE, and stops on a failed key instead of listing it forever and then reassigning a slot that still held keys
 - AOF durability (#449):
   - under `appendfsync everysec` and `no`, writes now reach the OS before their replies, so a crash of the process loses no acknowledged write. before, they sat in an 8 KB buffer, and `no` never flushed it
   - a write that fails part way, as on a full disk, no longer leaves a partial record that made the rest of the file unreadable

README.md

@@ -293,7 +293,7 @@ make cluster-clean  # stop + delete ./data/cluster/
 | `CLUSTER SETSLOT slot IMPORTING\|MIGRATING\|NODE\|STABLE ...` | manage slot migration state |
 | `CLUSTER COUNTKEYSINSLOT slot` | count keys owned by a slot on this node |
 | `CLUSTER GETKEYSINSLOT slot count` | list keys in a slot (used during migration) |
-| `MIGRATE host port key db timeout [COPY] [REPLACE]` | move a key to another node |
+| `MIGRATE host port key db timeout [COPY] [REPLACE] [AUTH pw \| AUTH2 user pw]` | move a key to another node |
 | `RESTORE key ttl payload` | receive a migrated key |
 | `ASKING` | tell the server to honor an ASK redirect for the next command |
 

crates/ember-cli/src/cluster.rs

@@ -1083,8 +1083,12 @@ async fn run_cluster_reshard(
             break;
         }
 
-        // 3. migrate keys
-        loop {
+        // 3. migrate keys. a key that fails to move stays in the slot, so
+        // the next batch would list it again: stop instead of looping, and
+        // leave the slot with the source so the key isn't orphaned
+        let target_port_s = target_port.to_string();
+        let mut stuck = false;
+        while !stuck {
             let keys_frame = match source_conn
                 .send(&["CLUSTER", "GETKEYSINSLOT", &slot.to_string(), "100"])
                 .await
@@ -1095,6 +1099,7 @@ async fn run_cluster_reshard(
                         "{}",
                         format!("GETKEYSINSLOT failed for slot {slot}: {e}").red()
                     );
+                    stuck = true;
                     break;
                 }
             };
@@ -1105,18 +1110,25 @@ async fn run_cluster_reshard(
             }
 
             for key in &keys {
-                let r = source_conn
-                    .send(&[
-                        "MIGRATE",
-                        &target_host,
-                        &target_port.to_string(),
-                        key,
-                        "0",
-                        "5000",
-                        "REPLACE",
-                    ])
-                    .await;
-                if !matches!(&r, Ok(f) if is_ok(f)) {
+                let mut args = vec![
+                    "MIGRATE",
+                    &target_host,
+                    &target_port_s,
+                    key,
+                    "0",
+                    "5000",
+                    "REPLACE",
+                ];
+                if let Some(pw) = password {
+                    args.extend(["AUTH", pw]);
+                }
+                let r = source_conn.send(&args).await;
+                // NOKEY: the key expired or was deleted since it was listed
+                let moved = matches!(
+                    &r,
+                    Ok(ember_protocol::types::Frame::Simple(s)) if s == "OK" || s == "NOKEY"
+                );
+                if !moved {
                     eprintln!(
                         "{}",
                         format!(
@@ -1125,10 +1137,19 @@ async fn run_cluster_reshard(
                         )
                         .red()
                     );
-                    // continue with other keys
+                    stuck = true;
+                    break;
                 }
             }
         }
+        if stuck {
+            eprintln!(
+                "{}",
+                format!("slot {slot} is still migrating; fix the error and run reshard again")
+                    .red()
+            );
+            break;
+        }
 
         // 4. SETSLOT NODE on both source and target (and the node we're connected to)
         let slot_s = slot.to_string();

crates/ember-cluster/README.md

@@ -61,7 +61,7 @@ the following CLUSTER commands are supported at the protocol layer:
 | `CLUSTER COUNTKEYSINSLOT <slot>` | key count in slot |
 | `CLUSTER GETKEYSINSLOT <slot> <count>` | get keys in slot |
 | `ASKING` | follow ASK redirect |
-| `MIGRATE <host> <port> <key> <db> <timeout>` | migrate a key |
+| `MIGRATE <host> <port> <key> <db> <timeout> [COPY] [REPLACE] [AUTH pw \| AUTH2 user pw]` | migrate a key |
 
 ## design notes
 

crates/ember-cluster/src/migration.rs

@@ -180,8 +180,6 @@ pub struct MigrationManager {
     outgoing: HashMap<u16, Migration>,
     /// Migrations where this node is the target (slot is importing in).
     incoming: HashMap<u16, Migration>,
-    /// Keys that have been migrated but not yet confirmed.
-    pending_keys: HashMap<u16, HashSet<Vec<u8>>>,
     /// Keys a MIGRATE is copying to another node right now.
     in_flight: HashSet<Vec<u8>>,
 }
@@ -230,7 +228,6 @@ impl MigrationManager {
 
         let migration = Migration::new_importing(slot, source, local_id);
         self.incoming.insert(slot, migration);
-        self.pending_keys.insert(slot, HashSet::new());
         self.incoming
             .get(&slot)
             .ok_or(MigrationError::NoMigrationInProgress { slot })
@@ -254,17 +251,13 @@ impl MigrationManager {
 
         let migration = Migration::new_migrating(slot, local_id, target);
         self.outgoing.insert(slot, migration);
-        self.pending_keys.insert(slot, HashSet::new());
         self.outgoing
             .get(&slot)
             .ok_or(MigrationError::NoMigrationInProgress { slot })
     }
 
-    /// Record that a key has been migrated.
-    pub fn key_migrated(&mut self, slot: u16, key: Vec<u8>) {
-        if let Some(keys) = self.pending_keys.get_mut(&slot) {
-            keys.insert(key);
-        }
+    /// Counts a key that MIGRATE moved out of `slot`.
+    pub fn key_migrated(&mut self, slot: u16) {
         if let Some(migration) = self.outgoing.get_mut(&slot) {
             migration.record_migrated(1);
         }
@@ -286,16 +279,8 @@ impl MigrationManager {
         self.in_flight.contains(key)
     }
 
-    /// Check if a specific key has been migrated.
-    pub fn is_key_migrated(&self, slot: u16, key: &[u8]) -> bool {
-        self.pending_keys
-            .get(&slot)
-            .is_some_and(|keys| keys.contains(key))
-    }
-
     /// Complete a migration and clean up state.
     pub fn complete_migration(&mut self, slot: u16) -> Option<Migration> {
-        self.pending_keys.remove(&slot);
         // Try outgoing first, then incoming
         self.outgoing
             .remove(&slot)
@@ -308,7 +293,6 @@ impl MigrationManager {
 
     /// Abort a migration and clean up state.
     pub fn abort_migration(&mut self, slot: u16) -> Option<Migration> {
-        self.pending_keys.remove(&slot);
         self.outgoing
             .remove(&slot)
             .or_else(|| self.incoming.remove(&slot))
@@ -615,18 +599,12 @@ mod tests {
     }
 
     #[test]
-    fn manager_key_tracking() {
+    fn manager_counts_migrated_keys() {
         let mut manager = MigrationManager::new();
-        let local = node_id();
-        let target = node_id();
-
-        manager.start_migrate(100, local, target).unwrap();
-
-        assert!(!manager.is_key_migrated(100, b"key1"));
-
-        manager.key_migrated(100, b"key1".to_vec());
-        assert!(manager.is_key_migrated(100, b"key1"));
-        assert!(!manager.is_key_migrated(100, b"key2"));
+        manager.start_migrate(100, node_id(), node_id()).unwrap();
+        manager.key_migrated(100);
+        manager.key_migrated(100);
+        assert_eq!(manager.get_outgoing(100).unwrap().keys_migrated, 2);
     }
 
     #[test]
@@ -636,15 +614,14 @@ mod tests {
         let target = node_id();
 
         manager.start_migrate(100, local, target).unwrap();
-        manager.key_migrated(100, b"key1".to_vec());
+        manager.key_migrated(100);
 
         let completed = manager.complete_migration(100);
         assert!(completed.is_some());
         assert_eq!(completed.unwrap().state, MigrationState::Complete);
 
         // State should be cleaned up
         assert!(!manager.is_migrating(100));
-        assert!(!manager.is_key_migrated(100, b"key1"));
     }
 
     #[test]

crates/ember-protocol/src/command/mod.rs

@@ -19,6 +19,14 @@ pub enum SetExpire {
     PxAt(u64),
 }
 
+/// Credentials MIGRATE sends to the target with AUTH.
+#[derive(Debug, Clone, PartialEq)]
+pub struct MigrateAuth {
+    /// Set by AUTH2. AUTH uses the default user.
+    pub username: Option<String>,
+    pub password: String,
+}
+
 /// A parsed client command, ready for execution.
 #[derive(Debug, Clone, PartialEq)]
 pub enum Command {
@@ -629,7 +637,8 @@ pub enum Command {
     /// CLUSTER GETKEYSINSLOT `slot` `count`. Return keys in a slot.
     ClusterGetKeysInSlot { slot: u16, count: u32 },
 
-    /// MIGRATE `host` `port` `key` `db` `timeout` \[COPY\] \[REPLACE\] \[KEYS key...\].
+    /// MIGRATE `host` `port` `key` `db` `timeout` \[COPY\] \[REPLACE\]
+    /// \[AUTH password | AUTH2 username password\].
     /// Migrate a key to another node.
     Migrate {
         host: String,
@@ -639,6 +648,8 @@ pub enum Command {
         timeout_ms: u64,
         copy: bool,
         replace: bool,
+        /// Credentials for the target, sent with AUTH before the key.
+        auth: Option<MigrateAuth>,
     },
 
     /// RESTORE `key` `ttl` `serialized-value` \[REPLACE\].

crates/ember-protocol/src/command/parse/cluster.rs

@@ -230,6 +230,7 @@ pub(super) fn parse_cluster_setslot(args: &[Frame]) -> Result<Command, ProtocolE
 
 pub(super) fn parse_migrate(args: &[Frame]) -> Result<Command, ProtocolError> {
     // MIGRATE host port key db timeout [COPY] [REPLACE]
+    //         [AUTH password | AUTH2 username password]
     if args.len() < 5 {
         return Err(wrong_arity("MIGRATE"));
     }
@@ -246,13 +247,32 @@ pub(super) fn parse_migrate(args: &[Frame]) -> Result<Command, ProtocolError> {
 
     let mut copy = false;
     let mut replace = false;
+    let mut auth = None;
 
-    for arg in &args[5..] {
+    let mut rest = args[5..].iter();
+    while let Some(arg) = rest.next() {
         let mut kw = [0u8; MAX_KEYWORD_LEN];
         let opt = uppercase_arg(arg, &mut kw)?;
+        let mut next = || {
+            rest.next()
+                .ok_or_else(|| wrong_arity("MIGRATE"))
+                .and_then(extract_string)
+        };
         match opt {
             "COPY" => copy = true,
             "REPLACE" => replace = true,
+            "AUTH" => {
+                auth = Some(MigrateAuth {
+                    username: None,
+                    password: next()?,
+                })
+            }
+            "AUTH2" => {
+                auth = Some(MigrateAuth {
+                    username: Some(next()?),
+                    password: next()?,
+                })
+            }
             _ => {
                 return Err(ProtocolError::InvalidCommandFrame(format!(
                     "unknown MIGRATE option '{opt}'"
@@ -269,5 +289,6 @@ pub(super) fn parse_migrate(args: &[Frame]) -> Result<Command, ProtocolError> {
         timeout_ms,
         copy,
         replace,
+        auth,
     })
 }

crates/ember-protocol/src/command/parse/mod.rs

@@ -9,7 +9,9 @@ use bytes::Bytes;
 use crate::error::ProtocolError;
 use crate::types::Frame;
 
-use super::{BitOpKind, BitRange, BitRangeUnit, Command, ScoreBound, SetExpire, ZAddFlags};
+use super::{
+    BitOpKind, BitRange, BitRangeUnit, Command, MigrateAuth, ScoreBound, SetExpire, ZAddFlags,
+};
 
 mod cluster;
 mod hashes;

crates/ember-protocol/src/command/tests.rs

@@ -1830,6 +1830,7 @@ fn migrate_basic() {
             timeout_ms: 5000,
             copy: false,
             replace: false,
+            auth: None,
         },
     );
 }
@@ -1856,10 +1857,40 @@ fn migrate_with_options() {
             timeout_ms: 10000,
             copy: true,
             replace: true,
+            auth: None,
         },
     );
 }
 
+#[test]
+fn migrate_with_auth() {
+    let parse = |extra: &[&str]| {
+        let mut args = vec!["MIGRATE", "h", "1", "k", "0", "10"];
+        args.extend_from_slice(extra);
+        match Command::from_frame(cmd(&args)) {
+            Ok(Command::Migrate { auth, .. }) => Ok(auth),
+            Ok(other) => panic!("{other:?}"),
+            Err(e) => Err(e),
+        }
+    };
+    assert_eq!(
+        parse(&["AUTH", "pw"]).unwrap(),
+        Some(MigrateAuth {
+            username: None,
+            password: "pw".into()
+        })
+    );
+    assert_eq!(
+        parse(&["REPLACE", "AUTH2", "u", "pw"]).unwrap(),
+        Some(MigrateAuth {
+            username: Some("u".into()),
+            password: "pw".into()
+        })
+    );
+    assert!(parse(&["AUTH"]).is_err());
+    assert!(parse(&["AUTH2", "u"]).is_err());
+}
+
 #[test]
 fn migrate_wrong_arity() {
     let err = Command::from_frame(cmd(&["MIGRATE", "host", "port", "key"])).unwrap_err();

crates/ember-protocol/src/lib.rs

@@ -27,7 +27,7 @@ mod serialize;
 pub mod slots;
 pub mod types;
 
-pub use command::{Command, SetExpire, ZAddFlags};
+pub use command::{Command, MigrateAuth, SetExpire, ZAddFlags};
 pub use error::ProtocolError;
 pub use parse::{parse_frame, parse_frame_bytes, parse_request};
 pub use types::wire;