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,18 @@ public enum Lambda {
         var logger = logger
         do {
             while !Task.isCancelled {
+
+                if let runtimeClient = runtimeClient as? LambdaRuntimeClient,
+                    let futureConnectionClosed = await runtimeClient.futureConnectionClosed
+                {
+                    // Wait for the futureConnectionClosed to complete,
+                    // which will happen when the Lambda HTTP Server (or MockServer) closes the connection
+                    // This allows us to exit the run loop gracefully.
+                    // The futureConnectionClosed is always an error, let it throw to terminate the Lambda run loop.
+                    let _ = try await futureConnectionClosed.get()
+                }
+
+                logger.trace("Waiting for next invocation")
                 let (invocation, writer) = try await runtimeClient.nextInvocation()
                 logger[metadataKey: "aws-request-id"] = "\(invocation.metadata.requestID)"
 
@@ -76,14 +88,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+ServiceLifecycle.swift

@@ -18,7 +18,20 @@ import ServiceLifecycle
 extension LambdaRuntime: Service {
     public func run() async throws {
         try await cancelWhenGracefulShutdown {
-            try await self._run()
+            do {
+                try await self._run()
+            } 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
+                }
+            }
         }
     }
 }

Sources/AWSLambdaRuntime/LambdaRuntime.swift

@@ -59,7 +59,20 @@ public final class LambdaRuntime<Handler>: Sendable where Handler: StreamingLamb
     #if !ServiceLifecycleSupport
     @inlinable
     internal func run() async throws {
-        try await _run()
+        do {
+            try await _run()
+        } 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
+            }
+        }
     }
     #endif
 

Sources/AWSLambdaRuntime/LambdaRuntimeClient.swift

@@ -92,6 +92,9 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
         case closed
     }
 
+    @usableFromInline
+    var futureConnectionClosed: EventLoopFuture<LambdaRuntimeError>? = nil
+
     private let eventLoop: any EventLoop
     private let logger: Logger
     private let configuration: Configuration
@@ -118,10 +121,8 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
         } catch {
             result = .failure(error)
         }
-
         await runtime.close()
 
-        //try? await runtime.close()
         return try result.get()
     }
 
@@ -330,6 +331,8 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
                     try channel.pipeline.syncOperations.addHTTPClientHandlers()
                     // Lambda quotas... An invocation payload is maximal 6MB in size:
                     //   https://docs.aws.amazon.com/lambda/latest/dg/gettingstarted-limits.html
+                    // TODO: should we enforce this here ?  What about streaming functions that
+                    //       support up to 20Mb responses ?
                     try channel.pipeline.syncOperations.addHandler(
                         NIOHTTPClientResponseAggregator(maxContentLength: 6 * 1024 * 1024)
                     )
@@ -364,6 +367,16 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
             channel.closeFuture.whenComplete { result in
                 self.assumeIsolated { runtimeClient in
                     runtimeClient.channelClosed(channel)
+
+                    // at this stage, we lost the connection to the Lambda Service,
+                    // this is very unlikely to happen when running in a lambda function deployed in the cloud
+                    // however, this happens when performance testing against the MockServer
+                    // shutdown this runtime.
+                    // The Lambda service will create a new runtime environment anyway
+                    runtimeClient.logger.trace("Connection to Lambda Service HTTP Server lost, exiting")
+                    runtimeClient.futureConnectionClosed = runtimeClient.eventLoop.makeFailedFuture(
+                        LambdaRuntimeError(code: .connectionToControlPlaneLost)
+                    )
                 }
             }
 
@@ -382,6 +395,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
                 return handler
             }
         } catch {
+
             switch self.connectionState {
             case .disconnected, .connected:
                 fatalError("Unexpected state: \(self.connectionState)")

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
         }