[core] Implement Lambda streaming with custom HTTP headers (#521)

Fix https://github.com/swift-server/swift-aws-lambda-runtime/issues/520

Examples/Streaming/README.md

@@ -13,15 +13,55 @@ The sample code creates a `SendNumbersWithPause` struct that conforms to the `St
 
 The `handle(...)` method of this protocol receives incoming events as a Swift NIO `ByteBuffer` and returns the output as a `ByteBuffer`.
 
-The response is streamed through the `LambdaResponseStreamWriter`, which is passed as an argument in the `handle` function.  The code calls the `write(_:)` function of the `LambdaResponseStreamWriter` with partial data repeatedly written before
-finally closing the response stream by calling `finish()`. Developers can also choose to return the entire output and not
-stream the response by calling `writeAndFinish(_:)`.
+The response is streamed through the `LambdaResponseStreamWriter`, which is passed as an argument in the `handle` function. 
+
+### Setting HTTP Status Code and Headers
+
+Before streaming the response body, you can set the HTTP status code and headers using the `writeStatusAndHeaders(_:)` method:
+
+```swift
+try await responseWriter.writeStatusAndHeaders(
+    StreamingLambdaStatusAndHeadersResponse(
+        statusCode: 200,
+        headers: [
+            "Content-Type": "text/plain",
+            "x-my-custom-header": "streaming-example"
+        ]
+    )
+)
+```
+
+The `StreamingLambdaStatusAndHeadersResponse` structure allows you to specify:
+- **statusCode**: HTTP status code (e.g., 200, 404, 500)
+- **headers**: Dictionary of single-value HTTP headers (optional)
+
+### Streaming the Response Body
+
+After setting headers, you can stream the response body by calling the `write(_:)` function of the `LambdaResponseStreamWriter` with partial data repeatedly before finally closing the response stream by calling `finish()`. Developers can also choose to return the entire output and not stream the response by calling `writeAndFinish(_:)`.
+
+```swift
+// Stream data in chunks
+for i in 1...3 {
+    try await responseWriter.write(ByteBuffer(string: "Number: \(i)\n"))
+    try await Task.sleep(for: .milliseconds(1000))
+}
+
+// Close the response stream
+try await responseWriter.finish()
+```
 
 An error is thrown if `finish()` is called multiple times or if it is called after having called `writeAndFinish(_:)`.
 
+### Example Usage Patterns
+
+The example includes two handler implementations:
+
+1. **SendNumbersWithPause**: Demonstrates basic streaming with headers, sending numbers with delays
+2. **ConditionalStreamingHandler**: Shows how to handle different response scenarios, including error responses with appropriate status codes
+
 The `handle(...)` method is marked as `mutating` to allow handlers to be implemented with a `struct`.
 
-Once the struct is created and the `handle(...)` method is defined, the sample code creates a `LambdaRuntime` struct and initializes it with the handler just created.  Then, the code calls `run()` to start the interaction with the AWS Lambda control plane.
+Once the struct is created and the `handle(...)` method is defined, the sample code creates a `LambdaRuntime` struct and initializes it with the handler just created. Then, the code calls `run()` to start the interaction with the AWS Lambda control plane.
 
 ## Build & Package 
 
@@ -68,7 +108,7 @@ aws lambda create-function \
 ```
 
 > [!IMPORTANT] 
-> The timeout value must be bigger than the time it takes for your function to stream its output. Otherwise, the Lambda control plane will terminate the execution environment before your code has a chance to finish writing the stream. Here, the sample function stream responses during 10 seconds and we set the timeout for 15 seconds.
+> The timeout value must be bigger than the time it takes for your function to stream its output. Otherwise, the Lambda control plane will terminate the execution environment before your code has a chance to finish writing the stream. Here, the sample function stream responses during 3 seconds and we set the timeout for 5 seconds.
 
 The `--architectures` flag is only required when you build the binary on an Apple Silicon machine (Apple M1 or more recent). It defaults to `x64`.
 
@@ -139,13 +179,7 @@ This should output the following result, with a one-second delay between each nu
 1
 2
 3
-4
-5
-6
-7
-8
-9
-10
+Streaming complete!
 ```
 
 ### Undeploy

Examples/Streaming/Sources/main.swift

@@ -15,22 +15,46 @@
 import AWSLambdaRuntime
 import NIOCore
 
+#if canImport(FoundationEssentials)
+import FoundationEssentials
+#else
+import Foundation
+#endif
+
 struct SendNumbersWithPause: StreamingLambdaHandler {
     func handle(
         _ event: ByteBuffer,
         responseWriter: some LambdaResponseStreamWriter,
         context: LambdaContext
     ) async throws {
-        for i in 1...10 {
+
+        // Send HTTP status code and headers before streaming the response body
+        try await responseWriter.writeStatusAndHeaders(
+            StreamingLambdaStatusAndHeadersResponse(
+                statusCode: 418,  // I'm a tea pot
+                headers: [
+                    "Content-Type": "text/plain",
+                    "x-my-custom-header": "streaming-example",
+                ]
+            )
+        )
+
+        // Stream numbers with pauses to demonstrate streaming functionality
+        for i in 1...3 {
             // Send partial data
-            try await responseWriter.write(ByteBuffer(string: "\(i)\n"))
-            // Perform some long asynchronous work
+            try await responseWriter.write(ByteBuffer(string: "Number: \(i)\n"))
+
+            // Perform some long asynchronous work to simulate processing
             try await Task.sleep(for: .milliseconds(1000))
         }
+
+        // Send final message
+        try await responseWriter.write(ByteBuffer(string: "Streaming complete!\n"))
+
         // All data has been sent. Close off the response stream.
         try await responseWriter.finish()
     }
 }
 
-let runtime = LambdaRuntime.init(handler: SendNumbersWithPause())
+let runtime = LambdaRuntime(handler: SendNumbersWithPause())
 try await runtime.run()

Examples/Streaming/template.yaml

@@ -8,7 +8,7 @@ Resources:
     Type: AWS::Serverless::Function
     Properties:
       CodeUri: .build/plugins/AWSLambdaPackager/outputs/AWSLambdaPackager/StreamingNumbers/StreamingNumbers.zip
-      Timeout: 15
+      Timeout: 5  # Must be bigger than the time it takes to stream the output
       Handler: swift.bootstrap  # ignored by the Swift runtime
       Runtime: provided.al2
       MemorySize: 128
@@ -17,6 +17,9 @@ Resources:
       FunctionUrlConfig:
         AuthType: AWS_IAM
         InvokeMode: RESPONSE_STREAM
+      Environment:
+        Variables:
+          LOG_LEVEL: trace
 
 Outputs:
   # print Lambda function URL