Merge branch 'main' into sebsto/localserver+streaming

sebsto · sebsto · commit 5a4415ccf22d · 2025-07-10T19:09:17.000+01:00
diff --git a/Examples/HelloJSON/Package.swift b/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: [
diff --git a/Sources/AWSLambdaRuntime/Lambda+LocalServer.swift b/Sources/AWSLambdaRuntime/Lambda+LocalServer.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
@@ -279,9 +284,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)
         }
diff --git a/Sources/AWSLambdaRuntime/LambdaRuntime.swift b/Sources/AWSLambdaRuntime/LambdaRuntime.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,7 +40,7 @@ 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
@@ -62,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)
         }
@@ -100,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),
diff --git a/Sources/AWSLambdaRuntime/LambdaRuntimeError.swift b/Sources/AWSLambdaRuntime/LambdaRuntimeError.swift
@@ -16,6 +16,7 @@
 package struct LambdaRuntimeError: Error {
     @usableFromInline
     package enum Code: Sendable {
+        /// internal error codes for LambdaRuntimeClient
         case closingRuntimeClient
 
         case connectionToControlPlaneLost
diff --git a/Sources/AWSLambdaRuntime/Utils.swift b/Sources/AWSLambdaRuntime/Utils.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
+        }
+    }
+}
diff --git a/Sources/MockServer/MockHTTPServer.swift b/Sources/MockServer/MockHTTPServer.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")
     }
 
diff --git a/Tests/AWSLambdaRuntimeTests/LambdaRuntimeClientTests.swift b/Tests/AWSLambdaRuntimeTests/LambdaRuntimeClientTests.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 {
diff --git a/Tests/AWSLambdaRuntimeTests/LambdaRuntimeTests.swift b/Tests/AWSLambdaRuntimeTests/LambdaRuntimeTests.swift
@@ -0,0 +1,88 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftAWSLambdaRuntime open source project
+//
+// Copyright (c) 2025 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 Foundation
+import Logging
+import NIOCore
+import Synchronization
+import Testing
+
+@testable import AWSLambdaRuntime
+
+@Suite("LambdaRuntimeTests")
+struct LambdaRuntimeTests {
+
+    @Test("LambdaRuntime can only be run once")
+    func testLambdaRuntimerunOnce() async throws {
+
+        // First runtime
+        let runtime1 = LambdaRuntime(
+            handler: MockHandler(),
+            eventLoop: Lambda.defaultEventLoop,
+            logger: Logger(label: "LambdaRuntimeTests.Runtime1")
+        )
+
+        // Second runtime
+        let runtime2 = LambdaRuntime(
+            handler: MockHandler(),
+            eventLoop: Lambda.defaultEventLoop,
+            logger: Logger(label: "LambdaRuntimeTests.Runtime2")
+        )
+
+        try await withThrowingTaskGroup(of: Void.self) { taskGroup in
+            // start the first runtime
+            taskGroup.addTask {
+                // ChannelError will be thrown when we cancel the task group
+                await #expect(throws: ChannelError.self) {
+                    try await runtime1.run()
+                }
+            }
+
+            // wait a small amount to ensure runtime1 task is started
+            try await Task.sleep(for: .seconds(1))
+
+            // Running the second runtime should trigger LambdaRuntimeError
+            await #expect(throws: LambdaRuntimeError.self) {
+                try await runtime2.run()
+            }
+
+            // cancel runtime 1 / task 1
+            taskGroup.cancelAll()
+        }
+
+        // Running the second runtime should work now
+        try await withThrowingTaskGroup(of: Void.self) { taskGroup in
+            taskGroup.addTask {
+                // ChannelError will be thrown when we cancel the task group
+                await #expect(throws: ChannelError.self) {
+                    try await runtime2.run()
+                }
+            }
+
+            // Set timeout and cancel the runtime 2
+            try await Task.sleep(for: .seconds(2))
+            taskGroup.cancelAll()
+        }
+    }
+}
+
+struct MockHandler: StreamingLambdaHandler {
+    mutating func handle(
+        _ event: NIOCore.ByteBuffer,
+        responseWriter: some AWSLambdaRuntime.LambdaResponseStreamWriter,
+        context: AWSLambdaRuntime.LambdaContext
+    ) async throws {
+
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -123,6 +123,7 @@ struct HttpServer {`
`123`	`123`	`}`
`124`	`124`	`}`
`125`	`125`	`}`
	`126`	+ // it's ok to keep this at `info` level because it is only used for local testing and unit tests
`126`	`127`	`logger.info("Server shutting down")`
`127`	`128`	`}`
`128`	`129`