Propagate Connection Closed Information up to top-level (fix #465)

Sources/AWSLambdaRuntime/Lambda+LocalServer.swift

@@ -452,7 +452,7 @@ internal struct LambdaHTTPServer {
             await self.responsePool.push(
                 LocalServerResponse(
                     id: requestId,
-                    status: .ok,
+                    status: .accepted,
                     // the local server has no mecanism to collect headers set by the lambda function
                     headers: HTTPHeaders(),
                     body: body,

Sources/AWSLambdaRuntime/Lambda.swift

@@ -41,6 +41,8 @@ public enum Lambda {
         var logger = logger
         do {
             while !Task.isCancelled {
+
+                logger.trace("Waiting for next invocation")
                 let (invocation, writer) = try await runtimeClient.nextInvocation()
                 logger[metadataKey: "aws-request-id"] = "\(invocation.metadata.requestID)"
 
@@ -76,14 +78,18 @@ public enum Lambda {
                             logger: logger
                         )
                     )
+                    logger.trace("Handler finished processing invocation")
                 } catch {
+                    logger.trace("Handler failed processing invocation", metadata: ["Handler error": "\(error)"])
                     try await writer.reportError(error)
                     continue
                 }
+                logger.handler.metadata.removeValue(forKey: "aws-request-id")
             }
         } catch is CancellationError {
             // don't allow cancellation error to propagate further
         }
+
     }
 
     /// The default EventLoop the Lambda is scheduled on.

Sources/AWSLambdaRuntime/LambdaRuntime.swift

@@ -94,16 +94,29 @@ public final class LambdaRuntime<Handler>: Sendable where Handler: StreamingLamb
             let ip = String(ipAndPort[0])
             guard let port = Int(ipAndPort[1]) else { throw LambdaRuntimeError(code: .invalidPort) }
 
-            try await LambdaRuntimeClient.withRuntimeClient(
-                configuration: .init(ip: ip, port: port),
-                eventLoop: self.eventLoop,
-                logger: self.logger
-            ) { runtimeClient in
-                try await Lambda.runLoop(
-                    runtimeClient: runtimeClient,
-                    handler: handler,
+            do {
+                try await LambdaRuntimeClient.withRuntimeClient(
+                    configuration: .init(ip: ip, port: port),
+                    eventLoop: self.eventLoop,
                     logger: self.logger
-                )
+                ) { runtimeClient in
+                    try await Lambda.runLoop(
+                        runtimeClient: runtimeClient,
+                        handler: handler,
+                        logger: self.logger
+                    )
+                }
+            } catch {
+                // catch top level errors that have not been handled until now
+                // this avoids the runtime to crash and generate a backtrace
+                self.logger.error("LambdaRuntime.run() failed with error", metadata: ["error": "\(error)"])
+                if let error = error as? LambdaRuntimeError,
+                    error.code != .connectionToControlPlaneLost
+                {
+                    // if the error is a LambdaRuntimeError but not a connection error,
+                    // we rethrow it to preserve existing behaviour
+                    throw error
+                }
             }
 
         } else {

Sources/AWSLambdaRuntime/LambdaRuntimeClient.swift

@@ -67,10 +67,26 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
         NIOLoopBound<LambdaChannelHandler<LambdaRuntimeClient>>, any Error
     >
 
-    private enum ConnectionState {
+    private enum ConnectionState: Equatable {
         case disconnected
+        case lostConnection
         case connecting([ConnectionContinuation])
         case connected(Channel, LambdaChannelHandler<LambdaRuntimeClient>)
+
+        static func == (lhs: ConnectionState, rhs: ConnectionState) -> Bool {
+            switch (lhs, rhs) {
+            case (.disconnected, .disconnected):
+                return true
+            case (.connecting, .connecting):
+                return true
+            case (.connected, .connected):
+                return true
+            case (.lostConnection, .lostConnection):
+                return true
+            default:
+                return false
+            }
+        }
     }
 
     enum LambdaState {
@@ -97,6 +113,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
     private let configuration: Configuration
 
     private var connectionState: ConnectionState = .disconnected
+
     private var lambdaState: LambdaState = .idle(previousRequestID: nil)
     private var closingState: ClosingState = .notClosing
 
@@ -118,10 +135,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
         } catch {
             result = .failure(error)
         }
-
         await runtime.close()
-
-        //try? await runtime.close()
         return try result.get()
     }
 
@@ -157,18 +171,27 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
 
             case .connected(let channel, _):
                 channel.close(mode: .all, promise: nil)
+
+            case .lostConnection:
+                continuation.resume()
             }
         }
     }
 
     @usableFromInline
     func nextInvocation() async throws -> (Invocation, Writer) {
-        try await withTaskCancellationHandler {
+
+        if self.connectionState == .lostConnection {
+            throw LambdaRuntimeError(code: .connectionToControlPlaneLost)
+        }
+
+        return try await withTaskCancellationHandler {
             switch self.lambdaState {
             case .idle:
                 self.lambdaState = .waitingForNextInvocation
                 let handler = try await self.makeOrGetConnection()
                 let invocation = try await handler.nextInvocation()
+
                 guard case .waitingForNextInvocation = self.lambdaState else {
                     fatalError("Invalid state: \(self.lambdaState)")
                 }
@@ -262,7 +285,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
 
     private func channelClosed(_ channel: any Channel) {
         switch (self.connectionState, self.closingState) {
-        case (_, .closed):
+        case (_, .closed), (.lostConnection, _):
             fatalError("Invalid state: \(self.connectionState), \(self.closingState)")
 
         case (.disconnected, .notClosing):
@@ -283,7 +306,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
         case (.connecting(let array), .notClosing):
             self.connectionState = .disconnected
             for continuation in array {
-                continuation.resume(throwing: LambdaRuntimeError(code: .lostConnectionToControlPlane))
+                continuation.resume(throwing: LambdaRuntimeError(code: .connectionToControlPlaneLost))
             }
 
         case (.connecting(let array), .closing(let continuation)):
@@ -322,6 +345,9 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
             return loopBound.value
         case .connected(_, let handler):
             return handler
+
+        case .lostConnection:
+            throw LambdaRuntimeError(code: .connectionToControlPlaneLost)
         }
 
         let bootstrap = ClientBootstrap(group: self.eventLoop)
@@ -363,7 +389,19 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
             )
             channel.closeFuture.whenComplete { result in
                 self.assumeIsolated { runtimeClient in
+
+                    // resume any pending continuation on the handler
+                    if case .connected(_, let handler) = runtimeClient.connectionState {
+                        if case .connected(_, let lambdaState) = handler.state {
+                            if case .waitingForNextInvocation(let continuation) = lambdaState {
+                                continuation.resume(throwing: LambdaRuntimeError(code: .connectionToControlPlaneLost))
+                            }
+                        }
+                    }
+
+                    // close the channel
                     runtimeClient.channelClosed(channel)
+                    runtimeClient.connectionState = .lostConnection
                 }
             }
 
@@ -380,8 +418,12 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
                     }
                 }
                 return handler
+            case .lostConnection:
+                // this should never happen
+                fatalError("Lost connection to Lambda service")
             }
         } catch {
+
             switch self.connectionState {
             case .disconnected, .connected:
                 fatalError("Unexpected state: \(self.connectionState)")
@@ -394,6 +436,9 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
                     }
                 }
                 throw error
