Add support for sending additional command in handshake

- Add ignore case to ValkeyPromise for commands whose results we don't care about
- Add deque of pending commands to ConnectedState
- Push remaining pending commands to active state once we receive the hello

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

Author: adam-fowler

Sources/Valkey/Connection/ValkeyChannelHandler+stateMachine.swift

@@ -65,11 +65,11 @@ extension ValkeyChannelHandler {
         @usableFromInline
         struct ConnectedState {
             let context: Context
-            var pendingHelloCommand: PendingCommand
+            var pendingCommands: Deque<PendingCommand>
 
-            func cancel(requestID: Int) -> PendingCommand? {
-                if pendingHelloCommand.requestID == requestID {
-                    return pendingHelloCommand
+            func cancel(requestID: Int) -> Deque<PendingCommand>? {
+                if pendingCommands.first?.requestID == requestID {
+                    return pendingCommands
                 }
                 return nil
             }
@@ -85,11 +85,11 @@ extension ValkeyChannelHandler {
 
         /// handler has become active
         @usableFromInline
-        mutating func setConnected(context: Context, pendingHelloCommand: PendingCommand) {
+        mutating func setConnected(context: Context, pendingCommands: Deque<PendingCommand>) {
             switch consume self.state {
             case .initialized:
                 self = .connected(
-                    .init(context: context, pendingHelloCommand: pendingHelloCommand)
+                    .init(context: context, pendingCommands: pendingCommands)
                 )
             case .connected:
                 preconditionFailure("Cannot set connected state when state is connected")
@@ -155,15 +155,18 @@ extension ValkeyChannelHandler {
             switch consume self.state {
             case .initialized:
                 preconditionFailure("Cannot send command when initialized")
-            case .connected(let state):
+            case .connected(var state):
+                guard let helloCommand = state.pendingCommands.popFirst() else {
+                    preconditionFailure("Cannot be in connected state with no pending commands")
+                }
                 switch token.value {
                 case .bulkError(let message), .simpleError(let message):
                     let error = ValkeyClientError(.commandError, message: String(buffer: message))
                     self = .closed(error)
-                    return .respondAndClose(state.pendingHelloCommand, error)
+                    return .respondAndClose(helloCommand, error)
                 default:
-                    self = .active(.init(context: state.context, pendingCommands: .init()))
-                    return .respond(state.pendingHelloCommand, .cancel)
+                    self = .active(.init(context: state.context, pendingCommands: state.pendingCommands))
+                    return .respond(helloCommand, .cancel)
                 }
             case .active(var state):
                 guard let command = state.pendingCommands.popFirst() else {
@@ -204,8 +207,12 @@ extension ValkeyChannelHandler {
                     self = .closed(nil)
                     return .respondAndClose(command, nil)
                 }
-            case .closed:
-                preconditionFailure("Cannot receive command on closed connection")
+            case .closed(let error):
+                guard let error else {
+                    preconditionFailure("Cannot receive command on closed connection with no error")
+                }
+                self = .closed(error)
+                return .closeWithError(error)
             }
         }
 
@@ -221,19 +228,22 @@ extension ValkeyChannelHandler {
             case .initialized:
                 preconditionFailure("Cannot wait until connection has succeeded")
             case .connected(let state):
-                switch state.pendingHelloCommand.promise {
+                guard let helloCommand = state.pendingCommands.first else {
+                    preconditionFailure("Cannot be in connected state with no pending commands")
+                }
+                switch helloCommand.promise {
                 case .nio(let promise):
                     self = .connected(state)
                     return .waitForPromise(promise)
-                case .swift:
+                case .swift, .ignore:
                     preconditionFailure("Connected state cannot be setup with a Swift continuation")
                 }
             case .active(let state):
                 self = .active(state)
                 return .done
             case .closing(let state):
                 self = .closing(state)
-                return .done
+                return .reportedClosed(nil)
             case .closed(let error):
                 self = .closed(error)
                 return .reportedClosed(error)
@@ -253,12 +263,15 @@ extension ValkeyChannelHandler {
             case .initialized:
                 preconditionFailure("Cannot cancel when initialized")
             case .connected(let state):
-                if state.pendingHelloCommand.deadline <= now {
+                guard let helloCommand = state.pendingCommands.first else {
+                    preconditionFailure("Cannot be in connected state with no pending commands")
+                }
+                if helloCommand.deadline <= now {
                     self = .closed(ValkeyClientError(.timeout))
-                    return .failPendingCommandsAndClose(state.context, [state.pendingHelloCommand])
+                    return .failPendingCommandsAndClose(state.context, state.pendingCommands)
                 } else {
                     self = .connected(state)
-                    return .reschedule(state.pendingHelloCommand.deadline)
+                    return .reschedule(helloCommand.deadline)
                 }
             case .active(let state):
                 if let firstCommand = state.pendingCommands.first {
@@ -303,11 +316,11 @@ extension ValkeyChannelHandler {
             case .initialized:
                 preconditionFailure("Cannot cancel when initialized")
             case .connected(let state):
-                if let command = state.cancel(requestID: requestID) {
+                if let commands = state.cancel(requestID: requestID) {
                     self = .closed(CancellationError())
                     return .failPendingCommandsAndClose(
                         state.context,
-                        cancel: [command],
+                        cancel: .init(commands),
                         closeConnectionDueToCancel: []
                     )
                 } else {
@@ -360,7 +373,7 @@ extension ValkeyChannelHandler {
                 self = .closed(nil)
                 return .doNothing
             case .connected(let state):
-                self = .closing(.init(context: state.context, pendingCommands: [state.pendingHelloCommand]))
+                self = .closing(.init(context: state.context, pendingCommands: state.pendingCommands))
                 return .waitForPendingCommands(state.context)
             case .active(let state):
                 if state.pendingCommands.count > 0 {
@@ -393,7 +406,7 @@ extension ValkeyChannelHandler {
                 return .doNothing
             case .connected(let state):
                 self = .closed(nil)
-                return .failPendingCommandsAndClose(state.context, [state.pendingHelloCommand])
+                return .failPendingCommandsAndClose(state.context, state.pendingCommands)
             case .active(let state):
                 self = .closed(nil)
                 return .failPendingCommandsAndClose(state.context, state.pendingCommands)
@@ -421,7 +434,7 @@ extension ValkeyChannelHandler {
                 return .doNothing
             case .connected(let state):
                 self = .closed(nil)
-                return .failPendingCommandsAndSubscriptions([state.pendingHelloCommand])
+                return .failPendingCommandsAndSubscriptions(state.pendingCommands)
             case .active(let state):
                 self = .closed(nil)
                 return .failPendingCommandsAndSubscriptions(state.pendingCommands)

Sources/Valkey/Connection/ValkeyChannelHandler.swift

@@ -20,13 +20,16 @@ import NIOCore
 enum ValkeyPromise<T: Sendable>: Sendable {
     case nio(EventLoopPromise<T>)
     case swift(CheckedContinuation<T, any Error>)
+    case ignore
 
     func succeed(_ t: T) {
         switch self {
         case .nio(let eventLoopPromise):
             eventLoopPromise.succeed(t)
         case .swift(let checkedContinuation):
             checkedContinuation.resume(returning: t)
+        case .ignore:
+            break
         }
     }
 
@@ -36,6 +39,8 @@ enum ValkeyPromise<T: Sendable>: Sendable {
             eventLoopPromise.fail(e)
         case .swift(let checkedContinuation):
             checkedContinuation.resume(throwing: e)
+        case .ignore:
+            break
         }
     }
 }
@@ -275,25 +280,35 @@ final class ValkeyChannelHandler: ChannelInboundHandler {
     @usableFromInline
     func setConnected(context: ChannelHandlerContext) {
         // Send initial HELLO command
-        let command = HELLO(
+        let helloCommand = HELLO(
             arguments: .init(
                 protover: 3,
                 auth: configuration.authentication.map { .init(username: $0.username, password: $0.password) },
                 clientname: configuration.clientName
             )
         )
+        // set client info
+        let clientInfoLibName = CLIENT.SETINFO(attr: .libname("valkey-swift"))
+        let clientInfoLibVersion = CLIENT.SETINFO(attr: .libver("0.1.0"))
+
         self.encoder.reset()
-        command.encode(into: &self.encoder)
-        let buffer = self.encoder.buffer
+        helloCommand.encode(into: &self.encoder)
+        clientInfoLibName.encode(into: &self.encoder)
+        clientInfoLibVersion.encode(into: &self.encoder)
 
         let promise = eventLoop.makePromise(of: RESPToken.self)
+
         let deadline = .now() + self.configuration.commandTimeout
-        context.writeAndFlush(self.wrapOutboundOut(buffer), promise: nil)
+        context.writeAndFlush(self.wrapOutboundOut(self.encoder.buffer), promise: nil)
         scheduleDeadlineCallback(deadline: deadline)
 
         self.stateMachine.setConnected(
             context: context,
-            pendingHelloCommand: .init(promise: .nio(promise), requestID: 0, deadline: deadline)
+            pendingCommands: [
+                .init(promise: .nio(promise), requestID: 0, deadline: deadline),  // HELLO
+                .init(promise: .ignore, requestID: 0, deadline: deadline),  // CLIENT.SETINFO
+                .init(promise: .ignore, requestID: 0, deadline: deadline),  // CLIENT.SETINFO
+            ]
         )
     }
 

Tests/ValkeyTests/Utils/NIOAsyncTestingChannel+hello.swift

@@ -21,7 +21,15 @@ import Testing
 extension NIOAsyncTestingChannel {
     func processHello() async throws {
         let hello = try await self.waitForOutboundWrite(as: ByteBuffer.self)
-        #expect(hello == RESPToken(.array([.bulkString("HELLO"), .bulkString("3")])).base)
+        var expectedBuffer = ByteBuffer()
+        expectedBuffer.writeImmutableBuffer(RESPToken(.array([.bulkString("HELLO"), .bulkString("3")])).base)
+        expectedBuffer.writeImmutableBuffer(
+            RESPToken(.array([.bulkString("CLIENT"), .bulkString("SETINFO"), .bulkString("lib-name"), .bulkString("valkey-swift")])).base
+        )
+        expectedBuffer.writeImmutableBuffer(
+            RESPToken(.array([.bulkString("CLIENT"), .bulkString("SETINFO"), .bulkString("lib-ver"), .bulkString("0.1.0")])).base
+        )
+        #expect(hello == expectedBuffer)
         try await self.writeInbound(
             RESPToken(
                 .map([
@@ -35,5 +43,7 @@ extension NIOAsyncTestingChannel {
                 ])
             ).base
         )
+        try await self.writeInbound(RESPToken.ok.base)
+        try await self.writeInbound(RESPToken.ok.base)
     }
 }

Tests/ValkeyTests/ValkeyChannelHandlerStateMachineTests.swift

@@ -183,7 +183,12 @@ struct ValkeyChannelHandlerStateMachineTests {
         }
         expect(
             stateMachine.state
-                == .closing(.init(context: "testGracefulShutdown", pendingCommands: [.init(promise: .nio(promise), requestID: 23, deadline: .now())]))
+                == .closing(
+                    .init(
+                        context: "testGracefulShutdown",
+                        pendingCommands: [.init(promise: .nio(promise), requestID: 23, deadline: .now())]
+                    )
+                )
         )
         switch stateMachine.receivedResponse(token: .ok) {
         case .respondAndClose(let command, let error):
@@ -218,7 +223,10 @@ struct ValkeyChannelHandlerStateMachineTests {
         expect(
             stateMachine.state
                 == .closing(
-                    .init(context: "testClosedClosingState", pendingCommands: [.init(promise: .nio(promise), requestID: 17, deadline: .now())])
+                    .init(
+                        context: "testClosedClosingState",
+                        pendingCommands: [.init(promise: .nio(promise), requestID: 17, deadline: .now())]
+                    )
                 )
         )
         switch stateMachine.setClosed() {
@@ -460,7 +468,7 @@ extension ValkeyChannelHandler.StateMachine<String>.State {
         case .connected(let lhs):
             switch rhs {
             case .connected(let rhs):
-                return lhs.context == rhs.context && lhs.pendingHelloCommand.requestID == rhs.pendingHelloCommand.requestID
+                return lhs.context == rhs.context && lhs.pendingCommands.map { $0.requestID } == rhs.pendingCommands.map { $0.requestID }
             default:
                 return false
             }
@@ -535,7 +543,7 @@ extension ValkeyChannelHandler.StateMachine {
         let promise = EmbeddedEventLoop().makePromise(of: RESPToken.self)
         self.setConnected(
             context: context,
-            pendingHelloCommand: .init(promise: .nio(promise), requestID: 0, deadline: .now() + .seconds(30))
+            pendingCommands: [.init(promise: .nio(promise), requestID: 0, deadline: .now() + .seconds(30))]
         )
     }
 

Tests/ValkeyTests/ValkeyConnectionTests.swift

@@ -46,8 +46,9 @@ struct ConnectionTests {
         let logger = Logger(label: "test")
         _ = try await ValkeyConnection.setupChannelAndConnect(channel, configuration: .init(), logger: logger)
 
-        let outbound = try await channel.waitForOutboundWrite(as: ByteBuffer.self)
-        #expect(outbound == RESPToken(.command(["HELLO", "3"])).base)
+        var outbound = try await channel.waitForOutboundWrite(as: ByteBuffer.self)
+        let hello3 = RESPToken(.command(["HELLO", "3"])).base
+        #expect(outbound.readSlice(length: hello3.readableBytes) == hello3)
     }
 
     @Test
@@ -57,8 +58,9 @@ struct ConnectionTests {
         let logger = Logger(label: "test")
         _ = try await ValkeyConnection.setupChannelAndConnect(channel, configuration: .init(), logger: logger)
 
-        let outbound = try await channel.waitForOutboundWrite(as: ByteBuffer.self)
-        #expect(outbound == RESPToken(.command(["HELLO", "3"])).base)
+        var outbound = try await channel.waitForOutboundWrite(as: ByteBuffer.self)
+        let hello3 = RESPToken(.command(["HELLO", "3"])).base
+        #expect(outbound.readSlice(length: hello3.readableBytes) == hello3)
         await #expect(throws: ValkeyClientError(.commandError, message: "Not supported")) {
             try await channel.writeInbound(RESPToken(.bulkError("Not supported")).base)
         }
@@ -79,8 +81,9 @@ struct ConnectionTests {
             logger: logger
         )
 
-        let outbound = try await channel.waitForOutboundWrite(as: ByteBuffer.self)
-        #expect(outbound == RESPToken(.command(["HELLO", "3", "AUTH", "john", "smith"])).base)
+        var outbound = try await channel.waitForOutboundWrite(as: ByteBuffer.self)
+        let hello3 = RESPToken(.command(["HELLO", "3", "AUTH", "john", "smith"])).base
+        #expect(outbound.readSlice(length: hello3.readableBytes) == hello3)
     }
 
     @Test
@@ -95,8 +98,9 @@ struct ConnectionTests {
             logger: logger
         )
 
-        let outbound = try await channel.waitForOutboundWrite(as: ByteBuffer.self)
-        #expect(outbound == RESPToken(.command(["HELLO", "3", "SETNAME", "Testing"])).base)
+        var outbound = try await channel.waitForOutboundWrite(as: ByteBuffer.self)
+        let hello3 = RESPToken(.command(["HELLO", "3", "SETNAME", "Testing"])).base
+        #expect(outbound.readSlice(length: hello3.readableBytes) == hello3)
     }
 
     @Test