terminator handler

motivation: make it simpler to register shutdown hooks

changes:
* introduce Terminaotr helper that allow registering and de-registaring shutdown handlers
* expose the new terminator hanler on the InitializationContext and deprecate ShutdownContext
* deprecate the Handler ::shutdown protocol requirment
* update the runtime code to use the new terminator instead of calling shutdown on the handler
* add and adjust tests

Author: tomerd

Sources/AWSLambdaRuntimeCore/LambdaContext.swift

@@ -38,10 +38,14 @@ extension Lambda {
         /// `ByteBufferAllocator` to allocate `ByteBuffer`
         public let allocator: ByteBufferAllocator
 
-        init(logger: Logger, eventLoop: EventLoop, allocator: ByteBufferAllocator) {
+        /// `Terminator` to register shutdown operations
+        public let terminator: Terminator
+
+        init(logger: Logger, eventLoop: EventLoop, allocator: ByteBufferAllocator, terminator: Terminator) {
             self.eventLoop = eventLoop
             self.logger = logger
             self.allocator = allocator
+            self.terminator = terminator
         }
 
         /// This interface is not part of the public API and must not be used by adopters. This API is not part of semver versioning.
@@ -52,7 +56,8 @@ extension Lambda {
             InitializationContext(
                 logger: logger,
                 eventLoop: eventLoop,
-                allocator: ByteBufferAllocator()
+                allocator: ByteBufferAllocator(),
+                terminator: Lambda.Terminator()
             )
         }
     }
@@ -208,6 +213,7 @@ public struct LambdaContext: CustomDebugStringConvertible {
 
 // MARK: - ShutdownContext
 
+/*
 extension Lambda {
     /// Lambda runtime shutdown context.
     /// The Lambda runtime generates and passes the `ShutdownContext` to the Lambda handler as an argument.
@@ -229,3 +235,4 @@ extension Lambda {
         }
     }
 }
+*/

Sources/AWSLambdaRuntimeCore/LambdaHandler.swift

@@ -182,14 +182,15 @@ 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 shutdown(context: Lambda.ShutdownContext) -> EventLoopFuture<Void>
+    //func shutdown(context: Lambda.ShutdownContext) -> EventLoopFuture<Void>
 }
 
+/*
 extension ByteBufferLambdaHandler {
     public func shutdown(context: Lambda.ShutdownContext) -> EventLoopFuture<Void> {
         context.eventLoop.makeSucceededFuture(())
     }
-}
+}*/
 
 extension ByteBufferLambdaHandler {
     /// Initializes and runs the lambda function.

Sources/AWSLambdaRuntimeCore/LambdaRunner.swift

@@ -34,13 +34,16 @@ extension Lambda {
         /// Run the user provided initializer. This *must* only be called once.
         ///
         /// - Returns: An `EventLoopFuture<LambdaHandler>` fulfilled with the outcome of the initialization.
-        func initialize<Handler: ByteBufferLambdaHandler>(logger: Logger, handlerType: Handler.Type) -> EventLoopFuture<Handler> {
+        func initialize<Handler: ByteBufferLambdaHandler>(logger: Logger, terminator: Terminator, handlerType: Handler.Type) -> EventLoopFuture<Handler> {
             logger.debug("initializing lambda")
             // 1. create the handler from the factory
-            // 2. report initialization error if one occured
-            let context = InitializationContext(logger: logger,
-                                                eventLoop: self.eventLoop,
-                                                allocator: self.allocator)
+            // 2. report initialization error if one occurred
+            let context = InitializationContext(
+                logger: logger,
+                eventLoop: self.eventLoop,
+                allocator: self.allocator,
+                terminator: terminator
+            )
             return Handler.makeHandler(context: context)
                 // Hopping back to "our" EventLoop is important in case the factory returns a future
                 // that originated from a foreign EventLoop/EventLoopGroup.

Sources/AWSLambdaRuntimeCore/LambdaRuntime.swift

@@ -74,9 +74,10 @@ public final class LambdaRuntime<Handler: ByteBufferLambdaHandler> {
 
         var logger = self.logger
         logger[metadataKey: "lifecycleId"] = .string(self.configuration.lifecycle.id)
+        let terminator = Lambda.Terminator()
         let runner = Lambda.Runner(eventLoop: self.eventLoop, configuration: self.configuration)
 
-        let startupFuture = runner.initialize(logger: logger, handlerType: Handler.self)
+        let startupFuture = runner.initialize(logger: logger, terminator: terminator, handlerType: Handler.self)
         startupFuture.flatMap { handler -> EventLoopFuture<(Handler, Result<Int, Error>)> in
             // after the startup future has succeeded, we have a handler that we can use
             // to `run` the lambda.
@@ -88,16 +89,29 @@ public final class LambdaRuntime<Handler: ByteBufferLambdaHandler> {
         .flatMap { handler, runnerResult -> EventLoopFuture<Int> in
             // after the lambda finishPromise has succeeded or failed we need to
             // shutdown the handler
-            let shutdownContext = Lambda.ShutdownContext(logger: logger, eventLoop: self.eventLoop)
-            return handler.shutdown(context: shutdownContext).flatMapErrorThrowing { error in
+            //let shutdownContext = Lambda.ShutdownContext(logger: logger, eventLoop: self.eventLoop)
+            /*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)")
                 throw Lambda.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
+            }*/
+
+            // after the lambda finishPromise has succeeded or failed we need to
+            // shutdown the handler
+            return terminator.terminate(eventLoop: self.eventLoop).flatMapErrorThrowing { error in
+                // if, we had an error shutting down the handler, we want to concatenate it with
+                // the runner result
+                logger.error("Error shutting down handler: \(error)")
+                throw Lambda.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()

Sources/AWSLambdaRuntimeCore/Terminator.swift

@@ -0,0 +1,84 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftAWSLambdaRuntime open source project
+//
+// Copyright (c) 2022 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 NIOConcurrencyHelpers
+import NIOCore
+
+extension Lambda {
+    /// Lambda terminator.
+    /// Helper utility to manage the lambda shutdown sequence
+    #warning("API docs, impl, async")
+    public final class Terminator {
+        public typealias Handler =  (EventLoop) -> EventLoopFuture<Void>
+        public typealias RegistrationKey = String
+
+        private let lock: Lock
+        private var handlers: [String: (name: String, handler: Handler)]
+
+        public init() {
+            self.lock = .init()
+            self.handlers = [:]
+        }
+
+        #warning("API docs")
+        @discardableResult
+        public func register(name: String, handler: @escaping Handler) -> RegistrationKey {
+            let key = LambdaRequestID().uuidString // UUID basically
+            self.lock.withLock {
+                self.handlers[key] = (name: name, handler: handler)
+            }
+            return key
+        }
+
+        #warning("API docs")
+        public func deregister(_ key: RegistrationKey) {
+            self.lock.withLock {
+                self.handlers[key] = nil
+            }
+        }
+
+        #warning("API docs")
+        internal func terminate(eventLoop: EventLoop) -> EventLoopFuture<Void> {
+            let handlers = self.lock.withLock {
+                self.handlers
+            }
+
+            #warning("FIXME: cascade in reverse order?")
+            return EventLoopFuture.whenAllComplete(handlers.values.map{ $0.handler(eventLoop) }, on: eventLoop).flatMap{ results in
+                let errors = results.compactMap { $0.error }
+                if errors.isEmpty {
+                    return eventLoop.makeSucceededVoidFuture()
+                } else {
+                    return eventLoop.makeFailedFuture(TerminationError(underlying: errors))
+                }
+            }
+        }
+    }
+
+    struct TerminationError: Error {
+        let underlying: [Error]
+    }
+}
+
+private extension Result {
+    var error: Error? {
+        switch self {
+        case .failure(let error):
+            return error
+        case .success:
+            return .none
+        }
+    }
+}

Tests/AWSLambdaRuntimeCoreTests/LambdaRuntimeTest.swift

@@ -66,23 +66,34 @@ class LambdaRuntimeTest: XCTestCase {
         let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
         defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
 
-        struct ShutdownError: Error {}
+        struct ShutdownError: Error {
+            let description: String
+        }
 
         struct ShutdownErrorHandler: EventLoopLambdaHandler {
             typealias Event = String
             typealias Output = Void
 
             static func makeHandler(context: Lambda.InitializationContext) -> EventLoopFuture<ShutdownErrorHandler> {
-                context.eventLoop.makeSucceededFuture(ShutdownErrorHandler())
+                // register shutdown operation
+                context.terminator.register(name: "test 1", handler: { eventLoop in
+                    return eventLoop.makeFailedFuture(ShutdownError(description: "1"))
+                })
+                context.terminator.register(name: "test 2", handler: { eventLoop in
+                    return eventLoop.makeSucceededVoidFuture()
+                })
+                context.terminator.register(name: "test 3", handler: { eventLoop in
+                    return eventLoop.makeFailedFuture(ShutdownError(description: "2"))
+                })
+                context.terminator.register(name: "test 3", handler: { eventLoop in
+                    return eventLoop.makeSucceededVoidFuture()
+                })
+                return context.eventLoop.makeSucceededFuture(ShutdownErrorHandler())
             }
 
             func handle(_ event: String, context: LambdaContext) -> EventLoopFuture<Void> {
                 context.eventLoop.makeSucceededVoidFuture()
             }
-
-            func shutdown(context: Lambda.ShutdownContext) -> EventLoopFuture<Void> {
-                context.eventLoop.makeFailedFuture(ShutdownError())
-            }
         }
 
         let eventLoop = eventLoopGroup.next()
@@ -95,7 +106,7 @@ class LambdaRuntimeTest: XCTestCase {
                 XCTFail("Unexpected error: \(error)"); return
             }
 
-            XCTAssert(shutdownError is ShutdownError)
+            XCTAssertEqual(shutdownError as? Lambda.TerminationError, Lambda.TerminationError(underlying: [ShutdownError(description: "1"), ShutdownError(description: "2")]))
             XCTAssertEqual(runtimeError as? Lambda.RuntimeError, .badStatusCode(.internalServerError))
         }
     }

Tests/AWSLambdaRuntimeCoreTests/Utils.swift

@@ -23,10 +23,11 @@ func runLambda<Handler: ByteBufferLambdaHandler>(behavior: LambdaServerBehavior,
     defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
     let logger = Logger(label: "TestLogger")
     let configuration = Lambda.Configuration(runtimeEngine: .init(requestTimeout: .milliseconds(100)))
+    let terminator = Lambda.Terminator()
     let runner = Lambda.Runner(eventLoop: eventLoopGroup.next(), configuration: configuration)
     let server = try MockLambdaServer(behavior: behavior).start().wait()
     defer { XCTAssertNoThrow(try server.stop().wait()) }
-    try runner.initialize(logger: logger, handlerType: handlerType).flatMap { handler in
+    try runner.initialize(logger: logger, terminator: terminator, handlerType: handlerType).flatMap { handler in
         runner.run(logger: logger, handler: handler)
     }.wait()
 }
@@ -66,3 +67,13 @@ extension Lambda.RuntimeError: Equatable {
         String(describing: lhs) == String(describing: rhs)
     }
 }
+
+extension Lambda.TerminationError: Equatable {
+    public static func == (lhs: Lambda.TerminationError, rhs: Lambda.TerminationError) -> Bool {
+        guard lhs.underlying.count == rhs.underlying.count else {
+            return false
+        }
+        // technically incorrect, but good enough for our tests
+        return String(describing: lhs) == String(describing: rhs)
+    }
+}

Tests/AWSLambdaRuntimeTests/Lambda+CodableTest.swift

@@ -172,7 +172,8 @@ class CodableLambdaTest: XCTestCase {
         Lambda.InitializationContext(
             logger: Logger(label: "test"),
             eventLoop: self.eventLoopGroup.next(),
-            allocator: ByteBufferAllocator()
+            allocator: ByteBufferAllocator(),
+            terminator: Lambda.Terminator()
         )
     }
 }