Skip to content

Commit 437a60d

Browse files
fabianfettfabianfett
authored
Added shutdown() -> EventLoopFuture<Void> to the ByteBufferLambdaHandler (#122)
* Added `syncShutdown() throws` to the `ByteBufferLambdaHandler` * Updated to use a `ShutdownContext` * Review comments addressed.
1 parent 7728066 commit 437a60d

File tree

6 files changed

+210
-9
lines changed

6 files changed

+210
-9
lines changed

‎Sources/AWSLambdaRuntimeCore/LambdaContext.swift

+24
Original file line numberDiff line numberDiff line change
@@ -122,3 +122,27 @@ extension Lambda {
122122
}
123123
}
124124
}
125+
126+
// MARK: - ShutdownContext
127+
128+
extension Lambda {
129+
/// Lambda runtime shutdown context.
130+
/// The Lambda runtime generates and passes the `ShutdownContext` to the Lambda handler as an argument.
131+
public final class ShutdownContext {
132+
/// `Logger` to log with
133+
///
134+
/// - note: The `LogLevel` can be configured using the `LOG_LEVEL` environment variable.
135+
public let logger: Logger
136+
137+
/// The `EventLoop` the Lambda is executed on. Use this to schedule work with.
138+
///
139+
/// - note: The `EventLoop` is shared with the Lambda runtime engine and should be handled with extra care.
140+
/// Most importantly the `EventLoop` must never be blocked.
141+
public let eventLoop: EventLoop
142+
143+
internal init(logger: Logger, eventLoop: EventLoop) {
144+
self.eventLoop = eventLoop
145+
self.logger = logger
146+
}
147+
}
148+
}

‎Sources/AWSLambdaRuntimeCore/LambdaHandler.swift

+13
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,19 @@ public protocol ByteBufferLambdaHandler {
164164
/// - Returns: An `EventLoopFuture` to report the result of the Lambda back to the runtime engine.
165165
/// The `EventLoopFuture` should be completed with either a response encoded as `ByteBuffer` or an `Error`
166166
func handle(context: Lambda.Context, event: ByteBuffer) -> EventLoopFuture<ByteBuffer?>
167+
168+
/// The method to clean up your resources.
169+
/// Concrete Lambda handlers implement this method to shutdown their `HTTPClient`s and database connections.
170+
///
171+
/// - Note: In case your Lambda fails while creating your LambdaHandler in the `HandlerFactory`, this method
172+
/// **is not invoked**. In this case you must cleanup the created resources immediately in the `HandlerFactory`.
173+
func shutdown(context: Lambda.ShutdownContext) -> EventLoopFuture<Void>
174+
}
175+
176+
public extension ByteBufferLambdaHandler {
177+
func shutdown(context: Lambda.ShutdownContext) -> EventLoopFuture<Void> {
178+
context.eventLoop.makeSucceededFuture(Void())
179+
}
167180
}
168181

