Cluster client transaction support

Signed-off-by: Adam Fowler <[email protected]>

Author: adam-fowler

Sources/Valkey/Cluster/ValkeyClusterClient.swift

@@ -303,6 +303,62 @@ public final class ValkeyClusterClient: Sendable {
         }
     }
 
+    /// Pipeline a series of commands as a transaction to Valkey connection
+    ///
+    /// Another client will never be served in the middle of the execution of these
+    /// commands. See https://valkey.io/topics/transactions/ for more information.
+    ///
+    /// EXEC and MULTI commands are added to the pipelined commands and the output
+    /// of the EXEC command is transformed into an array of RESPToken Results, one for
+    /// each command.
+    ///
+    /// This is an alternative version of the transaction function ``ValkeyConnection/transaction(_:)->(_,_)``
+    /// that allows for a collection of ValkeyCommands. It provides more flexibility but the command
+    /// responses are returned as ``RESPToken`` instead of the response type for the command.
+    ///
+    /// - Parameter commands: Collection of ValkeyCommands
+    /// - Returns: Array holding the RESPToken responses of all the commands
+    /// - Throws: ValkeyTransactionError when EXEC aborts
+    @inlinable
+    public func transaction<Commands: Collection & Sendable>(
+        _ commands: Commands
+    ) async throws -> [Result<RESPToken, Error>] where Commands.Element == any ValkeyCommand {
+        let hashSlot = try self.hashSlot(for: commands.flatMap { $0.keysAffected })
+
+        var clientSelector: () async throws -> ValkeyNodeClient = {
+            try await self.nodeClient(for: hashSlot.map { [$0] } ?? [])
+        }
+
+        var asking = false
+        var attempt = 0
+        while !Task.isCancelled {
+            do {
+                let client = try await clientSelector()
+                if asking {
+                    return try await client.transactionWithAsk(commands)
+                } else {
+                    return try await client.transaction(commands)
+                }
+            } catch let error as ValkeyClusterError where error == .noNodeToTalkTo {
+                // TODO: Rerun node discovery!
+            } catch {
+                let retryAction = self.getTransactionRetryAction(from: error)
+                switch retryAction {
+                case .redirect(let redirectError):
+                    clientSelector = { try await self.nodeClient(for: redirectError) }
+                    asking = (redirectError.redirection == .ask)
+                case .tryAgain:
+                    let wait = self.clientConfiguration.retryParameters.calculateWaitTime(retry: attempt)
+                    try await Task.sleep(for: wait)
+                    attempt += 1
+                case .dontRetry:
+                    throw error
+                }
+            }
+        }
+        throw ValkeyClusterError.clientRequestCancelled
+    }
+
     struct Redirection {
         let node: ValkeyNodeClient
         let ask: Bool
@@ -565,6 +621,33 @@ public final class ValkeyClusterClient: Sendable {
         }
     }
 
+    @usableFromInline
+    /* private */ func getTransactionRetryAction(from error: some Error) -> RetryAction {
+        switch error {
+        case let transactionError as ValkeyTransactionError:
+            switch transactionError {
+            case .transactionErrors(let results, let execError):
+                for result in results {
+                    if case .failure(let queuedError) = result {
+                        let queuedAction = self.getRetryAction(from: queuedError)
+                        guard case .dontRetry = queuedAction else {
+                            return queuedAction
+                        }
+                    }
+                }
+                let execAction = self.getRetryAction(from: execError)
+                guard case .dontRetry = execAction else {
+                    return execAction
+                }
+                return .dontRetry
+            case .transactionAborted:
+                return .dontRetry
+            }
+        default:
+            return .dontRetry
+        }
+    }
+
     private func queueAction(_ action: RunAction) {
         self.actionStreamContinuation.yield(action)
     }

Sources/Valkey/Connection/ValkeyConnection.swift

@@ -386,41 +386,9 @@ public final actor ValkeyConnection: ValkeyClientProtocol, Sendable {
         return try await _execute(
             buffer: encoder.buffer,
             promises: promises,
-            valkeyPromises: promises.map { .nio($0) }
-        ) { promises -> Result<[Result<RESPToken, Error>], any Error> in
-            let responses: EXEC.Response
-            do {
-                let execFutureResult = promises.last!.futureResult
-                responses = try await execFutureResult.get().decode(as: EXEC.Response.self)
-            } catch let error as ValkeyClientError where error.errorCode == .commandError {
-                // we received an error while running the EXEC command. Extract queuing
-                // results and throw error
-                var results: [Result<RESPToken, Error>] = .init()
-                results.reserveCapacity(promises.count - 2)
-                for promise in promises[1..<(promises.count - 1)] {
-                    results.append(await promise.futureResult._result())
-                }
-                return .failure(ValkeyTransactionError.transactionErrors(queuedResults: results, execError: error))
-            } catch {
-                return .failure(error)
-            }
-            // If EXEC returned nil then transaction was aborted because a
-            // WATCHed variable changed
-            guard let responses else {
-                return .failure(ValkeyTransactionError.transactionAborted)
-            }
-            // We convert all the RESP errors in the response from EXEC to Result.failure
-            return .success(
-                responses.map {
-                    switch $0.identifier {
-                    case .simpleError, .bulkError:
-                        .failure(ValkeyClientError(.commandError, message: $0.errorString.map { Swift.String(buffer: $0) }))
-                    default:
-                        .success($0)
-                    }
-                }
-            )
-        }.get()
+            valkeyPromises: promises.map { .nio($0) },
+            processResults: self._processTransactionPromises
+        ).get()
     }
 
     /// Pipeline a series of commands to Valkey connection and precede each command with an ASKING
