Add support for ASK errors being returned by Valkey cluster (#199)

Author: adam-fowler

Sources/Valkey/Cluster/HashSlotShardMap.swift

@@ -181,7 +181,8 @@ package struct HashSlotShardMap: Sendable {
     ///
     /// 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 {
+    package mutating func updateSlots(with movedError: ValkeyClusterRedirectionError) -> UpdateSlotsResult {
+        assert(movedError.redirection == .move)
         if let shardIndex = self.slotToShardID[Int(movedError.slot.rawValue)].value {
             // if the slot had a shard assignment before
             var shard = self.shardIDToShard[shardIndex]

Sources/Valkey/Cluster/ValkeyClusterClient.swift

@@ -148,18 +148,28 @@ public final class ValkeyClusterClient: Sendable {
             try await self.nodeClient(for: hashSlot.map { [$0] } ?? [])
         }
 
+        var asking = false
         while !Task.isCancelled {
             do {
                 let client = try await clientSelector()
-                return try await client.execute(command)
+                if asking {
+                    // if asking we need to call ASKING beforehand otherwise we will get a MOVE error
+                    return try await client.execute(
+                        ASKING(),
+                        command
+                    ).1.get()
+                } else {
+                    return try await client.execute(command)
+                }
             } catch ValkeyClusterError.noNodeToTalkTo {
                 // TODO: Rerun node discovery!
             } catch let error as ValkeyClientError where error.errorCode == .commandError {
-                guard let errorMessage = error.message, let movedError = ValkeyMovedError(errorMessage) else {
+                guard let errorMessage = error.message, let redirectError = ValkeyClusterRedirectionError(errorMessage) else {
                     throw error
                 }
-                self.logger.trace("Received move error", metadata: ["error": "\(movedError)"])
-                clientSelector = { try await self.nodeClient(for: movedError) }
+                self.logger.trace("Received redirect error", metadata: ["error": "\(redirectError)"])
+                clientSelector = { try await self.nodeClient(for: redirectError) }
+                asking = (redirectError.redirection == .ask)
             }
         }
         throw CancellationError()
@@ -366,25 +376,29 @@ public final class ValkeyClusterClient: Sendable {
     /// This internal method is used when handling cluster topology changes indicated by
     /// MOVED responses from Valkey nodes.
     ///
-    /// - Parameter moveError: The MOVED error response from a Valkey node.
+    /// - Parameter redirectError: The MOVED/ASK error response from a Valkey node.
     /// - Returns: A ``ValkeyNode`` connected to the node that can handle the request.
     /// - Throws:
     ///   - `ValkeyClusterError.waitedForDiscoveryAfterMovedErrorThreeTimes` if unable to resolve
     ///     the MOVED error after multiple attempts
     ///   - `ValkeyClusterError.clientRequestCancelled` if the request is cancelled
     @usableFromInline
-    /* private */ func nodeClient(for moveError: ValkeyMovedError) async throws -> ValkeyNodeClient {
+    /* private */ func nodeClient(for redirectError: ValkeyClusterRedirectionError) async throws -> ValkeyNodeClient {
         var counter = 0
         while counter < 3 {
             defer { counter += 1 }
-            let action = try self.stateLock.withLock { stateMachine throws(ValkeyClusterError) -> StateMachine.PoolForMovedErrorAction in
-                try stateMachine.poolFastPath(for: moveError)
+            let action = try self.stateLock.withLock { stateMachine throws(ValkeyClusterError) -> StateMachine.PoolForRedirectErrorAction in
+                try stateMachine.poolFastPath(for: redirectError)
             }
 
             switch action {
             case .connectionPool(let node):
                 return node
 
+            case .runAndUseConnectionPool(let node):
+                self.queueAction(.runClient(node))
+                return node
+
             case .waitForDiscovery:
                 break
 
@@ -432,7 +446,7 @@ public final class ValkeyClusterClient: Sendable {
     ///   - `ValkeyClusterError.clusterIsUnavailable` if no healthy nodes are available
     ///   - `ValkeyClusterError.clusterIsMissingSlotAssignment` if the slot assignment cannot be determined
     @inlinable
-    func nodeClient(for slots: [HashSlot]) async throws -> ValkeyNodeClient {
+    package func nodeClient(for slots: some (Collection<HashSlot> & Sendable)) async throws -> ValkeyNodeClient {
         var retries = 0
         while retries < 3 {
             defer { retries += 1 }
@@ -492,7 +506,7 @@ public final class ValkeyClusterClient: Sendable {
     /// Handles the transition to a degraded state when a moved error is received.
     ///
     /// - Parameter action: The action containing operations for degraded mode.
-    private func runMovedToDegraded(_ action: StateMachine.PoolForMovedErrorAction.MoveToDegraded) {
+    private func runMovedToDegraded(_ action: StateMachine.PoolForRedirectErrorAction.MoveToDegraded) {
         if let cancelToken = action.runDiscoveryAndCancelTimer {
             cancelToken.yield()
             self.queueAction(.runClusterDiscovery(runNodeDiscovery: false))

Sources/Valkey/Cluster/ValkeyClusterClientStateMachine.swift

@@ -559,8 +559,9 @@ package struct ValkeyClusterClientStateMachine<
     }
 
     @usableFromInline
-    package enum PoolForMovedErrorAction {
+    package enum PoolForRedirectErrorAction {
         case connectionPool(ConnectionPool)
+        case runAndUseConnectionPool(ConnectionPool)
         case moveToDegraded(MoveToDegraded)
         case waitForDiscovery
 
@@ -577,7 +578,29 @@ package struct ValkeyClusterClientStateMachine<
     }
 
     @usableFromInline
-    package mutating func poolFastPath(for movedError: ValkeyMovedError) throws(ValkeyClusterError) -> PoolForMovedErrorAction {
+    package mutating func poolFastPath(for redirectError: ValkeyClusterRedirectionError) throws(ValkeyClusterError) -> PoolForRedirectErrorAction {
+        switch redirectError.redirection {
+        case .ask:
+            return try poolFastPath(forAskError: redirectError)
+        case .move:
+            return try poolFastPath(forMovedError: redirectError)
+        }
+    }
+
+    @usableFromInline
+    package mutating func poolFastPath(forAskError askError: ValkeyClusterRedirectionError) throws(ValkeyClusterError) -> PoolForRedirectErrorAction {
+        switch self.runningClients.addNode(ValkeyNodeDescription(redirectionError: askError)) {
+        case .useExistingPool(let connectionPool):
+            return .connectionPool(connectionPool)
+        case .runAndUsePool(let connectionPool):
+            return .runAndUseConnectionPool(connectionPool)
+        }
+    }
+
+    @usableFromInline
+    package mutating func poolFastPath(
+        forMovedError movedError: ValkeyClusterRedirectionError
+    ) throws(ValkeyClusterError) -> PoolForRedirectErrorAction {
         switch self.clusterState {
         case .unavailable(let unavailableContext):
             if unavailableContext.start.advanced(by: self.configuration.circuitBreakerDuration) > self.clock.now {
@@ -606,7 +629,6 @@ package struct ValkeyClusterClientStateMachine<
                     return .connectionPool(pool)
                 }
             }
-
             let circuitBreakerTimerID = self.nextTimerID()
 
             self.clusterState = .degraded(

Sources/Valkey/Cluster/ValkeyMovedError.swift renamed to Sources/Valkey/Cluster/ValkeyClusterRedirectionError.swift

@@ -7,16 +7,31 @@
 //
 import NIOCore
 
-/// Represents a `MOVED` redirection error from a Valkey cluster node.
+/// Represents a redirection error from a Valkey cluster node.
 ///
 /// When a client sends a command to a Valkey cluster node that doesn't own
 /// the hash slot for the specified key, the node responds with a `MOVED` error
 /// containing information about which node actually owns that slot.
 ///
+/// When a client sends a command to a Valkey cluster node for a hash slot that
+/// is currently migrating for a key that does not exist the node responds
+/// with a `ASK` error containing information about which node is importing
+/// that hash slot
+///
 /// This error provides the necessary information for clients to redirect their
 /// request to the correct node in the cluster.
 @usableFromInline
-package struct ValkeyMovedError: Hashable, Sendable {
+package struct ValkeyClusterRedirectionError: Hashable, Sendable {
+    @usableFromInline
+    package enum Redirection: Sendable {
+        case move
+        case ask
+    }
+
+    /// Request type
+    @usableFromInline
+    package var redirection: Redirection
+
     /// The hash slot number that triggered the redirection.
     package var slot: HashSlot
 
@@ -26,7 +41,8 @@ package struct ValkeyMovedError: Hashable, Sendable {
     /// The port number of the node that owns the requested hash slot.
     package var port: Int
 
-    package init(slot: HashSlot, endpoint: String, port: Int) {
+    package init(request: Redirection, slot: HashSlot, endpoint: String, port: Int) {
+        self.redirection = request
         self.slot = slot
         self.endpoint = endpoint
         self.port = port
@@ -38,25 +54,33 @@ package struct ValkeyMovedError: Hashable, Sendable {
     }
 }
 
-extension ValkeyMovedError {
+extension ValkeyClusterRedirectionError {
     static let movedPrefix = "MOVED "
+    static let askPrefix = "ASK "
 
-    /// Attempts to parse a Valkey MOVED error from a String.
+    /// Attempts to parse a Valkey MOVED/ASK error from a String.
     ///
     /// This method extracts the hash slot, endpoint, and port information from the string
-    /// if it represents a Valkey MOVED error. MOVED errors are returned by Valkey cluster
-    /// nodes when a client attempts to access a key that belongs to a different node.
+    /// if it represents a Valkey MOVED/ASK error. Redirection errors are returned by Valkey
+    /// cluster nodes when a client attempts to access a key that belongs to a different node
+    /// or the hashslot is currently migrating.
     ///
     /// The error format is expected to be: `"MOVED <slot> <endpoint>:<port>"`
     ///
-    /// - Returns: A `ValkeyMovedError` if the token represents a valid MOVED error, or `nil` otherwise.
+    /// - Returns: A `ValkeyClusterRedirectionError` if the token represents a valid MOVED/ASK error, or `nil` otherwise.
     @usableFromInline
     init?(_ errorMessage: String) {
-        guard errorMessage.hasPrefix(Self.movedPrefix) else {
+        let msg: String.SubSequence
+        let request: Redirection
+        if errorMessage.hasPrefix(Self.movedPrefix) {
+            msg = errorMessage.dropFirst(Self.movedPrefix.count)
+            request = .move
+        } else if errorMessage.hasPrefix(Self.askPrefix) {
+            msg = errorMessage.dropFirst(Self.askPrefix.count)
+            request = .ask
+        } else {
             return nil
         }
-
-        let msg = errorMessage.dropFirst(Self.movedPrefix.count)
         guard let spaceAfterSlotIndex = msg.firstIndex(where: { $0 == " " }) else {
             return nil
         }
@@ -79,6 +103,6 @@ extension ValkeyMovedError {
             return nil
         }
 
-        self = ValkeyMovedError(slot: slot, endpoint: Swift.String(endpoint), port: port)
+        self = ValkeyClusterRedirectionError(request: request, slot: slot, endpoint: Swift.String(endpoint), port: port)
     }
 }

Sources/Valkey/Cluster/ValkeyNodeDescription.swift

@@ -80,6 +80,18 @@ package struct ValkeyNodeDescription: Identifiable, Hashable, Sendable {
         self.port = description.tlsPort ?? description.port ?? 6379
     }
 
+    /// Creates a node description from a redirection error.
+    ///
+    /// This initializer converts a `ValkeyClusterRedirectionError` to a `ValkeyNodeDescription`.
+    ///
+    /// - Parameter redirectionError: A `ValkeyClusterRedirectionError` instance.
+    package init(redirectionError: ValkeyClusterRedirectionError) {
+        self.host = nil
+        self.ip = nil
+        self.endpoint = redirectionError.endpoint
+        self.port = redirectionError.port
+    }
+
     /// Determines whether this node description matches a given cluster node description.
     ///
     /// This method compares the essential connection properties of this node with

Sources/Valkey/Commands/ClusterCommands.swift

@@ -284,8 +284,6 @@ public enum CLUSTER {
     /// Returns the key names in a hash slot.
     @_documentation(visibility: internal)
     public struct GETKEYSINSLOT: ValkeyCommand {
-        public typealias Response = RESPToken.Array
-
         @inlinable public static var name: String { "CLUSTER GETKEYSINSLOT" }
 
         public var slot: Int
@@ -462,8 +460,6 @@ public enum CLUSTER {
     /// Returns the ID of a node.
     @_documentation(visibility: internal)
     public struct MYID: ValkeyCommand {
-        public typealias Response = ByteBuffer
-
         @inlinable public static var name: String { "CLUSTER MYID" }
 
         @inlinable public init() {
@@ -477,8 +473,6 @@ public enum CLUSTER {
     /// Returns the shard ID of a node.
     @_documentation(visibility: internal)
     public struct MYSHARDID: ValkeyCommand {
-        public typealias Response = ByteBuffer
-
         @inlinable public static var name: String { "CLUSTER MYSHARDID" }
 
         @inlinable public init() {
@@ -1014,7 +1008,7 @@ extension ValkeyClientProtocol {
     /// - Response: [Array]: An array with up to count elements.
     @inlinable
     @discardableResult
-    public func clusterGetkeysinslot(slot: Int, count: Int) async throws -> RESPToken.Array {
+    public func clusterGetkeysinslot(slot: Int, count: Int) async throws -> CLUSTER.GETKEYSINSLOT.Response {
         try await execute(CLUSTER.GETKEYSINSLOT(slot: slot, count: count))
     }
 
@@ -1108,7 +1102,7 @@ extension ValkeyClientProtocol {
     /// - Response: [String]: The node id.
     @inlinable
     @discardableResult
-    public func clusterMyid() async throws -> ByteBuffer {
+    public func clusterMyid() async throws -> CLUSTER.MYID.Response {
         try await execute(CLUSTER.MYID())
     }
 
@@ -1120,7 +1114,7 @@ extension ValkeyClientProtocol {
     /// - Response: [String]: The node's shard id.
     @inlinable
     @discardableResult
-    public func clusterMyshardid() async throws -> ByteBuffer {
+    public func clusterMyshardid() async throws -> CLUSTER.MYSHARDID.Response {
         try await execute(CLUSTER.MYSHARDID())
     }
 

Sources/Valkey/Commands/Custom/ClusterCustomCommands.swift

@@ -7,6 +7,18 @@
 //
 import NIOCore
 
+extension CLUSTER.GETKEYSINSLOT {
+    public typealias Response = [ValkeyKey]
+}
+
+extension CLUSTER.MYID {
+    public typealias Response = String
+}
+
+extension CLUSTER.MYSHARDID {
+    public typealias Response = String
+}
+
 extension CLUSTER.SHARDS {
     public typealias Response = ValkeyClusterDescription
 }

Sources/Valkey/Commands/GenericCommands.swift

@@ -352,15 +352,15 @@ public struct MIGRATE<Host: RESPStringRenderable>: ValkeyCommand {
         public var respEntries: Int {
             switch self {
             case .key(let key): key.respEntries
-            case .emptyString: "\"\"".respEntries
+            case .emptyString: "".respEntries
             }
         }
 
         @inlinable
         public func encode(into commandEncoder: inout ValkeyCommandEncoder) {
             switch self {
             case .key(let key): key.encode(into: &commandEncoder)
-            case .emptyString: "\"\"".encode(into: &commandEncoder)
+            case .emptyString: "".encode(into: &commandEncoder)
             }
         }
     }

Sources/Valkey/Node/ValkeyNodeClient.swift

@@ -35,17 +35,17 @@ public final class ValkeyNodeClient: Sendable {
         ContinuousClock
     >
     /// Server address
-    let serverAddress: ValkeyServerAddress
+    public let serverAddress: ValkeyServerAddress
     /// Connection pool
     let connectionPool: Pool
 
     let connectionFactory: ValkeyConnectionFactory
     /// configuration
-    var configuration: ValkeyClientConfiguration { self.connectionFactory.configuration }
+    public var configuration: ValkeyClientConfiguration { self.connectionFactory.configuration }
     /// EventLoopGroup to use
-    let eventLoopGroup: any EventLoopGroup
+    public let eventLoopGroup: any EventLoopGroup
     /// Logger
-    let logger: Logger
+    public let logger: Logger
 
     package init(
         _ address: ValkeyServerAddress,

Sources/Valkey/Node/ValkeyRunningClients.swift

@@ -89,6 +89,22 @@ struct ValkeyRunningClientsStateMachine<
         )
     }
 
+    enum AddNodeAction {
+        case useExistingPool(ConnectionPool)
+        case runAndUsePool(ConnectionPool)
+    }
+
+    mutating func addNode(
+        _ node: ValkeyNodeDescription
+    ) -> AddNodeAction {
+        if let pool = self.clientMap[node.id] {
+            return .useExistingPool(pool.pool)
+        }
+        let newPool = self.makePool(for: node)
+        self.clientMap[node.id] = NodeBundle(pool: newPool, nodeDescription: node)
+        return .runAndUsePool(newPool)
+    }
+
     @inlinable
     subscript(_ index: ValkeyNodeID) -> NodeBundle? {
         self.clientMap[index]