make LambdaRuntime a singleton without breaking the API

sebsto · sebsto · commit 456569a51c82 · 2025-03-14T13:49:07.000+01:00
diff --git a/Sources/AWSLambdaRuntime/ControlPlaneRequest.swift b/Sources/AWSLambdaRuntime/ControlPlaneRequest.swift
@@ -36,19 +36,19 @@ package struct InvocationMetadata: Hashable {
     package let clientContext: String?
     package let cognitoIdentity: String?
 
-    package init(headers: HTTPHeaders) throws(LambdaRuntimeError) {
+    package init(headers: HTTPHeaders) throws(LambdaRuntimeClientError) {
         guard let requestID = headers.first(name: AmazonHeaders.requestID), !requestID.isEmpty else {
-            throw LambdaRuntimeError(code: .nextInvocationMissingHeaderRequestID)
+            throw LambdaRuntimeClientError(code: .nextInvocationMissingHeaderRequestID)
         }
 
         guard let deadline = headers.first(name: AmazonHeaders.deadline),
             let unixTimeInMilliseconds = Int64(deadline)
         else {
-            throw LambdaRuntimeError(code: .nextInvocationMissingHeaderDeadline)
+            throw LambdaRuntimeClientError(code: .nextInvocationMissingHeaderDeadline)
         }
 
         guard let invokedFunctionARN = headers.first(name: AmazonHeaders.invokedFunctionARN) else {
-            throw LambdaRuntimeError(code: .nextInvocationMissingHeaderInvokeFuctionARN)
+            throw LambdaRuntimeClientError(code: .nextInvocationMissingHeaderInvokeFuctionARN)
         }
 
         self.requestID = requestID
diff --git a/Sources/AWSLambdaRuntime/FoundationSupport/Lambda+JSON.swift b/Sources/AWSLambdaRuntime/FoundationSupport/Lambda+JSON.swift
@@ -108,7 +108,11 @@ extension LambdaRuntime {
             handler: LambdaHandlerAdapter(handler: ClosureHandler(body: body))
         )
 
-        self.init(handler: handler)
+        do {
+            try self.init(handler: handler)
+        } catch {
+            fatalError("Failed to initialize LambdaRuntime: \(error)")
+        }
     }
 
     /// Initialize an instance with a `LambdaHandler` defined in the form of a closure **with a `Void` return type**.
@@ -132,7 +136,11 @@ extension LambdaRuntime {
             handler: LambdaHandlerAdapter(handler: ClosureHandler(body: body))
         )
 
-        self.init(handler: handler)
+        do {
+            try self.init(handler: handler)
+        } catch {
+            fatalError("Failed to initialize LambdaRuntime: \(error)")
+        }
     }
 }
 #endif  // trait: FoundationJSONSupport
diff --git a/Sources/AWSLambdaRuntime/LambdaHandlers.swift b/Sources/AWSLambdaRuntime/LambdaHandlers.swift
@@ -179,7 +179,11 @@ extension LambdaRuntime {
     public convenience init(
         body: @Sendable @escaping (ByteBuffer, LambdaResponseStreamWriter, LambdaContext) async throws -> Void
     ) where Handler == StreamingClosureHandler {
-        self.init(handler: StreamingClosureHandler(body: body))
+        do {
+            try self.init(handler: StreamingClosureHandler(body: body))
+        } catch {
+            fatalError("Failed to initialize LambdaRuntime: \(error)")
+        }
     }
 
     /// Initialize an instance with a ``LambdaHandler`` defined in the form of a closure **with a non-`Void` return type**, an encoder, and a decoder.
@@ -213,8 +217,11 @@ extension LambdaRuntime {
             decoder: decoder,
             handler: streamingAdapter
         )
-
-        self.init(handler: codableWrapper)
+        do {
+            try self.init(handler: codableWrapper)
+        } catch {
+            fatalError("Failed to initialize LambdaRuntime: \(error)")
+        }
     }
 
     /// Initialize an instance with a ``LambdaHandler`` defined in the form of a closure **with a `Void` return type**, an encoder, and a decoder.
@@ -238,7 +245,10 @@ extension LambdaRuntime {
             decoder: decoder,
             handler: LambdaHandlerAdapter(handler: ClosureHandler(body: body))
         )
-
-        self.init(handler: handler)
+        do {
+            try self.init(handler: handler)
+        } catch {
+            fatalError("Failed to initialize LambdaRuntime: \(error)")
+        }
     }
 }
