Introduce LambdaRuntimeClientProtocol (#344)

Author: aryan-25

Sources/AWSLambdaRuntimeCore/ControlPlaneRequest.swift

@@ -23,20 +23,20 @@ enum ControlPlaneRequest: Hashable {
 }
 
 enum ControlPlaneResponse: Hashable {
-    case next(Invocation, ByteBuffer)
+    case next(InvocationMetadata, ByteBuffer)
     case accepted
     case error(ErrorResponse)
 }
 
-struct Invocation: Hashable {
-    let requestID: String
-    let deadlineInMillisSinceEpoch: Int64
-    let invokedFunctionARN: String
-    let traceID: String
-    let clientContext: String?
-    let cognitoIdentity: String?
+package struct InvocationMetadata: Hashable {
+    package let requestID: String
+    package let deadlineInMillisSinceEpoch: Int64
+    package let invokedFunctionARN: String
+    package let traceID: String
+    package let clientContext: String?
+    package let cognitoIdentity: String?
 
-    init(headers: HTTPHeaders) throws {
+    package init(headers: HTTPHeaders) throws {
         guard let requestID = headers.first(name: AmazonHeaders.requestID), !requestID.isEmpty else {
             throw LambdaRuntimeError.invocationMissingHeader(AmazonHeaders.requestID)
         }

Sources/AWSLambdaRuntimeCore/LambdaRunner.swift

@@ -132,7 +132,7 @@ final class LambdaRunner {
 }
 
 extension LambdaContext {
-    init(logger: Logger, eventLoop: EventLoop, allocator: ByteBufferAllocator, invocation: Invocation) {
+    init(logger: Logger, eventLoop: EventLoop, allocator: ByteBufferAllocator, invocation: InvocationMetadata) {
         self.init(
             requestID: invocation.requestID,
             traceID: invocation.traceID,

Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift

@@ -32,14 +32,14 @@ struct LambdaRuntimeClient {
     }
 
     /// Requests invocation from the control plane.
-    func getNextInvocation(logger: Logger) -> EventLoopFuture<(Invocation, ByteBuffer)> {
+    func getNextInvocation(logger: Logger) -> EventLoopFuture<(InvocationMetadata, ByteBuffer)> {
         let url = Consts.invocationURLPrefix + Consts.getNextInvocationURLSuffix
         logger.debug("requesting work from lambda runtime engine using \(url)")
         return self.httpClient.get(url: url, headers: LambdaRuntimeClient.defaultHeaders).flatMapThrowing { response in
             guard response.status == .ok else {
                 throw LambdaRuntimeError.badStatusCode(response.status)
             }
-            let invocation = try Invocation(headers: response.headers)
+            let invocation = try InvocationMetadata(headers: response.headers)
             guard let event = response.body else {
                 throw LambdaRuntimeError.noBody
             }
@@ -59,7 +59,7 @@ struct LambdaRuntimeClient {
     /// Reports a result to the Runtime Engine.
     func reportResults(
         logger: Logger,
-        invocation: Invocation,
+        invocation: InvocationMetadata,
         result: Result<ByteBuffer?, Error>
     ) -> EventLoopFuture<Void> {
         var url = Consts.invocationURLPrefix + "/" + invocation.requestID

Sources/AWSLambdaRuntimeCore/LambdaRuntimeClientProtocol.swift

@@ -0,0 +1,38 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftAWSLambdaRuntime open source project
+//
+// Copyright (c) 2024 Apple Inc. and the SwiftAWSLambdaRuntime project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of SwiftAWSLambdaRuntime project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import NIOCore
+
+package protocol LambdaResponseStreamWriter {
+    mutating func write(_ buffer: ByteBuffer) async throws
+    func finish() async throws
+    func writeAndFinish(_ buffer: ByteBuffer) async throws
+    func reportError(_ error: any Error) async throws
+}
+
+package protocol LambdaRuntimeClientProtocol {
+    associatedtype Writer: LambdaResponseStreamWriter
+
+    func nextInvocation() async throws -> (Invocation, Writer)
+}
+
+package struct Invocation {
+    package var metadata: InvocationMetadata
+    package var event: ByteBuffer
+
+    package init(metadata: InvocationMetadata, event: ByteBuffer) {
+        self.metadata = metadata
+        self.event = event
+    }
+}

Tests/AWSLambdaRuntimeCoreTests/ControlPlaneRequestTests.swift

@@ -25,9 +25,9 @@ class InvocationTest: XCTestCase {
             (AmazonHeaders.invokedFunctionARN, "arn:aws:lambda:us-east-1:123456789012:function:custom-runtime"),
         ])
 
-        var invocation: Invocation?
+        var invocation: InvocationMetadata?
 
-        XCTAssertNoThrow(invocation = try Invocation(headers: headers))
+        XCTAssertNoThrow(invocation = try InvocationMetadata(headers: headers))
         XCTAssertNotNil(invocation)
 
         guard !invocation!.traceID.isEmpty else {

Tests/AWSLambdaRuntimeCoreTests/LambdaRuntimeClientTest.swift

@@ -279,8 +279,8 @@ class LambdaRuntimeClientTest: XCTestCase {
             (AmazonHeaders.invokedFunctionARN, "arn:aws:lambda:us-east-1:123456789012:function:custom-runtime"),
             (AmazonHeaders.traceID, "Root=\(AmazonHeaders.generateXRayTraceID());Sampled=1"),
         ])
-        var inv: Invocation?
-        XCTAssertNoThrow(inv = try Invocation(headers: header))
+        var inv: InvocationMetadata?
+        XCTAssertNoThrow(inv = try InvocationMetadata(headers: header))
         guard let invocation = inv else { return }
 
         let result = client.reportResults(
@@ -332,8 +332,8 @@ class LambdaRuntimeClientTest: XCTestCase {
             (AmazonHeaders.invokedFunctionARN, "arn:aws:lambda:us-east-1:123456789012:function:custom-runtime"),
             (AmazonHeaders.traceID, "Root=\(AmazonHeaders.generateXRayTraceID());Sampled=1"),
         ])
-        var inv: Invocation?
-        XCTAssertNoThrow(inv = try Invocation(headers: header))
+        var inv: InvocationMetadata?
+        XCTAssertNoThrow(inv = try InvocationMetadata(headers: header))
         guard let invocation = inv else { return }
 
         let result = client.reportResults(logger: logger, invocation: invocation, result: Result.success(nil))