LambdaRuntimeClient works with Invocation

WIP: This is a change that I would like to talk about. I don't want to make all the changes necessary for this to correctly work before we haven't talked about if this is something we want to pursue.

### Motivation:

- We want to store different entities that are needed when executing a handler within the LambdaContext (Logger, EventLoop, ByteBufferAllocator, …)
- Currently the LambdaRuntimeClient creates a LambdaContext. Having the LambdaContext though as a wrapper around EventLoopGroup the API becomes messy.
- Conceptionally the Lambda control plane api call is “get next Invocation” (API naming)

### Changes:

- LambdaContext has now a reference to the EventLoop (needed for changes down the road)
- LambdaContext properties `traceId`, `invokedFunctionArn`, `deadline` are not optional anymore since they will be always set when executing a lambda
- LambdaRuntimeClient responds with an Invocation and does not use the LambdaContext at all anymore.
- Added an LambdaRuntimeClientError.invocationMissingHeader(String) error

### Open ends:

- tests fail, since I made some properties that are always set during the actual Lambda invocation non optional
- we will need to build some kind of Deadline into the context
- Invocation is currently a struct and I'd assume copied four times with six internal Strings. Maybe we can make this faster by using an internal Storage class. (potentially 4 ref counts vs 24).
- I don’t know if Context really is a value type? grpc-swift has the context as class:
https://github.com/grpc/grpc-swift/blob/nio/Sources/GRPC/ServerCallContexts/ServerCallContext.swift
- What do you think about namespacing LambdaContext within Lambda so that we only have Context.

Author: fabianfett

Sources/MockServer/main.swift

@@ -107,7 +107,12 @@ internal final class HTTPHandler: ChannelInboundHandler {
             case .json:
                 responseBody = "{ \"body\": \"\(requestId)\" }"
             }
