Update internal slot to shard map based on MOVED error (#100)

Author: fabianfett

Notice.txt

@@ -78,3 +78,13 @@ This product contains a version of the ConnectionPoolModule from postgres-nio
     * https://github.com/vapor/postgres-nio/blob/main/LICENSE
   * HOMEPAGE:
     * https://github.com/vapor/postgres-nio
+
+---
+
+This product was influenced by valkey-glide.
+  - It adapted valkey-glide's hashslot update logic after a MOVED error. 
+
+  * LICENSE (Apache License 2.0)
+    * https://github.com/valkey-io/valkey-glide/blob/main/LICENSE
+  * HOMEPAGE:
+    * https://github.com/valkey-io/valkey-glide

Sources/Valkey/Cluster/HashSlotShardMap.swift

@@ -51,7 +51,7 @@ extension ValkeyShardNodeIDs: ExpressibleByArrayLiteral {
 
 /// This object allows us to efficiently look up the Valkey shard given a hash slot.
 ///
-/// The `HashSlotShardMap` maintains an internal array where each element corresponds to one hash slot (0-16383).
+/// The ``HashSlotShardMap`` maintains an internal array where each element corresponds to one hash slot (0-16383).
 /// This makes looking up the shard as efficient as a simple array access operation.
 ///
 /// Hash slots are assigned to shards in a Valkey cluster, and each key is mapped to a specific slot
@@ -167,6 +167,71 @@ package struct HashSlotShardMap: Sendable {
         }
     }
 
+    @usableFromInline
+    package enum UpdateSlotsResult: Equatable {
+        case updatedSlotToExistingNode
+        case updatedSlotToUnknownNode
+    }
+
+    /// Handles MOVED errors by updating the client's slot and node mappings based on the new primary's role:
+    ///
+    /// 1. **No Change**: If the new primary is already the current slot owner, no updates are needed.
+    /// 2. **Failover**: If the new primary is a replica within the same shard (indicating a failover),
+    ///    the slot ownership is updated by promoting the replica to the primary in the existing shard addresses.
+    /// 3. **Slot Migration**: If the new primary is an existing primary in another shard, this indicates a slot migration,
+    ///    and the slot mapping is updated to point to the new shard addresses.
+    /// 4. **Replica Moved to a Different Shard**: If the new primary is a replica in a different shard, it can be due to:
+    ///    - The replica became the primary of its shard after a failover, with new slots migrated to it.
+    ///    - The replica has moved to a different shard as the primary.
+    ///      Since further information is unknown, the replica is removed from its original shard and added as the primary of a new shard.
+    /// 5. **New Node**: If the new primary is unknown, it is added as a new node in a new shard, possibly indicating scale-out.
+    ///
+    /// This logic was first implemented in `valkey-glide` (see `Notice.txt`) and adopted for Swift here.
+    @usableFromInline
+    package mutating func updateSlots(with movedError: ValkeyMovedError) -> UpdateSlotsResult {
+        if let shardIndex = self.slotToShardID[Int(movedError.slot.rawValue)].value {
+            // if the slot had a shard assignment before
+            var shard = self.shardIDToShard[shardIndex]
+
+            // 1. No change
+            if shard.master == movedError.nodeID {
+                return .updatedSlotToExistingNode
+            }
+
+            // 2. Failover
+            if shard.replicas.contains(movedError.nodeID) {
+                // lets promote the replica to be the primary and remove the old primary for now
+                shard.master = movedError.nodeID
+                shard.replicas.removeAll { $0 == movedError.nodeID }
+                self.shardIDToShard[shardIndex] = shard
+                return .updatedSlotToExistingNode
+            }
+        }
+
+        // 3. Slot migration to an existing primary
+        if let newShardIndex = self.shardIDToShard.firstIndex(where: { $0.master == movedError.nodeID }) {
+            self.slotToShardID[Int(movedError.slot.rawValue)] = .init(newShardIndex)
+            return .updatedSlotToExistingNode
+        }
+
+        // 4. Replica moved to a different shard
+        if let ogShardIndexOfNewPrimary = self.shardIDToShard.firstIndex(where: { $0.replicas.contains(movedError.nodeID) }) {
+            // remove replica from its og shard
+            self.shardIDToShard[ogShardIndexOfNewPrimary].replicas.removeAll(where: { $0 == movedError.nodeID })
+            // create a new shard with the replica
+            let newShardIndex = self.shardIDToShard.endIndex
+            self.shardIDToShard.append(.init(master: movedError.nodeID))
+            self.slotToShardID[Int(movedError.slot.rawValue)] = .init(newShardIndex)
+            return .updatedSlotToExistingNode
+        }
+
+        // 5. totally new node
+        let newShardIndex = self.shardIDToShard.endIndex
+        self.shardIDToShard.append(.init(master: movedError.nodeID))
+        self.slotToShardID[Int(movedError.slot.rawValue)] = .init(newShardIndex)
+        return .updatedSlotToUnknownNode
+    }
+
     /// An internal type representing an optional shard ID with efficient storage.
     ///
     /// This type uses a special sentinel value to represent a missing shard ID

Sources/Valkey/Cluster/ValkeyMovedError.swift

@@ -23,18 +23,24 @@ import NIOCore
 /// This error provides the necessary information for clients to redirect their
 /// request to the correct node in the cluster.
 @usableFromInline
-struct ValkeyMovedError: Hashable, Sendable {
+package struct ValkeyMovedError: Hashable, Sendable {
     /// The hash slot number that triggered the redirection.
-    var slot: HashSlot
-    
+    package var slot: HashSlot
+
     /// The hostname or IP address of the node that owns the requested hash slot.
-    var endpoint: String
-    
+    package var endpoint: String
+
     /// The port number of the node that owns the requested hash slot.
-    var port: Int
+    package var port: Int
+
+    package init(slot: HashSlot, endpoint: String, port: Int) {
+        self.slot = slot
+        self.endpoint = endpoint
+        self.port = port
+    }
 
     @usableFromInline
-    var nodeID: ValkeyNodeID {
+    package var nodeID: ValkeyNodeID {
         ValkeyNodeID(endpoint: self.endpoint, port: self.port)
     }
 }

Tests/ValkeyTests/Cluster/HashSlotShardMapTests.swift

@@ -667,4 +667,172 @@ struct HashSlotShardMapTests {
         #expect(shardNodes.replicas.contains(expectedReplica1))
         #expect(shardNodes.replicas.contains(expectedReplica2))
     }
+
+    func makeExampleCusterWithNShardsAndMReplicasPerShard(shards: Int, replicas: Int) -> ValkeyClusterDescription {
+        let defaultRangeSize = Int(HashSlot.max.rawValue + 1) / shards
+        var range: ClosedRange<HashSlot> = 0...0
+
+        var result = [ValkeyClusterDescription.Shard]()
+        var nodeIndex = 1
+
+        for i in 0..<shards {
+            if i == 0 {
+                if shards == 1 {
+                    range = HashSlot.min...HashSlot.max
+                } else {
+                    range = HashSlot.min...(range.upperBound.advanced(by: defaultRangeSize - 1))
+                }
+            } else if i == shards - 1 {
+                range = (range.upperBound.advanced(by: 1))...HashSlot.max
+            } else {
+                range = (range.upperBound.advanced(by: 1))...(range.upperBound.advanced(by: defaultRangeSize - 1))
+            }
+
+            var shard = ValkeyClusterDescription.Shard(slots: [range], nodes: [])
+            for _ in 0..<(replicas + 1) {
+                defer { nodeIndex += 1 }
+                shard.nodes.append(
+                    .init(
+                        id: "node-\(nodeIndex)",
+                        port: nil,
+                        tlsPort: 6379,
+                        ip: "192.168.64.\(nodeIndex)",
+                        hostname: "node-\(nodeIndex).valkey.io",
+                        endpoint: "node-\(nodeIndex).valkey.io",
+                        role: .replica,
+                        replicationOffset: 14,
+                        health: .online
+                    )
+                )
+            }
+            let primaryIndex = shard.nodes.indices.randomElement()!
+            shard.nodes[primaryIndex].role = .master
+
+            result.append(shard)
+        }
+
+        return ValkeyClusterDescription(result)
+    }
+
+    @Test("Case 1: MovedError specifies the already exisiting shard primary node")
+    func movedErrorSpecifiesTheAlreadyExisitingShardPrimaryNode() throws {
+        let clusterDescription = self.makeExampleCusterWithNShardsAndMReplicasPerShard(shards: 3, replicas: 1)
+
+        var map = HashSlotShardMap()
+        map.updateCluster(clusterDescription.shards)
+
+        let ogShard = try map.nodeID(for: CollectionOfOne(2))
+        let update = map.updateSlots(with: ValkeyMovedError(slot: 2, endpoint: ogShard.master.endpoint, port: ogShard.master.port))
+        #expect(update == .updatedSlotToExistingNode)
+        let updatedShard = try map.nodeID(for: CollectionOfOne(2))
+        #expect(updatedShard == ogShard)
+    }
+
+    @Test("Case 2: MovedError specifies a previous shard replica node")
+    func movedErrorSpecifiesAPreviousShardReplicaNode() throws {
+        let clusterDescription = self.makeExampleCusterWithNShardsAndMReplicasPerShard(shards: 3, replicas: 3)
+
+        var map = HashSlotShardMap()
+        map.updateCluster(clusterDescription.shards)
+
+        let ogShard = try map.nodeID(for: CollectionOfOne(2))
+        let luckyReplica = ogShard.replicas.randomElement()!
+
+        let update = map.updateSlots(with: ValkeyMovedError(slot: 2, endpoint: luckyReplica.endpoint, port: luckyReplica.port))
+        #expect(update == .updatedSlotToExistingNode)
+        let updatedShard = try map.nodeID(for: CollectionOfOne(2))
+        #expect(updatedShard.master == luckyReplica)
+        #expect(updatedShard != ogShard)
+
+        // test neighboring hashes have seen an update as well
+        let updatedShard1 = try map.nodeID(for: CollectionOfOne(1))
+        let updatedShard3 = try map.nodeID(for: CollectionOfOne(3))
+
+        #expect(updatedShard == updatedShard1)
+        #expect(updatedShard == updatedShard3)
+    }
+
+    @Test("Case 3: MovedError specifies another shards primary node")
+    func movedErrorSpecifiesOtherShardPrimaryNode() throws {
+        let clusterDescription = self.makeExampleCusterWithNShardsAndMReplicasPerShard(shards: 3, replicas: 3)
+
+        var map = HashSlotShardMap()
+        map.updateCluster(clusterDescription.shards)
+
+        let ogShard = try map.nodeID(for: CollectionOfOne(2))
+        let otherShard = try map.nodeID(for: CollectionOfOne(.max))
+        let newPrimary = otherShard.master
+
+        let update = map.updateSlots(with: ValkeyMovedError(slot: 2, endpoint: newPrimary.endpoint, port: newPrimary.port))
+        #expect(update == .updatedSlotToExistingNode)
+        let updatedShard = try map.nodeID(for: CollectionOfOne(2))
+        #expect(updatedShard == otherShard)
+
+        // test neighboring hashes have not been updated
+        let updatedShard1 = try map.nodeID(for: CollectionOfOne(1))
+        let updatedShard3 = try map.nodeID(for: CollectionOfOne(3))
+
+        #expect(ogShard == updatedShard1)
+        #expect(ogShard == updatedShard3)
+    }
+
+    @Test("Case 4: MovedError specifies another shards replica node")
+    func movedErrorSpecifiesOtherShardReplicaNode() throws {
+        let clusterDescription = self.makeExampleCusterWithNShardsAndMReplicasPerShard(shards: 3, replicas: 3)
+
+        var map = HashSlotShardMap()
+        map.updateCluster(clusterDescription.shards)
+
+        let ogShard = try map.nodeID(for: CollectionOfOne(2))
+        let otherShard = try map.nodeID(for: CollectionOfOne(.max))
+        let newPrimary = otherShard.replicas.randomElement()!
+
+        let update = map.updateSlots(with: ValkeyMovedError(slot: 2, endpoint: newPrimary.endpoint, port: newPrimary.port))
+        #expect(update == .updatedSlotToExistingNode)
+        let updatedShard = try map.nodeID(for: CollectionOfOne(2))
+        #expect(updatedShard.master == newPrimary)
+        #expect(updatedShard.replicas.isEmpty)
+        #expect(updatedShard != ogShard)
+
+        // test neighboring hashes have not been updated
+        let updatedShard1 = try map.nodeID(for: CollectionOfOne(1))
+        let updatedShard3 = try map.nodeID(for: CollectionOfOne(3))
+
+        #expect(ogShard == updatedShard1)
+        #expect(ogShard == updatedShard3)
+
+        // test other shard has been updated and new primary replica has been removed there
+        let otherShardUpdated = try map.nodeID(for: CollectionOfOne(.max))
+        #expect(!otherShardUpdated.replicas.contains(newPrimary))
+    }
+
+    @Test("Case 5: MovedError specifies previously unknown node")
+    func movedErrorSpecifiesPreviouslyUnknownNode() throws {
+        let clusterDescription = self.makeExampleCusterWithNShardsAndMReplicasPerShard(shards: 3, replicas: 3)
+
+        var map = HashSlotShardMap()
+        map.updateCluster(clusterDescription.shards)
+
+        let ogShard = try map.nodeID(for: CollectionOfOne(2))
+        let otherShard = try map.nodeID(for: CollectionOfOne(.max))
+        let newPrimary = ValkeyNodeID(endpoint: "new.valkey.io", port: 6379)
+
+        let update = map.updateSlots(with: ValkeyMovedError(slot: 2, endpoint: newPrimary.endpoint, port: newPrimary.port))
+        #expect(update == .updatedSlotToUnknownNode)
+        let updatedShard = try map.nodeID(for: CollectionOfOne(2))
+        #expect(updatedShard.master == newPrimary)
+        #expect(updatedShard.replicas.isEmpty)
+        #expect(updatedShard != ogShard)
+
+        // test neighboring hashes have not been updated
+        let updatedShard1 = try map.nodeID(for: CollectionOfOne(1))
+        let updatedShard3 = try map.nodeID(for: CollectionOfOne(3))
+
+        #expect(ogShard == updatedShard1)
+        #expect(ogShard == updatedShard3)
+
+        // test other shard has been updated and new primary replica has been removed there
+        let otherShardUpdated = try map.nodeID(for: CollectionOfOne(.max))
+        #expect(!otherShardUpdated.replicas.contains(newPrimary))
+    }
 }