Fix missing host header, implement runtime's run function, fix defaultEventLoop

Author: aryan-25

Sources/AWSLambdaRuntimeCore/NewLambda.swift

@@ -47,9 +47,5 @@ extension Lambda {
     }
 
     /// The default EventLoop the Lambda is scheduled on.
-    package static var defaultEventLoop: any EventLoop {
-        get {
-            NIOSingletons.posixEventLoopGroup.next()
-        }
-    }
+    package static var defaultEventLoop: any EventLoop = NIOSingletons.posixEventLoopGroup.next()
 }

Sources/AWSLambdaRuntimeCore/NewLambdaRuntime.swift

@@ -19,7 +19,7 @@ import NIOCore
 import Synchronization
 
 package final class NewLambdaRuntime<Handler>: Sendable where Handler: StreamingLambdaHandler {
-    let handlerMutex: Mutex<Handler>
+    let handlerMutex: Mutex<Handler?>
     let logger: Logger
     let eventLoop: EventLoop
 
@@ -34,5 +34,35 @@ package final class NewLambdaRuntime<Handler>: Sendable where Handler: Streaming
     }
 
     package func run() async throws {
+        guard let runtimeEndpoint = Lambda.env("AWS_LAMBDA_RUNTIME_API") else {
+            throw NewLambdaRuntimeError(code: .cannotStartLambdaRuntime)
+        }
+
+        let ipAndPort = runtimeEndpoint.split(separator: ":", maxSplits: 1)
+        let ip = String(ipAndPort[0])
+        let port = Int(ipAndPort[1])!
+
+        let handler = self.handlerMutex.withLock { maybeHandler in
+            defer {
+                maybeHandler = nil
+            }
+            return maybeHandler
+        }
+
+        guard let handler else {
+            throw NewLambdaRuntimeError(code: .runtimeCanOnlyBeStartedOnce)
+        }
+
+        try await NewLambdaRuntimeClient.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
+            )
+        }
     }
 }

Sources/AWSLambdaRuntimeCore/NewLambdaRuntimeClient.swift

@@ -303,7 +303,7 @@ final actor NewLambdaRuntimeClient: LambdaRuntimeClientProtocol {
                         NIOHTTPClientResponseAggregator(maxContentLength: 6 * 1024 * 1024)
                     )
                     try channel.pipeline.syncOperations.addHandler(
-                        LambdaChannelHandler(delegate: self, logger: self.logger)
+                        LambdaChannelHandler(delegate: self, logger: self.logger, configuration: self.configuration)
                     )
                     return channel.eventLoop.makeSucceededFuture(())
                 } catch {
@@ -433,10 +433,32 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
     private var reusableErrorBuffer: ByteBuffer?
     private let logger: Logger
     private let delegate: Delegate
+    private let configuration: NewLambdaRuntimeClient.Configuration
 
-    init(delegate: Delegate, logger: Logger) {
+    let defaultHeaders: HTTPHeaders
+    /// These headers must be sent along an invocation or initialization error report
+    let errorHeaders: HTTPHeaders
+    /// These headers must be sent along an invocation or initialization error report
+    let streamingHeaders: HTTPHeaders
+
+    init(delegate: Delegate, logger: Logger, configuration: NewLambdaRuntimeClient.Configuration) {
         self.delegate = delegate
         self.logger = logger
+        self.configuration = configuration
+        self.defaultHeaders = [
+            "host": "\(self.configuration.ip):\(self.configuration.port)",
+            "user-agent": "Swift-Lambda/Unknown",
+        ]
+        self.errorHeaders = [
+            "host": "\(self.configuration.ip):\(self.configuration.port)",
+            "user-agent": "Swift-Lambda/Unknown",
+            "lambda-runtime-function-error-type": "Unhandled",
+        ]
+        self.streamingHeaders = [
+            "host": "\(self.configuration.ip):\(self.configuration.port)",
+            "user-agent": "Swift-Lambda/Unknown",
+            "transfer-encoding": "streaming",
+        ]
     }
 
     func nextInvocation(isolation: isolated (any Actor)? = #isolation) async throws -> Invocation {
@@ -578,7 +600,7 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
                 version: .http1_1,
                 method: .POST,
                 uri: url,
-                headers: NewLambdaRuntimeClient.streamingHeaders
+                headers: self.streamingHeaders
             )
 
             context.write(self.wrapOutboundOut(.head(httpRequest)), promise: nil)
@@ -604,11 +626,12 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
             let headers: HTTPHeaders =
                 if byteBuffer?.readableBytes ?? 0 < 6_000_000 {
                     [
+                        "host": "\(self.configuration.ip):\(self.configuration.port)",
                         "user-agent": "Swift-Lambda/Unknown",
                         "content-length": "\(byteBuffer?.readableBytes ?? 0)",
                     ]
                 } else {
-                    NewLambdaRuntimeClient.streamingHeaders
+                    self.streamingHeaders
                 }
 
             let httpRequest = HTTPRequestHead(
@@ -634,7 +657,7 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
             version: .http1_1,
             method: .GET,
             uri: self.nextInvocationPath,
-            headers: NewLambdaRuntimeClient.defaultHeaders
+            headers: self.defaultHeaders
         )
 
         context.write(self.wrapOutboundOut(.head(httpRequest)), promise: nil)
@@ -650,7 +673,7 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
             version: .http1_1,
             method: .POST,
             uri: url,
-            headers: NewLambdaRuntimeClient.errorHeaders
+            headers: self.errorHeaders
         )
 
         if self.reusableErrorBuffer == nil {
@@ -797,22 +820,3 @@ extension LambdaChannelHandler: ChannelInboundHandler {
         context.fireChannelInactive()
     }
 }
-
-extension NewLambdaRuntimeClient {
-    static let defaultHeaders: HTTPHeaders = [
-        "user-agent": "Swift-Lambda/Unknown"
-    ]
-
-    /// These headers must be sent along an invocation or initialization error report
-    static let errorHeaders: HTTPHeaders = [
-        "user-agent": "Swift-Lambda/Unknown",
-        "lambda-runtime-function-error-type": "Unhandled",
-    ]
-
-    /// These headers must be sent along an invocation or initialization error report
-    static let streamingHeaders: HTTPHeaders = [
-        "user-agent": "Swift-Lambda/Unknown",
-        "transfer-encoding": "streaming",
-    ]
-
-}

Sources/AWSLambdaRuntimeCore/NewLambdaRuntimeError.swift

@@ -29,6 +29,8 @@ package struct NewLambdaRuntimeError: Error {
         case nextInvocationMissingHeaderDeadline
         case nextInvocationMissingHeaderInvokeFuctionARN
 
+        case cannotStartLambdaRuntime
+        case runtimeCanOnlyBeStartedOnce
     }
 
     package init(code: Code, underlying: (any Error)? = nil) {