Add an HTTP/1.1 connection pool (#105)

motivation: Better performance thanks to connection reuse

changes:
- Added a connection pool for HTTP/1.1
- All requests automatically use the connection pool
- Up to 8 parallel connections per (scheme, host, port)
- Multiple additional unit tests

Sources/AsyncHTTPClient/HTTPClient.swift

@@ -18,6 +18,7 @@ import NIOConcurrencyHelpers
 import NIOHTTP1
 import NIOHTTPCompression
 import NIOSSL
+import NIOTLS
 
 /// HTTPClient class provides API for request execution.
 ///
@@ -48,7 +49,10 @@ public class HTTPClient {
     public let eventLoopGroup: EventLoopGroup
     let eventLoopGroupProvider: EventLoopGroupProvider
     let configuration: Configuration
-    let isShutdown = NIOAtomic<Bool>.makeAtomic(value: false)
+    let pool: ConnectionPool
+    var state: State
+    private var tasks = [UUID: TaskProtocol]()
+    private let stateLock = Lock()
 
     /// Create an `HTTPClient` with specified `EventLoopGroup` provider and configuration.
     ///
@@ -64,24 +68,79 @@ public class HTTPClient {
             self.eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
         }
         self.configuration = configuration
+        self.pool = ConnectionPool(configuration: configuration)
+        self.state = .upAndRunning
     }
 
     deinit {
-        assert(self.isShutdown.load(), "Client not shut down before the deinit. Please call client.syncShutdown() when no longer needed.")
+        assert(self.pool.connectionProviderCount == 0)
+        assert(self.state == .shutDown, "Client not shut down before the deinit. Please call client.syncShutdown() when no longer needed.")
     }
 
     /// Shuts down the client and `EventLoopGroup` if it was created by the client.
     public func syncShutdown() throws {
-        switch self.eventLoopGroupProvider {
-        case .shared:
-            self.isShutdown.store(true)
-            return
-        case .createNew:
-            if self.isShutdown.compareAndExchange(expected: false, desired: true) {
-                try self.eventLoopGroup.syncShutdownGracefully()
-            } else {
+        try self.syncShutdown(requiresCleanClose: false)
+    }
+
+    /// Shuts down the client and `EventLoopGroup` if it was created by the client.
+    ///
+    /// - parameters:
+    ///     - requiresCleanClose: Determine if the client should throw when it is shutdown in a non-clean state
+    ///
+    /// - Note:
+    /// The `requiresCleanClose` will let the client do additional checks about its internal consistency on shutdown and
+    /// throw the appropriate error if needed. For instance, if its internal connection pool has any non-released connections,
+    /// this indicate shutdown was called too early before tasks were completed or explicitly canceled.
+    /// In general, setting this parameter to `true` should make it easier and faster to catch related programming errors.
+    internal func syncShutdown(requiresCleanClose: Bool) throws {
+        var closeError: Error?
+
+        let tasks = try self.stateLock.withLock { () -> Dictionary<UUID, TaskProtocol>.Values in
+            if self.state != .upAndRunning {
                 throw HTTPClientError.alreadyShutdown
             }
+            self.state = .shuttingDown
+            return self.tasks.values
+        }
+
+        self.pool.prepareForClose()
+
+        if !tasks.isEmpty, requiresCleanClose {
+            closeError = HTTPClientError.uncleanShutdown
+        }
+
+        for task in tasks {
+            task.cancel()
+        }
+
+        try? EventLoopFuture.andAllComplete((tasks.map { $0.completion }), on: self.eventLoopGroup.next()).wait()
+
+        self.pool.syncClose()
+
+        do {
+            try self.stateLock.withLock {
+                switch self.eventLoopGroupProvider {
+                case .shared:
+                    self.state = .shutDown
+                    return
+                case .createNew:
+                    switch self.state {
+                    case .shuttingDown:
+                        self.state = .shutDown
+                        try self.eventLoopGroup.syncShutdownGracefully()
+                    case .shutDown, .upAndRunning:
+                        assertionFailure("The only valid state at this point is \(State.shutDown)")
+                    }
+                }
+            }
+        } catch {
+            if closeError == nil {
+                closeError = error
+            }
+        }
+
+        if let closeError = closeError {
+            throw closeError
         }
     }
 
@@ -188,8 +247,7 @@ public class HTTPClient {
     public func execute<Delegate: HTTPClientResponseDelegate>(request: Request,
                                                               delegate: Delegate,
                                                               deadline: NIODeadline? = nil) -> Task<Delegate.Response> {
-        let eventLoop = self.eventLoopGroup.next()
-        return self.execute(request: request, delegate: delegate, eventLoop: eventLoop, deadline: deadline)
+        return self.execute(request: request, delegate: delegate, eventLoop: .indifferent, deadline: deadline)
     }
 
     /// Execute arbitrary HTTP request and handle response processing using provided delegate.
@@ -201,31 +259,35 @@ public class HTTPClient {
     ///     - deadline: Point in time by which the request must complete.
     public func execute<Delegate: HTTPClientResponseDelegate>(request: Request,
                                                               delegate: Delegate,
-                                                              eventLoop: EventLoopPreference,
+                                                              eventLoop eventLoopPreference: EventLoopPreference,
                                                               deadline: NIODeadline? = nil) -> Task<Delegate.Response> {
-        switch eventLoop.preference {
+        let taskEL: EventLoop
+        switch eventLoopPreference.preference {
         case .indifferent:
-            return self.execute(request: request, delegate: delegate, eventLoop: self.eventLoopGroup.next(), deadline: deadline)
+            taskEL = self.pool.associatedEventLoop(for: ConnectionPool.Key(request)) ?? self.eventLoopGroup.next()
         case .delegate(on: let eventLoop):
             precondition(self.eventLoopGroup.makeIterator().contains { $0 === eventLoop }, "Provided EventLoop must be part of clients EventLoopGroup.")
-            return self.execute(request: request, delegate: delegate, eventLoop: eventLoop, deadline: deadline)
+            taskEL = eventLoop
         case .delegateAndChannel(on: let eventLoop):
             precondition(self.eventLoopGroup.makeIterator().contains { $0 === eventLoop }, "Provided EventLoop must be part of clients EventLoopGroup.")
-            return self.execute(request: request, delegate: delegate, eventLoop: eventLoop, deadline: deadline)
-        case .testOnly_exact(channelOn: let channelEL, delegateOn: let delegateEL):
-            return self.execute(request: request,
-                                delegate: delegate,
-                                eventLoop: delegateEL,
-                                channelEL: channelEL,
-                                deadline: deadline)
+            taskEL = eventLoop
+        case .testOnly_exact(_, delegateOn: let delegateEL):
+            taskEL = delegateEL
+        }
+
+        let failedTask: Task<Delegate.Response>? = self.stateLock.withLock {
+            switch state {
+            case .upAndRunning:
+                return nil
+            case .shuttingDown, .shutDown:
+                return Task<Delegate.Response>.failedTask(eventLoop: taskEL, error: HTTPClientError.alreadyShutdown)
+            }
+        }
+
+        if let failedTask = failedTask {
+            return failedTask
         }
-    }
 
-    private func execute<Delegate: HTTPClientResponseDelegate>(request: Request,
-                                                               delegate: Delegate,
-                                                               eventLoop delegateEL: EventLoop,
-                                                               channelEL: EventLoop? = nil,
-                                                               deadline: NIODeadline? = nil) -> Task<Delegate.Response> {
         let redirectHandler: RedirectHandler<Delegate.Response>?
         switch self.configuration.redirectConfiguration.configuration {
         case .follow(let max, let allowCycles):
@@ -236,72 +298,73 @@ public class HTTPClient {
             redirectHandler = RedirectHandler<Delegate.Response>(request: request) { newRequest in
                 self.execute(request: newRequest,
                              delegate: delegate,
-                             eventLoop: delegateEL,
-                             channelEL: channelEL,
+                             eventLoop: eventLoopPreference,
                              deadline: deadline)
             }
         case .disallow:
             redirectHandler = nil
         }
 
-        let task = Task<Delegate.Response>(eventLoop: delegateEL)
+        let task = Task<Delegate.Response>(eventLoop: taskEL)
+        self.stateLock.withLock {
+            self.tasks[task.id] = task
+        }
+        let promise = task.promise
 
-        var bootstrap = ClientBootstrap(group: channelEL ?? delegateEL)
-            .channelOption(ChannelOptions.socket(SocketOptionLevel(IPPROTO_TCP), TCP_NODELAY), value: 1)
-            .channelInitializer { channel in
-                let encoder = HTTPRequestEncoder()
-                let decoder = ByteToMessageHandler(HTTPResponseDecoder(leftOverBytesStrategy: .forwardBytes))
-                return channel.pipeline.addHandlers([encoder, decoder], position: .first).flatMap {
-                    switch self.configuration.proxy {
-                    case .none:
-                        return channel.pipeline.addSSLHandlerIfNeeded(for: request, tlsConfiguration: self.configuration.tlsConfiguration)
-                    case .some(let proxy):
-                        return channel.pipeline.addProxyHandler(for: request, decoder: decoder, encoder: encoder, tlsConfiguration: self.configuration.tlsConfiguration, proxy: proxy)
-                    }
-                }.flatMap {
-                    switch self.configuration.decompression {
-                    case .disabled:
-                        return channel.eventLoop.makeSucceededFuture(())
-                    case .enabled(let limit):
-                        return channel.pipeline.addHandler(NIOHTTPResponseDecompressor(limit: limit))
-                    }
-                }.flatMap {
-                    if let timeout = self.resolve(timeout: self.configuration.timeout.read, deadline: deadline) {
-                        return channel.pipeline.addHandler(IdleStateHandler(readTimeout: timeout))
-                    } else {
-                        return channel.eventLoop.makeSucceededFuture(())
-                    }
-                }.flatMap {
-                    let taskHandler = TaskHandler(task: task,
-                                                  kind: request.kind,
-                                                  delegate: delegate,
-                                                  redirectHandler: redirectHandler,
-                                                  ignoreUncleanSSLShutdown: self.configuration.ignoreUncleanSSLShutdown)
-                    return channel.pipeline.addHandler(taskHandler)
-                }
+        promise.futureResult.whenComplete { _ in
+            self.stateLock.withLock {
+                self.tasks[task.id] = nil
+            }
+        }
+
+        let connection = self.pool.getConnection(for: request, preference: eventLoopPreference, on: taskEL, deadline: deadline)
+
+        connection.flatMap { connection -> EventLoopFuture<Void> in
+            let channel = connection.channel
+            let addedFuture: EventLoopFuture<Void>
+
+            switch self.configuration.decompression {
+            case .disabled:
+                addedFuture = channel.eventLoop.makeSucceededFuture(())
+            case .enabled(let limit):
+                let decompressHandler = NIOHTTPResponseDecompressor(limit: limit)
+                addedFuture = channel.pipeline.addHandler(decompressHandler)
             }
 
-        if let timeout = self.resolve(timeout: self.configuration.timeout.connect, deadline: deadline) {
-            bootstrap = bootstrap.connectTimeout(timeout)
-        }
+            return addedFuture.flatMap {
+                if let timeout = self.resolve(timeout: self.configuration.timeout.read, deadline: deadline) {
+                    return channel.pipeline.addHandler(IdleStateHandler(readTimeout: timeout))
+                } else {
+                    return channel.eventLoop.makeSucceededFuture(())
+                }
+            }.flatMap {
+                let taskHandler = TaskHandler(task: task,
+                                              kind: request.kind,
+                                              delegate: delegate,
+                                              redirectHandler: redirectHandler,
+                                              ignoreUncleanSSLShutdown: self.configuration.ignoreUncleanSSLShutdown)
+                return channel.pipeline.addHandler(taskHandler)
+            }.flatMap {
+                task.setConnection(connection)
 
-        let eventLoopChannel: EventLoopFuture<Channel>
-        switch request.kind {
-        case .unixSocket:
-            let socketPath = request.url.baseURL?.path ?? request.url.path
-            eventLoopChannel = bootstrap.connect(unixDomainSocketPath: socketPath)
-        case .host:
-            let address = self.resolveAddress(request: request, proxy: self.configuration.proxy)
-            eventLoopChannel = bootstrap.connect(host: address.host, port: address.port)
-        }
+                let isCancelled = task.lock.withLock {
+                    task.cancelled
+                }
 
-        eventLoopChannel.map { channel in
-            task.setChannel(channel)
-        }
-        .flatMap { channel in
-            channel.writeAndFlush(request)
-        }
-        .cascadeFailure(to: task.promise)
+                if !isCancelled {
+                    return channel.writeAndFlush(request).flatMapError { _ in
+                        // At this point the `TaskHandler` will already be present
+                        // to handle the failure and pass it to the `promise`
+                        channel.eventLoop.makeSucceededFuture(())
+                    }
+                } else {
+                    return channel.eventLoop.makeSucceededFuture(())
+                }
+            }.flatMapError { error in
+                connection.release()
+                return channel.eventLoop.makeFailedFuture(error)
+            }
+        }.cascadeFailure(to: promise)
 
         return task
     }
@@ -319,10 +382,10 @@ public class HTTPClient {
         }
     }
 
-    private func resolveAddress(request: Request, proxy: Configuration.Proxy?) -> (host: String, port: Int) {
-        switch self.configuration.proxy {
+    static func resolveAddress(host: String, port: Int, proxy: Configuration.Proxy?) -> (host: String, port: Int) {
+        switch proxy {
         case .none:
-            return (request.host, request.port)
+            return (host, port)
         case .some(let proxy):
             return (proxy.host, proxy.port)
         }
@@ -436,6 +499,12 @@ public class HTTPClient {
         /// Decompression is enabled.
         case enabled(limit: NIOHTTPDecompression.DecompressionLimit)
     }
+
+    enum State {
+        case upAndRunning
+        case shuttingDown
+        case shutDown
+    }
 }
 
 extension HTTPClient.Configuration {
@@ -490,37 +559,84 @@ extension HTTPClient.Configuration {
     }
 }
 
-private extension ChannelPipeline {
-    func addProxyHandler(for request: HTTPClient.Request, decoder: ByteToMessageHandler<HTTPResponseDecoder>, encoder: HTTPRequestEncoder, tlsConfiguration: TLSConfiguration?, proxy: HTTPClient.Configuration.Proxy?) -> EventLoopFuture<Void> {
-        let handler = HTTPClientProxyHandler(host: request.host, port: request.port, authorization: proxy?.authorization, onConnect: { channel in
-            channel.pipeline.removeHandler(decoder).flatMap {
-                channel.pipeline.addHandler(
-                    ByteToMessageHandler(HTTPResponseDecoder(leftOverBytesStrategy: .forwardBytes)),
-                    position: .after(encoder)
-                )
-            }.flatMap {
-                channel.pipeline.addSSLHandlerIfNeeded(for: request, tlsConfiguration: tlsConfiguration)
+extension ChannelPipeline {
+    func addProxyHandler(host: String, port: Int, authorization: HTTPClient.Authorization?) -> EventLoopFuture<Void> {
+        let encoder = HTTPRequestEncoder()
+        let decoder = ByteToMessageHandler(HTTPResponseDecoder(leftOverBytesStrategy: .forwardBytes))
+        let handler = HTTPClientProxyHandler(host: host, port: port, authorization: authorization) { channel in
+            let encoderRemovePromise = self.eventLoop.next().makePromise(of: Void.self)
+            channel.pipeline.removeHandler(encoder, promise: encoderRemovePromise)
+            return encoderRemovePromise.futureResult.flatMap {
+                channel.pipeline.removeHandler(decoder)
             }
-        })
-        return self.addHandler(handler)
+        }
+        return addHandlers([encoder, decoder, handler])
     }
 
-    func addSSLHandlerIfNeeded(for request: HTTPClient.Request, tlsConfiguration: TLSConfiguration?) -> EventLoopFuture<Void> {
-        guard request.useTLS else {
+    func addSSLHandlerIfNeeded(for key: ConnectionPool.Key, tlsConfiguration: TLSConfiguration?, handshakePromise: EventLoopPromise<Void>) -> EventLoopFuture<Void> {
+        guard key.scheme == .https else {
+            handshakePromise.succeed(())
             return self.eventLoop.makeSucceededFuture(())
         }
 
         do {
             let tlsConfiguration = tlsConfiguration ?? TLSConfiguration.forClient()
             let context = try NIOSSLContext(configuration: tlsConfiguration)
-            return self.addHandler(try NIOSSLClientHandler(context: context, serverHostname: request.host.isIPAddress ? nil : request.host),
-                                   position: .first)
+            let handlers: [ChannelHandler] = [
+                try NIOSSLClientHandler(context: context, serverHostname: key.host.isIPAddress ? nil : key.host),
+                TLSEventsHandler(completionPromise: handshakePromise),
+            ]
+
+            return self.addHandlers(handlers)
         } catch {
             return self.eventLoop.makeFailedFuture(error)
         }
     }
 }
 
+class TLSEventsHandler: ChannelInboundHandler, RemovableChannelHandler {
+    typealias InboundIn = NIOAny
+
+    var completionPromise: EventLoopPromise<Void>?
+
+    init(completionPromise: EventLoopPromise<Void>) {
+        self.completionPromise = completionPromise
+    }
+
+    func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) {
+        if let tlsEvent = event as? TLSUserEvent {
+            switch tlsEvent {
+            case .handshakeCompleted:
+                self.completionPromise?.succeed(())
+                self.completionPromise = nil
+                context.pipeline.removeHandler(self, promise: nil)
+            case .shutdownCompleted:
+                break
+            }
+        }
+        context.fireUserInboundEventTriggered(event)
+    }
+
+    func errorCaught(context: ChannelHandlerContext, error: Error) {
+        if let sslError = error as? NIOSSLError {
+            switch sslError {
+            case .handshakeFailed:
+                self.completionPromise?.fail(error)
+                self.completionPromise = nil
+                context.pipeline.removeHandler(self, promise: nil)
+            default:
+                break
+            }
+        }
+        context.fireErrorCaught(error)
+    }
+
+    func handlerRemoved(context: ChannelHandlerContext) {
+        struct NoResult: Error {}
+        self.completionPromise?.fail(NoResult())
+    }
+}
+
 /// Possible client errors.
 public struct HTTPClientError: Error, Equatable, CustomStringConvertible {
     private enum Code: Equatable {
@@ -539,6 +655,7 @@ public struct HTTPClientError: Error, Equatable, CustomStringConvertible {
         case proxyAuthenticationRequired
         case redirectLimitReached
         case redirectCycleDetected
+        case uncleanShutdown
     }
 
     private var code: Code
@@ -581,4 +698,6 @@ public struct HTTPClientError: Error, Equatable, CustomStringConvertible {
     public static let redirectLimitReached = HTTPClientError(code: .redirectLimitReached)
     /// Redirect Cycle detected.
     public static let redirectCycleDetected = HTTPClientError(code: .redirectCycleDetected)
+    /// Unclean shutdown
+    public static let uncleanShutdown = HTTPClientError(code: .uncleanShutdown)
 }