refactor

motivation: safer concurrency, better tests

changes:
* make http client thread safe
* add concurrency tests
* refactor test
* make mock server more ribust

Author: tomerd

Sources/MockServer/main.swift

@@ -64,41 +64,42 @@ internal final class HTTPHandler: ChannelInboundHandler {
     private let mode: Mode
     private let keepAlive: Bool
 
-    private var requestHead: HTTPRequestHead!
-    private var requestBody: ByteBuffer?
+    private var pending = CircularBuffer<(head: HTTPRequestHead, body: ByteBuffer?)>()
 
     public init(logger: Logger, keepAlive: Bool, mode: Mode) {
         self.logger = logger
-        self.mode = mode
         self.keepAlive = keepAlive
+        self.mode = mode
     }
 
     func channelRead(context: ChannelHandlerContext, data: NIOAny) {
         let requestPart = unwrapInboundIn(data)
 
         switch requestPart {
         case .head(let head):
-            self.requestHead = head
-            self.requestBody?.clear()
+            self.pending.append((head: head, body: nil))
         case .body(var buffer):
-            if self.requestBody == nil {
-                self.requestBody = buffer
+            var request = self.pending.removeFirst()
+            if request.body == nil {
+                request.body = buffer
             } else {
-                self.requestBody!.writeBuffer(&buffer)
+                request.body!.writeBuffer(&buffer)
             }
+            self.pending.prepend(request)
         case .end:
-            self.processRequest(context: context)
+            let request = self.pending.removeFirst()
+            self.processRequest(context: context, request: request)
         }
     }
 
-    func processRequest(context: ChannelHandlerContext) {
-        self.logger.debug("\(self) processing \(self.requestHead.uri)")
+    func processRequest(context: ChannelHandlerContext, request: (head: HTTPRequestHead, body: ByteBuffer?)) {
+        self.logger.debug("\(self) processing \(request.head.uri)")
 
         var responseStatus: HTTPResponseStatus
         var responseBody: String?
         var responseHeaders: [(String, String)]?
 
-        if self.requestHead.uri.hasSuffix("/next") {
+        if request.head.uri.hasSuffix("/next") {
             let requestId = UUID().uuidString
             responseStatus = .ok
             switch self.mode {
@@ -108,7 +109,7 @@ internal final class HTTPHandler: ChannelInboundHandler {
                 responseBody = "{ \"body\": \"\(requestId)\" }"
             }
             responseHeaders = [(AmazonHeaders.requestID, requestId)]
-        } else if self.requestHead.uri.hasSuffix("/response") {
+        } else if request.head.uri.hasSuffix("/response") {
             responseStatus = .accepted
         } else {
             responseStatus = .notFound

Sources/SwiftAwsLambda/HttpClient.swift

@@ -17,13 +17,13 @@ import NIOConcurrencyHelpers
 import NIOHTTP1
 
 /// A barebone HTTP client to interact with AWS Runtime Engine which is an HTTP server.
-internal class HTTPClient {
+internal final class HTTPClient {
     private let eventLoop: EventLoop
     private let configuration: Lambda.Configuration.RuntimeEngine
     private let targetHost: String
 
     private var state = State.disconnected
-    private let lock = Lock()
+    private let stateLock = Lock()
 
     init(eventLoop: EventLoop, configuration: Lambda.Configuration.RuntimeEngine) {
         self.eventLoop = eventLoop
@@ -47,37 +47,45 @@ internal class HTTPClient {
     }
 
     private func execute(_ request: Request) -> EventLoopFuture<Response> {
-        self.lock.lock()
+        self.stateLock.lock()
         switch self.state {
+        case .disconnected:
+            let promise = self.eventLoop.makePromise(of: Response.self)
+            self.state = .connecting(promise.futureResult)
+            self.stateLock.unlock()
+            self.connect().flatMap { channel -> EventLoopFuture<Response> in
+                self.stateLock.withLock {
+                    guard case .connecting = self.state else {
+                        preconditionFailure("invalid state \(self.state)")
+                    }
+                    self.state = .connected(channel)
+                }
+                return self.execute(request)
+            }.cascade(to: promise)
+            return promise.futureResult
+        case .connecting(let future):
+            let future = future.flatMap { _ in
+                self.execute(request)
+            }
+            self.state = .connecting(future)
+            self.stateLock.unlock()
+            return future
         case .connected(let channel):
             guard channel.isActive else {
-                // attempt to reconnect
                 self.state = .disconnected
-                self.lock.unlock()
+                self.stateLock.unlock()
                 return self.execute(request)
             }
-            self.lock.unlock()
+            self.stateLock.unlock()
             let promise = channel.eventLoop.makePromise(of: Response.self)
             let wrapper = HTTPRequestWrapper(request: request, promise: promise)
-            return channel.writeAndFlush(wrapper).flatMap {
-                promise.futureResult
-            }
-        case .disconnected:
-            return self.connect().flatMap {
-                self.lock.unlock()
-                return self.execute(request)
-            }
-        default:
-            preconditionFailure("invalid state \(self.state)")
+            channel.writeAndFlush(wrapper).cascadeFailure(to: promise)
+            return promise.futureResult
         }
     }
 
-    private func connect() -> EventLoopFuture<Void> {
-        guard case .disconnected = self.state else {
-            preconditionFailure("invalid state \(self.state)")
-        }
-        self.state = .connecting
-        let bootstrap = ClientBootstrap(group: eventLoop)
+    private func connect() -> EventLoopFuture<Channel> {
+        let bootstrap = ClientBootstrap(group: self.eventLoop)
             .channelInitializer { channel in
                 channel.pipeline.addHTTPClientHandlers().flatMap {
                     channel.pipeline.addHandlers([HTTPHandler(keepAlive: self.configuration.keepAlive),
@@ -88,9 +96,7 @@ internal class HTTPClient {
         do {
             // connect directly via socket address to avoid happy eyeballs (perf)
             let address = try SocketAddress(ipAddress: self.configuration.ip, port: self.configuration.port)
-            return bootstrap.connect(to: address).flatMapThrowing { channel in
-                self.state = .connected(channel)
-            }
+            return bootstrap.connect(to: address)
         } catch {
             return self.eventLoop.makeFailedFuture(error)
         }
@@ -126,13 +132,13 @@ internal class HTTPClient {
     }
 
     private enum State {
-        case connecting
-        case connected(Channel)
         case disconnected
+        case connecting(EventLoopFuture<Response>)
+        case connected(Channel)
     }
 }
 
-private class HTTPHandler: ChannelDuplexHandler {
+private final class HTTPHandler: ChannelDuplexHandler {
     typealias OutboundIn = HTTPClient.Request
     typealias InboundOut = HTTPClient.Response
     typealias InboundIn = HTTPClientResponsePart
@@ -207,15 +213,15 @@ private class HTTPHandler: ChannelDuplexHandler {
     }
 }
 
-private class UnaryHandler: ChannelInboundHandler, ChannelOutboundHandler {
+private final class UnaryHandler: ChannelInboundHandler, ChannelOutboundHandler {
     typealias OutboundIn = HTTPRequestWrapper
     typealias InboundIn = HTTPClient.Response
     typealias OutboundOut = HTTPClient.Request
 
     private let keepAlive: Bool
 
     private let lock = Lock()
-    private var pendingResponses = CircularBuffer<(EventLoopPromise<HTTPClient.Response>, Scheduled<Void>?)>()
+    private var pendingResponses = CircularBuffer<(promise: EventLoopPromise<HTTPClient.Response>, timeout: Scheduled<Void>?)>()
     private var lastError: Error?
 
     init(keepAlive: Bool) {
@@ -231,19 +237,21 @@ private class UnaryHandler: ChannelInboundHandler, ChannelOutboundHandler {
                 }
             }
         }
-        self.lock.withLockVoid { pendingResponses.append((wrapper.promise, timeoutTask)) }
+        self.lock.withLockVoid { pendingResponses.append((promise: wrapper.promise, timeout: timeoutTask)) }
         context.writeAndFlush(wrapOutboundOut(wrapper.request), promise: promise)
     }
 
     func channelRead(context: ChannelHandlerContext, data: NIOAny) {
         let response = unwrapInboundIn(data)
         if let pending = (self.lock.withLock { self.pendingResponses.popFirst() }) {
             let serverKeepAlive = response.headers["connection"].first?.lowercased() == "keep-alive"
-            let future = self.keepAlive && serverKeepAlive ? context.eventLoop.makeSucceededFuture(()) : context.channel.close()
-            future.whenComplete { _ in
-                pending.1?.cancel()
-                pending.0.succeed(response)
+            if !self.keepAlive || !serverKeepAlive {
+                pending.promise.futureResult.whenComplete { _ in
+                    _ = context.channel.close()
+                }
             }
+            pending.timeout?.cancel()
+            pending.promise.succeed(response)
         }
     }
 

Sources/SwiftAwsLambda/Lambda.swift

@@ -105,7 +105,7 @@ public enum Lambda {
             }
             set {
                 self.stateLock.withLockVoid {
-                    precondition(newValue.rawValue > _state.rawValue, "invalid state \(newValue) after \(_state)")
+                    precondition(newValue.rawValue > _state.rawValue, "invalid state \(newValue) after \(self._state)")
                     self._state = newValue
                 }
             }
@@ -124,12 +124,12 @@ public enum Lambda {
         }
 
         func stop() {
-            self.logger.info("lambda lifecycle stopping")
+            self.logger.debug("lambda lifecycle stopping")
             self.state = .stopping
         }
 
         func shutdown() {
-            self.logger.info("lambda lifecycle shutdown")
+            self.logger.debug("lambda lifecycle shutdown")
             self.state = .shutdown
         }
 

Sources/SwiftAwsLambda/LambdaRunner.swift

@@ -36,7 +36,7 @@ internal struct LambdaRunner {
     ///
     /// - Returns: An `EventLoopFuture<Void>` fulfilled with the outcome of the initialization.
     func initialize(logger: Logger) -> EventLoopFuture<Void> {
-        logger.info("initializing lambda")
+        logger.debug("initializing lambda")
         // We need to use `flatMap` instead of `whenFailure` to ensure we complete reporting the result before stopping.
         return self.lambdaHandler.initialize(eventLoop: self.eventLoop,
                                              lifecycleId: self.lifecycleId,
@@ -69,7 +69,7 @@ internal struct LambdaRunner {
             }
         }.always { result in
             // we are done!
-            logger.log(level: result.successful ? .info : .warning, "lambda invocation sequence completed \(result.successful ? "successfully" : "with failure")")
+            logger.log(level: result.successful ? .debug : .warning, "lambda invocation sequence completed \(result.successful ? "successfully" : "with failure")")
         }
     }
 }

Sources/SwiftAwsLambda/LambdaRuntimeClient.swift

@@ -145,7 +145,7 @@ internal struct JsonCodecError: Error, Equatable {
     }
 
     static func == (lhs: JsonCodecError, rhs: JsonCodecError) -> Bool {
-        return lhs.cause.localizedDescription == rhs.cause.localizedDescription
+        return String(describing: lhs.cause) == String(describing: rhs.cause)
     }
 }
 

Tests/SwiftAwsLambdaTests/Lambda+CodeableTest.swift

@@ -16,39 +16,47 @@
 import XCTest
 
 class CodableLambdaTest: XCTestCase {
-    func testSuceess() throws {
+    func testSuceess() {
+        let server = MockLambdaServer(behavior: GoodBehavior())
+        XCTAssertNoThrow(try server.start().wait())
+        defer { XCTAssertNoThrow(try server.stop().wait()) }
+
         let maxTimes = Int.random(in: 1 ... 10)
         let configuration = Lambda.Configuration(lifecycle: .init(maxTimes: maxTimes))
-        let server = try MockLambdaServer(behavior: GoodBehavior()).start().wait()
         let result = Lambda.run(handler: CodableEchoHandler(), configuration: configuration)
-        try server.stop().wait()
         assertLambdaLifecycleResult(result: result, shoudHaveRun: maxTimes)
     }
 
-    func testFailure() throws {
-        let server = try MockLambdaServer(behavior: BadBehavior()).start().wait()
+    func testFailure() {
+        let server = MockLambdaServer(behavior: BadBehavior())
+        XCTAssertNoThrow(try server.start().wait())
+        defer { XCTAssertNoThrow(try server.stop().wait()) }
+
         let result = Lambda.run(handler: CodableEchoHandler())
-        try server.stop().wait()
         assertLambdaLifecycleResult(result: result, shouldFailWithError: LambdaRuntimeClientError.badStatusCode(.internalServerError))
     }
 
-    func testClosureSuccess() throws {
+    func testClosureSuccess() {
+        let server = MockLambdaServer(behavior: GoodBehavior())
+        XCTAssertNoThrow(try server.start().wait())
+        defer { XCTAssertNoThrow(try server.stop().wait()) }
+
         let maxTimes = Int.random(in: 1 ... 10)
         let configuration = Lambda.Configuration(lifecycle: .init(maxTimes: maxTimes))
-        let server = try MockLambdaServer(behavior: GoodBehavior()).start().wait()
         let result = Lambda.run(configuration: configuration) { (_, payload: Request, callback) in
             callback(.success(Response(requestId: payload.requestId)))
         }
-        try server.stop().wait()
         assertLambdaLifecycleResult(result: result, shoudHaveRun: maxTimes)
     }
 
-    func testClosureFailure() throws {
-        let server = try MockLambdaServer(behavior: BadBehavior()).start().wait()
+    func testClosureFailure() {
+        let server = MockLambdaServer(behavior: BadBehavior())
+        XCTAssertNoThrow(try server.start().wait())
+        defer { XCTAssertNoThrow(try server.stop().wait()) }
+
         let result: LambdaLifecycleResult = Lambda.run { (_, payload: Request, callback) in
             callback(.success(Response(requestId: payload.requestId)))
         }
-        try server.stop().wait()
         assertLambdaLifecycleResult(result: result, shouldFailWithError: LambdaRuntimeClientError.badStatusCode(.internalServerError))
     }
 }

Tests/SwiftAwsLambdaTests/Lambda+StringTest.swift

@@ -16,39 +16,47 @@
 import XCTest
 
 class StringLambdaTest: XCTestCase {
-    func testSuceess() throws {
+    func testSuceess() {
+        let server = MockLambdaServer(behavior: GoodBehavior())
+        XCTAssertNoThrow(try server.start().wait())
+        defer { XCTAssertNoThrow(try server.stop().wait()) }
+
         let maxTimes = Int.random(in: 1 ... 10)
         let configuration = Lambda.Configuration(lifecycle: .init(maxTimes: maxTimes))
-        let server = try MockLambdaServer(behavior: GoodBehavior()).start().wait()
         let result = Lambda.run(handler: StringEchoHandler(), configuration: configuration)
-        try server.stop().wait()
         assertLambdaLifecycleResult(result: result, shoudHaveRun: maxTimes)
     }
 
-    func testFailure() throws {
-        let server = try MockLambdaServer(behavior: BadBehavior()).start().wait()
+    func testFailure() {
+        let server = MockLambdaServer(behavior: BadBehavior())
+        XCTAssertNoThrow(try server.start().wait())
+        defer { XCTAssertNoThrow(try server.stop().wait()) }
+
         let result = Lambda.run(handler: StringEchoHandler())
-        try server.stop().wait()
         assertLambdaLifecycleResult(result: result, shouldFailWithError: LambdaRuntimeClientError.badStatusCode(.internalServerError))
     }
 
-    func testClosureSuccess() throws {
+    func testClosureSuccess() {
+        let server = MockLambdaServer(behavior: GoodBehavior())
+        XCTAssertNoThrow(try server.start().wait())
+        defer { XCTAssertNoThrow(try server.stop().wait()) }
+
         let maxTimes = Int.random(in: 1 ... 10)
         let configuration = Lambda.Configuration(lifecycle: .init(maxTimes: maxTimes))
-        let server = try MockLambdaServer(behavior: GoodBehavior()).start().wait()
         let result = Lambda.run(configuration: configuration) { (_, payload: String, callback) in
             callback(.success(payload))
         }
-        try server.stop().wait()
         assertLambdaLifecycleResult(result: result, shoudHaveRun: maxTimes)
     }
 
-    func testClosureFailure() throws {
-        let server = try MockLambdaServer(behavior: BadBehavior()).start().wait()
+    func testClosureFailure() {
+        let server = MockLambdaServer(behavior: BadBehavior())
+        XCTAssertNoThrow(try server.start().wait())
+        defer { XCTAssertNoThrow(try server.stop().wait()) }
+
         let result: LambdaLifecycleResult = Lambda.run { (_, payload: String, callback) in
             callback(.success(payload))
         }
-        try server.stop().wait()
         assertLambdaLifecycleResult(result: result, shouldFailWithError: LambdaRuntimeClientError.badStatusCode(.internalServerError))
     }
 }

Tests/SwiftAwsLambdaTests/LambdaRunnerTest+XCTest.swift

@@ -27,6 +27,7 @@ extension LambdaRunnerTest {
         return [
             ("testSuccess", testSuccess),
             ("testFailure", testFailure),
+            ("testConcurrency", testConcurrency),
         ]
     }
 }