diff --git a/Sources/AWSLambdaRuntime/LambdaRuntime.swift b/Sources/AWSLambdaRuntime/LambdaRuntime.swift index 147f56d6..5f66df6f 100644 --- a/Sources/AWSLambdaRuntime/LambdaRuntime.swift +++ b/Sources/AWSLambdaRuntime/LambdaRuntime.swift @@ -82,7 +82,7 @@ public final class LambdaRuntime: Sendable where Handler: StreamingLamb // 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) + throw LambdaRuntimeError(code: .handlerCanOnlyBeGetOnce) } // are we running inside an AWS Lambda runtime environment ? diff --git a/Sources/AWSLambdaRuntime/LambdaRuntimeError.swift b/Sources/AWSLambdaRuntime/LambdaRuntimeError.swift index 499ed600..a9c0cbca 100644 --- a/Sources/AWSLambdaRuntime/LambdaRuntimeError.swift +++ b/Sources/AWSLambdaRuntime/LambdaRuntimeError.swift @@ -34,6 +34,7 @@ package struct LambdaRuntimeError: Error { case missingLambdaRuntimeAPIEnvironmentVariable case runtimeCanOnlyBeStartedOnce + case handlerCanOnlyBeGetOnce case invalidPort } diff --git a/Tests/AWSLambdaRuntimeTests/LambdaRuntime+ServiceLifeCycle.swift b/Tests/AWSLambdaRuntimeTests/LambdaRuntime+ServiceLifeCycle.swift new file mode 100644 index 00000000..7103ea8d --- /dev/null +++ b/Tests/AWSLambdaRuntimeTests/LambdaRuntime+ServiceLifeCycle.swift @@ -0,0 +1,46 @@ +//===----------------------------------------------------------------------===// +// +// This source file is part of the SwiftAWSLambdaRuntime open source project +// +// Copyright (c) 2024 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 +// +//===----------------------------------------------------------------------===// + +#if ServiceLifecycleSupport +@testable import AWSLambdaRuntime +import ServiceLifecycle +import Testing +import Logging + +@Suite +struct LambdaRuntimeServiceLifecycleTests { + @Test + func testLambdaRuntimeGracefulShutdown() async throws { + let runtime = LambdaRuntime { + (event: String, context: LambdaContext) in + "Hello \(event)" + } + + let serviceGroup = ServiceGroup( + services: [runtime], + gracefulShutdownSignals: [.sigterm, .sigint], + logger: Logger(label: "TestLambdaRuntimeGracefulShutdown") + ) + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + try await serviceGroup.run() + } + // wait a small amount to ensure we are waiting for continuation + try await Task.sleep(for: .milliseconds(100)) + + await serviceGroup.triggerGracefulShutdown() + } + } +} +#endif diff --git a/Tests/AWSLambdaRuntimeTests/LambdaRuntimeClientTests.swift b/Tests/AWSLambdaRuntimeTests/LambdaRuntimeClientTests.swift index cc901461..ba24b2cc 100644 --- a/Tests/AWSLambdaRuntimeTests/LambdaRuntimeClientTests.swift +++ b/Tests/AWSLambdaRuntimeTests/LambdaRuntimeClientTests.swift @@ -15,7 +15,6 @@ import Logging import NIOCore import NIOPosix -import ServiceLifecycle import Testing import struct Foundation.UUID @@ -90,7 +89,7 @@ struct LambdaRuntimeClientTests { } @Test - func testCancellation() async throws { + func testRuntimeClientCancellation() async throws { struct HappyBehavior: LambdaServerBehavior { let requestId = UUID().uuidString let event = "hello" @@ -140,28 +139,4 @@ struct LambdaRuntimeClientTests { } } } - #if ServiceLifecycleSupport - @Test - func testLambdaRuntimeGracefulShutdown() async throws { - let runtime = LambdaRuntime { - (event: String, context: LambdaContext) in - "Hello \(event)" - } - - let serviceGroup = ServiceGroup( - services: [runtime], - gracefulShutdownSignals: [.sigterm, .sigint], - logger: Logger(label: "TestLambdaRuntimeGracefulShutdown") - ) - try await withThrowingTaskGroup(of: Void.self) { group in - group.addTask { - try await serviceGroup.run() - } - // wait a small amount to ensure we are waiting for continuation - try await Task.sleep(for: .milliseconds(100)) - - await serviceGroup.triggerGracefulShutdown() - } - } - #endif } diff --git a/Tests/AWSLambdaRuntimeTests/LambdaRuntimeTests.swift b/Tests/AWSLambdaRuntimeTests/LambdaRuntimeTests.swift index 76057ddf..3f37bf09 100644 --- a/Tests/AWSLambdaRuntimeTests/LambdaRuntimeTests.swift +++ b/Tests/AWSLambdaRuntimeTests/LambdaRuntimeTests.swift @@ -50,17 +50,24 @@ struct LambdaRuntimeTests { } // wait a small amount to ensure runtime1 task is started - try await Task.sleep(for: .seconds(0.5)) + // on GH Actions, it might take a bit longer to start the runtime + try await Task.sleep(for: .seconds(2)) // Running the second runtime should trigger LambdaRuntimeError - await #expect(throws: LambdaRuntimeError.self) { - try await runtime2.run() + // start the first runtime + taskGroup.addTask { + await #expect(throws: LambdaRuntimeError.self) { + try await runtime2.run() + } } // cancel runtime 1 / task 1 taskGroup.cancelAll() } + // wait a small amount to ensure everything is cancelled and cleanup + try await Task.sleep(for: .seconds(1)) + // Running the second runtime should work now try await withThrowingTaskGroup(of: Void.self) { taskGroup in taskGroup.addTask { @@ -75,6 +82,53 @@ struct LambdaRuntimeTests { taskGroup.cancelAll() } } + @Test("run() must be cancellable") + func testLambdaRuntimeCancellable() async throws { + + let logger = Logger(label: "LambdaRuntimeTests.RuntimeCancellable") + // create a runtime + let runtime = LambdaRuntime( + handler: MockHandler(), + eventLoop: Lambda.defaultEventLoop, + logger: logger + ) + + // Running the runtime with structured concurrency + // Task group returns when all tasks are completed. + // Even cancelled tasks must cooperatlivly complete + await #expect(throws: Never.self) { + try await withThrowingTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask { + logger.trace("--- launching runtime ----") + try await runtime.run() + } + + // Add a timeout task to the group + taskGroup.addTask { + logger.trace("--- launching timeout task ----") + try await Task.sleep(for: .seconds(5)) + if Task.isCancelled { return } + logger.trace("--- throwing timeout error ----") + throw TestError.timeout // Fail the test if the timeout triggers + } + + do { + // Wait for the runtime to start + logger.trace("--- waiting for runtime to start ----") + try await Task.sleep(for: .seconds(1)) + + // Cancel all tasks, this should not throw an error + // and should allow the runtime to complete gracefully + logger.trace("--- cancel all tasks ----") + taskGroup.cancelAll() // Cancel all tasks + } catch { + logger.error("--- catch an error: \(error)") + throw error // Propagate the error to fail the test + } + } + } + + } } struct MockHandler: StreamingLambdaHandler { @@ -86,3 +140,15 @@ struct MockHandler: StreamingLambdaHandler { } } + +// Define a custom error for timeout +enum TestError: Error, CustomStringConvertible { + case timeout + + var description: String { + switch self { + case .timeout: + return "Test timed out waiting for the task to complete." + } + } +} diff --git a/Tests/AWSLambdaRuntimeTests/PoolTests.swift b/Tests/AWSLambdaRuntimeTests/PoolTests.swift index cfeea6a2..15d54a73 100644 --- a/Tests/AWSLambdaRuntimeTests/PoolTests.swift +++ b/Tests/AWSLambdaRuntimeTests/PoolTests.swift @@ -37,7 +37,7 @@ struct PoolTests { } @Test - func testCancellation() async throws { + func testPoolCancellation() async throws { let pool = LambdaHTTPServer.Pool() // Create a task that will be cancelled