propagate the connection closed info through a Future

Author: sebsto

Sources/AWSLambdaRuntime/Lambda.swift

@@ -41,6 +41,17 @@ 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 finish the run loop.
+                    let _ = try await futureConnectionClosed.get()
+                }
+
                 let (invocation, writer) = try await runtimeClient.nextInvocation()
                 logger[metadataKey: "aws-request-id"] = "\(invocation.metadata.requestID)"
 
@@ -84,6 +95,7 @@ public enum Lambda {
         } catch is CancellationError {
             // don't allow cancellation error to propagate further
         }
+
     }
 
     /// The default EventLoop the Lambda is scheduled on.

Sources/AWSLambdaRuntime/LambdaRuntime+ServiceLifecycle.swift

@@ -17,8 +17,14 @@ import ServiceLifecycle
 
 extension LambdaRuntime: Service {
     public func run() async throws {
-        try await cancelWhenGracefulShutdown {
-            try await self._run()
+        await cancelWhenGracefulShutdown {
+            do {
+                try await self._run()
+            } catch {
+                // catch top level error that have not been handled before
+                // this avoids the runtime to crash and generate a backtrace
+                self.logger.error("LambdaRuntime.run() failed with error", metadata: ["error": "\(error)"])
+            }
         }
     }
 }

Sources/AWSLambdaRuntime/LambdaRuntime.swift

@@ -59,7 +59,13 @@ 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 error that have not been handled before
+            // this avoids the runtime to crash and generate a backtrace
+            self.logger.error("LambdaRuntime.run() failed with error", metadata: ["error": "\(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
@@ -372,8 +375,10 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
                     // 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 API lost, exiting")
-                    exit(-1)
+                    runtimeClient.logger.trace("Connection to Lambda API. lost, exiting")
+                    runtimeClient.futureConnectionClosed = runtimeClient.eventLoop.makeFailedFuture(
+                        LambdaRuntimeError(code: .connectionToControlPlaneLost)
+                    )
                 }
             }
 
@@ -392,6 +397,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
                 return handler
             }
         } catch {
+
             switch self.connectionState {
             case .disconnected, .connected:
                 fatalError("Unexpected state: \(self.connectionState)")