[core] Only one LambdaRuntime.run() can be called at a time (fix #507)

Sources/AWSLambdaRuntime/LambdaRuntime.swift

@@ -15,16 +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 running at the time
+// We use an Atomic here to ensure thread safety
+private let _isRunning = Atomic<Bool>(false)
+
 // 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.
+// sadly crashes the compiler today (on Linux).
 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
     let handlerMutex: NIOLockedValueBox<Handler?>
@@ -48,7 +53,25 @@ public final class LambdaRuntime<Handler>: @unchecked Sendable where Handler: St
         self.logger.debug("LambdaRuntime initialized")
     }
 
+    /// Make sure only one run() is called at a time
     public func run() async throws {
+
+        // 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: .moreThanOneLambdaRuntimeInstance)
+        }
+
+        defer {
+            _isRunning.store(false, ordering: .releasing)
+        }
+
+        try await self._run()
+    }
+
+    private func _run() async throws {
         let handler = self.handlerMutex.withLockedValue { handler in
             let result = handler
             handler = nil

Sources/AWSLambdaRuntime/LambdaRuntimeError.swift

@@ -12,8 +12,10 @@
 //
 //===----------------------------------------------------------------------===//
 
-package struct LambdaRuntimeError: Error {
-    package enum Code {
+public struct LambdaRuntimeError: Error {
+    package enum Code: Sendable {
+
+        /// internal error codes for LambdaRuntimeClient
         case closingRuntimeClient
 
         case connectionToControlPlaneLost
@@ -32,14 +34,17 @@ package struct LambdaRuntimeError: Error {
         case missingLambdaRuntimeAPIEnvironmentVariable
         case runtimeCanOnlyBeStartedOnce
         case invalidPort
+
+        /// public error codes for LambdaRuntime
+        case moreThanOneLambdaRuntimeInstance
     }
 
     package init(code: Code, underlying: (any Error)? = nil) {
         self.code = code
         self.underlying = underlying
     }
 
-    package var code: Code
-    package var underlying: (any Error)?
+    var code: Code
+    public var underlying: (any Error)?
 
 }

Tests/AWSLambdaRuntimeTests/LambdaRuntimeTests.swift

@@ -0,0 +1,84 @@
+//===----------------------------------------------------------------------===//
+//
+// 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: "Runtime1")
+        )
+
+        // Second runtime
+        let runtime2 = LambdaRuntime(
+            handler: MockHandler(),
+            eventLoop: Lambda.defaultEventLoop,
+            logger: Logger(label: "Runtime1")
+        )
+
+        try await withThrowingTaskGroup(of: Void.self) { taskGroup in
+            // start the first runtime
+            taskGroup.addTask {
+                await #expect(throws: Never.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 {
+                await #expect(throws: Never.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 {
+
+    }
+}