Merge branch 'swift-server:main' into sebsto/servicelifecycle

Author: sebsto

Examples/HelloJSON/Package.swift

@@ -15,10 +15,7 @@ let package = Package(
         // during CI, the dependency on local version of swift-aws-lambda-runtime is added dynamically below
         .package(
             url: "https://github.com/swift-server/swift-aws-lambda-runtime.git",
-            branch: "ff-package-traits",
-            traits: [
-                .trait(name: "FoundationJSONSupport")
-            ]
+            branch: "main"
         )
     ],
     targets: [

Plugins/Documentation.docc/Proposals/0001-v2-plugins.md

@@ -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/FoundationSupport/Lambda+JSON.swift

@@ -48,11 +48,9 @@ extension Lambda {
     @usableFromInline
     static func withLocalServer(
         invocationEndpoint: String? = nil,
+        logger: Logger,
         _ body: sending @escaping () async throws -> Void
     ) async throws {
-        var logger = Logger(label: "LocalServer")
-        logger.logLevel = Lambda.env("LOG_LEVEL").flatMap(Logger.Level.init) ?? .info
-
         try await LambdaHTTPServer.withLocalServer(
             invocationEndpoint: invocationEndpoint,
             logger: logger
@@ -93,7 +91,7 @@ internal struct LambdaHTTPServer {
         case serverReturned(Swift.Result<Void, any Error>)
     }
 
-    struct UnsafeTransferBox<Value>: @unchecked Sendable {
+    fileprivate struct UnsafeTransferBox<Value>: @unchecked Sendable {
         let value: Value
 
         init(value: sending Value) {
@@ -133,6 +131,7 @@ internal struct LambdaHTTPServer {
                 }
             }
 
+        // it's ok to keep this at `info` level because it is only used for local testing and unit tests
         logger.info(
             "Server started and listening",
             metadata: [
@@ -202,12 +201,18 @@ internal struct LambdaHTTPServer {
                 return result
 
             case .serverReturned(let result):
-                logger.error(
-                    "Server shutdown before closure completed",
-                    metadata: [
-                        "error": "\(result.maybeError != nil ? "\(result.maybeError!)" : "none")"
-                    ]
-                )
+
+                if result.maybeError is CancellationError {
+                    logger.trace("Server's task cancelled")
+                } else {
+                    logger.error(
+                        "Server shutdown before closure completed",
+                        metadata: [
+                            "error": "\(result.maybeError != nil ? "\(result.maybeError!)" : "none")"
+                        ]
+                    )
+                }
+
                 switch await group.next()! {
                 case .closureResult(let result):
                     return result
@@ -265,9 +270,12 @@ internal struct LambdaHTTPServer {
                         }
                     }
                 }
+            } catch let error as CancellationError {
+                logger.trace("The task was cancelled", metadata: ["error": "\(error)"])
             } catch {
                 logger.error("Hit error: \(error)")
             }
+
         } onCancel: {
             channel.channel.close(promise: nil)
         }

Sources/AWSLambdaRuntime/Lambda+LocalServer.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.
@@ -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/LambdaHandlers.swift

@@ -13,22 +13,23 @@
 //===----------------------------------------------------------------------===//
 
 import Logging
-import NIOConcurrencyHelpers
 import NIOCore
+import Synchronization
 
 #if canImport(FoundationEssentials)
 import FoundationEssentials
 #else
 import Foundation
 #endif
 
-// We need `@unchecked` Sendable here, as `NIOLockedValueBox` does not understand `sending` today.
-// We don't want to use `NIOLockedValueBox` here anyway. We would love to use Mutex here, but this
-// sadly crashes the compiler today.
-public final class LambdaRuntime<Handler>: @unchecked Sendable where Handler: StreamingLambdaHandler {
-    // TODO: We want to change this to Mutex as soon as this doesn't crash the Swift compiler on Linux anymore
+// This is our guardian to ensure only one LambdaRuntime is running at the time
+// We use an Atomic here to ensure thread safety
+private let _isRunning = Atomic<Bool>(false)
+
+public final class LambdaRuntime<Handler>: Sendable where Handler: StreamingLambdaHandler {
     @usableFromInline
-    let handlerMutex: NIOLockedValueBox<Handler?>
+    /// we protect the handler behind a Mutex to ensure that we only ever have one copy of it
+    let handlerStorage: SendingStorage<Handler>
     @usableFromInline
     let logger: Logger
     @usableFromInline
@@ -39,14 +40,18 @@ public final class LambdaRuntime<Handler>: @unchecked Sendable where Handler: St
         eventLoop: EventLoop = Lambda.defaultEventLoop,
         logger: Logger = Logger(label: "LambdaRuntime")
     ) {
-        self.handlerMutex = NIOLockedValueBox(handler)
+        self.handlerStorage = SendingStorage(handler)
         self.eventLoop = eventLoop
 
         // by setting the log level here, we understand it can not be changed dynamically at runtime
         // 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")
     }
@@ -58,14 +63,24 @@ public final class LambdaRuntime<Handler>: @unchecked Sendable where Handler: St
     }
     #endif
 
-    @inlinable
+    /// Make sure only one run() is called at a time
+    // @inlinable
     internal func _run() async throws {
-        let handler = self.handlerMutex.withLockedValue { handler in
-            let result = handler
-            handler = nil
-            return result
+
+        // we use an atomic global variable to ensure only one LambdaRuntime is running at the time
+        let (_, original) = _isRunning.compareExchange(expected: false, desired: true, ordering: .acquiringAndReleasing)
+
+        // if the original value was already true, run() is already running
+        if original {
+            throw LambdaRuntimeError(code: .runtimeCanOnlyBeStartedOnce)
         }
 
+        defer {
+            _isRunning.store(false, ordering: .releasing)
+        }
+
+        // The handler can be non-sendable, we want to ensure we only ever have one copy of it
+        let handler = try? self.handlerStorage.get()
         guard let handler else {
             throw LambdaRuntimeError(code: .runtimeCanOnlyBeStartedOnce)
         }
@@ -96,8 +111,10 @@ public final class LambdaRuntime<Handler>: @unchecked Sendable where Handler: St
             #if LocalServerSupport
             // we're not running on Lambda and we're compiled in DEBUG mode,
             // let's start a local server for testing
-            try await Lambda.withLocalServer(invocationEndpoint: Lambda.env("LOCAL_LAMBDA_SERVER_INVOCATION_ENDPOINT"))
-            {
+            try await Lambda.withLocalServer(
+                invocationEndpoint: Lambda.env("LOCAL_LAMBDA_SERVER_INVOCATION_ENDPOINT"),
+                logger: self.logger
+            ) {
 
                 try await LambdaRuntimeClient.withRuntimeClient(
                     configuration: .init(ip: "127.0.0.1", port: 7000),

Sources/AWSLambdaRuntime/LambdaRuntime.swift

@@ -16,6 +16,7 @@
 package struct LambdaRuntimeError: Error {
     @usableFromInline
     package enum Code: Sendable {
+        /// internal error codes for LambdaRuntimeClient
         case closingRuntimeClient
 
         case connectionToControlPlaneLost

Sources/AWSLambdaRuntime/LambdaRuntimeError.swift

@@ -13,8 +13,11 @@
 //===----------------------------------------------------------------------===//
 
 import Dispatch
+import NIOConcurrencyHelpers
 import NIOPosix
 
+// import Synchronization
+
 enum Consts {
     static let apiPrefix = "/2018-06-01"
     static let invocationURLPrefix = "\(apiPrefix)/runtime/invocation"
@@ -36,28 +39,6 @@ enum AmazonHeaders {
     static let invokedFunctionARN = "Lambda-Runtime-Invoked-Function-Arn"
 }
 
-/// Helper function to trap signals
-func trap(signal sig: Signal, handler: @escaping (Signal) -> Void) -> DispatchSourceSignal {
-    let signalSource = DispatchSource.makeSignalSource(signal: sig.rawValue, queue: DispatchQueue.global())
-    signal(sig.rawValue, SIG_IGN)
-    signalSource.setEventHandler(handler: {
-        signalSource.cancel()
-        handler(sig)
-    })
-    signalSource.resume()
-    return signalSource
-}
-
-enum Signal: Int32 {
-    case HUP = 1
-    case INT = 2
-    case QUIT = 3
-    case ABRT = 6
-    case KILL = 9  // ignore-unacceptable-language
-    case ALRM = 14
-    case TERM = 15
-}
-
 extension DispatchWallTime {
     @usableFromInline
     init(millisSinceEpoch: Int64) {
@@ -132,3 +113,34 @@ extension AmazonHeaders {
         return "\(version)-\(datePadding)\(dateValue)-\(identifier)"
     }
 }
+
+/// Temporary storage for value being sent from one isolation domain to another
+// use NIOLockedValueBox instead of Mutex to avoid compiler crashes on 6.0
+// see https://github.com/swiftlang/swift/issues/78048
+@usableFromInline
+struct SendingStorage<Value>: ~Copyable, @unchecked Sendable {
+    @usableFromInline
+    struct ValueAlreadySentError: Error {
+        @usableFromInline
+        init() {}
+    }
+
+    @usableFromInline
+    // let storage: Mutex<Value?>
+    let storage: NIOLockedValueBox<Value?>
+
+    @inlinable
+    init(_ value: sending Value) {
+        self.storage = .init(value)
+    }
+
+    @inlinable
+    func get() throws -> Value {
+        // try self.storage.withLock {
+        try self.storage.withLockedValue {
+            guard let value = $0 else { throw ValueAlreadySentError() }
+            $0 = nil
+            return value
+        }
+    }
+}

Sources/AWSLambdaRuntime/Utils.swift

@@ -123,6 +123,7 @@ struct HttpServer {
                 }
             }
         }
+        // it's ok to keep this at `info` level because it is only used for local testing and unit tests
         logger.info("Server shutting down")
     }
 

Sources/MockServer/MockHTTPServer.swift

@@ -147,12 +147,11 @@ struct LambdaRuntimeClientTests {
             (event: String, context: LambdaContext) in
             "Hello \(event)"
         }
-        var logger = Logger(label: "LambdaRuntime")
-        logger.logLevel = .debug
+
         let serviceGroup = ServiceGroup(
             services: [runtime],
             gracefulShutdownSignals: [.sigterm, .sigint],
-            logger: logger
+            logger: Logger(label: "TestLambdaRuntimeGracefulShutdown")
         )
         try await withThrowingTaskGroup(of: Void.self) { group in
             group.addTask {