Skip to content

Commit feb35c7

Browse files
fabianfettfabianfett
committed
Added syncShutdown() throws to the ByteBufferLambdaHandler 
1 parent a5e2fd0 commit feb35c7

File tree

3 files changed

+145
-9
lines changed

3 files changed

+145
-9
lines changed

‎Sources/AWSLambdaRuntimeCore/LambdaHandler.swift

+11
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,17 @@ 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 syncShutdown() throws
174+
}
175+
176+
public extension ByteBufferLambdaHandler {
177+
func syncShutdown() throws {}
167178
}
168179

169180
private enum CodecError: Error {

‎Sources/AWSLambdaRuntimeCore/LambdaLifecycle.swift

+24-9
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,37 @@ 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 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)
89-
}
90+
return finishedPromise.futureResult.always { _ in
91+
// If the lambda is terminated (e.g. LocalServer shutdown), we make sure
92+
// developers have the chance to cleanup their resources.
93+
do {
94+
try handler.syncShutdown()
95+
} catch {
96+
logger.error("Error shutting down handler: \(error)")
97+
}
98+
}
99+
}.always { _ in
100+
// triggered when the Lambda has finished its last run or has a startup failure.
101+
self.markShutdown()
102+
}.cascade(to: self.shutdownPromise)
103+
104+
return startupFuture.map { _ in }
90105
}
91106

92107
// MARK: - Private

‎Tests/AWSLambdaRuntimeCoreTests/LambdaLifecycleTest.swift

+110
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
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.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+
func testSyncShutdownIsCalledWhenLambdaShutsdown() {
48+
struct CallbackLambdaHandler: ByteBufferLambdaHandler {
49+
let handler: (Lambda.Context, ByteBuffer) -> (EventLoopFuture<ByteBuffer?>)
50+
let shutdown: () throws -> Void
51+
52+
init(_ handler: @escaping (Lambda.Context, ByteBuffer) -> (EventLoopFuture<ByteBuffer?>), shutdown: @escaping () throws -> Void) {
53+
self.handler = handler
54+
self.shutdown = shutdown
55+
}
56+
57+
func handle(context: Lambda.Context, event: ByteBuffer) -> EventLoopFuture<ByteBuffer?> {
58+
self.handler(context, event)
59+
}
60+
61+
func syncShutdown() throws {
62+
try self.shutdown()
63+
}
64+
}
65+
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) }) {
74+
count += 1
75+
}
76+
77+
let eventLoop = eventLoopGroup.next()
78+
let logger = Logger(label: "TestLogger")
79+
let lifecycle = Lambda.Lifecycle(eventLoop: eventLoop, logger: logger, factory: {
80+
$0.makeSucceededFuture(handler)
81+
})
82+
83+
XCTAssertNoThrow(_ = try eventLoop.flatSubmit { lifecycle.start() }.wait())
84+
XCTAssertThrowsError(_ = try lifecycle.shutdownFuture.wait()) { error in
85+
XCTAssertEqual(.badStatusCode(HTTPResponseStatus.internalServerError), error as? Lambda.RuntimeError)
86+
}
87+
XCTAssertEqual(count, 1)
88+
}
89+
}
90+
91+
struct BadBehavior: LambdaServerBehavior {
92+
func getInvocation() -> GetInvocationResult {
93+
.failure(.internalServerError)
94+
}
95+
96+
func processResponse(requestId: String, response: String?) -> Result<Void, ProcessResponseError> {
97+
XCTFail("should not report a response")
98+
return .failure(.internalServerError)
99+
}
100+
101+
func processError(requestId: String, error: ErrorResponse) -> Result<Void, ProcessErrorError> {
102+
XCTFail("should not report an error")
103+
return .failure(.internalServerError)
104+
}
105+
106+
func processInitError(error: ErrorResponse) -> Result<Void, ProcessErrorError> {
107+
XCTFail("should not report an error")
108+
return .failure(.internalServerError)
109+
}
110+
}

0 commit comments

Comments
 (0)