From 44dea91ddae61556e80e1caf3c51f2d111194f39 Mon Sep 17 00:00:00 2001
From: tomer doron <tomer@apple.com>
Date: Tue, 10 Oct 2023 14:25:25 -0700
Subject: [PATCH 1/5] allow custom initialisation of the HandlerType of the
 LambdaRuntime

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 https://github.com/swift-server/swift-aws-lambda-runtime/pull/308
---
 Sources/AWSLambdaRuntimeCore/Lambda.swift     |  23 +++-
 .../AWSLambdaRuntimeCore/LambdaRunner.swift   |   8 +-
 .../AWSLambdaRuntimeCore/LambdaRuntime.swift  | 103 +++++++++++++++++-
 .../LambdaRunnerTest.swift                    |  63 +++++++++++
 Tests/AWSLambdaRuntimeCoreTests/Utils.swift   |  35 +++++-
 5 files changed, 221 insertions(+), 11 deletions(-)

diff --git a/Sources/AWSLambdaRuntimeCore/Lambda.swift b/Sources/AWSLambdaRuntimeCore/Lambda.swift
index ae225d01..55340b69 100644
--- 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
index 898fc3e3..56995bc5 100644
--- 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
index 96b77489..d35002de 100644
--- 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
index a561dbf8..2cecfcd2 100644
--- 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
index aecd3186..40db0e22 100644
--- 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)
+    }
+}

From d9e47ef912d71e4375f7d09344c8fafe13cade17 Mon Sep 17 00:00:00 2001
From: tomer doron <tomer@apple.com>
Date: Sat, 25 Nov 2023 20:44:12 -0800
Subject: [PATCH 2/5] introduce CoreByteBufferLambdaHandler

---
 Sources/AWSLambdaRuntimeCore/Lambda.swift     |   8 +-
 .../AWSLambdaRuntimeCore/LambdaHandler.swift  |  17 ++-
 .../AWSLambdaRuntimeCore/LambdaRunner.swift   |   4 +-
 .../AWSLambdaRuntimeCore/LambdaRuntime.swift  | 100 ++++++++----------
 4 files changed, 67 insertions(+), 62 deletions(-)

