fix unit tests

Author: sebsto

Sources/AWSLambdaRuntime/Lambda.swift

@@ -52,6 +52,7 @@ public enum Lambda {
                     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)"
 
@@ -87,10 +88,13 @@ 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

Sources/AWSLambdaRuntime/LambdaRuntime+ServiceLifecycle.swift

@@ -17,13 +17,20 @@ import ServiceLifecycle
 
 extension LambdaRuntime: Service {
     public func run() async throws {
-        await cancelWhenGracefulShutdown {
+        try await cancelWhenGracefulShutdown {
             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

@@ -65,6 +65,13 @@ public final class LambdaRuntime<Handler>: Sendable where Handler: StreamingLamb
             // 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

@@ -121,10 +121,8 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
         } catch {
             result = .failure(error)
         }
-
         await runtime.close()
 
-        //try? await runtime.close()
         return try result.get()
     }
 
@@ -375,7 +373,7 @@ 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")
+                    runtimeClient.logger.trace("Connection to Lambda Service HTTP Server lost, exiting")
                     runtimeClient.futureConnectionClosed = runtimeClient.eventLoop.makeFailedFuture(
                         LambdaRuntimeError(code: .connectionToControlPlaneLost)
                     )