Updated to use a ShutdownContext

Author: fabianfett

Sources/AWSLambdaRuntimeCore/LambdaContext.swift

@@ -122,3 +122,27 @@ extension Lambda {
         }
     }
 }
+
+// MARK: - ShutdownContext
+
+extension Lambda {
+    /// Lambda runtime shutdown context.
+    /// The Lambda runtime generates and passes the `ShutdownContext` to the Lambda handler as an argument.
+    public final class ShutdownContext {
+        /// `Logger` to log with
+        ///
+        /// - note: The `LogLevel` can be configured using the `LOG_LEVEL` environment variable.
+        public let logger: Logger
+
+        /// The `EventLoop` the Lambda is executed on. Use this to schedule work with.
+        ///
+        /// - note: The `EventLoop` is shared with the Lambda runtime engine and should be handled with extra care.
+        ///         Most importantly the `EventLoop` must never be blocked.
+        public let eventLoop: EventLoop
+
+        internal init(logger: Logger, eventLoop: EventLoop) {
+            self.eventLoop = eventLoop
+            self.logger = logger
+        }
+    }
+}

Sources/AWSLambdaRuntimeCore/LambdaHandler.swift

@@ -170,11 +170,13 @@ public protocol ByteBufferLambdaHandler {
     ///
     /// - Note: In case your Lambda fails while creating your LambdaHandler in the `HandlerFactory`, this method
     ///         **is not invoked**. In this case you must cleanup the created resources immediately in the `HandlerFactory`.
-    func syncShutdown() throws
+    func shutdown(context: Lambda.ShutdownContext) -> EventLoopFuture<Void>
 }
 
 public extension ByteBufferLambdaHandler {
-    func syncShutdown() throws {}
+    func shutdown(context: Lambda.ShutdownContext) -> EventLoopFuture<Void> {
+        context.eventLoop.makeSucceededFuture(Void())
+    }
 }
 
 private enum CodecError: Error {

Sources/AWSLambdaRuntimeCore/LambdaLifecycle.swift

@@ -81,21 +81,19 @@ extension Lambda {
             let runner = Runner(eventLoop: self.eventLoop, configuration: self.configuration)
 
             let startupFuture = runner.initialize(logger: logger, factory: self.factory)
-            startupFuture.flatMap { handler in
+            startupFuture.flatMap { handler -> EventLoopFuture<(ByteBufferLambdaHandler, Result<Int, Error>)> in
                 // after the startup future has succeeded, we have a handler that we can use
                 // to `run` the lambda.
                 let finishedPromise = self.eventLoop.makePromise(of: Int.self)
                 self.state = .active(runner, handler)
                 self.run(promise: finishedPromise)
-                return finishedPromise.futureResult.always { _ in
-                    // If the lambda is terminated (e.g. LocalServer shutdown), we make sure
-                    // developers have the chance to cleanup their resources.
-                    do {
-                        try handler.syncShutdown()
-                    } catch {
-                        logger.error("Error shutting down handler: \(error)")
-                    }
-                }
+                return finishedPromise.futureResult.mapResult { (handler, $0) }
+            }
+            .flatMap { (handler, result) -> EventLoopFuture<Int> in
+                let shutdownContext = ShutdownContext(logger: logger, eventLoop: self.eventLoop)
+                return handler.shutdown(context: shutdownContext).recover { error in
+                    logger.error("Error shutting down handler: \(error)")
+                }.flatMapResult { _ in result }
             }.always { _ in
                 // triggered when the Lambda has finished its last run or has a startup failure.
                 self.markShutdown()

Sources/AWSLambdaRuntimeCore/LambdaRunner.swift

@@ -115,7 +115,7 @@ private extension Lambda.Context {
 }
 
 // TODO: move to nio?
-private extension EventLoopFuture {
+extension EventLoopFuture {
     // callback does not have side effects, failing with original result
     func peekError(_ callback: @escaping (Error) -> Void) -> EventLoopFuture<Value> {
         self.flatMapError { error in

Tests/AWSLambdaRuntimeCoreTests/LambdaLifecycleTest.swift

@@ -47,9 +47,9 @@ class LambdaLifecycleTest: XCTestCase {
     func testSyncShutdownIsCalledWhenLambdaShutsdown() {
         struct CallbackLambdaHandler: ByteBufferLambdaHandler {
             let handler: (Lambda.Context, ByteBuffer) -> (EventLoopFuture<ByteBuffer?>)
-            let shutdown: () throws -> Void
+            let shutdown: (Lambda.ShutdownContext) -> EventLoopFuture<Void>
 
-            init(_ handler: @escaping (Lambda.Context, ByteBuffer) -> (EventLoopFuture<ByteBuffer?>), shutdown: @escaping () throws -> Void) {
+            init(_ handler: @escaping (Lambda.Context, ByteBuffer) -> (EventLoopFuture<ByteBuffer?>), shutdown: @escaping (Lambda.ShutdownContext) -> EventLoopFuture<Void>) {
                 self.handler = handler
                 self.shutdown = shutdown
             }
@@ -58,8 +58,8 @@ class LambdaLifecycleTest: XCTestCase {
                 self.handler(context, event)
             }
 
-            func syncShutdown() throws {
-                try self.shutdown()
+            func shutdown(context: Lambda.ShutdownContext) -> EventLoopFuture<Void> {
+                self.shutdown(context)
             }
         }
 
@@ -70,8 +70,9 @@ class LambdaLifecycleTest: XCTestCase {
         defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
 
         var count = 0
-        let handler = CallbackLambdaHandler({ XCTFail("Should not be reached"); return $0.eventLoop.makeSucceededFuture($1) }) {
+        let handler = CallbackLambdaHandler({ XCTFail("Should not be reached"); return $0.eventLoop.makeSucceededFuture($1) }) { context in
             count += 1
+            return context.eventLoop.makeSucceededFuture(Void())
         }
 
         let eventLoop = eventLoopGroup.next()