Merge branch 'main' into sebsto/new-plugin-proposal

Author: sebsto

.github/workflows/pull_request.yml

@@ -23,7 +23,7 @@ jobs:
     with:
       linux_5_9_enabled: false
       linux_5_10_enabled: false
-      linux_nightly_6_0_arguments_override: "--explicit-target-dependency-import-check error"
+      linux_nightly_next_arguments_override: "--explicit-target-dependency-import-check error"
       linux_nightly_main_arguments_override: "--explicit-target-dependency-import-check error"
 
   integration-tests:

.gitignore

@@ -11,4 +11,5 @@ Package.resolved
 .serverless
 .vscode
 Makefile
-.devcontainer
+.devcontainer
+.amazonq

Package@swift-6.0.swift

@@ -56,6 +56,11 @@ let package = Package(
                 .byName(name: "AWSLambdaRuntime"),
                 .product(name: "NIOTestUtils", package: "swift-nio"),
                 .product(name: "NIOFoundationCompat", package: "swift-nio"),
+            ],
+            swiftSettings: [
+                .define("FoundationJSONSupport"),
+                .define("ServiceLifecycleSupport"),
+                .define("LocalServerSupport"),
             ]
         ),
         // for perf testing

Sources/AWSLambdaRuntime/FoundationSupport/Lambda+JSON.swift

@@ -23,6 +23,8 @@ import class Foundation.JSONDecoder
 import class Foundation.JSONEncoder
 #endif
 
+import Logging
+
 public struct LambdaJSONEventDecoder: LambdaEventDecoder {
     @usableFromInline let jsonDecoder: JSONDecoder
 
@@ -87,10 +89,12 @@ extension LambdaRuntime {
     /// - Parameters:
     ///   - decoder: The decoder object that will be used to decode the incoming `ByteBuffer` event into the generic `Event` type. `JSONDecoder()` used as default.
     ///   - encoder: The encoder object that will be used to encode the generic `Output` into a `ByteBuffer`. `JSONEncoder()` used as default.
+    ///   - logger: The logger to use for the runtime. Defaults to a logger with label "LambdaRuntime".
     ///   - body: The handler in the form of a closure.
     public convenience init<Event: Decodable, Output>(
         decoder: JSONDecoder = JSONDecoder(),
         encoder: JSONEncoder = JSONEncoder(),
+        logger: Logger = Logger(label: "LambdaRuntime"),
         body: sending @escaping (Event, LambdaContext) async throws -> Output
     )
     where
@@ -108,14 +112,16 @@ extension LambdaRuntime {
             handler: LambdaHandlerAdapter(handler: ClosureHandler(body: body))
         )
 
-        self.init(handler: handler)
+        self.init(handler: handler, logger: logger)
     }
 
     /// Initialize an instance with a `LambdaHandler` defined in the form of a closure **with a `Void` return type**.
     /// - Parameter body: The handler in the form of a closure.
     /// - Parameter decoder: The decoder object that will be used to decode the incoming `ByteBuffer` event into the generic `Event` type. `JSONDecoder()` used as default.
+    /// - Parameter logger: The logger to use for the runtime. Defaults to a logger with label "LambdaRuntime".
     public convenience init<Event: Decodable>(
         decoder: JSONDecoder = JSONDecoder(),
+        logger: Logger = Logger(label: "LambdaRuntime"),
         body: sending @escaping (Event, LambdaContext) async throws -> Void
     )
     where
@@ -132,7 +138,7 @@ extension LambdaRuntime {
             handler: LambdaHandlerAdapter(handler: ClosureHandler(body: body))
         )
 
-        self.init(handler: handler)
+        self.init(handler: handler, logger: logger)
     }
 }
 #endif  // trait: FoundationJSONSupport

Sources/AWSLambdaRuntime/Lambda+LocalServer.swift

@@ -2,7 +2,7 @@
 //
 // This source file is part of the SwiftAWSLambdaRuntime open source project
 //
