Add a minimum connections configuration to the ConnectionPool

Sources/GRPC/ClientConnection.swift

@@ -131,6 +131,7 @@ public final class ClientConnection: Sendable {
     self.connectionManager = ConnectionManager(
       configuration: configuration,
       connectivityDelegate: monitor,
+      idleBehavior: .closeWhenIdleTimeout,
       logger: configuration.backgroundActivityLogger
     )
   }

Sources/GRPC/ConnectionManager.swift

@@ -25,6 +25,14 @@ import NIOHTTP2
 // event loop is being used.
 @usableFromInline
 internal final class ConnectionManager: @unchecked Sendable {
+
+  /// Whether the connection managed by this manager should be allowed to go idle and be closed, or
+  /// if it should remain open indefinitely even when there are no active streams.
+  internal enum IdleBehavior {
+    case closeWhenIdleTimeout
+    case neverGoIdle
+  }
+
   internal enum Reconnect {
     case none
     case after(TimeInterval)
@@ -324,6 +332,9 @@ internal final class ConnectionManager: @unchecked Sendable {
   /// attempts should be made at all.
   private let connectionBackoff: ConnectionBackoff?
 
+  /// Whether this connection should be allowed to go idle (and thus be closed when the idle timer fires).
+  internal let idleBehavior: IdleBehavior
+
   /// A logger.
   internal var logger: Logger
 
@@ -356,12 +367,14 @@ internal final class ConnectionManager: @unchecked Sendable {
     configuration: ClientConnection.Configuration,
     channelProvider: ConnectionManagerChannelProvider? = nil,
     connectivityDelegate: ConnectionManagerConnectivityDelegate?,
+    idleBehavior: IdleBehavior,
     logger: Logger
   ) {
     self.init(
       eventLoop: configuration.eventLoopGroup.next(),
       channelProvider: channelProvider ?? DefaultChannelProvider(configuration: configuration),
       callStartBehavior: configuration.callStartBehavior.wrapped,
+      idleBehavior: idleBehavior,
       connectionBackoff: configuration.connectionBackoff,
       connectivityDelegate: connectivityDelegate,
       http2Delegate: nil,
@@ -373,6 +386,7 @@ internal final class ConnectionManager: @unchecked Sendable {
     eventLoop: EventLoop,
     channelProvider: ConnectionManagerChannelProvider,
     callStartBehavior: CallStartBehavior.Behavior,
+    idleBehavior: IdleBehavior,
     connectionBackoff: ConnectionBackoff?,
     connectivityDelegate: ConnectionManagerConnectivityDelegate?,
     http2Delegate: ConnectionManagerHTTP2Delegate?,
@@ -392,6 +406,7 @@ internal final class ConnectionManager: @unchecked Sendable {
     self.connectionBackoff = connectionBackoff
     self.connectivityDelegate = connectivityDelegate
     self.http2Delegate = http2Delegate
+    self.idleBehavior = idleBehavior
 
     self.connectionID = connectionID
     self.channelNumber = channelNumber
@@ -799,7 +814,7 @@ internal final class ConnectionManager: @unchecked Sendable {
       // Yes, after some time.
       case let .after(delay):
         let error = GRPCStatus(code: .unavailable, message: "Connection closed while connecting")
-        // Fail the candidate mux promise. KEep the 'readyChannelMuxPromise' as we'll try again.
+        // Fail the candidate mux promise. Keep the 'readyChannelMuxPromise' as we'll try again.
         connecting.candidateMuxPromise.fail(error)
 
         let scheduled = self.eventLoop.scheduleTask(in: .seconds(timeInterval: delay)) {

Sources/GRPC/ConnectionPool/ConnectionPool.swift

@@ -85,6 +85,11 @@ internal final class ConnectionPool {
   @usableFromInline
   internal let maxWaiters: Int
 
+  /// The number of connections in the pool that should always be kept open (i.e. they won't go idle).
+  /// In other words, it's the number of connections for which we should ignore idle timers.
+  @usableFromInline
+  internal let minConnections: Int
+
   /// Configuration for backoff between subsequence connection attempts.
   @usableFromInline
   internal let connectionBackoff: ConnectionBackoff
@@ -157,6 +162,7 @@ internal final class ConnectionPool {
   init(
     eventLoop: EventLoop,
     maxWaiters: Int,
+    minConnections: Int,
     reservationLoadThreshold: Double,
     assumedMaxConcurrentStreams: Int,
     connectionBackoff: ConnectionBackoff,
@@ -176,6 +182,7 @@ internal final class ConnectionPool {
 
     self._connections = [:]
     self.maxWaiters = maxWaiters
+    self.minConnections = minConnections
     self.waiters = CircularBuffer(initialCapacity: 16)
 
     self.eventLoop = eventLoop
@@ -201,17 +208,25 @@ internal final class ConnectionPool {
       ]
     )
     self._connections.reserveCapacity(connections)
+    var numberOfKeepOpenConnections = self.minConnections
     while self._connections.count < connections {
-      self.addConnectionToPool()
+      // If we have less than the minimum number of connections, don't let
+      // the new connection close when idle.
+      let idleBehavior =
+        numberOfKeepOpenConnections > 0
+        ? ConnectionManager.IdleBehavior.neverGoIdle : .closeWhenIdleTimeout
+      numberOfKeepOpenConnections -= 1
+      self.addConnectionToPool(idleBehavior: idleBehavior)
     }
   }
 
   /// Make and add a new connection to the pool.
-  private func addConnectionToPool() {
+  private func addConnectionToPool(idleBehavior: ConnectionManager.IdleBehavior) {
     let manager = ConnectionManager(
       eventLoop: self.eventLoop,
       channelProvider: self.channelProvider,
       callStartBehavior: .waitsForConnectivity,
+      idleBehavior: idleBehavior,
       connectionBackoff: self.connectionBackoff,
       connectivityDelegate: self,
       http2Delegate: self,
@@ -220,6 +235,19 @@ internal final class ConnectionPool {
     let id = manager.id
     self._connections[id] = PerConnectionState(manager: manager)
     self.delegate?.connectionAdded(id: .init(id))
+
+    // If it's one of the connections that should be kept open, then connect
+    // straight away.
+    switch idleBehavior {
+    case .neverGoIdle:
+      self.eventLoop.execute {
+        if manager.sync.isIdle {
+          manager.sync.startConnecting()
+        }
+      }
+    case .closeWhenIdleTimeout:
+      ()
+    }
   }
 
   // MARK: - Called from the pool manager
@@ -689,8 +717,9 @@ extension ConnectionPool: ConnectionManagerConnectivityDelegate {
     // Grab the number of reserved streams (before invalidating the index by adding a connection).
     let reservedStreams = self._connections.values[index].reservedStreams
 
-    // Replace the connection with a new idle one.
-    self.addConnectionToPool()
+    // Replace the connection with a new idle one. Keep the idle behavior, so that
+    // if it's a connection that should be kept alive, we maintain it.
+    self.addConnectionToPool(idleBehavior: manager.idleBehavior)
 
     // Since we're removing this connection from the pool (and no new streams can be created on
     // the connection), the pool manager can ignore any streams reserved against this connection.
@@ -881,6 +910,22 @@ extension ConnectionPool {
       return self.pool._connections.values.reduce(0) { $0 &+ ($1.manager.sync.isIdle ? 1 : 0) }
     }
 
+    /// The number of active (i.e. connecting or ready) connections in the pool.
+    internal var activeConnections: Int {
+      self.pool.eventLoop.assertInEventLoop()
+      return self.pool._connections.values.reduce(0) {
+        $0 &+ (($1.manager.sync.isReady || $1.manager.sync.isConnecting) ? 1 : 0)
+      }
+    }
+
+    /// The number of connections in the pool in transient failure state.
+    internal var transientFailureConnections: Int {
+      self.pool.eventLoop.assertInEventLoop()
+      return self.pool._connections.values.reduce(0) {
+        $0 &+ ($1.manager.sync.isTransientFailure ? 1 : 0)
+      }
+    }
+
     /// The number of streams currently available to reserve across all connections in the pool.
     internal var availableStreams: Int {
       self.pool.eventLoop.assertInEventLoop()

Sources/GRPC/ConnectionPool/GRPCChannelPool.swift

@@ -275,6 +275,10 @@ extension GRPCChannelPool.Configuration {
     /// Defaults to 100.
     public var maxWaitersPerEventLoop: Int = 100
 
+    /// The minimum number of connections to keep open in this pool, per EventLoop.
+    /// This number of connections per EventLoop will never go idle and be closed.
+    public var minConnections: Int = 0
+
     /// The maximum amount of time a caller is willing to wait for a stream for before timing out.
     ///
     /// Defaults to 30 seconds.

Sources/GRPC/ConnectionPool/PoolManager.swift

@@ -33,6 +33,11 @@ internal final class PoolManager {
     @usableFromInline
     var maxWaiters: Int
 
+    /// The minimum number of connections to keep open per pool.
+    /// This number of connections will never go idle and be closed.
+    @usableFromInline
+    var minConnections: Int
+
     /// A load threshold in the range `0.0 ... 1.0` beyond which another connection will be started
     /// (assuming there is a connection available to start).
     @usableFromInline
@@ -62,6 +67,7 @@ internal final class PoolManager {
     internal init(
       maxConnections: Int,
       maxWaiters: Int,
+      minConnections: Int,
       loadThreshold: Double,
       assumedMaxConcurrentStreams: Int = 100,
       connectionBackoff: ConnectionBackoff,
@@ -70,6 +76,7 @@ internal final class PoolManager {
     ) {
       self.maxConnections = maxConnections
       self.maxWaiters = maxWaiters
+      self.minConnections = minConnections
       self.loadThreshold = loadThreshold
       self.assumedMaxConcurrentStreams = assumedMaxConcurrentStreams
       self.connectionBackoff = connectionBackoff
@@ -225,6 +232,7 @@ internal final class PoolManager {
       return ConnectionPool(
         eventLoop: eventLoop,
         maxWaiters: configuration.maxWaiters,
+        minConnections: configuration.minConnections,
         reservationLoadThreshold: configuration.loadThreshold,
         assumedMaxConcurrentStreams: configuration.assumedMaxConcurrentStreams,
         connectionBackoff: configuration.connectionBackoff,

Sources/GRPC/ConnectionPool/PooledChannel.swift

@@ -96,6 +96,7 @@ internal final class PooledChannel: GRPCChannel {
       perPoolConfiguration: .init(
         maxConnections: configuration.connectionPool.connectionsPerEventLoop,
         maxWaiters: configuration.connectionPool.maxWaitersPerEventLoop,
+        minConnections: configuration.connectionPool.minConnections,
         loadThreshold: configuration.connectionPool.reservationLoadThreshold,
         assumedMaxConcurrentStreams: 100,
         connectionBackoff: configuration.connectionBackoff,

Sources/GRPC/GRPCIdleHandler.swift

@@ -24,7 +24,8 @@ internal final class GRPCIdleHandler: ChannelInboundHandler {
   typealias OutboundOut = HTTP2Frame
 
   /// The amount of time to wait before closing the channel when there are no active streams.
-  private let idleTimeout: TimeAmount
+  /// If nil, then we shouldn't schedule idle tasks.
+  private let idleTimeout: TimeAmount?
 
   /// The ping handler.
   private var pingHandler: PingHandler
@@ -78,7 +79,12 @@ internal final class GRPCIdleHandler: ChannelInboundHandler {
     logger: Logger
   ) {
     self.mode = .client(connectionManager, multiplexer)
-    self.idleTimeout = idleTimeout
+    switch connectionManager.idleBehavior {
+    case .neverGoIdle:
+      self.idleTimeout = nil
+    case .closeWhenIdleTimeout:
+      self.idleTimeout = idleTimeout
+    }
     self.stateMachine = .init(role: .client, logger: logger)
     self.pingHandler = PingHandler(
       pingCode: 5,
@@ -135,7 +141,7 @@ internal final class GRPCIdleHandler: ChannelInboundHandler {
     }
 
     // Handle idle timeout creation/cancellation.
-    if let idleTask = operations.idleTask {
+    if let idleTimeout = self.idleTimeout, let idleTask = operations.idleTask {
       switch idleTask {
       case let .cancel(task):
         self.stateMachine.logger.debug("idle timeout task cancelled")
@@ -145,9 +151,9 @@ internal final class GRPCIdleHandler: ChannelInboundHandler {
         if self.idleTimeout != .nanoseconds(.max), let context = self.context {
           self.stateMachine.logger.debug(
             "scheduling idle timeout task",
-            metadata: [MetadataKey.delayMs: "\(self.idleTimeout.milliseconds)"]
+            metadata: [MetadataKey.delayMs: "\(idleTimeout.milliseconds)"]
           )
-          let task = context.eventLoop.scheduleTask(in: self.idleTimeout) {
+          let task = context.eventLoop.scheduleTask(in: idleTimeout) {
             self.stateMachine.logger.debug("idle timeout task fired")
             self.idleTimeoutFired()
           }

Tests/GRPCTests/ConnectionManagerTests.swift

@@ -59,6 +59,7 @@ class ConnectionManagerTests: GRPCTestCase {
       configuration: configuration,
       channelProvider: channelProvider.map { HookedChannelProvider($0) },
       connectivityDelegate: self.monitor,
+      idleBehavior: .closeWhenIdleTimeout,
       logger: self.logger
     )
   }
@@ -948,6 +949,7 @@ extension ConnectionManagerTests {
         return loop.makeFailedFuture(DoomedChannelError())
       },
       connectivityDelegate: nil,
+      idleBehavior: .closeWhenIdleTimeout,
       logger: self.logger
     )
     let candidate = manager.getHTTP2Multiplexer()
@@ -1207,6 +1209,7 @@ extension ConnectionManagerTests {
         return eventLoop.makeSucceededFuture(channel)
       },
       callStartBehavior: .waitsForConnectivity,
+      idleBehavior: .closeWhenIdleTimeout,
       connectionBackoff: ConnectionBackoff(),
       connectivityDelegate: nil,
       http2Delegate: http2,
@@ -1383,6 +1386,7 @@ extension ConnectionManagerTests {
       configuration: configuration,
       channelProvider: Provider(),
       connectivityDelegate: self.monitor,
+      idleBehavior: .closeWhenIdleTimeout,
       logger: self.logger
     )