+            case .lostConnection:
+                // this should never happen
+                fatalError("Lost connection to Lambda service")
             }
         }
     }
@@ -431,6 +476,9 @@ extension LambdaRuntimeClient: LambdaChannelHandlerDelegate {
 
                 isolated.connectionState = .disconnected
 
+            case .lostConnection:
+                // this should never happen
+                fatalError("Lost connection to Lambda service")
             }
         }
     }
@@ -463,7 +511,7 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
         }
     }
 
-    private var state: State = .disconnected
+    var state: State = .disconnected
     private var lastError: Error?
     private var reusableErrorBuffer: ByteBuffer?
     private let logger: Logger
@@ -885,6 +933,7 @@ extension LambdaChannelHandler: ChannelInboundHandler {
         // fail any pending responses with last error or assume peer disconnected
         switch self.state {
         case .connected(_, .waitingForNextInvocation(let continuation)):
+            self.state = .disconnected
             continuation.resume(throwing: self.lastError ?? ChannelError.ioOnClosedChannel)
         default:
             break

Sources/AWSLambdaRuntime/LambdaRuntimeError.swift

@@ -25,7 +25,6 @@ package struct LambdaRuntimeError: Error {
 
         case writeAfterFinishHasBeenSent
         case finishAfterFinishHasBeenSent
-        case lostConnectionToControlPlane
         case unexpectedStatusCodeForRequest
 
         case nextInvocationMissingHeaderRequestID

Sources/MockServer/MockHTTPServer.swift

@@ -216,7 +216,7 @@ struct HttpServer {
         } else if requestHead.uri.hasSuffix("/response") {
             responseStatus = .accepted
         } else if requestHead.uri.hasSuffix("/error") {
-            responseStatus = .ok
+            responseStatus = .accepted
         } else {
             responseStatus = .notFound
         }