Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 37 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
37 commits
Select commit Hold shift + click to select a range
456569a
make LambdaRuntime a singleton without breaking the API
sebsto Mar 14, 2025
a63a639
fix license header
sebsto Mar 14, 2025
40b6127
convert Mutex to NIOLockedValueBox
sebsto Mar 14, 2025
d474eba
Replace NIOLockedValueBox with Mutex
sebsto Mar 15, 2025
46e76f3
revert replacing NIOLockedValueBox by Mutex
sebsto Mar 15, 2025
299b9bc
remove typed throw (workaround for https://github.com/swiftlang/swift…
sebsto Mar 15, 2025
e3a3851
fix integration tests
sebsto Mar 15, 2025
4ebf24f
Replace NIOLockedValueBox with Mutex
sebsto Mar 15, 2025
f3e65ac
Merge branch 'main' into sebsto/fix_507
sebsto Mar 18, 2025
3aee427
use Atomic instead of Mutex.
sebsto Mar 19, 2025
e445f90
Merge branch 'sebsto/fix_507' of github.com:sebsto/swift-aws-lambda-r…
sebsto Mar 19, 2025
0509eaa
revert `try` on `runtime.init()` in doc
sebsto Mar 19, 2025
aa6395f
revert unwanted change
sebsto Mar 19, 2025
066ab76
revert unwanted change
sebsto Mar 19, 2025
e01b25d
swift-format
sebsto Mar 19, 2025
23c4b64
Merge branch 'main' into sebsto/fix_507
sebsto Mar 20, 2025
57e2396
Update Sources/AWSLambdaRuntime/LambdaRuntime.swift
sebsto Mar 21, 2025
85c988d
Update Sources/AWSLambdaRuntime/LambdaRuntime.swift
sebsto Mar 21, 2025
ab80580
Update Sources/AWSLambdaRuntime/LambdaRuntimeError.swift
sebsto Mar 21, 2025
10304ce
Update Sources/AWSLambdaRuntime/LambdaRuntimeError.swift
sebsto Mar 21, 2025
c1e68fc
Update Tests/AWSLambdaRuntimeTests/LambdaRuntimeTests.swift
sebsto Mar 21, 2025
918492a
Update Tests/AWSLambdaRuntimeTests/LambdaRuntimeTests.swift
sebsto Mar 21, 2025
6ada041
Update Tests/AWSLambdaRuntimeTests/LambdaRuntimeTests.swift
sebsto Mar 21, 2025
06df831
add package visibility to Code
sebsto Mar 21, 2025
50b9df8
swift format
sebsto Mar 24, 2025
fa0fd88
remove print statement
sebsto Mar 24, 2025
4b2b32d
Merge branch 'main' into sebsto/fix_507
sebsto Jun 1, 2025
6aeafbf
Merge branch 'main' into sebsto/fix_507
sebsto Jun 28, 2025
9fb86a7
Merge branch 'sebsto/fix_507' of github.com:sebsto/swift-aws-lambda-r…
sebsto Jun 28, 2025
f7ad7c0
Merge branch 'main' into sebsto/fix_507
sebsto Jun 28, 2025
b60d6e5
apply swift format
sebsto Jun 28, 2025
fbdf503
make LambdaRuntimeError package to be inlinable
sebsto Jun 28, 2025
6353d96
fix example dependencies
sebsto Jun 29, 2025
ee296b4
Merge branch 'main' into sebsto/fix_507
sebsto Jun 29, 2025
e00c1c4
improve logging when task is cancelled
sebsto Jun 29, 2025
b33b6ed
let the lambda runtime set the debug level
sebsto Jun 29, 2025
c5eb4ae
adjust test for task cancellation
sebsto Jun 29, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 1 addition & 4 deletions Examples/HelloJSON/Package.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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: [
Expand Down
26 changes: 17 additions & 9 deletions Sources/AWSLambdaRuntime/Lambda+LocalServer.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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: [
Expand Down Expand Up @@ -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 as? CancellationError) != nil {
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
Expand Down Expand Up @@ -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)
}
Expand Down
29 changes: 25 additions & 4 deletions Sources/AWSLambdaRuntime/LambdaRuntime.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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
@usableFromInline
Expand Down Expand Up @@ -58,8 +63,22 @@ 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 {

// 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)
}

let handler = self.handlerMutex.withLockedValue { handler in
let result = handler
handler = nil
Expand Down Expand Up @@ -96,8 +115,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),
Expand Down
4 changes: 4 additions & 0 deletions Sources/AWSLambdaRuntime/LambdaRuntimeError.swift
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
package struct LambdaRuntimeError: Error {
@usableFromInline
package enum Code: Sendable {
/// internal error codes for LambdaRuntimeClient
case closingRuntimeClient

case connectionToControlPlaneLost
Expand All @@ -34,6 +35,9 @@ package struct LambdaRuntimeError: Error {
case missingLambdaRuntimeAPIEnvironmentVariable
case runtimeCanOnlyBeStartedOnce
case invalidPort

/// public error codes for LambdaRuntime
case moreThanOneLambdaRuntimeInstance
}

@usableFromInline
Expand Down
1 change: 1 addition & 0 deletions Sources/MockServer/MockHTTPServer.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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")
}

Expand Down
5 changes: 2 additions & 3 deletions Tests/AWSLambdaRuntimeTests/LambdaRuntimeClientTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
88 changes: 88 additions & 0 deletions Tests/AWSLambdaRuntimeTests/LambdaRuntimeTests.swift
Original file line number Diff line number Diff line change
@@ -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 {

}
}
Loading