Review comments addressed.

swift-server · Jun 17, 2020 · a772032 · a772032
1 parent 42d6dec
commit a772032
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 18 deletions.
diff --git 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<Int> in
+            .flatMap { (handler, runnerResult) -> EventLoopFuture<Int> 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<Int, Error> 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
@@ -129,6 +129,7 @@ internal extension Lambda {
         case invocationMissingHeader(String)
         case noBody
         case json(Error)
+        case shutdownError(shutdownError: Error, runnerResult: Result<Int, Error>)
     }
 }
 

diff --git 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<ByteBuffer?>)
-            let shutdown: (Lambda.ShutdownContext) -> EventLoopFuture<Void>
-
-            init(_ handler: @escaping (Lambda.Context, ByteBuffer) -> (EventLoopFuture<ByteBuffer?>), shutdown: @escaping (Lambda.ShutdownContext) -> EventLoopFuture<Void>) {
-                self.handler = handler
-                self.shutdown = shutdown
-            }
+    struct CallbackLambdaHandler: ByteBufferLambdaHandler {
+        let handler: (Lambda.Context, ByteBuffer) -> (EventLoopFuture<ByteBuffer?>)
+        let shutdown: (Lambda.ShutdownContext) -> EventLoopFuture<Void>
 
-            func handle(context: Lambda.Context, event: ByteBuffer) -> EventLoopFuture<ByteBuffer?> {
-                self.handler(context, event)
-            }
+        init(_ handler: @escaping (Lambda.Context, ByteBuffer) -> (EventLoopFuture<ByteBuffer?>), shutdown: @escaping (Lambda.ShutdownContext) -> EventLoopFuture<Void>) {
+            self.handler = handler
+            self.shutdown = shutdown
+        }
 
-            func shutdown(context: Lambda.ShutdownContext) -> EventLoopFuture<Void> {
-                self.shutdown(context)
-            }
+        func handle(context: Lambda.Context, event: ByteBuffer) -> EventLoopFuture<ByteBuffer?> {
+            self.handler(context, event)
         }
 
+        func shutdown(context: Lambda.ShutdownContext) -> EventLoopFuture<Void> {
+            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 {