Skip to content

Commit e7ed5fd

Browse files
tomerdtomerd
committed
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 #308
1 parent f30a585 commit e7ed5fd

File tree

5 files changed

+221
-11
lines changed

5 files changed

+221
-11
lines changed

‎Sources/AWSLambdaRuntimeCore/Lambda.swift

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,22 @@ public enum Lambda {
8585
internal static func run(
8686
configuration: LambdaConfiguration = .init(),
8787
handlerType: (some ByteBufferLambdaHandler).Type
88+
) -> Result<Int, Error> {
89+
Self.run(configuration: configuration, handlerProvider: handlerType.makeHandler(context:))
90+
}
91+
92+
/// Run a Lambda defined by implementing the ``ByteBufferLambdaHandler`` protocol.
93+
/// The Runtime will manage the Lambdas application lifecycle automatically. It will invoke the
94+
/// ``ByteBufferLambdaHandler/makeHandler(context:)`` to create a new Handler.
95+
///
96+
/// - parameters:
97+
/// - configuration: A Lambda runtime configuration object
98+
/// - handlerProvider: A provider of the Handler to invoke.
99+
///
100+
/// - note: This is a blocking operation that will run forever, as its lifecycle is managed by the AWS Lambda Runtime Engine.
101+
internal static func run(
102+
configuration: LambdaConfiguration = .init(),
103+
handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<some ByteBufferLambdaHandler>
88104
) -> Result<Int, Error> {
89105
let _run = { (configuration: LambdaConfiguration) -> Result<Int, Error> in
90106
#if swift(<5.9)
@@ -95,7 +111,12 @@ public enum Lambda {
95111

96112
var result: Result<Int, Error>!
97113
MultiThreadedEventLoopGroup.withCurrentThreadAsEventLoop { eventLoop in
98-
let runtime = LambdaRuntime(handlerType: handlerType, eventLoop: eventLoop, logger: logger, configuration: configuration)
114+
let runtime = LambdaRuntime(
115+
handlerProvider: handlerProvider,
116+
eventLoop: eventLoop,
117+
logger: logger,
118+
configuration: configuration
119+
)
99120
#if DEBUG
100121
let signalSource = trap(signal: configuration.lifecycle.stopSignal) { signal in
101122
logger.info("intercepted signal: \(signal)")

‎Sources/AWSLambdaRuntimeCore/LambdaRunner.swift

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,11 @@ internal final class LambdaRunner {
3333
/// Run the user provided initializer. This *must* only be called once.
3434
///
3535
/// - Returns: An `EventLoopFuture<LambdaHandler>` fulfilled with the outcome of the initialization.
36-
func initialize<Handler: ByteBufferLambdaHandler>(handlerType: Handler.Type, logger: Logger, terminator: LambdaTerminator) -> EventLoopFuture<Handler> {
36+
func initialize<Handler: ByteBufferLambdaHandler>(
37+
handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<Handler>,
38+
logger: Logger,
39+
terminator: LambdaTerminator
40+
) -> EventLoopFuture<Handler> {
3741
logger.debug("initializing lambda")
3842
// 1. create the handler from the factory
3943
// 2. report initialization error if one occurred
@@ -44,7 +48,7 @@ internal final class LambdaRunner {
4448
terminator: terminator
4549
)
4650

47-
return handlerType.makeHandler(context: context)
51+
return handlerProvider(context)
4852
// Hopping back to "our" EventLoop is important in case the factory returns a future
4953
// that originated from a foreign EventLoop/EventLoopGroup.
5054
// This can happen if the factory uses a library (let's say a database client) that manages its own threads/loops

‎Sources/AWSLambdaRuntimeCore/LambdaRuntime.swift

Lines changed: 100 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ public final class LambdaRuntime<Handler: ByteBufferLambdaHandler> {
2525
private let logger: Logger
2626
private let configuration: LambdaConfiguration
2727

28+
private let handlerProvider: (LambdaInitializationContext) -> EventLoopFuture<Handler>
29+
2830
private var state = State.idle {
2931
willSet {
3032
self.eventLoop.assertInEventLoop()
@@ -39,14 +41,40 @@ public final class LambdaRuntime<Handler: ByteBufferLambdaHandler> {
3941
/// - eventLoop: An `EventLoop` to run the Lambda on.
4042
/// - logger: A `Logger` to log the Lambda events.
4143
public convenience init(_ handlerType: Handler.Type, eventLoop: EventLoop, logger: Logger) {
42-
self.init(handlerType: handlerType, eventLoop: eventLoop, logger: logger, configuration: .init())
44+
self.init(handlerProvider: handlerType.makeHandler(context:), eventLoop: eventLoop, logger: logger, configuration: .init())
4345
}
4446

45-
init(handlerType: Handler.Type, eventLoop: EventLoop, logger: Logger, configuration: LambdaConfiguration) {
47+
/// Create a new `LambdaRuntime`.
48+
///
49+
/// - parameters:
50+
/// - handlerProvider: A provider of the ``ByteBufferLambdaHandler`` the `LambdaRuntime` will manage.
51+
/// - eventLoop: An `EventLoop` to run the Lambda on.
52+
/// - logger: A `Logger` to log the Lambda events.
53+
public convenience init(
54+
handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<Handler>,
55+
eventLoop: EventLoop,
56+
logger: Logger
57+
) {
58+
self.init(
59+
handlerProvider: handlerProvider,
60+
eventLoop: eventLoop,
61+
logger: logger,
62+
configuration: .init()
63+
)
64+
}
65+
66+
init(
67+
handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<Handler>,
68+
eventLoop: EventLoop,
69+
logger: Logger,
70+
configuration: LambdaConfiguration
71+
) {
4672
self.eventLoop = eventLoop
4773
self.shutdownPromise = eventLoop.makePromise(of: Int.self)
4874
self.logger = logger
4975
self.configuration = configuration
76+
77+
self.handlerProvider = handlerProvider
5078
}
5179

5280
deinit {
@@ -85,7 +113,7 @@ public final class LambdaRuntime<Handler: ByteBufferLambdaHandler> {
85113
let terminator = LambdaTerminator()
86114
let runner = LambdaRunner(eventLoop: self.eventLoop, configuration: self.configuration)
87115

88-
let startupFuture = runner.initialize(handlerType: Handler.self, logger: logger, terminator: terminator)
116+
let startupFuture = runner.initialize(handlerProvider: self.handlerProvider, logger: logger, terminator: terminator)
89117
startupFuture.flatMap { handler -> EventLoopFuture<Result<Int, Error>> in
90118
// after the startup future has succeeded, we have a handler that we can use
91119
// to `run` the lambda.
@@ -229,6 +257,75 @@ public enum LambdaRuntimeFactory {
229257
public static func makeRuntime<H: EventLoopLambdaHandler>(_ handlerType: H.Type, eventLoop: any EventLoop, logger: Logger) -> LambdaRuntime<some ByteBufferLambdaHandler> {
230258
LambdaRuntime<CodableEventLoopLambdaHandler<H>>(CodableEventLoopLambdaHandler<H>.self, eventLoop: eventLoop, logger: logger)
231259
}
260+
261+
/// Create a new `LambdaRuntime`.
262+
///
263+
/// - parameters:
264+
/// - handlerProvider: A provider of the ``SimpleLambdaHandler`` the `LambdaRuntime` will manage.
265+
/// - eventLoop: An `EventLoop` to run the Lambda on.
266+
/// - logger: A `Logger` to log the Lambda events.
267+
@inlinable
268+
public static func makeRuntime<H: SimpleLambdaHandler>(
269+
handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<H>,
270+
eventLoop: any EventLoop,
271+
logger: Logger
272+
) -> LambdaRuntime<some ByteBufferLambdaHandler> {
273+
LambdaRuntime(
274+
handlerProvider: { context in
275+
handlerProvider(context).map {
276+
CodableSimpleLambdaHandler(handler: $0, allocator: context.allocator)
277+
}
278+
},
279+
eventLoop: eventLoop,
280+
logger: logger
281+
)
282+
}
283+
284+
/// Create a new `LambdaRuntime`.
285+
///
286+
/// - parameters:
287+
/// - handlerProvider: A provider of the ``LambdaHandler`` the `LambdaRuntime` will manage.
288+
/// - eventLoop: An `EventLoop` to run the Lambda on.
289+
/// - logger: A `Logger` to log the Lambda events.
290+
@inlinable
291+
public static func makeRuntime<H: LambdaHandler>(
292+
handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<H>,
293+
eventLoop: any EventLoop,
294+
logger: Logger
295+
) -> LambdaRuntime<some ByteBufferLambdaHandler> {
296+
LambdaRuntime(
297+
handlerProvider: { context in
298+
handlerProvider(context).map {
299+
CodableLambdaHandler(handler: $0, allocator: context.allocator)
300+
}
301+
},
302+
eventLoop: eventLoop,
303+
logger: logger
304+
)
305+
}
306+
307+
/// Create a new `LambdaRuntime`.
308+
///
309+
/// - parameters:
310+
/// - handlerProvider: A provider of the ``EventLoopLambdaHandler`` the `LambdaRuntime` will manage.
311+
/// - eventLoop: An `EventLoop` to run the Lambda on.
312+
/// - logger: A `Logger` to log the Lambda events.
313+
@inlinable
314+
public static func makeRuntime<H: EventLoopLambdaHandler>(
315+
handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<H>,
316+
eventLoop: any EventLoop,
317+
logger: Logger
318+
) -> LambdaRuntime<some ByteBufferLambdaHandler> {
319+
LambdaRuntime(
320+
handlerProvider: { context in
321+
handlerProvider(context).map {
322+
CodableEventLoopLambdaHandler(handler: $0, allocator: context.allocator)
323+
}
324+
},
325+
eventLoop: eventLoop,
326+
logger: logger
327+
)
328+
}
232329
}
233330

234331
/// This is safe since lambda runtime synchronizes by dispatching all methods to a single `EventLoop`

‎Tests/AWSLambdaRuntimeCoreTests/LambdaRunnerTest.swift

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
//===----------------------------------------------------------------------===//
1414

1515
@testable import AWSLambdaRuntimeCore
16+
import NIOCore
1617
import XCTest
1718

1819
class LambdaRunnerTest: XCTestCase {
@@ -68,4 +69,66 @@ class LambdaRunnerTest: XCTestCase {
6869
}
6970
XCTAssertNoThrow(try runLambda(behavior: Behavior(), handlerType: RuntimeErrorHandler.self))
7071
}
72+
73+
func testCustomProviderSuccess() {
74+
struct Behavior: LambdaServerBehavior {
75+
let requestId = UUID().uuidString
76+
let event = "hello"
77+
func getInvocation() -> GetInvocationResult {
78+
.success((self.requestId, self.event))
79+
}
80+
81+
func processResponse(requestId: String, response: String?) -> Result<Void, ProcessResponseError> {
82+
XCTAssertEqual(self.requestId, requestId, "expecting requestId to match")
83+
XCTAssertEqual(self.event, response, "expecting response to match")
84+
return .success(())
85+
}
86+
87+
func processError(requestId: String, error: ErrorResponse) -> Result<Void, ProcessErrorError> {
88+
XCTFail("should not report error")
89+
return .failure(.internalServerError)
90+
}
91+
92+
func processInitError(error: ErrorResponse) -> Result<Void, ProcessErrorError> {
93+
XCTFail("should not report init error")
94+
return .failure(.internalServerError)
95+
}
96+
}
97+
XCTAssertNoThrow(try runLambda(behavior: Behavior(), handlerProvider: { context in
98+
context.eventLoop.makeSucceededFuture(EchoHandler())
99+
}))
100+
}
101+
102+
func testCustomProviderFailure() {
103+
struct Behavior: LambdaServerBehavior {
104+
let requestId = UUID().uuidString
105+
let event = "hello"
106+
func getInvocation() -> GetInvocationResult {
107+
.success((self.requestId, self.event))
108+
}
109+
110+
func processResponse(requestId: String, response: String?) -> Result<Void, ProcessResponseError> {
111+
XCTFail("should not report processing")
112+
return .failure(.internalServerError)
113+
}
114+
115+
func processError(requestId: String, error: ErrorResponse) -> Result<Void, ProcessErrorError> {
116+
XCTFail("should not report error")
117+
return .failure(.internalServerError)
118+
}
119+
120+
func processInitError(error: ErrorResponse) -> Result<Void, ProcessErrorError> {
121+
XCTAssertEqual(String(describing: CustomError()), error.errorMessage, "expecting error to match")
122+
return .success(())
123+
}
124+
}
125+
126+
struct CustomError: Error {}
127+
128+
XCTAssertThrowsError(try runLambda(behavior: Behavior(), handlerProvider: { context -> EventLoopFuture<EchoHandler> in
129+
context.eventLoop.makeFailedFuture(CustomError())
130+
})) { error in
131+
XCTAssertNotNil(error as? CustomError, "expecting error to match")
132+
}
133+
}
71134
}

‎Tests/AWSLambdaRuntimeCoreTests/Utils.swift

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,18 +19,32 @@ import NIOPosix
1919
import XCTest
2020

2121
func runLambda<Handler: SimpleLambdaHandler>(behavior: LambdaServerBehavior, handlerType: Handler.Type) throws {
22-
try runLambda(behavior: behavior, handlerType: CodableSimpleLambdaHandler<Handler>.self)
22+
try runLambda(behavior: behavior, handlerProvider: CodableSimpleLambdaHandler<Handler>.makeHandler(context:))
2323
}
2424

2525
func runLambda<Handler: LambdaHandler>(behavior: LambdaServerBehavior, handlerType: Handler.Type) throws {
26-
try runLambda(behavior: behavior, handlerType: CodableLambdaHandler<Handler>.self)
26+
try runLambda(behavior: behavior, handlerProvider: CodableLambdaHandler<Handler>.makeHandler(context:))
2727
}
2828

2929
func runLambda<Handler: EventLoopLambdaHandler>(behavior: LambdaServerBehavior, handlerType: Handler.Type) throws {
30-
try runLambda(behavior: behavior, handlerType: CodableEventLoopLambdaHandler<Handler>.self)
30+
try runLambda(behavior: behavior, handlerProvider: CodableEventLoopLambdaHandler<Handler>.makeHandler(context:))
3131
}
3232

33-
func runLambda(behavior: LambdaServerBehavior, handlerType: (some ByteBufferLambdaHandler).Type) throws {
33+
func runLambda<Handler: EventLoopLambdaHandler>(
34+
behavior: LambdaServerBehavior,
35+
handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<Handler>
36+
) throws {
37+
try runLambda(behavior: behavior, handlerProvider: { context in
38+
handlerProvider(context).map {
39+
CodableEventLoopLambdaHandler(handler: $0, allocator: context.allocator)
40+
}
41+
})
42+
}
43+
44+
func runLambda(
45+
behavior: LambdaServerBehavior,
46+
handlerProvider: @escaping (LambdaInitializationContext) -> EventLoopFuture<some ByteBufferLambdaHandler>
47+
) throws {
3448
let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
3549
defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
3650
let logger = Logger(label: "TestLogger")
@@ -39,7 +53,7 @@ func runLambda(behavior: LambdaServerBehavior, handlerType: (some ByteBufferLamb
3953
let runner = LambdaRunner(eventLoop: eventLoopGroup.next(), configuration: configuration)
4054
let server = try MockLambdaServer(behavior: behavior).start().wait()
4155
defer { XCTAssertNoThrow(try server.stop().wait()) }
42-
try runner.initialize(handlerType: handlerType, logger: logger, terminator: terminator).flatMap { handler in
56+
try runner.initialize(handlerProvider: handlerProvider, logger: logger, terminator: terminator).flatMap { handler in
4357
runner.run(handler: handler, logger: logger)
4458
}.wait()
4559
}
@@ -89,3 +103,14 @@ extension LambdaTerminator.TerminationError: Equatable {
89103
return String(describing: lhs) == String(describing: rhs)
90104
}
91105
}
106+
107+
// for backward compatibility in tests
108+
extension LambdaRunner {
109+
func initialize<Handler: ByteBufferLambdaHandler>(
110+
handlerType: Handler.Type,
111+
logger: Logger,
112+
terminator: LambdaTerminator
113+
) -> EventLoopFuture<Handler> {
114+
self.initialize(handlerProvider: handlerType.makeHandler(context:), logger: logger, terminator: terminator)
115+
}
116+
}

0 commit comments

Comments
 (0)