fix connection closure

Author: fabianfett

Sources/AWSLambdaRuntimeCore/NewLambdaRuntimeClient.swift

@@ -114,6 +114,8 @@ final actor NewLambdaRuntimeClient: LambdaRuntimeClientProtocol {
     }
 
     private func close() async {
+        self.logger.trace("Close lambda runtime client")
+
         guard case .notClosing = self.closingState else {
             return
         }
@@ -229,10 +231,24 @@ final actor NewLambdaRuntimeClient: LambdaRuntimeClientProtocol {
 
     private func channelClosed(_ channel: any Channel) {
         switch (self.connectionState, self.closingState) {
-        case (.disconnected, _),
-            (_, .closed):
+        case (_, .closed):
             fatalError("Invalid state: \(self.connectionState), \(self.closingState)")
 
+        case (.disconnected, .notClosing):
+            if let index = self.closingConnections.firstIndex(where: { $0 === channel }) {
+                self.closingConnections.remove(at: index)
+            }
+
+        case (.disconnected, .closing(let continuation)):
+            if let index = self.closingConnections.firstIndex(where: { $0 === channel }) {
+                self.closingConnections.remove(at: index)
+            }
+
+            if self.closingConnections.isEmpty {
+                self.closingState = .closed
+                continuation.resume()
+            }
+
         case (.connecting(let array), .notClosing):
             self.connectionState = .disconnected
             for continuation in array {
@@ -303,8 +319,14 @@ final actor NewLambdaRuntimeClient: LambdaRuntimeClientProtocol {
             let handler = try channel.pipeline.syncOperations.handler(
                 type: LambdaChannelHandler<NewLambdaRuntimeClient>.self
             )
+            self.logger.trace(
+                "Connection to control plane created",
+                metadata: [
+                    "lambda_port": "\(self.configuration.port)",
+                    "lambda_ip": "\(self.configuration.ip)",
+                ]
+            )
             channel.closeFuture.whenComplete { result in
-                self.eventLoop.preconditionInEventLoop()
                 self.assumeIsolated { runtimeClient in
                     runtimeClient.channelClosed(channel)
                 }
@@ -754,6 +776,12 @@ extension LambdaChannelHandler: ChannelInboundHandler {
     }
 
     func errorCaught(context: ChannelHandlerContext, error: Error) {
+        self.logger.trace(
+            "Channel error caught",
+            metadata: [
+                "error": "\(error)"
+            ]
+        )
         // pending responses will fail with lastError in channelInactive since we are calling context.close
         self.delegate.connectionErrorHappened(error, channel: context.channel)
 

Tests/AWSLambdaRuntimeCoreTests/NewLambdaRuntimeClientTests.swift

@@ -24,11 +24,11 @@ import struct Foundation.UUID
 @Suite
 struct NewLambdaRuntimeClientTests {
 
-    let logger = Logger(label: "NewLambdaClientRuntimeTest")
-
-    init() {
-
-    }
+    let logger = {
+        var logger = Logger(label: "NewLambdaClientRuntimeTest")
+        logger.logLevel = .trace
+        return logger
+    }()
 
     @Test
     func testSimpleInvocations() async throws {
@@ -57,12 +57,12 @@ struct NewLambdaRuntimeClientTests {
             }
         }
 
-        try await self.withMockServer(behaviour: HappyBehavior()) { mockServer, eventLoopGroup in
-            let configuration = NewLambdaRuntimeClient.Configuration(ip: "127.0.0.1", port: 7000)
+        try await withMockServer(behaviour: HappyBehavior()) { port in
+            let configuration = NewLambdaRuntimeClient.Configuration(ip: "127.0.0.1", port: port)
 
             try await NewLambdaRuntimeClient.withRuntimeClient(
                 configuration: configuration,
-                eventLoop: eventLoopGroup.next(),
+                eventLoop: NIOSingletons.posixEventLoopGroup.next(),
                 logger: self.logger
             ) { runtimeClient in
                 do {
@@ -86,26 +86,4 @@ struct NewLambdaRuntimeClientTests {
             }
         }
     }
-
-    func withMockServer<Result>(
-        behaviour: some LambdaServerBehavior,
-        _ body: (MockLambdaServer, MultiThreadedEventLoopGroup) async throws -> Result
-    ) async throws -> Result {
-        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
-        let server = MockLambdaServer(behavior: behaviour)
-        _ = try await server.start().get()
-
-        let result: Swift.Result<Result, any Error>
-        do {
-            result = .success(try await body(server, eventLoopGroup))
-        } catch {
-            result = .failure(error)
-        }
-
-        try? await server.stop().get()
-        try? await eventLoopGroup.shutdownGracefully()
-
-        return try result.get()
-    }
-
 }