allow custom initialization of the HandlerType of the LambdaRuntime (#310)

Motivation:

Provide the flexibility for custom initialization 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 Handler Type of Handler Provider
* Add and update tests

Originally suggested and coded by @tachyonics in #308

Author: tomerd

Sources/AWSLambdaRuntimeCore/Lambda.swift

@@ -38,7 +38,7 @@ public enum Lambda {
         configuration: LambdaConfiguration = .init(),
         handlerType: Handler.Type
     ) -> Result<Int, Error> {
-        Self.run(configuration: configuration, handlerType: CodableSimpleLambdaHandler<Handler>.self)
+        Self.run(configuration: configuration, handlerProvider: CodableSimpleLambdaHandler<Handler>.makeHandler(context:))
     }
 
     /// Run a Lambda defined by implementing the ``LambdaHandler`` protocol.
@@ -54,7 +54,7 @@ public enum Lambda {
         configuration: LambdaConfiguration = .init(),
         handlerType: Handler.Type
     ) -> Result<Int, Error> {
-        Self.run(configuration: configuration, handlerType: CodableLambdaHandler<Handler>.self)
+        Self.run(configuration: configuration, handlerProvider: CodableLambdaHandler<Handler>.makeHandler(context:))
     }
 
     /// Run a Lambda defined by implementing the ``EventLoopLambdaHandler`` protocol.
@@ -70,7 +70,7 @@ public enum Lambda {
         configuration: LambdaConfiguration = .init(),
         handlerType: Handler.Type
     ) -> Result<Int, Error> {
-        Self.run(configuration: configuration, handlerType: CodableEventLoopLambdaHandler<Handler>.self)
+        Self.run(configuration: configuration, handlerProvider: CodableEventLoopLambdaHandler<Handler>.makeHandler(context:))
     }
 
     /// Run a Lambda defined by implementing the ``ByteBufferLambdaHandler`` protocol.
@@ -85,6 +85,19 @@ 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 ``LambdaRuntimeHandler`` protocol.
+    /// - parameters:
+    ///     - configuration: A Lambda runtime configuration object
+    ///     - handlerProvider: A provider of the ``LambdaRuntimeHandler`` 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 LambdaRuntimeHandler>
     ) -> Result<Int, Error> {
         let _run = { (configuration: LambdaConfiguration) -> Result<Int, Error> in
             #if swift(<5.9)
@@ -95,7 +108,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)")

Sources/AWSLambdaRuntimeCore/LambdaHandler.swift

@@ -398,7 +398,7 @@ extension EventLoopLambdaHandler {
 /// - note: This is a low level protocol designed to power the higher level ``EventLoopLambdaHandler`` and
 ///         ``LambdaHandler`` based APIs.
 ///         Most users are not expected to use this protocol.
-public protocol ByteBufferLambdaHandler {
+public protocol ByteBufferLambdaHandler: LambdaRuntimeHandler {
     /// Create a Lambda handler for the runtime.
     ///
     /// Use this to initialize all your resources that you want to cache between invocations. This could be database
@@ -433,6 +433,28 @@ extension ByteBufferLambdaHandler {
     }
 }
 
+// MARK: - LambdaRuntimeHandler
+
+/// An `EventLoopFuture` based processing protocol for a Lambda that takes a `ByteBuffer` and returns
+/// an optional `ByteBuffer` asynchronously.
+///
+/// - note: This is a low level protocol designed to enable use cases where a frameworks initializes the
+///         runtime with a handler outside the normal initialization of
+///         ``ByteBufferLambdaHandler``, ``EventLoopLambdaHandler`` and ``LambdaHandler`` based APIs.
+///         Most users are not expected to use this protocol.
+public protocol LambdaRuntimeHandler {
+    /// The Lambda handling method.
+    /// Concrete Lambda handlers implement this method to provide the Lambda functionality.
+    ///
+    /// - parameters:
+    ///     - context: Runtime ``LambdaContext``.
+    ///     - event: The event or input payload encoded as `ByteBuffer`.
+    ///
+    /// - Returns: An `EventLoopFuture` to report the result of the Lambda back to the runtime engine.
+    ///            The `EventLoopFuture` should be completed with either a response encoded as `ByteBuffer` or an `Error`.
+    func handle(_ buffer: ByteBuffer, context: LambdaContext) -> EventLoopFuture<ByteBuffer?>
+}
+
 // MARK: - Other
 
 @usableFromInline

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: LambdaRuntimeHandler>(
+        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
@@ -59,7 +63,7 @@ internal final class LambdaRunner {
             }
     }
 
-    func run(handler: some ByteBufferLambdaHandler, logger: Logger) -> EventLoopFuture<Void> {
+    func run(handler: some LambdaRuntimeHandler, logger: Logger) -> EventLoopFuture<Void> {
         logger.debug("lambda invocation sequence starting")
         // 1. request invocation from lambda runtime engine
         self.isGettingNextInvocation = true

Sources/AWSLambdaRuntimeCore/LambdaRuntime.swift

@@ -19,12 +19,14 @@ import NIOCore
 /// `LambdaRuntime` manages the Lambda process lifecycle.
 ///
 /// Use this API, if you build a higher level web framework which shall be able to run inside the Lambda environment.
-public final class LambdaRuntime<Handler: ByteBufferLambdaHandler> {
+public final class LambdaRuntime<Handler: LambdaRuntimeHandler> {
     private let eventLoop: EventLoop
     private let shutdownPromise: EventLoopPromise<Int>
     private let logger: Logger
     private let configuration: LambdaConfiguration
 
+    private let handlerProvider: (LambdaInitializationContext) -> EventLoopFuture<Handler>
+
     private var state = State.idle {
         willSet {
             self.eventLoop.assertInEventLoop()
@@ -35,18 +37,41 @@ public final class LambdaRuntime<Handler: ByteBufferLambdaHandler> {
     /// Create a new `LambdaRuntime`.
     ///
     /// - parameters:
-    ///     - handlerType: The ``ByteBufferLambdaHandler`` type the `LambdaRuntime` shall create and manage.
+    ///     - handlerProvider: A provider of the ``Handler`` the `LambdaRuntime` will manage.
     ///     - 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())
+    @usableFromInline
+    convenience init(
+        handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<Handler>,
+        eventLoop: EventLoop,
+        logger: Logger
+    ) {
+        self.init(
+            handlerProvider: handlerProvider,
+            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 ``Handler`` the `LambdaRuntime` will manage.
+    ///     - eventLoop: An `EventLoop` to run the Lambda on.
+    ///     - logger: A `Logger` to log the Lambda events.
+    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 +110,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.
@@ -175,7 +200,7 @@ public final class LambdaRuntime<Handler: ByteBufferLambdaHandler> {
     private enum State {
         case idle
         case initializing
-        case active(LambdaRunner, any ByteBufferLambdaHandler)
+        case active(LambdaRunner, any LambdaRuntimeHandler)
         case shuttingdown
         case shutdown
 
@@ -204,8 +229,16 @@ public enum LambdaRuntimeFactory {
     ///     - eventLoop: An `EventLoop` to run the Lambda on.
     ///     - logger: A `Logger` to log the Lambda events.
     @inlinable
-    public static func makeRuntime<H: SimpleLambdaHandler>(_ handlerType: H.Type, eventLoop: any EventLoop, logger: Logger) -> LambdaRuntime<some ByteBufferLambdaHandler> {
-        LambdaRuntime<CodableSimpleLambdaHandler<H>>(CodableSimpleLambdaHandler<H>.self, eventLoop: eventLoop, logger: logger)
+    public static func makeRuntime<Handler: SimpleLambdaHandler>(
+        _ handlerType: Handler.Type,
+        eventLoop: any EventLoop,
+        logger: Logger
+    ) -> LambdaRuntime<some ByteBufferLambdaHandler> {
+        LambdaRuntime<CodableSimpleLambdaHandler<Handler>>(
+            handlerProvider: CodableSimpleLambdaHandler<Handler>.makeHandler(context:),
+            eventLoop: eventLoop,
+            logger: logger
+        )
     }
 
     /// Create a new `LambdaRuntime`.
@@ -215,8 +248,16 @@ public enum LambdaRuntimeFactory {
     ///     - eventLoop: An `EventLoop` to run the Lambda on.
     ///     - logger: A `Logger` to log the Lambda events.
     @inlinable
-    public static func makeRuntime<H: LambdaHandler>(_ handlerType: H.Type, eventLoop: any EventLoop, logger: Logger) -> LambdaRuntime<some ByteBufferLambdaHandler> {
-        LambdaRuntime<CodableLambdaHandler<H>>(CodableLambdaHandler<H>.self, eventLoop: eventLoop, logger: logger)
+    public static func makeRuntime<Handler: LambdaHandler>(
+        _ handlerType: Handler.Type,
+        eventLoop: any EventLoop,
+        logger: Logger
+    ) -> LambdaRuntime<some LambdaRuntimeHandler> {
+        LambdaRuntime<CodableLambdaHandler<Handler>>(
+            handlerProvider: CodableLambdaHandler<Handler>.makeHandler(context:),
+            eventLoop: eventLoop,
+            logger: logger
+        )
     }
 
     /// Create a new `LambdaRuntime`.
@@ -226,8 +267,79 @@ public enum LambdaRuntimeFactory {
     ///     - eventLoop: An `EventLoop` to run the Lambda on.
     ///     - logger: A `Logger` to log the Lambda events.
     @inlinable
-    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)
+    public static func makeRuntime<Handler: EventLoopLambdaHandler>(
+        _ handlerType: Handler.Type,
+        eventLoop: any EventLoop,
+        logger: Logger
+    ) -> LambdaRuntime<some LambdaRuntimeHandler> {
+        LambdaRuntime<CodableEventLoopLambdaHandler<Handler>>(
+            handlerProvider: CodableEventLoopLambdaHandler<Handler>.makeHandler(context:),
+            eventLoop: eventLoop,
+            logger: logger
+        )
+    }
+
+    /// Create a new `LambdaRuntime`.
+    ///
+    /// - parameters:
+    ///     - handlerType: The ``ByteBufferLambdaHandler`` type the `LambdaRuntime` shall create and manage.
+    ///     - eventLoop: An `EventLoop` to run the Lambda on.
+    ///     - logger: A `Logger` to log the Lambda events.
+    @inlinable
+    public static func makeRuntime<Handler: ByteBufferLambdaHandler>(
+        _ handlerType: Handler.Type,
+        eventLoop: any EventLoop,
+        logger: Logger
+    ) -> LambdaRuntime<some LambdaRuntimeHandler> {
+        LambdaRuntime<Handler>(
+            handlerProvider: Handler.makeHandler(context:),
+            eventLoop: eventLoop,
+            logger: logger
+        )
+    }
+
+    /// Create a new `LambdaRuntime`.
+    ///
+    /// - parameters:
+    ///     - handlerProvider: A provider of the ``Handler`` 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<Handler: LambdaRuntimeHandler>(
+        handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<Handler>,
+        eventLoop: any EventLoop,
+        logger: Logger
+    ) -> LambdaRuntime<Handler> {
+        LambdaRuntime(
+            handlerProvider: handlerProvider,
+            eventLoop: eventLoop,
+            logger: logger
+        )
+    }
+
+    /// Create a new `LambdaRuntime`.
+    ///
+    /// - parameters:
+    ///     - handlerProvider: A provider of the ``Handler`` 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<Handler: LambdaRuntimeHandler>(
+        handlerProvider: @escaping (LambdaInitializationContext) async throws -> Handler,
+        eventLoop: any EventLoop,
+        logger: Logger
+    ) -> LambdaRuntime<Handler> {
+        LambdaRuntime(
+            handlerProvider: { context in
+                let promise = eventLoop.makePromise(of: Handler.self)
+                promise.completeWithTask {
+                    try await handlerProvider(context)
+                }
+                return promise.futureResult
+            },
+            eventLoop: eventLoop,
+            logger: logger
+        )
     }
 }