Add a control to limit connection reuses (#678)

Motivation:

Sometimes it can be helpful to limit the number of times a connection
can be used before discarding it. AHC has no such support for this at
the moment.

Modifications:

- Add a `maximumUsesPerConnection` configuration option which defaults
  to `nil` (i.e. no limit).
- For HTTP1 we count down uses in the state machine and close the
  connection if it hits zero.
- For HTTP2, each use maps to a stream so we count down remaining uses
  in the state machine which we combine with max concurrent streams to
  limit how many streams are available per connection. We also count
  remaining uses in the HTTP2 idle handler: we treat no remaining uses
  as receiving a GOAWAY frame and notify the pool which then drains the
  streams and replaces the connection.

Result:

Users can control how many times each connection can be used.

Author: glbrntt

Sources/AsyncHTTPClient/ConnectionPool/HTTP2/HTTP2Connection.swift

@@ -81,6 +81,7 @@ final class HTTP2Connection {
     private var openStreams = Set<ChannelBox>()
     let id: HTTPConnectionPool.Connection.ID
     let decompression: HTTPClient.Decompression
+    let maximumConnectionUses: Int?
 
     var closeFuture: EventLoopFuture<Void> {
         self.channel.closeFuture
@@ -89,11 +90,13 @@ final class HTTP2Connection {
     init(channel: Channel,
          connectionID: HTTPConnectionPool.Connection.ID,
          decompression: HTTPClient.Decompression,
+         maximumConnectionUses: Int?,
          delegate: HTTP2ConnectionDelegate,
          logger: Logger) {
         self.channel = channel
         self.id = connectionID
         self.decompression = decompression
+        self.maximumConnectionUses = maximumConnectionUses
         self.logger = logger
         self.multiplexer = HTTP2StreamMultiplexer(
             mode: .client,
@@ -120,12 +123,14 @@ final class HTTP2Connection {
         connectionID: HTTPConnectionPool.Connection.ID,
         delegate: HTTP2ConnectionDelegate,
         decompression: HTTPClient.Decompression,
+        maximumConnectionUses: Int?,
         logger: Logger
     ) -> EventLoopFuture<(HTTP2Connection, Int)> {
         let connection = HTTP2Connection(
             channel: channel,
             connectionID: connectionID,
             decompression: decompression,
+            maximumConnectionUses: maximumConnectionUses,
             delegate: delegate,
             logger: logger
         )
@@ -192,7 +197,7 @@ final class HTTP2Connection {
             let sync = self.channel.pipeline.syncOperations
 
             let http2Handler = NIOHTTP2Handler(mode: .client, initialSettings: nioDefaultSettings)
-            let idleHandler = HTTP2IdleHandler(delegate: self, logger: self.logger)
+            let idleHandler = HTTP2IdleHandler(delegate: self, logger: self.logger, maximumConnectionUses: self.maximumConnectionUses)
 
             try sync.addHandler(http2Handler, position: .last)
             try sync.addHandler(idleHandler, position: .last)

Sources/AsyncHTTPClient/ConnectionPool/HTTP2/HTTP2IdleHandler.swift

@@ -35,9 +35,10 @@ final class HTTP2IdleHandler<Delegate: HTTP2IdleHandlerDelegate>: ChannelDuplexH
     let logger: Logger
     let delegate: Delegate
 
-    private var state: StateMachine = .init()
+    private var state: StateMachine
 
-    init(delegate: Delegate, logger: Logger) {
+    init(delegate: Delegate, logger: Logger, maximumConnectionUses: Int? = nil) {
+        self.state = StateMachine(maximumUses: maximumConnectionUses)
         self.delegate = delegate
         self.logger = logger
     }
@@ -140,19 +141,23 @@ extension HTTP2IdleHandler {
         }
 
         enum State {
-            case initialized
-            case connected
-            case active(openStreams: Int, maxStreams: Int)
+            case initialized(maximumUses: Int?)
+            case connected(remainingUses: Int?)
+            case active(openStreams: Int, maxStreams: Int, remainingUses: Int?)
             case closing(openStreams: Int, maxStreams: Int)
             case closed
         }
 
-        var state: State = .initialized
+        var state: State
+
+        init(maximumUses: Int?) {
+            self.state = .initialized(maximumUses: maximumUses)
+        }
 
         mutating func channelActive() {
             switch self.state {
-            case .initialized:
-                self.state = .connected
+            case .initialized(let maximumUses):
+                self.state = .connected(remainingUses: maximumUses)
 
             case .connected, .active, .closing, .closed:
                 break
@@ -171,17 +176,17 @@ extension HTTP2IdleHandler {
             case .initialized:
                 preconditionFailure("Invalid state: \(self.state)")
 
-            case .connected:
+            case .connected(let remainingUses):
                 // a settings frame might have multiple entries for `maxConcurrentStreams`. We are
                 // only interested in the last value! If no `maxConcurrentStreams` is set, we assume
                 // the http/2 default of 100.
                 let maxStreams = settings.last(where: { $0.parameter == .maxConcurrentStreams })?.value ?? 100
-                self.state = .active(openStreams: 0, maxStreams: maxStreams)
+                self.state = .active(openStreams: 0, maxStreams: maxStreams, remainingUses: remainingUses)
                 return .notifyConnectionNewMaxStreamsSettings(maxStreams)
 
-            case .active(openStreams: let openStreams, maxStreams: let maxStreams):
+            case .active(openStreams: let openStreams, maxStreams: let maxStreams, remainingUses: let remainingUses):
                 if let newMaxStreams = settings.last(where: { $0.parameter == .maxConcurrentStreams })?.value, newMaxStreams != maxStreams {
-                    self.state = .active(openStreams: openStreams, maxStreams: newMaxStreams)
+                    self.state = .active(openStreams: openStreams, maxStreams: newMaxStreams, remainingUses: remainingUses)
                     return .notifyConnectionNewMaxStreamsSettings(newMaxStreams)
                 }
                 return .nothing
@@ -205,7 +210,7 @@ extension HTTP2IdleHandler {
                 self.state = .closing(openStreams: 0, maxStreams: 0)
                 return .notifyConnectionGoAwayReceived(close: true)
 
-            case .active(let openStreams, let maxStreams):
+            case .active(let openStreams, let maxStreams, _):
                 self.state = .closing(openStreams: openStreams, maxStreams: maxStreams)
                 return .notifyConnectionGoAwayReceived(close: openStreams == 0)
 
@@ -228,7 +233,7 @@ extension HTTP2IdleHandler {
                 self.state = .closing(openStreams: 0, maxStreams: 0)
                 return .close
 
-            case .active(let openStreams, let maxStreams):
+            case .active(let openStreams, let maxStreams, _):
                 if openStreams == 0 {
                     self.state = .closed
                     return .close
@@ -247,10 +252,19 @@ extension HTTP2IdleHandler {
             case .initialized, .connected:
                 preconditionFailure("Invalid state: \(self.state)")
 
-            case .active(var openStreams, let maxStreams):
+            case .active(var openStreams, let maxStreams, let remainingUses):
                 openStreams += 1
-                self.state = .active(openStreams: openStreams, maxStreams: maxStreams)
-                return .nothing
+                let remainingUses = remainingUses.map { $0 - 1 }
+                self.state = .active(openStreams: openStreams, maxStreams: maxStreams, remainingUses: remainingUses)
+
+                if remainingUses == 0 {
+                    // Treat running out of connection uses as if we received a GOAWAY frame. This
+                    // will notify the delegate (i.e. connection pool) that the connection can no
+                    // longer be used.
+                    return self.goAwayReceived()
+                } else {
+                    return .nothing
+                }
 
             case .closing(var openStreams, let maxStreams):
                 // A stream might be opened, while we are closing because of race conditions. For
@@ -271,10 +285,10 @@ extension HTTP2IdleHandler {
             case .initialized, .connected:
                 preconditionFailure("Invalid state: \(self.state)")
 
-            case .active(var openStreams, let maxStreams):
+            case .active(var openStreams, let maxStreams, let remainingUses):
                 openStreams -= 1
                 assert(openStreams >= 0)
-                self.state = .active(openStreams: openStreams, maxStreams: maxStreams)
+                self.state = .active(openStreams: openStreams, maxStreams: maxStreams, remainingUses: remainingUses)
                 return .notifyConnectionStreamClosed(currentlyAvailable: maxStreams - openStreams)
 
             case .closing(var openStreams, let maxStreams):

Sources/AsyncHTTPClient/ConnectionPool/HTTPConnectionPool+Factory.swift

@@ -84,6 +84,7 @@ extension HTTPConnectionPool.ConnectionFactory {
                     connectionID: connectionID,
                     delegate: http2ConnectionDelegate,
                     decompression: self.clientConfiguration.decompression,
+                    maximumConnectionUses: self.clientConfiguration.maximumUsesPerConnection,
                     logger: logger
                 ).whenComplete { result in
                     switch result {

Sources/AsyncHTTPClient/ConnectionPool/HTTPConnectionPool.swift

@@ -71,7 +71,8 @@ final class HTTPConnectionPool {
         self._state = StateMachine(
             idGenerator: idGenerator,
             maximumConcurrentHTTP1Connections: clientConfiguration.connectionPool.concurrentHTTP1ConnectionsPerHostSoftLimit,
-            retryConnectionEstablishment: clientConfiguration.connectionPool.retryConnectionEstablishment
+            retryConnectionEstablishment: clientConfiguration.connectionPool.retryConnectionEstablishment,
+            maximumConnectionUses: clientConfiguration.maximumUsesPerConnection
         )
     }