Merge branch 'main' into cli

Author: adam-fowler

Sources/Valkey/ValkeyClient.swift

@@ -45,6 +45,7 @@ public final class ValkeyClient: Sendable {
         case runNodeClient(ValkeyNodeClient)
         case runTimer(ValkeyTimer)
         case runRole
+        case verifyReplicas
     }
     let actionStream: AsyncStream<RunAction>
     let actionStreamContinuation: AsyncStream<RunAction>.Continuation
@@ -185,18 +186,30 @@ extension ValkeyClient {
 
         case .runTimer(let timer):
             await runTimerAction(timer)
+
+        case .verifyReplicas:
+            break
         }
     }
 
     /// Run ROLE command and update topology based on results
     private func runRoleAction() async {
+        struct NotReplicaError: Error {}
         var topology: StateMachine.TopologyRefreshInput
         let nodeClient = self.getNode(readOnly: false)
         do {
             let role = try await nodeClient.execute(ROLE())
             switch role {
             case .primary(let primaryState):
                 let replicas = primaryState.replicas.map { ValkeyServerAddress.hostname($0.ip, port: $0.port) }
+                // add nodes primary found
+                let action = self.stateMachine.withLock { $0.addNewNodesFromRefresh(replicas) }
+                for replicaClient in action.clientsToRun {
+                    self.queueAction(.runNodeClient(replicaClient))
+                }
+                // verify all replicas are replicas
+                guard try await self.verifyReplicas(action.clientsToRun) else { throw NotReplicaError() }
+
                 topology = .primary(replicas: replicas)
                 self.logger.debug("Found replicas \(replicas)")
 
@@ -277,7 +290,7 @@ extension ValkeyClient {
         }
     }
 
-    func runTimerFiredAction(_ action: StateMachine.TimerFiredAction) {
+    private func runTimerFiredAction(_ action: StateMachine.TimerFiredAction) {
         switch action {
         case .runRole:
             self.queueAction(.runRole)
@@ -286,6 +299,21 @@ extension ValkeyClient {
             break
         }
     }
+
+    private func verifyReplicas(_ replicas: [ValkeyNodeClient]) async throws -> Bool {
+        try await withThrowingTaskGroup(of: Bool.self) { group in
+            for replica in replicas {
+                group.addTask {
+                    let role = try await replica.execute(ROLE())
+                    guard case .replica = role else {
+                        return false
+                    }
+                    return true
+                }
+            }
+            return try await group.reduce(true) { $0 && $1 }
+        }
+    }
 }
 
 // MARK: ValkeyClientProtocol methods

Sources/Valkey/ValkeyClientStateMachine.swift

@@ -29,8 +29,13 @@ struct ValkeyClientStateMachine<
     /// Current state
     @usableFromInline
     enum State {
+        @usableFromInline
+        struct HealthyContext {
+            @usableFromInline
+            let nodes: ValkeyNodeIDs<ValkeyServerAddress>
+        }
         case uninitialized
-        case running(ValkeyNodeIDs<ValkeyServerAddress>)
+        case healthy(HealthyContext)
         case shutdown
     }
     /// Topology refresh state
@@ -77,10 +82,10 @@ struct ValkeyClientStateMachine<
 
     @inlinable
     func getNode(_ selection: ValkeyNodeSelection) -> ConnectionPool {
-        guard case .running(let nodes) = self.state else {
+        guard case .healthy(let healthyContext) = self.state else {
             preconditionFailure("Cannot get a node if the client statemachine isn't initialized")
         }
-        let nodeID = selection.select(nodeIDs: nodes)
+        let nodeID = selection.select(nodeIDs: healthyContext.nodes)
         if let pool = self.runningClients[nodeID]?.pool {
             return pool
         } else {
@@ -109,8 +114,8 @@ struct ValkeyClientStateMachine<
         let client: ConnectionPool?
         switch action {
         case .useExistingPool:
-            if case .running(let currentNodes) = self.state {
-                if currentNodes.primary == address {
+            if case .healthy(let healthyContext) = self.state {
+                if healthyContext.nodes.primary == address {
                     return .init()
                 }
             }
@@ -119,7 +124,7 @@ struct ValkeyClientStateMachine<
             client = node
         }
 
-        self.state = .running(nodes)
+        self.state = .healthy(.init(nodes: nodes))
         if self.configuration.findReplicas {
             switch self.topologyRefreshState {
             case .notRefreshing:
@@ -136,6 +141,17 @@ struct ValkeyClientStateMachine<
         }
     }
 
+    func getNodes() -> ValkeyNodeIDs<ValkeyServerAddress>? {
+        switch self.state {
+        case .uninitialized:
+            preconditionFailure("Invalid state: \(self.state)")
+        case .healthy(let healthContext):
+            return healthContext.nodes
+        case .shutdown:
+            return nil
+        }
+    }
+
     enum TimerFiredAction {
         case runRole
         case doNothing
@@ -147,7 +163,7 @@ struct ValkeyClientStateMachine<
             switch self.state {
             case .uninitialized:
                 preconditionFailure("Invalid state: \(self.state)")
-            case .running:
+            case .healthy:
                 break
             case .shutdown:
                 return .doNothing
@@ -175,7 +191,7 @@ struct ValkeyClientStateMachine<
             switch self.state {
             case .shutdown, .uninitialized:
                 return .cancelTimer(token)
-            case .running:
+            case .healthy:
                 break
             }
 
@@ -191,6 +207,28 @@ struct ValkeyClientStateMachine<
         }
     }
 
+    struct NewNodesAction {
+        var clientsToRun: [ConnectionPool] = []
+    }
+
+    mutating func addNewNodesFromRefresh(_ nodes: [ValkeyServerAddress]) -> NewNodesAction {
+        // If we are shutdown ignore
+        guard case .healthy = self.state else {
+            return .init()
+        }
+
+        switch self.topologyRefreshState {
+        case .notRefreshing, .waitingForNextRefresh:
+            preconditionFailure("Invalid state: \(self.topologyRefreshState)")
+
+        case .refreshing:
+            let nodeDescriptions = nodes.map { ValkeyClientNodeDescription(address: $0) }
+            let action = self.runningClients.updateNodes(nodeDescriptions, removeUnmentionedPools: false)
+            precondition(action.poolsToShutdown.isEmpty, "There should be no pools to shutdown.")
+            return .init(clientsToRun: action.poolsToRun.map { $0.0 })
+        }
+    }
+
     struct TopologyRefreshAction {
         enum NextAction: Equatable {
             case refreshTopology
@@ -215,7 +253,7 @@ struct ValkeyClientStateMachine<
     /// Depending on what node we are connected to we either get a primary node or a list of replicas
     mutating func topologyRefreshSucceeded(_ topology: TopologyRefreshInput) -> TopologyRefreshAction {
         // If we are shutdown ignore
-        guard case .running(let nodes) = self.state else {
+        guard case .healthy(let healthyContext) = self.state else {
             return .init()
         }
 
@@ -227,13 +265,13 @@ struct ValkeyClientStateMachine<
             switch topology {
             case .primary(let replicas):
                 var nodeDescriptions = [
-                    ValkeyClientNodeDescription(address: nodes.primary)
+                    ValkeyClientNodeDescription(address: healthyContext.nodes.primary)
                 ]
                 nodeDescriptions.append(
                     contentsOf: replicas.lazy.map { ValkeyClientNodeDescription(address: $0) }
                 )
-                let newNodes = ValkeyNodeIDs(primary: nodes.primary, replicas: replicas)
-                self.state = .running(newNodes)
+                let newNodes = ValkeyNodeIDs(primary: healthyContext.nodes.primary, replicas: replicas)
+                self.state = .healthy(.init(nodes: newNodes))
                 let action = self.runningClients.updateNodes(nodeDescriptions, removeUnmentionedPools: true)
                 let refreshTimerID = self.nextTimerID()
                 self.topologyRefreshState = .waitingForNextRefresh(.init(timer: .init(id: refreshTimerID)))
@@ -254,7 +292,7 @@ struct ValkeyClientStateMachine<
                     ValkeyClientNodeDescription(address: primary)
                 ]
                 let newNodes = ValkeyNodeIDs(primary: primary, replicas: [])
-                self.state = .running(newNodes)
+                self.state = .healthy(.init(nodes: newNodes))
                 let action = self.runningClients.updateNodes(nodeDescriptions, removeUnmentionedPools: false)
                 return .init(
                     nextAction: .refreshTopology,
@@ -271,7 +309,7 @@ struct ValkeyClientStateMachine<
 
     mutating func topologyRefreshFailed() -> TopologyRefreshFailedAction {
         // If we are shutdown ignore
-        guard case .running = self.state else {
+        guard case .healthy = self.state else {
             return .init(retryTimer: nil)
         }
         switch self.topologyRefreshState {

Tests/IntegrationTests/StandaloneReplicaIntegrationTests.swift

@@ -321,8 +321,8 @@ struct StandaloneReplicaIntegrationTests {
     func withFailover<Value>(_ client: ValkeyClient, operation: () async throws -> Value) async throws -> Value {
         // get primary address
         let (host, port) = client.stateMachine.withLock {
-            guard case .running(let nodes) = $0.state else { fatalError("Expected a running primary node") }
-            guard case .hostname(let host, let port) = nodes.primary.value else { fatalError("Expected a hostname") }
+            guard case .healthy(let healthContext) = $0.state else { fatalError("Expected a running primary node") }
+            guard case .hostname(let host, let port) = healthContext.nodes.primary.value else { fatalError("Expected a hostname") }
             return (host, port)
         }
         // extract replica address