diff --git a/Sources/AWSLambdaRuntimeCore/Lambda.swift b/Sources/AWSLambdaRuntimeCore/Lambda.swift
index 55340b69..ef37fd34 100644
--- a/Sources/AWSLambdaRuntimeCore/Lambda.swift
+++ b/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.
@@ -100,7 +100,7 @@ public enum Lambda {
     /// - 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>
+        handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<some CoreByteBufferLambdaHandler>
     ) -> Result<Int, Error> {
         let _run = { (configuration: LambdaConfiguration) -> Result<Int, Error> in
             #if swift(<5.9)
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift b/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
index fc3611ba..5ebaafc7 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
+++ b/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: CoreByteBufferLambdaHandler {
     /// 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,21 @@ extension ByteBufferLambdaHandler {
     }
 }
 
+// MARK: - FIXME
+
+public protocol CoreByteBufferLambdaHandler {
+    /// 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
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
index 56995bc5..bca0e30a 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
@@ -33,7 +33,7 @@ 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>(
+    func initialize<Handler: CoreByteBufferLambdaHandler>(
         handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<Handler>,
         logger: Logger,
         terminator: LambdaTerminator
@@ -63,7 +63,7 @@ internal final class LambdaRunner {
             }
     }
 
-    func run(handler: some ByteBufferLambdaHandler, logger: Logger) -> EventLoopFuture<Void> {
+    func run(handler: some CoreByteBufferLambdaHandler, logger: Logger) -> EventLoopFuture<Void> {
         logger.debug("lambda invocation sequence starting")
         // 1. request invocation from lambda runtime engine
         self.isGettingNextInvocation = true
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRuntime.swift b/Sources/AWSLambdaRuntimeCore/LambdaRuntime.swift
index d35002de..eeff2c08 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRuntime.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRuntime.swift
@@ -19,7 +19,7 @@ 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: CoreByteBufferLambdaHandler> {
     private let eventLoop: EventLoop
     private let shutdownPromise: EventLoopPromise<Int>
     private let logger: Logger
@@ -34,23 +34,14 @@ public final class LambdaRuntime<Handler: ByteBufferLambdaHandler> {
         }
     }
 
-    /// 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.
-    public convenience init(_ handlerType: Handler.Type, eventLoop: EventLoop, logger: Logger) {
-        self.init(handlerProvider: handlerType.makeHandler(context:), eventLoop: eventLoop, logger: logger, configuration: .init())
-    }
-
     /// 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(
+    @usableFromInline
+    convenience init(
         handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<Handler>,
         eventLoop: EventLoop,
         logger: Logger
@@ -63,6 +54,12 @@ public final class LambdaRuntime<Handler: ByteBufferLambdaHandler> {
         )
     }
 
+    /// 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.
     init(
         handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<Handler>,
         eventLoop: EventLoop,
@@ -203,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 CoreByteBufferLambdaHandler)
         case shuttingdown
         case shutdown
 
@@ -232,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`.
@@ -243,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 CoreByteBufferLambdaHandler> {
+        LambdaRuntime<CodableLambdaHandler<Handler>>(
+            handlerProvider: CodableLambdaHandler<Handler>.makeHandler(context:),
+            eventLoop: eventLoop,
+            logger: logger
+        )
     }
 
     /// Create a new `LambdaRuntime`.
@@ -254,28 +267,13 @@ 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)
-    }
-
-    /// 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>,
+    public static func makeRuntime<Handler: EventLoopLambdaHandler>(
+        _ handlerType: Handler.Type,
         eventLoop: any EventLoop,
         logger: Logger
-    ) -> LambdaRuntime<some ByteBufferLambdaHandler> {
-        LambdaRuntime(
-            handlerProvider: { context in
-                handlerProvider(context).map {
-                    CodableSimpleLambdaHandler(handler: $0, allocator: context.allocator)
-                }
-            },
+    ) -> LambdaRuntime<some CoreByteBufferLambdaHandler> {
+        LambdaRuntime<CodableEventLoopLambdaHandler<Handler>>(
+            handlerProvider: CodableEventLoopLambdaHandler<Handler>.makeHandler(context:),
             eventLoop: eventLoop,
             logger: logger
         )
@@ -284,21 +282,17 @@ public enum LambdaRuntimeFactory {
     /// Create a new `LambdaRuntime`.
     ///
     /// - parameters:
-    ///     - handlerProvider: A provider of the ``LambdaHandler``  the `LambdaRuntime` will manage.
+    ///     - handlerType: The ``EventLoopLambdaHandler`` 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<H: LambdaHandler>(
-        handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<H>,
+    public static func makeRuntime<Handler: ByteBufferLambdaHandler>(
+        _ handlerType: Handler.Type,
         eventLoop: any EventLoop,
         logger: Logger
-    ) -> LambdaRuntime<some ByteBufferLambdaHandler> {
-        LambdaRuntime(
-            handlerProvider: { context in
-                handlerProvider(context).map {
-                    CodableLambdaHandler(handler: $0, allocator: context.allocator)
-                }
-            },
+    ) -> LambdaRuntime<some CoreByteBufferLambdaHandler> {
+        LambdaRuntime<Handler>(
+            handlerProvider: Handler.makeHandler(context:),
             eventLoop: eventLoop,
             logger: logger
         )
@@ -311,17 +305,13 @@ 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>(
-        handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<H>,
+    public static func makeRuntime<Handler: CoreByteBufferLambdaHandler>(
+        handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<Handler>,
         eventLoop: any EventLoop,
         logger: Logger
-    ) -> LambdaRuntime<some ByteBufferLambdaHandler> {
+    ) -> LambdaRuntime<Handler> {
         LambdaRuntime(
-            handlerProvider: { context in
-                handlerProvider(context).map {
-                    CodableEventLoopLambdaHandler(handler: $0, allocator: context.allocator)
-                }
-            },
+            handlerProvider: handlerProvider,
             eventLoop: eventLoop,
             logger: logger
         )

From 729e0dc3ef05153f59e752ebe62c0901dd0badf5 Mon Sep 17 00:00:00 2001
From: tomer doron <tomer@apple.com>
Date: Thu, 7 Dec 2023 10:28:35 -0800
Subject: [PATCH 3/5] rename

---
 Sources/AWSLambdaRuntimeCore/Lambda.swift     |  2 +-
 .../AWSLambdaRuntimeCore/LambdaHandler.swift  | 13 +++-
 .../AWSLambdaRuntimeCore/LambdaRunner.swift   |  4 +-
 .../AWSLambdaRuntimeCore/LambdaRuntime.swift  | 41 +++++++++---
 .../LambdaRunnerTest.swift                    | 62 +++++++++++++++++++
 Tests/AWSLambdaRuntimeCoreTests/Utils.swift   | 27 ++++++++
 6 files changed, 135 insertions(+), 14 deletions(-)

diff --git a/Sources/AWSLambdaRuntimeCore/Lambda.swift b/Sources/AWSLambdaRuntimeCore/Lambda.swift
index ef37fd34..a0aaaf38 100644
--- a/Sources/AWSLambdaRuntimeCore/Lambda.swift
+++ b/Sources/AWSLambdaRuntimeCore/Lambda.swift
@@ -100,7 +100,7 @@ public enum Lambda {
     /// - 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 CoreByteBufferLambdaHandler>
+        handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<some NonInitializingByteBufferLambdaHandler>
     ) -> Result<Int, Error> {
         let _run = { (configuration: LambdaConfiguration) -> Result<Int, Error> in
             #if swift(<5.9)
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift b/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
index 5ebaafc7..95dc2ebc 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
+++ b/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: CoreByteBufferLambdaHandler {
+public protocol ByteBufferLambdaHandler: NonInitializingByteBufferLambdaHandler {
     /// 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,9 +433,16 @@ extension ByteBufferLambdaHandler {
     }
 }
 
-// MARK: - FIXME
+// MARK: - NonInitializingByteBufferLambdaHandler
 
-public protocol CoreByteBufferLambdaHandler {
+/// 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 NonInitializingByteBufferLambdaHandler {
     /// The Lambda handling method.
     /// Concrete Lambda handlers implement this method to provide the Lambda functionality.
     ///
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
index bca0e30a..72d04848 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
@@ -33,7 +33,7 @@ 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: CoreByteBufferLambdaHandler>(
+    func initialize<Handler: NonInitializingByteBufferLambdaHandler>(
         handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<Handler>,
         logger: Logger,
         terminator: LambdaTerminator
@@ -63,7 +63,7 @@ internal final class LambdaRunner {
             }
     }
 
-    func run(handler: some CoreByteBufferLambdaHandler, logger: Logger) -> EventLoopFuture<Void> {
+    func run(handler: some NonInitializingByteBufferLambdaHandler, logger: Logger) -> EventLoopFuture<Void> {
         logger.debug("lambda invocation sequence starting")
         // 1. request invocation from lambda runtime engine
         self.isGettingNextInvocation = true
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRuntime.swift b/Sources/AWSLambdaRuntimeCore/LambdaRuntime.swift
index eeff2c08..155690d9 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRuntime.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRuntime.swift
@@ -19,7 +19,7 @@ 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: CoreByteBufferLambdaHandler> {
+public final class LambdaRuntime<Handler: NonInitializingByteBufferLambdaHandler> {
     private let eventLoop: EventLoop
     private let shutdownPromise: EventLoopPromise<Int>
     private let logger: Logger
@@ -200,7 +200,7 @@ public final class LambdaRuntime<Handler: CoreByteBufferLambdaHandler> {
     private enum State {
         case idle
         case initializing
-        case active(LambdaRunner, any CoreByteBufferLambdaHandler)
+        case active(LambdaRunner, any NonInitializingByteBufferLambdaHandler)
         case shuttingdown
         case shutdown
 
@@ -252,7 +252,7 @@ public enum LambdaRuntimeFactory {
         _ handlerType: Handler.Type,
         eventLoop: any EventLoop,
         logger: Logger
-    ) -> LambdaRuntime<some CoreByteBufferLambdaHandler> {
+    ) -> LambdaRuntime<some NonInitializingByteBufferLambdaHandler> {
         LambdaRuntime<CodableLambdaHandler<Handler>>(
             handlerProvider: CodableLambdaHandler<Handler>.makeHandler(context:),
             eventLoop: eventLoop,
@@ -271,7 +271,7 @@ public enum LambdaRuntimeFactory {
         _ handlerType: Handler.Type,
         eventLoop: any EventLoop,
         logger: Logger
-    ) -> LambdaRuntime<some CoreByteBufferLambdaHandler> {
+    ) -> LambdaRuntime<some NonInitializingByteBufferLambdaHandler> {
         LambdaRuntime<CodableEventLoopLambdaHandler<Handler>>(
             handlerProvider: CodableEventLoopLambdaHandler<Handler>.makeHandler(context:),
             eventLoop: eventLoop,
@@ -282,7 +282,7 @@ public enum LambdaRuntimeFactory {
     /// Create a new `LambdaRuntime`.
     ///
     /// - parameters:
-    ///     - handlerType: The ``EventLoopLambdaHandler`` type the `LambdaRuntime` shall create and manage.
+    ///     - 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
@@ -290,7 +290,7 @@ public enum LambdaRuntimeFactory {
         _ handlerType: Handler.Type,
         eventLoop: any EventLoop,
         logger: Logger
-    ) -> LambdaRuntime<some CoreByteBufferLambdaHandler> {
+    ) -> LambdaRuntime<some NonInitializingByteBufferLambdaHandler> {
         LambdaRuntime<Handler>(
             handlerProvider: Handler.makeHandler(context:),
             eventLoop: eventLoop,
@@ -301,11 +301,11 @@ public enum LambdaRuntimeFactory {
     /// Create a new `LambdaRuntime`.
     ///
     /// - parameters:
-    ///     - handlerProvider: A provider of the ``EventLoopLambdaHandler``  the `LambdaRuntime` will manage.
+    ///     - handlerProvider: A provider of the ``CoreByteBufferLambdaHandler``  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: CoreByteBufferLambdaHandler>(
+    public static func makeRuntime<Handler: NonInitializingByteBufferLambdaHandler>(
         handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<Handler>,
         eventLoop: any EventLoop,
         logger: Logger
@@ -316,6 +316,31 @@ public enum LambdaRuntimeFactory {
             logger: logger
         )
     }
+
+    /// Create a new `LambdaRuntime`.
+    ///
+    /// - parameters:
+    ///     - handlerProvider: A provider of the ``CoreByteBufferLambdaHandler``  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: NonInitializingByteBufferLambdaHandler>(
+        handlerProvider: @escaping (LambdaInitializationContext) async throws -> Handler,
+        eventLoop: any EventLoop,
+        logger: Logger
+    ) -> LambdaRuntime<Handler> {
+        return LambdaRuntime(
+            handlerProvider: { context in
+                let promise = eventLoop.makePromise(of: Handler.self)
+                promise.completeWithTask {
+                    try await handlerProvider(context)
+                }
+                return promise.futureResult
+            },
+            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
index 2cecfcd2..1b74949d 100644
--- a/Tests/AWSLambdaRuntimeCoreTests/LambdaRunnerTest.swift
+++ b/Tests/AWSLambdaRuntimeCoreTests/LambdaRunnerTest.swift
@@ -131,4 +131,66 @@ class LambdaRunnerTest: XCTestCase {
             XCTAssertNotNil(error as? CustomError, "expecting error to match")
         }
     }
+
+    func testCustomAsyncProviderSuccess() {
+        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 async throws -> EchoHandler in
+            EchoHandler()
+        }))
+    }
+
+    func testCustomAsyncProviderFailure() {
+        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 async throws -> EchoHandler in
+            throw CustomError()
+        })) { error in
+            XCTAssertNotNil(error as? CustomError, "expecting error to match")
+        }
+    }
 }
diff --git a/Tests/AWSLambdaRuntimeCoreTests/Utils.swift b/Tests/AWSLambdaRuntimeCoreTests/Utils.swift
index 40db0e22..8e2d3c38 100644
--- a/Tests/AWSLambdaRuntimeCoreTests/Utils.swift
+++ b/Tests/AWSLambdaRuntimeCoreTests/Utils.swift
@@ -41,6 +41,33 @@ func runLambda<Handler: EventLoopLambdaHandler>(
     })
 }
 
+func runLambda<Handler: EventLoopLambdaHandler>(
+    behavior: LambdaServerBehavior,
+    handlerProvider: @escaping (LambdaInitializationContext) async throws -> Handler
+) throws {
+    try runLambda(behavior: behavior, handlerProvider: { context in
+        let handler = try await handlerProvider(context)
+        return CodableEventLoopLambdaHandler(handler: handler, allocator: context.allocator)
+    })
+}
+
+func runLambda<Handler: ByteBufferLambdaHandler>(
+    behavior: LambdaServerBehavior,
+    handlerProvider: @escaping (LambdaInitializationContext) async throws -> Handler
+) throws {
+    let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+    try runLambda(
+        behavior: behavior,
+        handlerProvider: { context in
+            let promise = eventLoopGroup.next().makePromise(of: Handler.self)
+            promise.completeWithTask {
+                try await handlerProvider(context)
+            }
+            return promise.futureResult
+        }
+    )
+}
+
 func runLambda(
     behavior: LambdaServerBehavior,
     handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<some ByteBufferLambdaHandler>

From a7eb20eb9af26639c0daba7a5e2f0e2e9eae81a8 Mon Sep 17 00:00:00 2001
From: tomer doron <tomer@apple.com>
Date: Fri, 8 Dec 2023 13:17:59 -0800
Subject: [PATCH 4/5] another name!

---
 Sources/AWSLambdaRuntimeCore/Lambda.swift     |  9 +++----
 .../AWSLambdaRuntimeCore/LambdaHandler.swift  | 10 ++++----
 .../AWSLambdaRuntimeCore/LambdaRunner.swift   |  4 ++--
 .../AWSLambdaRuntimeCore/LambdaRuntime.swift  | 24 +++++++++----------
 .../LambdaRunnerTest.swift                    |  4 ++--
 5 files changed, 24 insertions(+), 27 deletions(-)

diff --git a/Sources/AWSLambdaRuntimeCore/Lambda.swift b/Sources/AWSLambdaRuntimeCore/Lambda.swift
index a0aaaf38..5bd0adb9 100644
--- a/Sources/AWSLambdaRuntimeCore/Lambda.swift
+++ b/Sources/AWSLambdaRuntimeCore/Lambda.swift
@@ -89,18 +89,15 @@ public enum Lambda {
         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.
-    ///
+    /// Run a Lambda defined by implementing the ``LambdaRuntimeHandler`` protocol.
     /// - parameters:
     ///     - configuration: A Lambda runtime configuration object
-    ///     - handlerProvider: A provider of the Handler to invoke.
+    ///     - 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 NonInitializingByteBufferLambdaHandler>
+        handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<some LambdaRuntimeHandler>
     ) -> Result<Int, Error> {
         let _run = { (configuration: LambdaConfiguration) -> Result<Int, Error> in
             #if swift(<5.9)
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift b/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
index 95dc2ebc..3a7e3c27 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
+++ b/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: NonInitializingByteBufferLambdaHandler {
+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,16 +433,16 @@ extension ByteBufferLambdaHandler {
     }
 }
 
-// MARK: - NonInitializingByteBufferLambdaHandler
+// 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 
+/// - 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 NonInitializingByteBufferLambdaHandler {
+public protocol LambdaRuntimeHandler {
     /// The Lambda handling method.
     /// Concrete Lambda handlers implement this method to provide the Lambda functionality.
     ///
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
index 72d04848..9557f41f 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
@@ -33,7 +33,7 @@ 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: NonInitializingByteBufferLambdaHandler>(
+    func initialize<Handler: LambdaRuntimeHandler>(
         handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<Handler>,
         logger: Logger,
         terminator: LambdaTerminator
@@ -63,7 +63,7 @@ internal final class LambdaRunner {
             }
     }
 
-    func run(handler: some NonInitializingByteBufferLambdaHandler, 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
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRuntime.swift b/Sources/AWSLambdaRuntimeCore/LambdaRuntime.swift
index 155690d9..c570a0b3 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRuntime.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRuntime.swift
@@ -19,7 +19,7 @@ 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: NonInitializingByteBufferLambdaHandler> {
+public final class LambdaRuntime<Handler: LambdaRuntimeHandler> {
     private let eventLoop: EventLoop
     private let shutdownPromise: EventLoopPromise<Int>
     private let logger: Logger
@@ -37,7 +37,7 @@ public final class LambdaRuntime<Handler: NonInitializingByteBufferLambdaHandler
     /// Create a new `LambdaRuntime`.
     ///
     /// - parameters:
-    ///     - handlerProvider: A provider of the ``ByteBufferLambdaHandler``  the `LambdaRuntime` will 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.
     @usableFromInline
@@ -57,7 +57,7 @@ public final class LambdaRuntime<Handler: NonInitializingByteBufferLambdaHandler
     /// Create a new `LambdaRuntime`.
     ///
     /// - parameters:
-    ///     - handlerProvider: A provider of the ``ByteBufferLambdaHandler``  the `LambdaRuntime` will 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.
     init(
@@ -200,7 +200,7 @@ public final class LambdaRuntime<Handler: NonInitializingByteBufferLambdaHandler
     private enum State {
         case idle
         case initializing
-        case active(LambdaRunner, any NonInitializingByteBufferLambdaHandler)
+        case active(LambdaRunner, any LambdaRuntimeHandler)
         case shuttingdown
         case shutdown
 
@@ -252,7 +252,7 @@ public enum LambdaRuntimeFactory {
         _ handlerType: Handler.Type,
         eventLoop: any EventLoop,
         logger: Logger
-    ) -> LambdaRuntime<some NonInitializingByteBufferLambdaHandler> {
+    ) -> LambdaRuntime<some LambdaRuntimeHandler> {
         LambdaRuntime<CodableLambdaHandler<Handler>>(
             handlerProvider: CodableLambdaHandler<Handler>.makeHandler(context:),
             eventLoop: eventLoop,
@@ -271,7 +271,7 @@ public enum LambdaRuntimeFactory {
         _ handlerType: Handler.Type,
         eventLoop: any EventLoop,
         logger: Logger
-    ) -> LambdaRuntime<some NonInitializingByteBufferLambdaHandler> {
+    ) -> LambdaRuntime<some LambdaRuntimeHandler> {
         LambdaRuntime<CodableEventLoopLambdaHandler<Handler>>(
             handlerProvider: CodableEventLoopLambdaHandler<Handler>.makeHandler(context:),
             eventLoop: eventLoop,
@@ -290,7 +290,7 @@ public enum LambdaRuntimeFactory {
         _ handlerType: Handler.Type,
         eventLoop: any EventLoop,
         logger: Logger
-    ) -> LambdaRuntime<some NonInitializingByteBufferLambdaHandler> {
+    ) -> LambdaRuntime<some LambdaRuntimeHandler> {
         LambdaRuntime<Handler>(
             handlerProvider: Handler.makeHandler(context:),
             eventLoop: eventLoop,
@@ -301,11 +301,11 @@ public enum LambdaRuntimeFactory {
     /// Create a new `LambdaRuntime`.
     ///
     /// - parameters:
-    ///     - handlerProvider: A provider of the ``CoreByteBufferLambdaHandler``  the `LambdaRuntime` will 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.
     @inlinable
-    public static func makeRuntime<Handler: NonInitializingByteBufferLambdaHandler>(
+    public static func makeRuntime<Handler: LambdaRuntimeHandler>(
         handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<Handler>,
         eventLoop: any EventLoop,
         logger: Logger
@@ -320,16 +320,16 @@ public enum LambdaRuntimeFactory {
     /// Create a new `LambdaRuntime`.
     ///
     /// - parameters:
-    ///     - handlerProvider: A provider of the ``CoreByteBufferLambdaHandler``  the `LambdaRuntime` will 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.
     @inlinable
-    public static func makeRuntime<Handler: NonInitializingByteBufferLambdaHandler>(
+    public static func makeRuntime<Handler: LambdaRuntimeHandler>(
         handlerProvider: @escaping (LambdaInitializationContext) async throws -> Handler,
         eventLoop: any EventLoop,
         logger: Logger
     ) -> LambdaRuntime<Handler> {
-        return LambdaRuntime(
+        LambdaRuntime(
             handlerProvider: { context in
                 let promise = eventLoop.makePromise(of: Handler.self)
                 promise.completeWithTask {
diff --git a/Tests/AWSLambdaRuntimeCoreTests/LambdaRunnerTest.swift b/Tests/AWSLambdaRuntimeCoreTests/LambdaRunnerTest.swift
index 1b74949d..6fd91aec 100644
--- a/Tests/AWSLambdaRuntimeCoreTests/LambdaRunnerTest.swift
+++ b/Tests/AWSLambdaRuntimeCoreTests/LambdaRunnerTest.swift
@@ -156,7 +156,7 @@ class LambdaRunnerTest: XCTestCase {
                 return .failure(.internalServerError)
             }
         }
-        XCTAssertNoThrow(try runLambda(behavior: Behavior(), handlerProvider: { context async throws -> EchoHandler in
+        XCTAssertNoThrow(try runLambda(behavior: Behavior(), handlerProvider: { _ async throws -> EchoHandler in
             EchoHandler()
         }))
     }
@@ -187,7 +187,7 @@ class LambdaRunnerTest: XCTestCase {
 
         struct CustomError: Error {}
 
-        XCTAssertThrowsError(try runLambda(behavior: Behavior(), handlerProvider: { context async throws -> EchoHandler in
+        XCTAssertThrowsError(try runLambda(behavior: Behavior(), handlerProvider: { _ async throws -> EchoHandler in
             throw CustomError()
         })) { error in
             XCTAssertNotNil(error as? CustomError, "expecting error to match")

From 5283e93418543830fe30b3f428e25bfbc99621a5 Mon Sep 17 00:00:00 2001
From: tomer doron <tomer@apple.com>
Date: Fri, 8 Dec 2023 19:11:17 -0800
Subject: [PATCH 5/5] Update Sources/AWSLambdaRuntimeCore/Lambda.swift

---
 Sources/AWSLambdaRuntimeCore/Lambda.swift | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Sources/AWSLambdaRuntimeCore/Lambda.swift b/Sources/AWSLambdaRuntimeCore/Lambda.swift
index 5bd0adb9..f23dc57c 100644
--- a/Sources/AWSLambdaRuntimeCore/Lambda.swift
+++ b/Sources/AWSLambdaRuntimeCore/Lambda.swift
@@ -92,7 +92,7 @@ public enum Lambda {
     /// Run a Lambda defined by implementing the ``LambdaRuntimeHandler`` protocol.
     /// - parameters:
     ///     - configuration: A Lambda runtime configuration object
-    ///     - handlerProvider: A provider of the ``LambdaRuntimeHandler``` to invoke.
+    ///     - 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(