169182
private enum CodecError: Error {

‎Sources/AWSLambdaRuntimeCore/LambdaLifecycle.swift

+29-8
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ extension Lambda {
2929

3030
private var state = State.idle {
3131
willSet {
32-
assert(self.eventLoop.inEventLoop, "State may only be changed on the `Lifecycle`'s `eventLoop`")
32+
self.eventLoop.assertInEventLoop()
3333
precondition(newValue.order > self.state.order, "invalid state \(newValue) after \(self.state.order)")
3434
}
3535
}
@@ -71,22 +71,43 @@ extension Lambda {
7171
///
7272
/// - note: This method must be called on the `EventLoop` the `Lifecycle` has been initialized with.
7373
public func start() -> EventLoopFuture<Void> {
74-
assert(self.eventLoop.inEventLoop, "Start must be called on the `EventLoop` the `Lifecycle` has been initialized with.")
74+
self.eventLoop.assertInEventLoop()
7575

7676
logger.info("lambda lifecycle starting with \(self.configuration)")
7777
self.state = .initializing
78-
// triggered when the Lambda has finished its last run
79-
let finishedPromise = self.eventLoop.makePromise(of: Int.self)
80-
finishedPromise.futureResult.always { _ in
81-
self.markShutdown()
82-
}.cascade(to: self.shutdownPromise)
78+
8379
var logger = self.logger
8480
logger[metadataKey: "lifecycleId"] = .string(self.configuration.lifecycle.id)
8581
let runner = Runner(eventLoop: self.eventLoop, configuration: self.configuration)
86-
return runner.initialize(logger: logger, factory: self.factory).map { handler in
82+
83+
let startupFuture = runner.initialize(logger: logger, factory: self.factory)
84+
startupFuture.flatMap { handler -> EventLoopFuture<(ByteBufferLambdaHandler, Result<Int, Error>)> in
85+
// after the startup future has succeeded, we have a handler that we can use
86+
// to `run` the lambda.
87+
let finishedPromise = self.eventLoop.makePromise(of: Int.self)
8788
self.state = .active(runner, handler)
8889
self.run(promise: finishedPromise)
90+
return finishedPromise.futureResult.mapResult { (handler, $0) }
8991
}
92+
.flatMap { (handler, runnerResult) -> EventLoopFuture<Int> in
93+
// after the lambda finishPromise has succeeded or failed we need to
94+
// shutdown the handler
95+
let shutdownContext = ShutdownContext(logger: logger, eventLoop: self.eventLoop)
96+
return handler.shutdown(context: shutdownContext).flatMapErrorThrowing { error in
97+
// if, we had an error shuting down the lambda, we want to concatenate it with
98+
// the runner result
99+
logger.error("Error shutting down handler: \(error)")
100+
throw RuntimeError.shutdownError(shutdownError: error, runnerResult: runnerResult)
101+
}.flatMapResult { (_) -> Result<Int, Error> in
102+
// we had no error shutting down the lambda. let's return the runner's result
103+
runnerResult
104+
}
105+
}.always { _ in
106+
// triggered when the Lambda has finished its last run or has a startup failure.
107+
self.markShutdown()
108+
}.cascade(to: self.shutdownPromise)
109+
110+
return startupFuture.map { _ in }
90111
}
91112

92113
// MARK: - Private

‎Sources/AWSLambdaRuntimeCore/LambdaRunner.swift

+1-1
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,7 @@ private extension Lambda.Context {
115115
}
116116

117117
// TODO: move to nio?
118-
private extension EventLoopFuture {
118+
extension EventLoopFuture {
119119
// callback does not have side effects, failing with original result
120120
func peekError(_ callback: @escaping (Error) -> Void) -> EventLoopFuture<Value> {
121121
self.flatMapError { error in

‎Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift

+1
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,7 @@ internal extension Lambda {
133133
case invocationMissingHeader(String)
134134
case noBody
135135
case json(Error)
136+
case shutdownError(shutdownError: Error, runnerResult: Result<Int, Error>)
136137
}
137138
}
138139

‎Tests/AWSLambdaRuntimeCoreTests/LambdaLifecycleTest.swift

+142
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
//===----------------------------------------------------------------------===//
2+
//
3+
// This source file is part of the SwiftAWSLambdaRuntime open source project
4+
//
5+
// Copyright (c) 2017-2018 Apple Inc. and the SwiftAWSLambdaRuntime project authors
6+
// Licensed under Apache License v2.0
7+
//
8+
// See LICENSE.txt for license information
9+
// See CONTRIBUTORS.txt for the list of SwiftAWSLambdaRuntime project authors
10+
//
11+
// SPDX-License-Identifier: Apache-2.0
12+
//
13+
//===----------------------------------------------------------------------===//
14+
15+
@testable import AWSLambdaRuntimeCore
16+
import Logging
17+
import NIO
18+
import NIOHTTP1
19+
import XCTest
20+
21+
class LambdaLifecycleTest: XCTestCase {
22+
func testShutdownFutureIsFulfilledWithStartUpError() {
23+
let server = MockLambdaServer(behavior: FailedBootstrapBehavior())
24+
XCTAssertNoThrow(try server.start().wait())
25+
defer { XCTAssertNoThrow(try server.stop().wait()) }
26+
let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
27+
defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
28+
29+
let eventLoop = eventLoopGroup.next()
30+
let logger = Logger(label: "TestLogger")
31+
let testError = TestError("kaboom")
32+
let lifecycle = Lambda.Lifecycle(eventLoop: eventLoop, logger: logger, factory: {
33+
$0.eventLoop.makeFailedFuture(testError)
34+
})
35+
36+
// eventLoop.submit in this case returns an EventLoopFuture<EventLoopFuture<ByteBufferHandler>>
37+
// which is why we need `wait().wait()`
38+
XCTAssertThrowsError(_ = try eventLoop.flatSubmit { lifecycle.start() }.wait()) { error in
39+
XCTAssertEqual(testError, error as? TestError)
40+
}
41+
42+
XCTAssertThrowsError(_ = try lifecycle.shutdownFuture.wait()) { error in
43+
XCTAssertEqual(testError, error as? TestError)
44+
}
45+
}
46+
47+
struct CallbackLambdaHandler: ByteBufferLambdaHandler {
48+
let handler: (Lambda.Context, ByteBuffer) -> (EventLoopFuture<ByteBuffer?>)
49+
let shutdown: (Lambda.ShutdownContext) -> EventLoopFuture<Void>
50+
51+
init(_ handler: @escaping (Lambda.Context, ByteBuffer) -> (EventLoopFuture<ByteBuffer?>), shutdown: @escaping (Lambda.ShutdownContext) -> EventLoopFuture<Void>) {
52+
self.handler = handler
53+
self.shutdown = shutdown
54+
}
55+
56+
func handle(context: Lambda.Context, event: ByteBuffer) -> EventLoopFuture<ByteBuffer?> {
57+
self.handler(context, event)
58+
}
59+
60+
func shutdown(context: Lambda.ShutdownContext) -> EventLoopFuture<Void> {
61+
self.shutdown(context)
62+
}
63+
}
64+
65+
func testShutdownIsCalledWhenLambdaShutsdown() {
66+
let server = MockLambdaServer(behavior: BadBehavior())
67+
XCTAssertNoThrow(try server.start().wait())
68+
defer { XCTAssertNoThrow(try server.stop().wait()) }
69+
let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
70+
defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
71+
72+
var count = 0
73+
let handler = CallbackLambdaHandler({ XCTFail("Should not be reached"); return $0.eventLoop.makeSucceededFuture($1) }) { context in
74+
count += 1
75+
return context.eventLoop.makeSucceededFuture(Void())
76+
}
77+
78+
let eventLoop = eventLoopGroup.next()
79+
let logger = Logger(label: "TestLogger")
80+
let lifecycle = Lambda.Lifecycle(eventLoop: eventLoop, logger: logger, factory: {
81+
$0.eventLoop.makeSucceededFuture(handler)
82+
})
83+
84+
XCTAssertNoThrow(_ = try eventLoop.flatSubmit { lifecycle.start() }.wait())
85+
XCTAssertThrowsError(_ = try lifecycle.shutdownFuture.wait()) { error in
86+
XCTAssertEqual(.badStatusCode(HTTPResponseStatus.internalServerError), error as? Lambda.RuntimeError)
87+
}
88+
XCTAssertEqual(count, 1)
89+
}
90+
91+
func testLambdaResultIfShutsdownIsUnclean() {
92+
let server = MockLambdaServer(behavior: BadBehavior())
93+
XCTAssertNoThrow(try server.start().wait())
94+
defer { XCTAssertNoThrow(try server.stop().wait()) }
95+
let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
96+
defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
97+
98+
var count = 0
99+
let handler = CallbackLambdaHandler({ XCTFail("Should not be reached"); return $0.eventLoop.makeSucceededFuture($1) }) { context in
100+
count += 1
101+
return context.eventLoop.makeFailedFuture(TestError("kaboom"))
102+
}
103+
104+
let eventLoop = eventLoopGroup.next()
105+
let logger = Logger(label: "TestLogger")
106+
let lifecycle = Lambda.Lifecycle(eventLoop: eventLoop, logger: logger, factory: {
107+
$0.eventLoop.makeSucceededFuture(handler)
108+
})
109+
110+
XCTAssertNoThrow(_ = try eventLoop.flatSubmit { lifecycle.start() }.wait())
111+
XCTAssertThrowsError(_ = try lifecycle.shutdownFuture.wait()) { error in
112+
guard case Lambda.RuntimeError.shutdownError(let shutdownError, .failure(let runtimeError)) = error else {
113+
XCTFail("Unexpected error"); return
114+
}
115+
116+
XCTAssertEqual(shutdownError as? TestError, TestError("kaboom"))
117+
XCTAssertEqual(runtimeError as? Lambda.RuntimeError, .badStatusCode(.internalServerError))
118+
}
119+
XCTAssertEqual(count, 1)
120+
}
121+
}
122+
123+
struct BadBehavior: LambdaServerBehavior {
124+
func getInvocation() -> GetInvocationResult {
125+
.failure(.internalServerError)
126+
}
127+
128+
func processResponse(requestId: String, response: String?) -> Result<Void, ProcessResponseError> {
129+
XCTFail("should not report a response")
130+
return .failure(.internalServerError)
131+
}
132+
133+
func processError(requestId: String, error: ErrorResponse) -> Result<Void, ProcessErrorError> {
134+
XCTFail("should not report an error")
135+
return .failure(.internalServerError)
136+
}
137+
138+
func processInitError(error: ErrorResponse) -> Result<Void, ProcessErrorError> {
139+
XCTFail("should not report an error")
140+
return .failure(.internalServerError)
141+
}
142+
}

0 commit comments

Comments
 (0)