diff --git a/Sources/AWSLambdaRuntime/LambdaRuntime.swift b/Sources/AWSLambdaRuntime/LambdaRuntime.swift
@@ -15,13 +15,21 @@
 import Logging
 import NIOConcurrencyHelpers
 import NIOCore
+import Synchronization
 
 #if canImport(FoundationEssentials)
 import FoundationEssentials
 #else
 import Foundation
 #endif
 
+// This is our gardian to ensure only one LambdaRuntime is initialized
+// We use a Mutex here to ensure thread safety
+// We don't use LambdaRuntime<> as the type here, as we don't know the concrete type that will be used
+private let _singleton = Mutex<Bool>(false)
+public enum LambdaRuntimeError: Error {
+    case moreThanOneLambdaRuntimeInstance
+}
 // 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.
@@ -35,7 +43,22 @@ public final class LambdaRuntime<Handler>: @unchecked Sendable where Handler: St
         handler: sending Handler,
         eventLoop: EventLoop = Lambda.defaultEventLoop,
         logger: Logger = Logger(label: "LambdaRuntime")
-    ) {
+    ) throws(LambdaRuntimeError) {
+
+        do {
+            try _singleton.withLock {
+                let alreadyCreated = $0
+                guard alreadyCreated == false else {
+                    throw LambdaRuntimeError.moreThanOneLambdaRuntimeInstance
+                }
+                $0 = true
+            }
+        } catch _ as LambdaRuntimeError {
+            throw LambdaRuntimeError.moreThanOneLambdaRuntimeInstance
+        } catch {
+            fatalError("An unknown error occurred: \(error)")
+        }
+
         self.handlerMutex = NIOLockedValueBox(handler)
         self.eventLoop = eventLoop
 
@@ -56,7 +79,7 @@ public final class LambdaRuntime<Handler>: @unchecked Sendable where Handler: St
         }
 
         guard let handler else {
-            throw LambdaRuntimeError(code: .runtimeCanOnlyBeStartedOnce)
+            throw LambdaRuntimeClientError(code: .runtimeCanOnlyBeStartedOnce)
         }
 
         // are we running inside an AWS Lambda runtime environment ?
@@ -66,7 +89,7 @@ public final class LambdaRuntime<Handler>: @unchecked Sendable where Handler: St
 
             let ipAndPort = runtimeEndpoint.split(separator: ":", maxSplits: 1)
             let ip = String(ipAndPort[0])
-            guard let port = Int(ipAndPort[1]) else { throw LambdaRuntimeError(code: .invalidPort) }
+            guard let port = Int(ipAndPort[1]) else { throw LambdaRuntimeClientError(code: .invalidPort) }
 
             try await LambdaRuntimeClient.withRuntimeClient(
                 configuration: .init(ip: ip, port: port),
diff --git a/Sources/AWSLambdaRuntime/LambdaRuntimeClient.swift b/Sources/AWSLambdaRuntime/LambdaRuntimeClient.swift
@@ -134,7 +134,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
 
             case .connecting(let continuations):
                 for continuation in continuations {
-                    continuation.resume(throwing: LambdaRuntimeError(code: .closingRuntimeClient))
+                    continuation.resume(throwing: LambdaRuntimeClientError(code: .closingRuntimeClient))
                 }
                 self.connectionState = .connecting([])
 
@@ -173,7 +173,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
     private func write(_ buffer: NIOCore.ByteBuffer) async throws {
         switch self.lambdaState {
         case .idle, .sentResponse:
-            throw LambdaRuntimeError(code: .writeAfterFinishHasBeenSent)
+            throw LambdaRuntimeClientError(code: .writeAfterFinishHasBeenSent)
 
         case .waitingForNextInvocation:
             fatalError("Invalid state: \(self.lambdaState)")
@@ -194,7 +194,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
     private func writeAndFinish(_ buffer: NIOCore.ByteBuffer?) async throws {
         switch self.lambdaState {
         case .idle, .sentResponse:
-            throw LambdaRuntimeError(code: .finishAfterFinishHasBeenSent)
+            throw LambdaRuntimeClientError(code: .finishAfterFinishHasBeenSent)
 
         case .waitingForNextInvocation:
             fatalError("Invalid state: \(self.lambdaState)")
@@ -261,7 +261,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
         case (.connecting(let array), .notClosing):
             self.connectionState = .disconnected
             for continuation in array {
-                continuation.resume(throwing: LambdaRuntimeError(code: .lostConnectionToControlPlane))
+                continuation.resume(throwing: LambdaRuntimeClientError(code: .lostConnectionToControlPlane))
             }
 
         case (.connecting(let array), .closing(let continuation)):
@@ -394,7 +394,7 @@ extension LambdaRuntimeClient: LambdaChannelHandlerDelegate {
                 }
 
                 for continuation in continuations {
-                    continuation.resume(throwing: LambdaRuntimeError(code: .connectionToControlPlaneLost))
+                    continuation.resume(throwing: LambdaRuntimeClientError(code: .connectionToControlPlaneLost))
                 }
 
             case .connected(let stateChannel, _):
@@ -489,7 +489,7 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
             fatalError("Invalid state: \(self.state)")
 
         case .disconnected:
-            throw LambdaRuntimeError(code: .connectionToControlPlaneLost)
+            throw LambdaRuntimeClientError(code: .connectionToControlPlaneLost)
         }
     }
 
@@ -528,10 +528,10 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
             )
 
         case .disconnected:
-            throw LambdaRuntimeError(code: .connectionToControlPlaneLost)
+            throw LambdaRuntimeClientError(code: .connectionToControlPlaneLost)
 
         case .closing:
-            throw LambdaRuntimeError(code: .connectionToControlPlaneGoingAway)
+            throw LambdaRuntimeClientError(code: .connectionToControlPlaneGoingAway)
         }
     }
 
