diff --git a/Sources/AWSLambdaRuntime/LambdaRuntime.swift b/Sources/AWSLambdaRuntime/LambdaRuntime.swift index 6bc2403c..b0ebe0ca 100644 --- a/Sources/AWSLambdaRuntime/LambdaRuntime.swift +++ b/Sources/AWSLambdaRuntime/LambdaRuntime.swift @@ -15,6 +15,7 @@ import Logging import NIOConcurrencyHelpers import NIOCore +import Synchronization #if canImport(FoundationEssentials) import FoundationEssentials @@ -22,9 +23,13 @@ import FoundationEssentials 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(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: @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 @@ -48,7 +53,25 @@ public final class LambdaRuntime: @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 diff --git a/Sources/AWSLambdaRuntime/LambdaRuntimeError.swift b/Sources/AWSLambdaRuntime/LambdaRuntimeError.swift index a6b4ac66..c4166731 100644 --- a/Sources/AWSLambdaRuntime/LambdaRuntimeError.swift +++ b/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,6 +34,9 @@ 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) { @@ -39,7 +44,7 @@ package struct LambdaRuntimeError: Error { self.underlying = underlying } - package var code: Code - package var underlying: (any Error)? + var code: Code + public var underlying: (any Error)? } diff --git a/Tests/AWSLambdaRuntimeTests/LambdaRuntimeTests.swift b/Tests/AWSLambdaRuntimeTests/LambdaRuntimeTests.swift new file mode 100644 index 00000000..bd3ca6d3 --- /dev/null +++ b/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 { + + } +}