Add support for reading from readonly replicas

Signed-off-by: Adam Fowler <adamfowler71@gmail.com>

Author: adam-fowler

Sources/Valkey/Cluster/ValkeyClusterClient.swift

@@ -146,8 +146,9 @@ public final class ValkeyClusterClient: Sendable {
     @inlinable
     public func execute<Command: ValkeyCommand>(_ command: Command) async throws -> Command.Response {
         let hashSlot = try self.hashSlot(for: command.keysAffected)
+        let readOnly = self.clientConfiguration.useReadOnlyReplicas && command.isReadOnly
         var clientSelector: () async throws -> ValkeyNodeClient = {
-            try await self.nodeClient(for: hashSlot.map { [$0] } ?? [])
+            try await self.nodeClient(for: hashSlot.map { [$0] } ?? [], readOnly: readOnly)
         }
 
         var asking = false
@@ -464,11 +465,12 @@ public final class ValkeyClusterClient: Sendable {
     @inlinable
     public func withConnection<Value>(
         forKeys keys: some Collection<ValkeyKey>,
+        readOnly: Bool = false,
         isolation: isolated (any Actor)? = #isolation,
         operation: (ValkeyConnection) async throws -> sending Value
     ) async throws -> Value {
         let hashSlots = keys.compactMap { HashSlot(key: $0) }
-        let node = try await self.nodeClient(for: hashSlots)
+        let node = try await self.nodeClient(for: hashSlots, readOnly: self.clientConfiguration.useReadOnlyReplicas && readOnly)
         return try await node.withConnection(isolation: isolation, operation: operation)
     }
 
@@ -570,7 +572,7 @@ public final class ValkeyClusterClient: Sendable {
                 // Get hash slot for key and add all the commands you have iterated through so far to the
                 // node associated with that key and break out of loop
                 let hashSlot = try self.hashSlot(for: keysAffected)
-                let node = try await self.nodeClient(for: hashSlot.map { [$0] } ?? [])
+                let node = try await self.nodeClient(for: hashSlot.map { [$0] } ?? [], readOnly: false)
                 let address = node.serverAddress
                 let nodeAndCommands = NodeAndCommands(node: node, commandIndices: .init(commands.startIndex..<index))
                 nodeMap[address] = nodeAndCommands
@@ -586,7 +588,7 @@ public final class ValkeyClusterClient: Sendable {
                 if keysAffected.count > 0 {
                     // If command affects a key get hash slot for key and add command to the node associated with that key
                     let hashSlot = try self.hashSlot(for: keysAffected)
-                    let node = try await self.nodeClient(for: hashSlot.map { [$0] } ?? [])
+                    let node = try await self.nodeClient(for: hashSlot.map { [$0] } ?? [], readOnly: false)
                     prevAddress = node.serverAddress
                     nodeMap[prevAddress, default: .init(node: node, commandIndices: [])].commandIndices.append(index)
                 } else {
@@ -597,7 +599,7 @@ public final class ValkeyClusterClient: Sendable {
             }
         } else {
             // if none of the commands affect any keys then choose a random node
-            let node = try await self.nodeClient(for: [])
+            let node = try await self.nodeClient(for: [], readOnly: false)
             let address = node.serverAddress
             let nodeAndCommands = NodeAndCommands(node: node, commandIndices: .init(commands.startIndex..<index))
             nodeMap[address] = nodeAndCommands
@@ -854,14 +856,14 @@ public final class ValkeyClusterClient: Sendable {
     ///   - `ValkeyClusterError.clusterIsUnavailable` if no healthy nodes are available
     ///   - `ValkeyClusterError.clusterIsMissingSlotAssignment` if the slot assignment cannot be determined
     @inlinable
-    package func nodeClient(for slots: some (Collection<HashSlot> & Sendable)) async throws -> ValkeyNodeClient {
+    package func nodeClient(for slots: some (Collection<HashSlot> & Sendable), readOnly: Bool) async throws -> ValkeyNodeClient {
         var retries = 0
         while retries < 3 {
             defer { retries += 1 }
 
             do {
                 return try self.stateLock.withLock { state -> ValkeyNodeClient in
-                    try state.poolFastPath(for: slots)
+                    try state.poolFastPath(for: slots, readOnly: readOnly)
                 }
             } catch let error as ValkeyClusterError where error == .clusterIsUnavailable {
                 let waiterID = self.nextRequestID()

Sources/Valkey/Cluster/ValkeyClusterClientStateMachine.swift

@@ -532,14 +532,20 @@ where
     }
 
     @inlinable
-    package func poolFastPath(for slots: some Collection<HashSlot>) throws(ValkeyClusterError) -> ConnectionPool {
+    package func poolFastPath(for slots: some Collection<HashSlot>, readOnly: Bool) throws(ValkeyClusterError) -> ConnectionPool {
         switch self.clusterState {
         case .unavailable:
             throw ValkeyClusterError.clusterIsUnavailable
 
         case .degraded(let context):
             let shardID = try context.hashSlotShardMap.nodeID(for: slots)
-            if let pool = self.runningClients[shardID.primary]?.pool {
+            let nodeID =
+                if readOnly {
+                    shardID.replicas.randomElement() ?? shardID.primary
+                } else {
+                    shardID.primary
+                }
+            if let pool = self.runningClients[nodeID]?.pool {
                 return pool
             }
             // If we don't have a node for a shard, that means that this shard got created from
@@ -549,7 +555,13 @@ where
 
         case .healthy(let context):
             let shardID = try context.hashSlotShardMap.nodeID(for: slots)
-            if let pool = self.runningClients[shardID.primary]?.pool {
+            let nodeID =
+                if readOnly {
+                    shardID.replicas.randomElement() ?? shardID.primary
+                } else {
+                    shardID.primary
+                }
+            if let pool = self.runningClients[nodeID]?.pool {
                 return pool
             }
             // If we don't have a node for a shard, that means that this shard got created from

Sources/Valkey/Subscriptions/ValkeyClusterClient+subscribe.swift

@@ -20,7 +20,7 @@ extension ValkeyClusterClient {
         isolation: isolated (any Actor)? = #isolation,
         _ operation: (ValkeyConnection) async throws -> sending Value
     ) async throws -> sending Value {
-        let node = try await self.nodeClient(for: [])
+        let node = try await self.nodeClient(for: [], readOnly: false)
         let id = node.subscriptionConnectionIDGenerator.next()
 
         let connection = try await withTaskCancellationHandler {

Sources/Valkey/ValkeyClientConfiguration.swift

@@ -171,6 +171,14 @@ public struct ValkeyClientConfiguration: Sendable {
     /// The TLS to use for the Valkey connection.
     public var tls: TLS
 
+    /// Execute readonly commands on replicas.
+    ///
+    /// Cluster by default will redirect commands from replica nodes to the primary node.
+    /// Setting this flag will allow replicas to run readonly commands. This will reduce
+    /// load on your primary nodes but there is a chance you will receive stale data as
+    /// the replica is not up to date.
+    public var useReadOnlyReplicas: Bool
+
     #if DistributedTracingSupport
     /// The distributed tracing configuration to use for the Valkey connection.
     /// Defaults to using the globally bootstrapped tracer with OpenTelemetry semantic conventions.
@@ -187,14 +195,16 @@ public struct ValkeyClientConfiguration: Sendable {
     ///   - commandTimeout: The timeout for a connection response.
     ///   - blockingCommandTimeout: The timeout for a blocking command response.
     ///   - tls: The TLS configuration.
+    ///   - useReadOnlyReplicas: Execute readonly commands on replicas
     public init(
         authentication: Authentication? = nil,
         connectionPool: ConnectionPool = .init(),
         keepAliveBehavior: KeepAliveBehavior = .init(),
         retryParameters: RetryParameters = .init(),
         commandTimeout: Duration = .seconds(30),
         blockingCommandTimeout: Duration = .seconds(120),
-        tls: TLS = .disable
+        tls: TLS = .disable,
+        useReadOnlyReplicas: Bool = false
     ) {
         self.authentication = authentication
         self.connectionPool = connectionPool
@@ -203,5 +213,6 @@ public struct ValkeyClientConfiguration: Sendable {
         self.commandTimeout = commandTimeout
         self.blockingCommandTimeout = blockingCommandTimeout
         self.tls = tls
+        self.useReadOnlyReplicas = useReadOnlyReplicas
     }
 }

Tests/ClusterIntegrationTests/ClusterIntegrationTests.swift

@@ -54,6 +54,24 @@ struct ClusterIntegrationTests {
         }
     }
 
+    @Test
+    @available(valkeySwift 1.0, *)
+    func testWithReadOnlyConnection() async throws {
+        var logger = Logger(label: "ValkeyCluster")
+        logger.logLevel = .trace
+        let firstNodeHostname = clusterFirstNodeHostname!
+        let firstNodePort = clusterFirstNodePort ?? 6379
+        try await Self.withValkeyCluster([(host: firstNodeHostname, port: firstNodePort)], logger: logger) { client in
+            try await Self.withKey(connection: client) { key in
+                try await client.set(key, value: "Hello")
+                try await client.withConnection(forKeys: [key], readOnly: true) { connection in
+                    let response = try await connection.get(key)
+                    #expect(response.map { String(buffer: $0) } == "Hello")
+                }
+            }
+        }
+    }
+
     @Test
     @available(valkeySwift 1.0, *)
     func testFailover() async throws {
@@ -270,11 +288,11 @@ struct ClusterIntegrationTests {
         afterMigrate: () async throws -> Void = {},
         finished: () async throws -> Void = {}
     ) async throws {
-        let nodeAClient = try await client.nodeClient(for: [hashSlot])
+        let nodeAClient = try await client.nodeClient(for: [hashSlot], readOnly: false)
         // find another shard
         var nodeBClient: ValkeyNodeClient
         repeat {
-            nodeBClient = try await client.nodeClient(for: [HashSlot(rawValue: Int.random(in: 0..<16384))!])
+            nodeBClient = try await client.nodeClient(for: [HashSlot(rawValue: Int.random(in: 0..<16384))!], readOnly: false)
         } while nodeAClient === nodeBClient
 
         guard let (hostnameA, portA) = nodeAClient.serverAddress.getHostnameAndPort() else { return }
@@ -341,7 +359,7 @@ struct ClusterIntegrationTests {
             try await ClusterIntegrationTests.withValkeyCluster([(host: firstNodeHostname, port: firstNodePort)], logger: logger) {
                 client in
                 try await ClusterIntegrationTests.withKey(connection: client, suffix: "{foo}") { key in
-                    let node = try await client.nodeClient(for: [HashSlot(key: key)])
+                    let node = try await client.nodeClient(for: [HashSlot(key: key)], readOnly: false)
                     var commands: [any ValkeyCommand] = .init()
                     commands.append(SET(key, value: "cluster pipeline test"))
                     commands.append(GET(key))
@@ -363,13 +381,13 @@ struct ClusterIntegrationTests {
                 client in
                 try await ClusterIntegrationTests.withKey(connection: client, suffix: "{foo}") { key in
                     let hashSlot = HashSlot(key: key)
-                    let node = try await client.nodeClient(for: [hashSlot])
+                    let node = try await client.nodeClient(for: [hashSlot], readOnly: false)
                     // get a key from same node
                     let key2 = try await {
                         while true {
                             let key2 = ValkeyKey(UUID().uuidString)
                             let hashSlot2 = HashSlot(key: key2)
-                            let node2 = try await client.nodeClient(for: [hashSlot2])
+                            let node2 = try await client.nodeClient(for: [hashSlot2], readOnly: false)
                             if node2.serverAddress == node.serverAddress {
                                 return key2
                             }
@@ -400,12 +418,12 @@ struct ClusterIntegrationTests {
                 client in
                 try await ClusterIntegrationTests.withKey(connection: client, suffix: "{foo}") { key in
                     let hashSlot = HashSlot(key: key)
-                    let node = try await client.nodeClient(for: [hashSlot])
+                    let node = try await client.nodeClient(for: [hashSlot], readOnly: false)
                     let key2 = try await {
                         while true {
                             let key2 = ValkeyKey(UUID().uuidString)
                             let hashSlot2 = HashSlot(key: key2)
-                            let node2 = try await client.nodeClient(for: [hashSlot2])
+                            let node2 = try await client.nodeClient(for: [hashSlot2], readOnly: false)
                             if node2.serverAddress != node.serverAddress {
                                 return key2
                             }
@@ -435,12 +453,12 @@ struct ClusterIntegrationTests {
                 client in
                 try await ClusterIntegrationTests.withKey(connection: client, suffix: "{foo}") { key in
                     let hashSlot = HashSlot(key: key)
-                    let node = try await client.nodeClient(for: [hashSlot])
+                    let node = try await client.nodeClient(for: [hashSlot], readOnly: false)
                     let key2 = try await {
                         while true {
                             let key2 = ValkeyKey(UUID().uuidString)
                             let hashSlot2 = HashSlot(key: key2)
-                            let node2 = try await client.nodeClient(for: [hashSlot2])
+                            let node2 = try await client.nodeClient(for: [hashSlot2], readOnly: false)
                             if node2.serverAddress != node.serverAddress {
                                 return key2
                             }
@@ -470,7 +488,7 @@ struct ClusterIntegrationTests {
             try await ClusterIntegrationTests.withValkeyCluster([(host: firstNodeHostname, port: firstNodePort)], logger: logger) {
                 clusterClient in
                 try await ClusterIntegrationTests.withKey(connection: clusterClient) { key in
-                    let node = try await clusterClient.nodeClient(for: [HashSlot(key: key)])
+                    let node = try await clusterClient.nodeClient(for: [HashSlot(key: key)], readOnly: false)
                     try await clusterClient.set(key, value: "bar")
                     let cluster = try await clusterClient.clusterShards()
                     let shard = try #require(
@@ -512,7 +530,7 @@ struct ClusterIntegrationTests {
                 let keySuffix = "{\(UUID().uuidString)}"
                 try await ClusterIntegrationTests.withKey(connection: client, suffix: keySuffix) { key in
                     let hashSlot = HashSlot(key: key)
-                    let node = try await client.nodeClient(for: [hashSlot])
+                    let node = try await client.nodeClient(for: [hashSlot], readOnly: false)
                     try await client.set(key, value: "Testing before import")
 
                     try await ClusterIntegrationTests.testMigratingHashSlot(hashSlot, client: client) {
@@ -548,7 +566,7 @@ struct ClusterIntegrationTests {
                 try await ClusterIntegrationTests.withKey(connection: client, suffix: keySuffix) { key in
                     try await ClusterIntegrationTests.withKey(connection: client, suffix: keySuffix) { key2 in
                         let hashSlot = HashSlot(key: key)
-                        let node = try await client.nodeClient(for: [hashSlot])
+                        let node = try await client.nodeClient(for: [hashSlot], readOnly: false)
                         try await client.lpush(key, elements: ["testing1"])
 
                         try await ClusterIntegrationTests.testMigratingHashSlot(hashSlot, client: client) {
@@ -820,7 +838,7 @@ struct ClusterIntegrationTests {
     @available(valkeySwift 1.0, *)
     static func withValkeyCluster<T>(
         _ nodeAddresses: [(host: String, port: Int)],
-        nodeClientConfiguration: ValkeyClientConfiguration = .init(),
+        nodeClientConfiguration: ValkeyClientConfiguration = .init(useReadOnlyReplicas: true),
         logger: Logger,
         _ body: (ValkeyClusterClient) async throws -> sending T
     ) async throws -> T {

Tests/ValkeyTests/Cluster/ValkeyClusterClientStateMachineTests.swift

@@ -57,7 +57,7 @@ struct ValkeyClusterClientStateMachineTests {
 
         // try to get shard
         #expect(throws: ValkeyClusterError.clusterIsUnavailable) {
-            try stateMachine.poolFastPath(for: CollectionOfOne(HashSlot(rawValue: 100)!))
+            try stateMachine.poolFastPath(for: CollectionOfOne(HashSlot(rawValue: 100)!), readOnly: false)
         }
         // let's register for the slow path
         let successNotifier = SuccessNotifier()
@@ -94,7 +94,7 @@ struct ValkeyClusterClientStateMachineTests {
 
         // try to get shard
         #expect(throws: ValkeyClusterError.clusterIsUnavailable) {
-            try stateMachine.poolFastPath(for: CollectionOfOne(HashSlot(rawValue: 100)!))
+            try stateMachine.poolFastPath(for: CollectionOfOne(HashSlot(rawValue: 100)!), readOnly: false)
         }
         // let's register for the slow path
         let successNotifier = SuccessNotifier()