allow custom initialisation of the HandlerType of the LambdaRuntime

tomerd · tomerd · commit 6aafe7984a1c · 2023-10-10T14:36:44.000-07:00
Motivation: Provide the flexibility for custom initialisation of the HandlerType as this will often be required by higher level frameworks. Modifications: * Modify the LambdaRuntime type to accept a closure to provide the handler rather than requiring that it is provided by a static method on the Handler type * Update downstream code to use HandlerProvider * Update upstream code to support passing Hanlder Type of Handler Provider * Add and update tests Originally suggested and coded by @tachyonics in #308
diff --git a/Sources/AWSLambdaRuntimeCore/Lambda.swift b/Sources/AWSLambdaRuntimeCore/Lambda.swift
@@ -85,6 +85,22 @@ public enum Lambda {
     internal static func run(
         configuration: LambdaConfiguration = .init(),
         handlerType: (some ByteBufferLambdaHandler).Type
+    ) -> Result<Int, Error> {
+        Self.run(configuration: configuration, handlerProvider: handlerType.makeHandler(context:))
+    }
+
+    /// Run a Lambda defined by implementing the ``ByteBufferLambdaHandler`` protocol.
+    /// The Runtime will manage the Lambdas application lifecycle automatically. It will invoke the
+    /// ``ByteBufferLambdaHandler/makeHandler(context:)`` to create a new Handler.
+    ///
+    /// - parameters:
+    ///     - configuration: A Lambda runtime configuration object
+    ///     - handlerProvider: A provider of the Handler to invoke.
+    ///
+    /// - note: This is a blocking operation that will run forever, as its lifecycle is managed by the AWS Lambda Runtime Engine.
+    internal static func run(
+        configuration: LambdaConfiguration = .init(),
+        handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<some ByteBufferLambdaHandler>
     ) -> Result<Int, Error> {
         let _run = { (configuration: LambdaConfiguration) -> Result<Int, Error> in
             #if swift(<5.9)
@@ -95,7 +111,12 @@ public enum Lambda {
 
             var result: Result<Int, Error>!
             MultiThreadedEventLoopGroup.withCurrentThreadAsEventLoop { eventLoop in
-                let runtime = LambdaRuntime(handlerType: handlerType, eventLoop: eventLoop, logger: logger, configuration: configuration)
+                let runtime = LambdaRuntime(
+                    handlerProvider: handlerProvider,
+                    eventLoop: eventLoop,
+                    logger: logger,
+                    configuration: configuration
+                )
                 #if DEBUG
                 let signalSource = trap(signal: configuration.lifecycle.stopSignal) { signal in
                     logger.info("intercepted signal: \(signal)")
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
@@ -33,7 +33,11 @@ internal final class LambdaRunner {
     /// 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>(handlerType: Handler.Type, logger: Logger, terminator: LambdaTerminator) -> EventLoopFuture<Handler> {
+    func initialize<Handler: ByteBufferLambdaHandler>(
+        handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<Handler>,
+        logger: Logger,
+        terminator: LambdaTerminator
+    ) -> EventLoopFuture<Handler> {
         logger.debug("initializing lambda")
         // 1. create the handler from the factory
         // 2. report initialization error if one occurred
@@ -44,7 +48,7 @@ internal final class LambdaRunner {
             terminator: terminator
         )
 
-        return handlerType.makeHandler(context: context)
+        return handlerProvider(context)
             // Hopping back to "our" EventLoop is important in case the factory returns a future
             // that originated from a foreign EventLoop/EventLoopGroup.
             // This can happen if the factory uses a library (let's say a database client) that manages its own threads/loops
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRuntime.swift b/Sources/AWSLambdaRuntimeCore/LambdaRuntime.swift
@@ -25,6 +25,8 @@ public final class LambdaRuntime<Handler: ByteBufferLambdaHandler> {
     private let logger: Logger
     private let configuration: LambdaConfiguration
 
+    private let handlerProvider: (LambdaInitializationContext) -> EventLoopFuture<Handler>
+
     private var state = State.idle {
         willSet {
             self.eventLoop.assertInEventLoop()
@@ -39,14 +41,40 @@ public final class LambdaRuntime<Handler: ByteBufferLambdaHandler> {
     ///     - eventLoop: An `EventLoop` to run the Lambda on.
     ///     - logger: A `Logger` to log the Lambda events.
     public convenience init(_ handlerType: Handler.Type, eventLoop: EventLoop, logger: Logger) {
-        self.init(handlerType: handlerType, eventLoop: eventLoop, logger: logger, configuration: .init())
+        self.init(handlerProvider: handlerType.makeHandler(context:), eventLoop: eventLoop, logger: logger, configuration: .init())
     }
 
-    init(handlerType: Handler.Type, eventLoop: EventLoop, logger: Logger, configuration: LambdaConfiguration) {
+    /// Create a new `LambdaRuntime`.
+    ///
+    /// - parameters:
+    ///     - handlerProvider: A provider of the ``ByteBufferLambdaHandler``  the `LambdaRuntime` will manage.
+    ///     - eventLoop: An `EventLoop` to run the Lambda on.
+    ///     - logger: A `Logger` to log the Lambda events.
+    public convenience init(
+        handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<Handler>,
+        eventLoop: EventLoop,
+        logger: Logger
+    ) {
+        self.init(
+            handlerProvider: handlerProvider,
+            eventLoop: eventLoop,
+            logger: logger,
+            configuration: .init()
+        )
+    }
+
+    init(
+        handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<Handler>,
+        eventLoop: EventLoop,
+        logger: Logger,
+        configuration: LambdaConfiguration
+    ) {
         self.eventLoop = eventLoop
         self.shutdownPromise = eventLoop.makePromise(of: Int.self)
         self.logger = logger
         self.configuration = configuration
+
+        self.handlerProvider = handlerProvider
     }
 
     deinit {
@@ -85,7 +113,7 @@ public final class LambdaRuntime<Handler: ByteBufferLambdaHandler> {
         let terminator = LambdaTerminator()
         let runner = LambdaRunner(eventLoop: self.eventLoop, configuration: self.configuration)
 
-        let startupFuture = runner.initialize(handlerType: Handler.self, logger: logger, terminator: terminator)
+        let startupFuture = runner.initialize(handlerProvider: self.handlerProvider, logger: logger, terminator: terminator)
         startupFuture.flatMap { handler -> EventLoopFuture<Result<Int, Error>> in
             // after the startup future has succeeded, we have a handler that we can use
             // to `run` the lambda.
@@ -229,6 +257,75 @@ public enum LambdaRuntimeFactory {
     public static func makeRuntime<H: EventLoopLambdaHandler>(_ handlerType: H.Type, eventLoop: any EventLoop, logger: Logger) -> LambdaRuntime<some ByteBufferLambdaHandler> {
         LambdaRuntime<CodableEventLoopLambdaHandler<H>>(CodableEventLoopLambdaHandler<H>.self, eventLoop: eventLoop, logger: logger)
     }
+
+    /// Create a new `LambdaRuntime`.
+    ///
+    /// - parameters:
+    ///     - handlerProvider: A provider of the ``SimpleLambdaHandler``  the `LambdaRuntime` will manage.
+    ///     - eventLoop: An `EventLoop` to run the Lambda on.
+    ///     - logger: A `Logger` to log the Lambda events.
+    @inlinable
+    public static func makeRuntime<H: SimpleLambdaHandler>(
+        handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<H>,
+        eventLoop: any EventLoop,
+        logger: Logger
+    ) -> LambdaRuntime<some ByteBufferLambdaHandler> {
+        LambdaRuntime(
+            handlerProvider: { context in
+                handlerProvider(context).map {
+                    CodableSimpleLambdaHandler(handler: $0, allocator: context.allocator)
+                }
+            },
+            eventLoop: eventLoop,
+            logger: logger
+        )
+    }
+
+    /// Create a new `LambdaRuntime`.
+    ///
+    /// - parameters:
+    ///     - handlerProvider: A provider of the ``LambdaHandler``  the `LambdaRuntime` will manage.
+    ///     - eventLoop: An `EventLoop` to run the Lambda on.
+    ///     - logger: A `Logger` to log the Lambda events.
+    @inlinable
+    public static func makeRuntime<H: LambdaHandler>(
+        handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<H>,
+        eventLoop: any EventLoop,
+        logger: Logger
+    ) -> LambdaRuntime<some ByteBufferLambdaHandler> {
+        LambdaRuntime(
+            handlerProvider: { context in
+                handlerProvider(context).map {
+                    CodableLambdaHandler(handler: $0, allocator: context.allocator)
+                }
+            },
+            eventLoop: eventLoop,
+            logger: logger
+        )
+    }
+
+    /// Create a new `LambdaRuntime`.
+    ///
+    /// - parameters:
+    ///     - handlerProvider: A provider of the ``EventLoopLambdaHandler``  the `LambdaRuntime` will manage.
+    ///     - eventLoop: An `EventLoop` to run the Lambda on.
+    ///     - logger: A `Logger` to log the Lambda events.
+    @inlinable
+    public static func makeRuntime<H: EventLoopLambdaHandler>(
+        handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<H>,
+        eventLoop: any EventLoop,
+        logger: Logger
+    ) -> LambdaRuntime<some ByteBufferLambdaHandler> {
+        LambdaRuntime(
+            handlerProvider: { context in
+                handlerProvider(context).map {
+                    CodableEventLoopLambdaHandler(handler: $0, allocator: context.allocator)
+                }
+            },
+            eventLoop: eventLoop,
+            logger: logger
+        )
+    }
 }
 
 /// This is safe since lambda runtime synchronizes by dispatching all methods to a single `EventLoop`
diff --git a/Tests/AWSLambdaRuntimeCoreTests/LambdaRunnerTest.swift b/Tests/AWSLambdaRuntimeCoreTests/LambdaRunnerTest.swift
@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 @testable import AWSLambdaRuntimeCore
+import NIOCore
 import XCTest
 
 class LambdaRunnerTest: XCTestCase {
@@ -68,4 +69,66 @@ class LambdaRunnerTest: XCTestCase {
         }
         XCTAssertNoThrow(try runLambda(behavior: Behavior(), handlerType: RuntimeErrorHandler.self))
     }
+
+    func testCustomProviderSuccess() {
+        struct Behavior: LambdaServerBehavior {
+            let requestId = UUID().uuidString
+            let event = "hello"
+            func getInvocation() -> GetInvocationResult {
+                .success((self.requestId, self.event))
+            }
+
+            func processResponse(requestId: String, response: String?) -> Result<Void, ProcessResponseError> {
+                XCTAssertEqual(self.requestId, requestId, "expecting requestId to match")
+                XCTAssertEqual(self.event, response, "expecting response to match")
+                return .success(())
+            }
+
+            func processError(requestId: String, error: ErrorResponse) -> Result<Void, ProcessErrorError> {
+                XCTFail("should not report error")
+                return .failure(.internalServerError)
+            }
+
+            func processInitError(error: ErrorResponse) -> Result<Void, ProcessErrorError> {
+                XCTFail("should not report init error")
+                return .failure(.internalServerError)
+            }
+        }
+        XCTAssertNoThrow(try runLambda(behavior: Behavior(), handlerProvider: { context in
+            context.eventLoop.makeSucceededFuture(EchoHandler())
+        }))
+    }
+
+    func testCustomProviderFailure() {
+        struct Behavior: LambdaServerBehavior {
+            let requestId = UUID().uuidString
+            let event = "hello"
+            func getInvocation() -> GetInvocationResult {
+                .success((self.requestId, self.event))
+            }
+
+            func processResponse(requestId: String, response: String?) -> Result<Void, ProcessResponseError> {
+                XCTFail("should not report processing")
+                return .failure(.internalServerError)
+            }
+
+            func processError(requestId: String, error: ErrorResponse) -> Result<Void, ProcessErrorError> {
+                XCTFail("should not report error")
+                return .failure(.internalServerError)
+            }
+
+            func processInitError(error: ErrorResponse) -> Result<Void, ProcessErrorError> {
+                XCTAssertEqual(String(describing: CustomError()), error.errorMessage, "expecting error to match")
+                return .success(())
+            }
+        }
+
+        struct CustomError: Error {}
+
+        XCTAssertThrowsError(try runLambda(behavior: Behavior(), handlerProvider: { context -> EventLoopFuture<EchoHandler> in
+            context.eventLoop.makeFailedFuture(CustomError())
+        })) { error in
+            XCTAssertNotNil(error as? CustomError, "expecting error to match")
+        }
+    }
 }
diff --git a/Tests/AWSLambdaRuntimeCoreTests/Utils.swift b/Tests/AWSLambdaRuntimeCoreTests/Utils.swift
@@ -19,18 +19,32 @@ import NIOPosix
 import XCTest
 
 func runLambda<Handler: SimpleLambdaHandler>(behavior: LambdaServerBehavior, handlerType: Handler.Type) throws {
-    try runLambda(behavior: behavior, handlerType: CodableSimpleLambdaHandler<Handler>.self)
+    try runLambda(behavior: behavior, handlerProvider: CodableSimpleLambdaHandler<Handler>.makeHandler(context:))
 }
 
 func runLambda<Handler: LambdaHandler>(behavior: LambdaServerBehavior, handlerType: Handler.Type) throws {
-    try runLambda(behavior: behavior, handlerType: CodableLambdaHandler<Handler>.self)
+    try runLambda(behavior: behavior, handlerProvider: CodableLambdaHandler<Handler>.makeHandler(context:))
 }
 
 func runLambda<Handler: EventLoopLambdaHandler>(behavior: LambdaServerBehavior, handlerType: Handler.Type) throws {
-    try runLambda(behavior: behavior, handlerType: CodableEventLoopLambdaHandler<Handler>.self)
+    try runLambda(behavior: behavior, handlerProvider: CodableEventLoopLambdaHandler<Handler>.makeHandler(context:))
 }
 
-func runLambda(behavior: LambdaServerBehavior, handlerType: (some ByteBufferLambdaHandler).Type) throws {
+func runLambda<Handler: EventLoopLambdaHandler>(
+    behavior: LambdaServerBehavior,
+    handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<Handler>
+) throws {
+    try runLambda(behavior: behavior, handlerProvider: { context in
+        handlerProvider(context).map {
+            CodableEventLoopLambdaHandler(handler: $0, allocator: context.allocator)
+        }
+    })
+}
+
+func runLambda(
+    behavior: LambdaServerBehavior,
+    handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<some ByteBufferLambdaHandler>
+) throws {
     let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
     defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
     let logger = Logger(label: "TestLogger")
@@ -39,7 +53,7 @@ func runLambda(behavior: LambdaServerBehavior, handlerType: (some ByteBufferLamb
     let runner = LambdaRunner(eventLoop: eventLoopGroup.next(), configuration: configuration)
     let server = try MockLambdaServer(behavior: behavior).start().wait()
     defer { XCTAssertNoThrow(try server.stop().wait()) }
-    try runner.initialize(handlerType: handlerType, logger: logger, terminator: terminator).flatMap { handler in
+    try runner.initialize(handlerProvider: handlerProvider, logger: logger, terminator: terminator).flatMap { handler in
         runner.run(handler: handler, logger: logger)
     }.wait()
 }
@@ -89,3 +103,14 @@ extension LambdaTerminator.TerminationError: Equatable {
         return String(describing: lhs) == String(describing: rhs)
     }
 }
+
+// for backward compatibility in tests
+extension LambdaRunner {
+    func initialize<Handler: ByteBufferLambdaHandler>(
+        handlerType: Handler.Type,
+        logger: Logger,
+        terminator: LambdaTerminator
+    ) -> EventLoopFuture<Handler> {
+        self.initialize(handlerProvider: handlerType.makeHandler(context:), logger: logger, terminator: terminator)
+    }
+}