@@ -553,13 +553,13 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
 
         case .connected(_, .idle),
             .connected(_, .sentResponse):
-            throw LambdaRuntimeError(code: .writeAfterFinishHasBeenSent)
+            throw LambdaRuntimeClientError(code: .writeAfterFinishHasBeenSent)
 
         case .disconnected:
-            throw LambdaRuntimeError(code: .connectionToControlPlaneLost)
+            throw LambdaRuntimeClientError(code: .connectionToControlPlaneLost)
 
         case .closing:
-            throw LambdaRuntimeError(code: .connectionToControlPlaneGoingAway)
+            throw LambdaRuntimeClientError(code: .connectionToControlPlaneGoingAway)
         }
     }
 
@@ -586,13 +586,13 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
             }
 
         case .connected(_, .sentResponse):
-            throw LambdaRuntimeError(code: .finishAfterFinishHasBeenSent)
+            throw LambdaRuntimeClientError(code: .finishAfterFinishHasBeenSent)
 
         case .disconnected:
-            throw LambdaRuntimeError(code: .connectionToControlPlaneLost)
+            throw LambdaRuntimeClientError(code: .connectionToControlPlaneLost)
 
         case .closing:
-            throw LambdaRuntimeError(code: .connectionToControlPlaneGoingAway)
+            throw LambdaRuntimeClientError(code: .connectionToControlPlaneGoingAway)
         }
     }
 
@@ -759,7 +759,7 @@ extension LambdaChannelHandler: ChannelInboundHandler {
                 self.delegate.connectionWillClose(channel: context.channel)
                 context.close(promise: nil)
                 continuation.resume(
-                    throwing: LambdaRuntimeError(code: .invocationMissingMetadata, underlying: error)
+                    throwing: LambdaRuntimeClientError(code: .invocationMissingMetadata, underlying: error)
                 )
             }
 
@@ -769,7 +769,7 @@ extension LambdaChannelHandler: ChannelInboundHandler {
                 continuation.resume()
             } else {
                 self.state = .connected(context, .idle)
-                continuation.resume(throwing: LambdaRuntimeError(code: .unexpectedStatusCodeForRequest))
+                continuation.resume(throwing: LambdaRuntimeClientError(code: .unexpectedStatusCodeForRequest))
             }
 
         case .disconnected, .closing, .connected(_, _):
