initial commit

Author: sebsto

Examples/Streaming/Package.swift

@@ -1,7 +1,6 @@
-// swift-tools-version:6.0
+// swift-tools-version:6.1
 
 import PackageDescription
-
 // needed for CI to test the local version of the library
 import struct Foundation.URL
 

Examples/Streaming/Sources/main.swift

@@ -14,13 +14,24 @@
 
 import AWSLambdaRuntime
 import NIOCore
+import NIOHTTP1
 
 struct SendNumbersWithPause: StreamingLambdaHandler {
     func handle(
         _ event: ByteBuffer,
         responseWriter: some LambdaResponseStreamWriter,
         context: LambdaContext
     ) async throws {
+        context.logger.info("Received event: \(event)")
+        try await responseWriter.writeHeaders(
+            HTTPHeaders([
+                ("X-Example-Header", "This is an example header")
+            ])
+        )
+
+        try await responseWriter.write(
+            ByteBuffer(string: "Starting to send numbers with a pause...\n")
+        )
         for i in 1...10 {
             // Send partial data
             try await responseWriter.write(ByteBuffer(string: "\(i)\n"))

Examples/Streaming/template.yaml

@@ -14,6 +14,9 @@ Resources:
       MemorySize: 128
       Architectures:
         - arm64
+      Environment:
+        Variables:
+          LOG_LEVEL: trace
       FunctionUrlConfig:
         AuthType: AWS_IAM
         InvokeMode: RESPONSE_STREAM

Sources/AWSLambdaRuntime/Lambda+LocalServer.swift

@@ -248,6 +248,7 @@ internal struct LambdaHTTPServer {
                         case .end:
                             precondition(requestHead != nil, "Received .end without .head")
                             // process the request
+                            // FIXME: this do not support response streaming
                             let response = try await self.processRequest(
                                 head: requestHead,
                                 body: requestBody,

Sources/AWSLambdaRuntime/LambdaHandlers.swift

@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 import NIOCore
+import NIOHTTP1
 
 /// The base handler protocol that receives a `ByteBuffer` representing the incoming event and returns the response as a `ByteBuffer` too.
 /// This handler protocol supports response streaming. Bytes can be streamed outwards through the ``LambdaResponseStreamWriter``
@@ -46,6 +47,10 @@ public protocol StreamingLambdaHandler: _Lambda_SendableMetatype {
 /// A writer object to write the Lambda response stream into. The HTTP response is started lazily.
 /// before the first call to ``write(_:)`` or ``writeAndFinish(_:)``.
 public protocol LambdaResponseStreamWriter {
+    /// Write the headers parts of the stream. This allows client to set headers before the first response part is written.
+    /// - Parameter buffer: The buffer to write.
+    func writeHeaders(_ headers: HTTPHeaders) async throws
+
     /// Write a response part into the stream. Bytes written are streamed continually.
     /// - Parameter buffer: The buffer to write.
     func write(_ buffer: ByteBuffer) async throws

Sources/AWSLambdaRuntime/LambdaRuntimeClient.swift

@@ -42,6 +42,11 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
             self.runtimeClient = runtimeClient
         }
 
+        @usableFromInline
+        func writeHeaders(_ headers: HTTPHeaders) async throws {
+            try await self.runtimeClient.appendHeaders(headers)
+        }
+
         @usableFromInline
         func write(_ buffer: NIOCore.ByteBuffer) async throws {
             try await self.runtimeClient.write(buffer)
@@ -188,6 +193,13 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
         }
     }
 
+    // we can use a var here because we are always isolated to this actor
+    private var userHeaders = HTTPHeaders()
+    private func appendHeaders(_ headers: HTTPHeaders) async throws {
+        // buffer the data to send them when we will send the headers 
+        userHeaders.add(contentsOf: headers)
+    }
+    
     private func write(_ buffer: NIOCore.ByteBuffer) async throws {
         switch self.lambdaState {
         case .idle, .sentResponse:
@@ -205,7 +217,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
             guard case .sendingResponse(requestID) = self.lambdaState else {
                 fatalError("Invalid state: \(self.lambdaState)")
             }
-            return try await handler.writeResponseBodyPart(buffer, requestID: requestID)
+            return try await handler.writeResponseBodyPart(self.userHeaders, buffer, requestID: requestID)
         }
     }
 
@@ -555,6 +567,7 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
 
     func writeResponseBodyPart(
         isolation: isolated (any Actor)? = #isolation,
+        _ userHeaders: HTTPHeaders,
         _ byteBuffer: ByteBuffer,
         requestID: String
     ) async throws {
@@ -564,10 +577,10 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
 
         case .connected(let context, .waitingForResponse):
             self.state = .connected(context, .sendingResponse)
-            try await self.sendResponseBodyPart(byteBuffer, sendHeadWithRequestID: requestID, context: context)
+            try await self.sendResponseBodyPart(userHeaders, byteBuffer, sendHeadWithRequestID: requestID, context: context)
 
         case .connected(let context, .sendingResponse):
-            try await self.sendResponseBodyPart(byteBuffer, sendHeadWithRequestID: nil, context: context)
+            try await self.sendResponseBodyPart(userHeaders, byteBuffer, sendHeadWithRequestID: nil, context: context)
 
         case .connected(_, .idle),
             .connected(_, .sentResponse):
@@ -616,6 +629,7 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
 
     private func sendResponseBodyPart(
         isolation: isolated (any Actor)? = #isolation,
+        _ userHeaders: HTTPHeaders,
         _ byteBuffer: ByteBuffer,
         sendHeadWithRequestID: String?,
         context: ChannelHandlerContext
@@ -625,13 +639,16 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
             // TODO: This feels super expensive. We should be able to make this cheaper. requestIDs are fixed length
             let url = Consts.invocationURLPrefix + "/" + requestID + Consts.postResponseURLSuffix
 
+            var responseHeaders = self.streamingHeaders
+            responseHeaders.add(contentsOf: userHeaders)
+            logger.trace("sendResponseBodyPart : ========== Sending response headers: \(responseHeaders)")
             let httpRequest = HTTPRequestHead(
                 version: .http1_1,
                 method: .POST,
                 uri: url,
-                headers: self.streamingHeaders
+                headers: responseHeaders
             )
-
+            
             context.write(self.wrapOutboundOut(.head(httpRequest)), promise: nil)
         }
 
@@ -663,13 +680,13 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
                     self.streamingHeaders
                 }
 
+            logger.trace("sendResponseFinish : ========== Sending response headers: \(headers)")
             let httpRequest = HTTPRequestHead(
                 version: .http1_1,
                 method: .POST,
                 uri: url,
                 headers: headers
             )
-
             context.write(self.wrapOutboundOut(.head(httpRequest)), promise: nil)
         }