-            responseHeaders = [(AmazonHeaders.requestID, requestId)]
+            responseHeaders = [
+                (AmazonHeaders.requestID, requestId),
+                (AmazonHeaders.invokedFunctionARN, "arn:aws:lambda:us-east-1:123456789012:function:custom-runtime"),
+                (AmazonHeaders.traceID, "Root=1-5bef4de7-ad49b0e87f6ef6c87fc2e700;Parent=9a9197af755a6419;Sampled=1"),
+                (AmazonHeaders.deadline, String(Date(timeIntervalSinceNow: 60).timeIntervalSince1970 * 1000))
+            ]
         } else if self.requestHead.uri.hasSuffix("/response") {
             responseStatus = .accepted
         } else {

Sources/SwiftAwsLambda/Lambda.swift

@@ -271,22 +271,22 @@ public protocol InitializableLambdaHandler {
 public struct LambdaContext {
     // from aws
     public let requestId: String
-    public let traceId: String?
-    public let invokedFunctionArn: String?
+    public let traceId: String
+    public let invokedFunctionArn: String
+    public let deadline: String
     public let cognitoIdentity: String?
     public let clientContext: String?
-    public let deadline: String?
     // utliity
     public let eventLoop: EventLoop
     public let allocator: ByteBufferAllocator
     public let logger: Logger
 
     public init(requestId: String,
-                traceId: String? = nil,
-                invokedFunctionArn: String? = nil,
+                traceId: String,
+                invokedFunctionArn: String,
+                deadline: String,
                 cognitoIdentity: String? = nil,
                 clientContext: String? = nil,
-                deadline: String? = nil,
                 eventLoop: EventLoop,
                 logger: Logger) {
         self.requestId = requestId
@@ -301,9 +301,7 @@ public struct LambdaContext {
         // mutate logger with context
         var logger = logger
         logger[metadataKey: "awsRequestId"] = .string(requestId)
-        if let traceId = traceId {
-            logger[metadataKey: "awsTraceId"] = .string(traceId)
-        }
+        logger[metadataKey: "awsTraceId"] = .string(traceId)
         self.logger = logger
     }
 }

Sources/SwiftAwsLambda/LambdaRunner.swift

@@ -56,17 +56,18 @@ internal struct LambdaRunner {
         // 1. request work from lambda runtime engine
         return self.runtimeClient.requestWork(logger: logger).peekError { error in
             logger.error("could not fetch work from lambda runtime engine: \(error)")
-        }.flatMap { context, payload in
+        }.flatMap { invocation, payload in
             // 2. send work to handler
+            let context = LambdaContext(logger: logger, eventLoop: self.eventLoop, invocation: invocation)
             logger.debug("sending work to lambda handler \(self.lambdaHandler)")
             return self.lambdaHandler.handle(eventLoop: self.eventLoop,
                                              lifecycleId: self.lifecycleId,
                                              offload: self.offload,
                                              context: context,
-                                             payload: payload).mapResult { (context, $0) }
-        }.flatMap { context, result in
+                                             payload: payload).mapResult { (invocation, $0) }
+        }.flatMap { invocation, result in
             // 3. report results to runtime engine
-            self.runtimeClient.reportResults(logger: logger, context: context, result: result).peekError { error in
+            self.runtimeClient.reportResults(logger: logger, invocation: invocation, result: result).peekError { error in
                 logger.error("failed reporting results to lambda runtime engine: \(error)")
             }
         }.always { result in
@@ -117,6 +118,20 @@ private extension LambdaHandler {
     }
 }
 
+private extension LambdaContext {
+    init(logger: Logger, eventLoop: EventLoop, invocation: Invocation) {
+        
+        self = LambdaContext(requestId: invocation.requestId,
+                             traceId: invocation.traceId,
+                             invokedFunctionArn: invocation.invokedFunctionArn,
+                             deadline: invocation.deadlineDate,
+                             cognitoIdentity: invocation.cognitoIdentity,
+                             clientContext: invocation.clientContext,
+                             eventLoop: eventLoop,
+                             logger: logger)
+    }
+}
+
 // TODO: move to nio?
 private extension EventLoopFuture {
     // callback does not have side effects, failing with original result

Sources/SwiftAwsLambda/LambdaRuntimeClient.swift

@@ -33,20 +33,18 @@ internal struct LambdaRuntimeClient {
     }
 
     /// Requests work from the Runtime Engine.
-    func requestWork(logger: Logger) -> EventLoopFuture<(LambdaContext, ByteBuffer)> {
+    func requestWork(logger: Logger) -> EventLoopFuture<(Invocation, ByteBuffer)> {
         let url = Consts.invocationURLPrefix + Consts.requestWorkURLSuffix
         logger.debug("requesting work from lambda runtime engine using \(url)")
         return self.httpClient.get(url: url).flatMapThrowing { response in
             guard response.status == .ok else {
                 throw LambdaRuntimeClientError.badStatusCode(response.status)
             }
+            let invocation = try Invocation(headers: response.headers)
             guard let payload = response.body else {
                 throw LambdaRuntimeClientError.noBody
             }
-            guard let context = LambdaContext(eventLoop: self.eventLoop, logger: logger, response: response) else {
-                throw LambdaRuntimeClientError.noContext
-            }
-            return (context, payload)
+            return (invocation, payload)
         }.flatMapErrorThrowing { error in
             switch error {
             case HTTPClient.Errors.timeout:
@@ -60,8 +58,8 @@ internal struct LambdaRuntimeClient {
     }
 
     /// Reports a result to the Runtime Engine.
-    func reportResults(logger: Logger, context: LambdaContext, result: Result<ByteBuffer, Error>) -> EventLoopFuture<Void> {
-        var url = Consts.invocationURLPrefix + "/" + context.requestId
+    func reportResults(logger: Logger, invocation: Invocation, result: Result<ByteBuffer, Error>) -> EventLoopFuture<Void> {
+        var url = Consts.invocationURLPrefix + "/" + invocation.requestId
         var body: ByteBuffer
         switch result {
         case .success(let data):
@@ -131,8 +129,8 @@ internal struct LambdaRuntimeClient {
 internal enum LambdaRuntimeClientError: Error, Equatable {
     case badStatusCode(HTTPResponseStatus)
     case upstreamError(String)
+    case invocationMissingHeader(String)
     case noBody
-    case noContext
     case json(JsonCodecError)
 }
 
@@ -151,6 +149,7 @@ internal struct JsonCodecError: Error, Equatable {
 internal struct ErrorResponse: Codable {
     var errorType: String
     var errorMessage: String
+    
 }
 
 private extension ErrorResponse {
@@ -181,26 +180,38 @@ private extension HTTPClient.Response {
     }
 }
 
-private extension LambdaContext {
-    init?(eventLoop: EventLoop, logger: Logger, response: HTTPClient.Response) {
-        guard let requestId = response.headerValue(AmazonHeaders.requestID) else {
-            return nil
+internal struct Invocation {
+    let requestId         : String
+    let deadlineDate      : String
+    let invokedFunctionArn: String
+    let traceId           : String
+    let clientContext     : String?
+    let cognitoIdentity   : String?
+
+    init(headers: HTTPHeaders) throws {
+
+        guard let requestId = headers[AmazonHeaders.requestID].first, !requestId.isEmpty else {
+            throw LambdaRuntimeClientError.invocationMissingHeader(AmazonHeaders.requestID)
         }
-        if requestId.isEmpty {
-            return nil
+
+        guard let unixTimeMilliseconds = headers[AmazonHeaders.deadline].first else {
+            throw LambdaRuntimeClientError.invocationMissingHeader(AmazonHeaders.deadline)
         }
-        let traceId = response.headerValue(AmazonHeaders.traceID)
-        let invokedFunctionArn = response.headerValue(AmazonHeaders.invokedFunctionARN)
-        let cognitoIdentity = response.headerValue(AmazonHeaders.cognitoIdentity)
-        let clientContext = response.headerValue(AmazonHeaders.clientContext)
-        let deadline = response.headerValue(AmazonHeaders.deadline)
-        self = LambdaContext(requestId: requestId,
-                             traceId: traceId,
-                             invokedFunctionArn: invokedFunctionArn,
-                             cognitoIdentity: cognitoIdentity,
-                             clientContext: clientContext,
-                             deadline: deadline,
-                             eventLoop: eventLoop,
-                             logger: logger)
+
+        guard let invokedFunctionArn = headers[AmazonHeaders.invokedFunctionARN].first else {
+            throw LambdaRuntimeClientError.invocationMissingHeader(AmazonHeaders.invokedFunctionARN)
+        }
+
+        guard let traceId = headers[AmazonHeaders.traceID].first else {
+            throw LambdaRuntimeClientError.invocationMissingHeader(AmazonHeaders.traceID)
+        }
+
+        self.requestId          = requestId
+        self.deadlineDate       = unixTimeMilliseconds
+        self.invokedFunctionArn = invokedFunctionArn
+        self.traceId            = traceId
+        self.clientContext      = headers["Lambda-Runtime-Client-Context"].first
+        self.cognitoIdentity    = headers["Lambda-Runtime-Cognito-Identity"].first
     }
+
 }

Tests/SwiftAwsLambdaTests/LambdaRuntimeClientTest+XCTest.swift

@@ -27,7 +27,7 @@ extension LambdaRuntimeClientTest {
         return [
             ("testGetWorkServerInternalError", testGetWorkServerInternalError),
             ("testGetWorkServerNoBodyError", testGetWorkServerNoBodyError),
-            ("testGetWorkServerNoContextError", testGetWorkServerNoContextError),
+            ("testGetWorkServerMissingHeaderRequestIDError", testGetWorkServerMissingHeaderRequestIDError),
             ("testProcessResponseInternalServerError", testProcessResponseInternalServerError),
             ("testProcessErrorInternalServerError", testProcessErrorInternalServerError),
             ("testProcessInitErrorInternalServerError", testProcessInitErrorInternalServerError),

Tests/SwiftAwsLambdaTests/LambdaRuntimeClientTest.swift

@@ -68,7 +68,7 @@ class LambdaRuntimeClientTest: XCTestCase {
         }
     }
 
-    func testGetWorkServerNoContextError() throws {
+    func testGetWorkServerMissingHeaderRequestIDError() throws {
         struct Behavior: LambdaServerBehavior {
             func getWork() -> GetWorkResult {
                 // no request id -> no context
@@ -91,7 +91,7 @@ class LambdaRuntimeClientTest: XCTestCase {
             }
         }
         XCTAssertThrowsError(try runLambda(behavior: Behavior(), handler: EchoHandler())) { error in
-            XCTAssertEqual(error as? LambdaRuntimeClientError, LambdaRuntimeClientError.noContext)
+            XCTAssertEqual(error as? LambdaRuntimeClientError, LambdaRuntimeClientError.invocationMissingHeader(AmazonHeaders.requestID))
         }
     }
 

Tests/SwiftAwsLambdaTests/MockLambdaServer.swift

@@ -139,7 +139,12 @@ internal final class HTTPHandler: ChannelInboundHandler {
                 }
                 responseStatus = .ok
                 responseBody = result
-                responseHeaders = [(AmazonHeaders.requestID, requestId)]
+                responseHeaders = [
+                    (AmazonHeaders.requestID, requestId),
+                    (AmazonHeaders.invokedFunctionARN, "arn:aws:lambda:us-east-1:123456789012:function:custom-runtime"),
+                    (AmazonHeaders.traceID, "Root=1-5bef4de7-ad49b0e87f6ef6c87fc2e700;Parent=9a9197af755a6419;Sampled=1"),
+                    (AmazonHeaders.deadline, String(Date(timeIntervalSinceNow: 60).timeIntervalSince1970 * 1000))
+                ]
             case .failure(let error):
                 responseStatus = .init(statusCode: error.rawValue)
             }