-// Copyright (c) 2020 Apple Inc. and the SwiftAWSLambdaRuntime project authors
+// Copyright (c) 2025 Apple Inc. and the SwiftAWSLambdaRuntime project authors
 // Licensed under Apache License v2.0
 //
 // See LICENSE.txt for license information
@@ -76,7 +76,7 @@ extension Lambda {
 /// 1. POST /invoke - the client posts the event to the lambda function
 ///
 /// This server passes the data received from /invoke POST request to the lambda function (GET /next) and then forwards the response back to the client.
-private struct LambdaHTTPServer {
+internal struct LambdaHTTPServer {
     private let invocationEndpoint: String
 
     private let invocationPool = Pool<LocalServerInvocation>()
@@ -166,17 +166,21 @@ private struct LambdaHTTPServer {
                     // consumed by iterating the group or by exiting the group. Since, we are never consuming
                     // the results of the group we need the group to automatically discard them; otherwise, this
                     // would result in a memory leak over time.
-                    try await withThrowingDiscardingTaskGroup { taskGroup in
-                        try await channel.executeThenClose { inbound in
-                            for try await connectionChannel in inbound {
-
-                                taskGroup.addTask {
-                                    logger.trace("Handling a new connection")
-                                    await server.handleConnection(channel: connectionChannel, logger: logger)
-                                    logger.trace("Done handling the connection")
+                    try await withTaskCancellationHandler {
+                        try await withThrowingDiscardingTaskGroup { taskGroup in
+                            try await channel.executeThenClose { inbound in
+                                for try await connectionChannel in inbound {
+
+                                    taskGroup.addTask {
+                                        logger.trace("Handling a new connection")
+                                        await server.handleConnection(channel: connectionChannel, logger: logger)
+                                        logger.trace("Done handling the connection")
+                                    }
                                 }
                             }
                         }
+                    } onCancel: {
+                        channel.channel.close(promise: nil)
                     }
                     return .serverReturned(.success(()))
                 } catch {
@@ -230,38 +234,42 @@ private struct LambdaHTTPServer {
         // Note that this method is non-throwing and we are catching any error.
         // We do this since we don't want to tear down the whole server when a single connection
         // encounters an error.
-        do {
-            try await channel.executeThenClose { inbound, outbound in
-                for try await inboundData in inbound {
-                    switch inboundData {
-                    case .head(let head):
-                        requestHead = head
-
-                    case .body(let body):
-                        requestBody.setOrWriteImmutableBuffer(body)
-
-                    case .end:
-                        precondition(requestHead != nil, "Received .end without .head")
-                        // process the request
-                        let response = try await self.processRequest(
-                            head: requestHead,
-                            body: requestBody,
-                            logger: logger
-                        )
-                        // send the responses
-                        try await self.sendResponse(
-                            response: response,
-                            outbound: outbound,
-                            logger: logger
-                        )
-
-                        requestHead = nil
-                        requestBody = nil
+        await withTaskCancellationHandler {
+            do {
+                try await channel.executeThenClose { inbound, outbound in
+                    for try await inboundData in inbound {
+                        switch inboundData {
+                        case .head(let head):
+                            requestHead = head
+
+                        case .body(let body):
+                            requestBody.setOrWriteImmutableBuffer(body)
+
+                        case .end:
+                            precondition(requestHead != nil, "Received .end without .head")
+                            // process the request
+                            let response = try await self.processRequest(
+                                head: requestHead,
+                                body: requestBody,
+                                logger: logger
+                            )
+                            // send the responses
+                            try await self.sendResponse(
+                                response: response,
+                                outbound: outbound,
+                                logger: logger
+                            )
+
+                            requestHead = nil
+                            requestBody = nil
+                        }
                     }
                 }
+            } catch {
+                logger.error("Hit error: \(error)")
             }
-        } catch {
-            logger.error("Hit error: \(error)")
+        } onCancel: {
+            channel.channel.close(promise: nil)
         }
     }
 
@@ -426,7 +434,7 @@ private struct LambdaHTTPServer {
     /// A shared data structure to store the current invocation or response requests and the continuation objects.
     /// This data structure is shared between instances of the HTTPHandler
     /// (one instance to serve requests from the Lambda function and one instance to serve requests from the client invoking the lambda function).
-    private final class Pool<T>: AsyncSequence, AsyncIteratorProtocol, Sendable where T: Sendable {
+    internal final class Pool<T>: AsyncSequence, AsyncIteratorProtocol, Sendable where T: Sendable {
         typealias Element = T
 
         enum State: ~Copyable {
@@ -462,26 +470,38 @@ private struct LambdaHTTPServer {
                 return nil
             }
 
-            return try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<T, any Error>) in
-                let nextAction = self.lock.withLock { state -> T? in
-                    switch consume state {
-                    case .buffer(var buffer):
-                        if let first = buffer.popFirst() {
-                            state = .buffer(buffer)
-                            return first
-                        } else {
-                            state = .continuation(continuation)
-                            return nil
-                        }
+            return try await withTaskCancellationHandler {
+                try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<T, any Error>) in
+                    let nextAction = self.lock.withLock { state -> T? in
+                        switch consume state {
+                        case .buffer(var buffer):
+                            if let first = buffer.popFirst() {
+                                state = .buffer(buffer)
+                                return first
+                            } else {
+                                state = .continuation(continuation)
+                                return nil
+                            }
 
-                    case .continuation:
-                        fatalError("Concurrent invocations to next(). This is illegal.")
+                        case .continuation:
+                            fatalError("Concurrent invocations to next(). This is illegal.")
+                        }
                     }
-                }
 
-                guard let nextAction else { return }
+                    guard let nextAction else { return }
 
-                continuation.resume(returning: nextAction)
+                    continuation.resume(returning: nextAction)
+                }
+            } onCancel: {
+                self.lock.withLock { state in
+                    switch consume state {
+                    case .buffer(let buffer):
+                        state = .buffer(buffer)
+                    case .continuation(let continuation):
+                        continuation?.resume(throwing: CancellationError())
+                        state = .buffer([])
+                    }
+                }
             }
         }
 

Sources/AWSLambdaRuntime/LambdaHandlers.swift

@@ -12,6 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+import Logging
 import NIOCore
 
 /// The base handler protocol that receives a `ByteBuffer` representing the incoming event and returns the response as a `ByteBuffer` too.
@@ -20,7 +21,7 @@ import NIOCore
 /// Background work can also be executed after returning the response. After closing the response stream by calling
 /// ``LambdaResponseStreamWriter/finish()`` or ``LambdaResponseStreamWriter/writeAndFinish(_:)``,
 /// the ``handle(_:responseWriter:context:)`` function is free to execute any background work.
-public protocol StreamingLambdaHandler {
+public protocol StreamingLambdaHandler: _Lambda_SendableMetatype {
     /// The handler function -- implement the business logic of the Lambda function here.
     /// - Parameters:
     ///   - event: The invocation's input data.
@@ -175,17 +176,22 @@ public struct ClosureHandler<Event: Decodable, Output>: LambdaHandler {
 
 extension LambdaRuntime {
     /// Initialize an instance with a ``StreamingLambdaHandler`` in the form of a closure.
-    /// - Parameter body: The handler in the form of a closure.
+    /// - Parameter
+    ///   - logger: The logger to use for the runtime. Defaults to a logger with label "LambdaRuntime".
+    ///   - body: The handler in the form of a closure.
     public convenience init(
+        logger: Logger = Logger(label: "LambdaRuntime"),
         body: @Sendable @escaping (ByteBuffer, LambdaResponseStreamWriter, LambdaContext) async throws -> Void
+
     ) where Handler == StreamingClosureHandler {
-        self.init(handler: StreamingClosureHandler(body: body))
+        self.init(handler: StreamingClosureHandler(body: body), logger: logger)
     }
 
     /// Initialize an instance with a ``LambdaHandler`` defined in the form of a closure **with a non-`Void` return type**, an encoder, and a decoder.
     /// - Parameters:
     ///   - encoder: The encoder object that will be used to encode the generic `Output` into a `ByteBuffer`.
     ///   - decoder: The decoder object that will be used to decode the incoming `ByteBuffer` event into the generic `Event` type.
+    ///   - logger: The logger to use for the runtime. Defaults to a logger with label "LambdaRuntime".
     ///   - body: The handler in the form of a closure.
     public convenience init<
         Event: Decodable,
@@ -195,6 +201,7 @@ extension LambdaRuntime {
     >(
         encoder: sending Encoder,
         decoder: sending Decoder,
+        logger: Logger = Logger(label: "LambdaRuntime"),
         body: sending @escaping (Event, LambdaContext) async throws -> Output
     )
     where
@@ -214,15 +221,17 @@ extension LambdaRuntime {
             handler: streamingAdapter
         )
 
-        self.init(handler: codableWrapper)
+        self.init(handler: codableWrapper, logger: logger)
     }
 
     /// Initialize an instance with a ``LambdaHandler`` defined in the form of a closure **with a `Void` return type**, an encoder, and a decoder.
     /// - Parameters:
     ///   - decoder: The decoder object that will be used to decode the incoming `ByteBuffer` event into the generic `Event` type.
+    ///   - logger: The logger to use for the runtime. Defaults to a logger with label "LambdaRuntime".
     ///   - body: The handler in the form of a closure.
     public convenience init<Event: Decodable, Decoder: LambdaEventDecoder>(
         decoder: sending Decoder,
+        logger: Logger = Logger(label: "LambdaRuntime"),
         body: sending @escaping (Event, LambdaContext) async throws -> Void
     )
     where
@@ -239,6 +248,6 @@ extension LambdaRuntime {
             handler: LambdaHandlerAdapter(handler: ClosureHandler(body: body))
         )
 
-        self.init(handler: handler)
+        self.init(handler: handler, logger: logger)
     }
 }

Sources/AWSLambdaRuntime/LambdaRuntime+ServiceLifecycle.swift

@@ -15,5 +15,11 @@
 #if ServiceLifecycleSupport
 import ServiceLifecycle
 
-extension LambdaRuntime: Service {}
+extension LambdaRuntime: Service {
+    public func run() async throws {
+        try await cancelWhenGracefulShutdown {
+            try await self._run()
+        }
+    }
+}
 #endif

Sources/AWSLambdaRuntime/LambdaRuntime.swift

@@ -46,13 +46,24 @@ public final class LambdaRuntime<Handler>: @unchecked Sendable where Handler: St
         // developers have to wait for AWS Lambda to dispose and recreate a runtime environment to pickup a change
         // this approach is less flexible but more performant than reading the value of the environment variable at each invocation
         var log = logger
-        log.logLevel = Lambda.env("LOG_LEVEL").flatMap(Logger.Level.init) ?? .info
+
+        // use the LOG_LEVEL environment variable to set the log level.
+        // if the environment variable is not set, use the default log level from the logger provided
+        log.logLevel = Lambda.env("LOG_LEVEL").flatMap(Logger.Level.init) ?? logger.logLevel
+
         self.logger = log
         self.logger.debug("LambdaRuntime initialized")
     }
 
+    #if !ServiceLifecycleSupport
+    @inlinable
+    internal func run() async throws {
+        try await _run()
+    }
+    #endif
+
     @inlinable
-    public func run() async throws {
+    internal func _run() async throws {
         let handler = self.handlerMutex.withLockedValue { handler in
             let result = handler
             handler = nil