diff --git a/Sources/AWSLambdaRuntimeCore/LambdaLifecycle.swift b/Sources/AWSLambdaRuntimeCore/LambdaLifecycle.swift index 9639e179..ec609901 100644 --- a/Sources/AWSLambdaRuntimeCore/LambdaLifecycle.swift +++ b/Sources/AWSLambdaRuntimeCore/LambdaLifecycle.swift @@ -89,11 +89,19 @@ extension Lambda { self.run(promise: finishedPromise) return finishedPromise.futureResult.mapResult { (handler, $0) } } - .flatMap { (handler, result) -> EventLoopFuture in + .flatMap { (handler, runnerResult) -> EventLoopFuture in + // after the lambda finishPromise has succeeded or failed we need to + // shutdown the handler let shutdownContext = ShutdownContext(logger: logger, eventLoop: self.eventLoop) - return handler.shutdown(context: shutdownContext).recover { error in + return handler.shutdown(context: shutdownContext).flatMapErrorThrowing { error in + // if, we had an error shuting down the lambda, we want to concatenate it with + // the runner result logger.error("Error shutting down handler: \(error)") - }.flatMapResult { _ in result } + throw RuntimeError.shutdownError(shutdownError: error, runnerResult: runnerResult) + }.flatMapResult { (_) -> Result in + // we had no error shutting down the lambda. let's return the runner's result + runnerResult + } }.always { _ in // triggered when the Lambda has finished its last run or has a startup failure. self.markShutdown() diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift b/Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift index 2976a8c2..499b63b7 100644 --- a/Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift +++ b/Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift @@ -129,6 +129,7 @@ internal extension Lambda { case invocationMissingHeader(String) case noBody case json(Error) + case shutdownError(shutdownError: Error, runnerResult: Result) } } diff --git a/Tests/AWSLambdaRuntimeCoreTests/LambdaLifecycleTest.swift b/Tests/AWSLambdaRuntimeCoreTests/LambdaLifecycleTest.swift index 152bc08a..1453d768 100644 --- a/Tests/AWSLambdaRuntimeCoreTests/LambdaLifecycleTest.swift +++ b/Tests/AWSLambdaRuntimeCoreTests/LambdaLifecycleTest.swift @@ -44,25 +44,25 @@ class LambdaLifecycleTest: XCTestCase { } } - func testSyncShutdownIsCalledWhenLambdaShutsdown() { - struct CallbackLambdaHandler: ByteBufferLambdaHandler { - let handler: (Lambda.Context, ByteBuffer) -> (EventLoopFuture) - let shutdown: (Lambda.ShutdownContext) -> EventLoopFuture - - init(_ handler: @escaping (Lambda.Context, ByteBuffer) -> (EventLoopFuture), shutdown: @escaping (Lambda.ShutdownContext) -> EventLoopFuture) { - self.handler = handler - self.shutdown = shutdown - } + struct CallbackLambdaHandler: ByteBufferLambdaHandler { + let handler: (Lambda.Context, ByteBuffer) -> (EventLoopFuture) + let shutdown: (Lambda.ShutdownContext) -> EventLoopFuture - func handle(context: Lambda.Context, event: ByteBuffer) -> EventLoopFuture { - self.handler(context, event) - } + init(_ handler: @escaping (Lambda.Context, ByteBuffer) -> (EventLoopFuture), shutdown: @escaping (Lambda.ShutdownContext) -> EventLoopFuture) { + self.handler = handler + self.shutdown = shutdown + } - func shutdown(context: Lambda.ShutdownContext) -> EventLoopFuture { - self.shutdown(context) - } + func handle(context: Lambda.Context, event: ByteBuffer) -> EventLoopFuture { + self.handler(context, event) } + func shutdown(context: Lambda.ShutdownContext) -> EventLoopFuture { + self.shutdown(context) + } + } + + func testShutdownIsCalledWhenLambdaShutsdown() { let server = MockLambdaServer(behavior: BadBehavior()) XCTAssertNoThrow(try server.start().wait()) defer { XCTAssertNoThrow(try server.stop().wait()) } @@ -87,6 +87,37 @@ class LambdaLifecycleTest: XCTestCase { } XCTAssertEqual(count, 1) } + + func testLambdaResultIfShutsdownIsUnclean() { + let server = MockLambdaServer(behavior: BadBehavior()) + XCTAssertNoThrow(try server.start().wait()) + defer { XCTAssertNoThrow(try server.stop().wait()) } + let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) + defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) } + + var count = 0 + let handler = CallbackLambdaHandler({ XCTFail("Should not be reached"); return $0.eventLoop.makeSucceededFuture($1) }) { context in + count += 1 + return context.eventLoop.makeFailedFuture(TestError("kaboom")) + } + + let eventLoop = eventLoopGroup.next() + let logger = Logger(label: "TestLogger") + let lifecycle = Lambda.Lifecycle(eventLoop: eventLoop, logger: logger, factory: { + $0.makeSucceededFuture(handler) + }) + + XCTAssertNoThrow(_ = try eventLoop.flatSubmit { lifecycle.start() }.wait()) + XCTAssertThrowsError(_ = try lifecycle.shutdownFuture.wait()) { error in + guard case Lambda.RuntimeError.shutdownError(let shutdownError, .failure(let runtimeError)) = error else { + XCTFail("Unexpected error"); return + } + + XCTAssertEqual(shutdownError as? TestError, TestError("kaboom")) + XCTAssertEqual(runtimeError as? Lambda.RuntimeError, .badStatusCode(.internalServerError)) + } + XCTAssertEqual(count, 1) + } } struct BadBehavior: LambdaServerBehavior {