Added header when reporting failing invocations or initializations #116

Package.swift

@@ -35,6 +35,8 @@ let package = Package(
         ]),
         .testTarget(name: "AWSLambdaRuntimeCoreTests", dependencies: [
             .byName(name: "AWSLambdaRuntimeCore"),
+            .product(name: "NIOTestUtils", package: "swift-nio"),
+            .product(name: "NIOFoundationCompat", package: "swift-nio"),
         ]),
         .testTarget(name: "AWSLambdaRuntimeTests", dependencies: [
             .byName(name: "AWSLambdaRuntimeCore"),

Sources/AWSLambdaRuntimeCore/HTTPClient.swift

@@ -43,13 +43,16 @@ internal final class HTTPClient {
                              timeout: timeout ?? self.configuration.requestTimeout))
     }
 
-    func post(url: String, body: ByteBuffer?, timeout: TimeAmount? = nil) -> EventLoopFuture<Response> {
-        self.execute(Request(targetHost: self.targetHost,
-                             url: url,
-                             method: .POST,
-                             headers: HTTPClient.headers,
-                             body: body,
-                             timeout: timeout ?? self.configuration.requestTimeout))
+    func post(url: String, body: ByteBuffer?, timeout: TimeAmount? = nil, additionalHeaders: HTTPHeaders? = nil) -> EventLoopFuture<Response> {
+        var headers = HTTPClient.headers
+        additionalHeaders.flatMap { headers.add(contentsOf: $0) }
+
+        return self.execute(Request(targetHost: self.targetHost,
+                                    url: url,
+                                    method: .POST,
+                                    headers: headers,
+                                    body: body,
+                                    timeout: timeout ?? self.configuration.requestTimeout))
     }
 
     /// cancels the current request if there is one

Sources/AWSLambdaRuntimeCore/LambdaConfiguration.swift

@@ -68,7 +68,7 @@ extension Lambda {
             let requestTimeout: TimeAmount?
 
             init(baseURL: String? = nil, keepAlive: Bool? = nil, requestTimeout: TimeAmount? = nil) {
-                let ipPort = env("AWS_LAMBDA_RUNTIME_API")?.split(separator: ":") ?? ["127.0.0.1", "7000"]
+                let ipPort = (baseURL ?? env("AWS_LAMBDA_RUNTIME_API"))?.split(separator: ":") ?? ["127.0.0.1", "7000"]
                 guard ipPort.count == 2, let port = Int(ipPort[1]) else {
                     preconditionFailure("invalid ip+port configuration \(ipPort)")
                 }

Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift

@@ -27,6 +27,9 @@ extension Lambda {
         private let allocator = ByteBufferAllocator()
         private let httpClient: HTTPClient
 
+        /// Headers that must be sent along an invocation or initialization error report
+        internal static let errorHeaders = HTTPHeaders([("lambda-runtime-function-error-type", "Unhandled")])
+
         init(eventLoop: EventLoop, configuration: Configuration.RuntimeEngine) {
             self.eventLoop = eventLoop
             self.httpClient = HTTPClient(eventLoop: eventLoop, configuration: configuration)
@@ -61,6 +64,8 @@ extension Lambda {
         func reportResults(logger: Logger, invocation: Invocation, result: Result<ByteBuffer?, Error>) -> EventLoopFuture<Void> {
             var url = Consts.invocationURLPrefix + "/" + invocation.requestID
             var body: ByteBuffer?
+            var additionalHeaders: HTTPHeaders?
+
             switch result {
             case .success(let buffer):
                 url += Consts.postResponseURLSuffix
@@ -71,9 +76,10 @@ extension Lambda {
                 let bytes = errorResponse.toJSONBytes()
                 body = self.allocator.buffer(capacity: bytes.count)
                 body!.writeBytes(bytes)
+                additionalHeaders = RuntimeClient.errorHeaders
             }
             logger.debug("reporting results to lambda runtime engine using \(url)")
-            return self.httpClient.post(url: url, body: body).flatMapThrowing { response in
+            return self.httpClient.post(url: url, body: body, additionalHeaders: additionalHeaders).flatMapThrowing { response in
                 guard response.status == .accepted else {
                     throw RuntimeError.badStatusCode(response.status)
                 }
@@ -98,7 +104,7 @@ extension Lambda {
             var body = self.allocator.buffer(capacity: bytes.count)
             body.writeBytes(bytes)
             logger.warning("reporting initialization error to lambda runtime engine using \(url)")
-            return self.httpClient.post(url: url, body: body).flatMapThrowing { response in
+            return self.httpClient.post(url: url, body: body, additionalHeaders: RuntimeClient.errorHeaders).flatMapThrowing { response in
                 guard response.status == .accepted else {
                     throw RuntimeError.badStatusCode(response.status)
                 }

Tests/AWSLambdaRuntimeCoreTests/LambdaRuntimeClientTest.swift

@@ -13,6 +13,11 @@
 //===----------------------------------------------------------------------===//
 
 @testable import AWSLambdaRuntimeCore
+import Logging
+import NIO
+import NIOFoundationCompat
+@testable import NIOHTTP1
+import NIOTestUtils
 import XCTest
 
 class LambdaRuntimeClientTest: XCTestCase {
@@ -209,6 +214,110 @@ class LambdaRuntimeClientTest: XCTestCase {
         XCTAssertEqual(#"{"errorType":"error","errorMessage":"🥑👨‍👩‍👧‍👧👩‍👩‍👧‍👧👨‍👨‍👧"}"#, String(decoding: emojiBytes, as: Unicode.UTF8.self))
     }
 
+    func testInitializationErrorReport() {
+        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
+
+        let server = NIOHTTP1TestServer(group: eventLoopGroup)
+        defer { XCTAssertNoThrow(try server.stop()) }
+
+        let logger = Logger(label: "TestLogger")
+        let client = Lambda.RuntimeClient(eventLoop: eventLoopGroup.next(), configuration: .init(baseURL: "127.0.0.1:\(server.serverPort)"))
+        let result = client.reportInitializationError(logger: logger, error: TestError("boom"))
+
+        var inboundHeader: HTTPServerRequestPart?
+        XCTAssertNoThrow(inboundHeader = try server.readInbound())
+        guard case .head(let head) = try? XCTUnwrap(inboundHeader) else { XCTFail("Expected to get a head first"); return }
+        XCTAssertEqual(head.headers["lambda-runtime-function-error-type"], ["Unhandled"])
+        XCTAssertEqual(head.headers["user-agent"], ["Swift-Lambda/Unknown"])
+
+        var inboundBody: HTTPServerRequestPart?
+        XCTAssertNoThrow(inboundBody = try server.readInbound())
+        guard case .body(let body) = try? XCTUnwrap(inboundBody) else { XCTFail("Expected body after head"); return }
+        XCTAssertEqual(try JSONDecoder().decode(ErrorResponse.self, from: body).errorMessage, "boom")
+
+        XCTAssertEqual(try server.readInbound(), .end(nil))
+
+        XCTAssertNoThrow(try server.writeOutbound(.head(.init(version: .init(major: 1, minor: 1), status: .accepted))))
+        XCTAssertNoThrow(try server.writeOutbound(.end(nil)))
+        XCTAssertNoThrow(try result.wait())
+    }
+
+    func testInvocationErrorReport() {
+        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
+
+        let server = NIOHTTP1TestServer(group: eventLoopGroup)
+        defer { XCTAssertNoThrow(try server.stop()) }
+
+        let logger = Logger(label: "TestLogger")
+        let client = Lambda.RuntimeClient(eventLoop: eventLoopGroup.next(), configuration: .init(baseURL: "127.0.0.1:\(server.serverPort)"))
+
+        let header = HTTPHeaders([
+            (AmazonHeaders.requestID, "test"),
+            (AmazonHeaders.deadline, String(Date(timeIntervalSinceNow: 60).millisSinceEpoch)),
+            (AmazonHeaders.invokedFunctionARN, "arn:aws:lambda:us-east-1:123456789012:function:custom-runtime"),
+            (AmazonHeaders.traceID, "Root=1-5bef4de7-ad49b0e87f6ef6c87fc2e700;Parent=9a9197af755a6419;Sampled=1"),
+        ])
+        var inv: Lambda.Invocation?
+        XCTAssertNoThrow(inv = try Lambda.Invocation(headers: header))
+        guard let invocation = inv else { return }
+
+        let result = client.reportResults(logger: logger, invocation: invocation, result: Result.failure(TestError("boom")))
+
+        var inboundHeader: HTTPServerRequestPart?
+        XCTAssertNoThrow(inboundHeader = try server.readInbound())
+        guard case .head(let head) = try? XCTUnwrap(inboundHeader) else { XCTFail("Expected to get a head first"); return }
+        XCTAssertEqual(head.headers["lambda-runtime-function-error-type"], ["Unhandled"])
+        XCTAssertEqual(head.headers["user-agent"], ["Swift-Lambda/Unknown"])
+
+        var inboundBody: HTTPServerRequestPart?
+        XCTAssertNoThrow(inboundBody = try server.readInbound())
+        guard case .body(let body) = try? XCTUnwrap(inboundBody) else { XCTFail("Expected body after head"); return }
+        XCTAssertEqual(try JSONDecoder().decode(ErrorResponse.self, from: body).errorMessage, "boom")
+
+        XCTAssertEqual(try server.readInbound(), .end(nil))
+
+        XCTAssertNoThrow(try server.writeOutbound(.head(.init(version: .init(major: 1, minor: 1), status: .accepted))))
+        XCTAssertNoThrow(try server.writeOutbound(.end(nil)))
+        XCTAssertNoThrow(try result.wait())
+    }
+
+    func testSuccessHeaders() {
+        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
+
+        let server = NIOHTTP1TestServer(group: eventLoopGroup)
+        defer { XCTAssertNoThrow(try server.stop()) }
+
+        let logger = Logger(label: "TestLogger")
+        let client = Lambda.RuntimeClient(eventLoop: eventLoopGroup.next(), configuration: .init(baseURL: "127.0.0.1:\(server.serverPort)"))
+
+        let header = HTTPHeaders([
+            (AmazonHeaders.requestID, "test"),
+            (AmazonHeaders.deadline, String(Date(timeIntervalSinceNow: 60).millisSinceEpoch)),
+            (AmazonHeaders.invokedFunctionARN, "arn:aws:lambda:us-east-1:123456789012:function:custom-runtime"),
+            (AmazonHeaders.traceID, "Root=1-5bef4de7-ad49b0e87f6ef6c87fc2e700;Parent=9a9197af755a6419;Sampled=1"),
+        ])
+        var inv: Lambda.Invocation?
+        XCTAssertNoThrow(inv = try Lambda.Invocation(headers: header))
+        guard let invocation = inv else { return }
+
+        let result = client.reportResults(logger: logger, invocation: invocation, result: Result.success(nil))
+
+        var inboundHeader: HTTPServerRequestPart?
+        XCTAssertNoThrow(inboundHeader = try server.readInbound())
+        guard case .head(let head) = try? XCTUnwrap(inboundHeader) else { XCTFail("Expected to get a head first"); return }
+        XCTAssertFalse(head.headers.contains(name: "lambda-runtime-function-error-type"))
+        XCTAssertEqual(head.headers["user-agent"], ["Swift-Lambda/Unknown"])
+
+        XCTAssertEqual(try server.readInbound(), .end(nil))
+
+        XCTAssertNoThrow(try server.writeOutbound(.head(.init(version: .init(major: 1, minor: 1), status: .accepted))))
+        XCTAssertNoThrow(try server.writeOutbound(.end(nil)))
+        XCTAssertNoThrow(try result.wait())
+    }
+
     class Behavior: LambdaServerBehavior {
         var state = 0