diff --git a/Sources/AWSLambdaRuntime/LambdaRuntimeClientError.swift b/Sources/AWSLambdaRuntime/LambdaRuntimeClientError.swift
@@ -12,7 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-package struct LambdaRuntimeError: Error {
+package struct LambdaRuntimeClientError: Error {
     package enum Code {
         case closingRuntimeClient
 
diff --git a/Tests/AWSLambdaRuntimeTests/LambdaRuntimeTests.swift b/Tests/AWSLambdaRuntimeTests/LambdaRuntimeTests.swift
@@ -0,0 +1,38 @@
+import Foundation
+import Logging
+import NIOCore
+import Synchronization
+import Testing
+
+@testable import AWSLambdaRuntime
+
+@Suite("LambdaRuntimeTests")
+final class LambdaRuntimeTests {
+
+    @Test("LambdaRuntime can only be initialized once")
+    func testLambdaRuntimeInitializationFatalError() throws {
+
+        // First initialization should succeed
+        try _ = LambdaRuntime(handler: MockHandler(), eventLoop: Lambda.defaultEventLoop, logger: Logger(label: "Test"))
+
+        // Second initialization should trigger LambdaRuntimeError
+        #expect(throws: LambdaRuntimeError.self) {
+            try _ = LambdaRuntime(
+                handler: MockHandler(),
+                eventLoop: Lambda.defaultEventLoop,
+                logger: Logger(label: "Test")
+            )
+        }
+
+    }
+}
+
+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
`@@ -108,7 +108,11 @@ extension LambdaRuntime {`
`108`	`108`	`handler: LambdaHandlerAdapter(handler: ClosureHandler(body: body))`
`109`	`109`	`)`
`110`	`110`
`111`		`- self.init(handler: handler)`
	`111`	`+ do {`
	`112`	`+ try self.init(handler: handler)`
	`113`	`+ } catch {`
	`114`	`+ fatalError("Failed to initialize LambdaRuntime: \(error)")`
	`115`	`+ }`
`112`	`116`	`}`
`113`	`117`
`114`	`118`	/// Initialize an instance with a `LambdaHandler` defined in the form of a closure with a `Void` return type.
`@@ -132,7 +136,11 @@ extension LambdaRuntime {`
`132`	`136`	`handler: LambdaHandlerAdapter(handler: ClosureHandler(body: body))`
`133`	`137`	`)`
`134`	`138`
`135`		`- self.init(handler: handler)`
	`139`	`+ do {`
	`140`	`+ try self.init(handler: handler)`
	`141`	`+ } catch {`
	`142`	`+ fatalError("Failed to initialize LambdaRuntime: \(error)")`
	`143`	`+ }`
`136`	`144`	`}`
`137`	`145`	`}`
`138`	`146`	`#endif // trait: FoundationJSONSupport`
Original file line number	Diff line number	Diff line change
`@@ -134,7 +134,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {`
`134`	`134`
`135`	`135`	`case .connecting(let continuations):`
`136`	`136`	`for continuation in continuations {`
`137`		`- continuation.resume(throwing: LambdaRuntimeError(code: .closingRuntimeClient))`
	`137`	`+ continuation.resume(throwing: LambdaRuntimeClientError(code: .closingRuntimeClient))`
`138`	`138`	`}`
`139`	`139`	`self.connectionState = .connecting([])`
`140`	`140`
`@@ -173,7 +173,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {`
`173`	`173`	`private func write(_ buffer: NIOCore.ByteBuffer) async throws {`
`174`	`174`	`switch self.lambdaState {`
`175`	`175`	`case .idle, .sentResponse:`
`176`		`- throw LambdaRuntimeError(code: .writeAfterFinishHasBeenSent)`
	`176`	`+ throw LambdaRuntimeClientError(code: .writeAfterFinishHasBeenSent)`
`177`	`177`
`178`	`178`	`case .waitingForNextInvocation:`
`179`	`179`	`fatalError("Invalid state: \(self.lambdaState)")`
`@@ -194,7 +194,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {`
`194`	`194`	`private func writeAndFinish(_ buffer: NIOCore.ByteBuffer?) async throws {`
`195`	`195`	`switch self.lambdaState {`
`196`	`196`	`case .idle, .sentResponse:`
`197`		`- throw LambdaRuntimeError(code: .finishAfterFinishHasBeenSent)`
	`197`	`+ throw LambdaRuntimeClientError(code: .finishAfterFinishHasBeenSent)`
`198`	`198`
`199`	`199`	`case .waitingForNextInvocation:`
`200`	`200`	`fatalError("Invalid state: \(self.lambdaState)")`
`@@ -261,7 +261,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {`
`261`	`261`	`case (.connecting(let array), .notClosing):`
`262`	`262`	`self.connectionState = .disconnected`
`263`	`263`	`for continuation in array {`
`264`		`- continuation.resume(throwing: LambdaRuntimeError(code: .lostConnectionToControlPlane))`
	`264`	`+ continuation.resume(throwing: LambdaRuntimeClientError(code: .lostConnectionToControlPlane))`
`265`	`265`	`}`
`266`	`266`
`267`	`267`	`case (.connecting(let array), .closing(let continuation)):`
`@@ -394,7 +394,7 @@ extension LambdaRuntimeClient: LambdaChannelHandlerDelegate {`
`394`	`394`	`}`
`395`	`395`
`396`	`396`	`for continuation in continuations {`
`397`		`- continuation.resume(throwing: LambdaRuntimeError(code: .connectionToControlPlaneLost))`
	`397`	`+ continuation.resume(throwing: LambdaRuntimeClientError(code: .connectionToControlPlaneLost))`
`398`	`398`	`}`
`399`	`399`
`400`	`400`	`case .connected(let stateChannel, _):`
`@@ -489,7 +489,7 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>`
`489`	`489`	`fatalError("Invalid state: \(self.state)")`
`490`	`490`
`491`	`491`	`case .disconnected:`
`492`		`- throw LambdaRuntimeError(code: .connectionToControlPlaneLost)`
	`492`	`+ throw LambdaRuntimeClientError(code: .connectionToControlPlaneLost)`
`493`	`493`	`}`
`494`	`494`	`}`
`495`	`495`
`@@ -528,10 +528,10 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>`
`528`	`528`	`)`
`529`	`529`
`530`	`530`	`case .disconnected:`
`531`		`- throw LambdaRuntimeError(code: .connectionToControlPlaneLost)`
	`531`	`+ throw LambdaRuntimeClientError(code: .connectionToControlPlaneLost)`
`532`	`532`
`533`	`533`	`case .closing:`
`534`		`- throw LambdaRuntimeError(code: .connectionToControlPlaneGoingAway)`
	`534`	`+ throw LambdaRuntimeClientError(code: .connectionToControlPlaneGoingAway)`
`535`	`535`	`}`
`536`	`536`	`}`
`537`	`537`
`@@ -553,13 +553,13 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>`
`553`	`553`
`554`	`554`	`case .connected(_, .idle),`
`555`	`555`	`.connected(_, .sentResponse):`
`556`		`- throw LambdaRuntimeError(code: .writeAfterFinishHasBeenSent)`
	`556`	`+ throw LambdaRuntimeClientError(code: .writeAfterFinishHasBeenSent)`
`557`	`557`
`558`	`558`	`case .disconnected:`
`559`		`- throw LambdaRuntimeError(code: .connectionToControlPlaneLost)`
	`559`	`+ throw LambdaRuntimeClientError(code: .connectionToControlPlaneLost)`
`560`	`560`
`561`	`561`	`case .closing:`
`562`		`- throw LambdaRuntimeError(code: .connectionToControlPlaneGoingAway)`
	`562`	`+ throw LambdaRuntimeClientError(code: .connectionToControlPlaneGoingAway)`
`563`	`563`	`}`
`564`	`564`	`}`
`565`	`565`
`@@ -586,13 +586,13 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>`
`586`	`586`	`}`
`587`	`587`
`588`	`588`	`case .connected(_, .sentResponse):`
`589`		`- throw LambdaRuntimeError(code: .finishAfterFinishHasBeenSent)`
	`589`	`+ throw LambdaRuntimeClientError(code: .finishAfterFinishHasBeenSent)`
`590`	`590`
`591`	`591`	`case .disconnected:`
`592`		`- throw LambdaRuntimeError(code: .connectionToControlPlaneLost)`
	`592`	`+ throw LambdaRuntimeClientError(code: .connectionToControlPlaneLost)`
`593`	`593`
`594`	`594`	`case .closing:`
`595`		`- throw LambdaRuntimeError(code: .connectionToControlPlaneGoingAway)`
	`595`	`+ throw LambdaRuntimeClientError(code: .connectionToControlPlaneGoingAway)`
`596`	`596`	`}`
`597`	`597`	`}`
`598`	`598`
`@@ -759,7 +759,7 @@ extension LambdaChannelHandler: ChannelInboundHandler {`
`759`	`759`	`self.delegate.connectionWillClose(channel: context.channel)`
`760`	`760`	`context.close(promise: nil)`
`761`	`761`	`continuation.resume(`
`762`		`- throwing: LambdaRuntimeError(code: .invocationMissingMetadata, underlying: error)`
	`762`	`+ throwing: LambdaRuntimeClientError(code: .invocationMissingMetadata, underlying: error)`
`763`	`763`	`)`
`764`	`764`	`}`
`765`	`765`
`@@ -769,7 +769,7 @@ extension LambdaChannelHandler: ChannelInboundHandler {`
`769`	`769`	`continuation.resume()`
`770`	`770`	`} else {`
`771`	`771`	`self.state = .connected(context, .idle)`
`772`		`- continuation.resume(throwing: LambdaRuntimeError(code: .unexpectedStatusCodeForRequest))`
	`772`	`+ continuation.resume(throwing: LambdaRuntimeClientError(code: .unexpectedStatusCodeForRequest))`
`773`	`773`	`}`
`774`	`774`
`775`	`775`	`case .disconnected, .closing, .connected(_, _):`