Don’t create an EventLoopFuture chain (#4)

motivation: prevent memory bloat

changes:
* recurse instead of event loop chain
* fix flaky test

Author: fabianfett

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
@@ -46,54 +46,60 @@ internal class HTTPClient {
                                     timeout: timeout ?? self.configuration.requestTimeout))
     }
 
-
+    // TODO: cap reconnect attempt
     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),
                                                   UnaryHandler(keepAlive: self.configuration.keepAlive)])
                 }
             }
-        
+
         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 eventLoop.makeFailedFuture(error)
+            return self.eventLoop.makeFailedFuture(error)
         }
     }
 
@@ -127,13 +133,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
@@ -208,15 +214,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) {
@@ -232,19 +238,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

@@ -18,11 +18,11 @@ import Glibc
 import Darwin.C
 #endif
 
-import Backtrace
+// import Backtrace
 import Logging
 import NIO
 import NIOConcurrencyHelpers
-import Dispatch
+// import Dispatch
 
 public enum Lambda {
     /// Run a Lambda defined by implementing the `LambdaClosure` closure.
@@ -63,17 +63,17 @@ public enum Lambda {
     }
 
     internal static func runAsync(eventLoopGroup: EventLoopGroup, handler: LambdaHandler, configuration: Configuration) -> EventLoopFuture<Int> {
-        Backtrace.install()
+        // Backtrace.install()
         var logger = Logger(label: "Lambda")
         logger.logLevel = configuration.general.logLevel
         let lifecycle = Lifecycle(eventLoop: eventLoopGroup.next(), logger: logger, configuration: configuration, handler: handler)
-        let signalSource = trap(signal: configuration.lifecycle.stopSignal) { signal in
-            logger.info("intercepted signal: \(signal)")
-            lifecycle.stop()
-        }
+        /* let signalSource = trap(signal: configuration.lifecycle.stopSignal) { signal in
+             logger.info("intercepted signal: \(signal)")
+             lifecycle.stop()
+         } */
         return lifecycle.start().always { _ in
             lifecycle.shutdown()
-            signalSource.cancel()
+            // signalSource.cancel()
         }
     }
 
@@ -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
                 }
             }
@@ -119,37 +119,51 @@ public enum Lambda {
             let runner = LambdaRunner(eventLoop: self.eventLoop, configuration: self.configuration, lambdaHandler: self.handler)
             return runner.initialize(logger: logger).flatMap { _ in
                 self.state = .active
-                return self.run(logger: logger, runner: runner, count: 0)
+                return self.run(runner: runner)
             }
         }
 
         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
         }
 
-        private func run(logger: Logger, runner: LambdaRunner, count: Int) -> EventLoopFuture<Int> {
-            switch self.state {
-            case .active:
-                if self.configuration.lifecycle.maxTimes > 0, count >= self.configuration.lifecycle.maxTimes {
-                    return self.eventLoop.makeSucceededFuture(count)
-                }
-                var logger = logger
-                logger[metadataKey: "lifecycleIteration"] = "\(count)"
-                return runner.run(logger: logger).flatMap { _ in
-                    // recursive! per aws lambda runtime spec the polling requests are to be done one at a time
-                    self.run(logger: logger, runner: runner, count: count + 1)
+        @inline(__always)
+        private func run(runner: LambdaRunner) -> EventLoopFuture<Int> {
+            let promise = self.eventLoop.makePromise(of: Int.self)
+
+            func _run(_ count: Int) {
+                switch self.state {
+                case .active:
+                    if self.configuration.lifecycle.maxTimes > 0, count >= self.configuration.lifecycle.maxTimes {
+                        return promise.succeed(count)
+                    }
+                    var logger = self.logger
+                    logger[metadataKey: "lifecycleIteration"] = "\(count)"
+                    runner.run(logger: logger).whenComplete { result in
+                        switch result {
+                        case .success:
+                            // recursive! per aws lambda runtime spec the polling requests are to be done one at a time
+                            _run(count + 1)
+                        case .failure(let error):
+                            promise.fail(error)
+                        }
+                    }
+                case .stopping, .shutdown:
+                    promise.succeed(count)
+                default:
+                    preconditionFailure("invalid run state: \(self.state)")
                 }
-            case .stopping, .shutdown:
-                return self.eventLoop.makeSucceededFuture(count)
-            default:
-                preconditionFailure("invalid run state: \(self.state)")
             }
+
+            _run(0)
+
+            return promise.futureResult
         }
     }
 
@@ -188,8 +202,8 @@ public enum Lambda {
             let stopSignal: Signal
 
             init(id: String? = nil, maxTimes: Int? = nil, stopSignal: Signal? = nil) {
-                self.id = id ?? "\(DispatchTime.now().uptimeNanoseconds)"
-                self.maxTimes = maxTimes ?? env("MAX_REQUESTS").flatMap(Int.init) ?? 0                
+                self.id = id ?? "lambda" // "\(DispatchTime.now().uptimeNanoseconds)"
+                self.maxTimes = maxTimes ?? env("MAX_REQUESTS").flatMap(Int.init) ?? 0
                 self.stopSignal = stopSignal ?? env("STOP_SIGNAL").flatMap(Int32.init).flatMap(Signal.init) ?? Signal.TERM
                 precondition(self.maxTimes >= 0, "maxTimes must be equal or larger than 0")
             }

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)
     }
 }