@@ -468,6 +436,48 @@ public final actor ValkeyConnection: ValkeyClientProtocol, Sendable {
         }
     }
 
+    /// Pipeline a series of commands as a transaction preceded with an ASKING command
+    ///
+    /// Once all the responses for the commands have been received the function returns
+    /// an array of RESPToken Results, one for each command.
+    ///
+    /// This is an internal function used by the cluster client
+    ///
+    /// - Parameter commands: Collection of ValkeyCommands
+    /// - Returns: Array holding the RESPToken responses of all the commands
+    @usableFromInline
+    func transactionWithAsk(
+        _ commands: some Collection<any ValkeyCommand>
+    ) async throws -> [Result<RESPToken, any Error>] {
+        // this currently allocates a promise for every command. We could collapse this down to one promise
+        var promises: [EventLoopPromise<RESPToken>] = []
+        promises.reserveCapacity(commands.count)
+        var valkeyPromises: [ValkeyPromise<RESPToken>] = []
+        valkeyPromises.reserveCapacity(commands.count + 3)
+        var encoder = ValkeyCommandEncoder()
+        ASKING().encode(into: &encoder)
+        MULTI().encode(into: &encoder)
+        promises.append(channel.eventLoop.makePromise(of: RESPToken.self))
+        valkeyPromises.append(.forget)
+        valkeyPromises.append(.nio(promises.last!))
+
+        for command in commands {
+            command.encode(into: &encoder)
+            promises.append(channel.eventLoop.makePromise(of: RESPToken.self))
+            valkeyPromises.append(.nio(promises.last!))
+        }
+        EXEC().encode(into: &encoder)
+        promises.append(channel.eventLoop.makePromise(of: RESPToken.self))
+        valkeyPromises.append(.nio(promises.last!))
+
+        return try await _execute(
+            buffer: encoder.buffer,
+            promises: promises,
+            valkeyPromises: valkeyPromises,
+            processResults: self._processTransactionPromises
+        ).get()
+    }
+
     /// Execute stream of commands written into buffer
     ///
     /// The function is provided with an array of EventLoopPromises for the responses of commands
@@ -498,6 +508,44 @@ public final actor ValkeyConnection: ValkeyClientProtocol, Sendable {
         }
     }
 
+    @usableFromInline
+    func _processTransactionPromises(
+        _ promises: [EventLoopPromise<RESPToken>]
+    ) async -> sending Result<[Result<RESPToken, Error>], any Error> {
+        let responses: EXEC.Response
+        do {
+            let execFutureResult = promises.last!.futureResult
+            responses = try await execFutureResult.get().decode(as: EXEC.Response.self)
+        } catch let error as ValkeyClientError where error.errorCode == .commandError {
+            // we received an error while running the EXEC command. Extract queuing
+            // results and throw error
+            var results: [Result<RESPToken, Error>] = .init()
+            results.reserveCapacity(promises.count - 2)
+            for promise in promises[1..<(promises.count - 1)] {
+                results.append(await promise.futureResult._result())
+            }
+            return .failure(ValkeyTransactionError.transactionErrors(queuedResults: results, execError: error))
+        } catch {
+            return .failure(error)
+        }
+        // If EXEC returned nil then transaction was aborted because a
+        // WATCHed variable changed
+        guard let responses else {
+            return .failure(ValkeyTransactionError.transactionAborted)
+        }
+        // We convert all the RESP errors in the response from EXEC to Result.failure
+        return .success(
+            responses.map {
+                switch $0.identifier {
+                case .simpleError, .bulkError:
+                    .failure(ValkeyClientError(.commandError, message: $0.errorString.map { Swift.String(buffer: $0) }))
+                default:
+                    .success($0)
+                }
+            }
+        )
+    }
+
     #if DistributedTracingSupport
     @usableFromInline
     func applyCommonAttributes(to attributes: inout SpanAttributes, commandName: String) {

Sources/Valkey/Node/ValkeyNodeClient.swift

@@ -272,6 +272,17 @@ extension ValkeyNodeClient {
             return .init(repeating: .failure(error), count: commands.count)
         }
     }
+
+    /// Internal command used by cluster client, that precedes each command with a ASKING
+    /// command
+    @usableFromInline
+    func transactionWithAsk<Commands: Collection & Sendable>(
+        _ commands: Commands
+    ) async throws -> [Result<RESPToken, any Error>] where Commands.Element == any ValkeyCommand {
+        try await self.withConnection { connection in
+            try await connection.transactionWithAsk(commands)
+        }
+    }
 }
 
 /// Extension that makes ``ValkeyNode`` conform to ``ValkeyNodeConnectionPool``.