Allow clients to set a max age for connections (#2235)

Motivation:

Some load balancers reject new streams when a client certificate has
expired but allow existing ones to continue. This is problematic because
there's no signal that the connection can't be used anymore so new RPCs
will continue to fail.

Modifications:

Allow clients to configure a max age for connections, after which time
the connection will shutdown gracefully. This allows new connections to
be established.

Result:

Client connections can be configured to age out

Author: glbrntt

Sources/GRPC/ClientConnection.swift

@@ -423,6 +423,14 @@ extension ClientConnection {
     /// Defaults to 30 minutes.
     public var connectionIdleTimeout: TimeAmount = .minutes(30)
 
+    /// The maximum allowed age of a connection.
+    ///
+    /// If set, no new RPCs will be started on the connection after the connection has been opened
+    /// for this period of time. Existing RPCs will be allowed to continue and the connection will
+    /// close once all RPCs on the connection have finished. If this isn't set then connections have
+    /// no limit on their lifetime.
+    public var connectionMaxAge: TimeAmount? = nil
+
     /// The behavior used to determine when an RPC should start. That is, whether it should wait for
     /// an active connection or fail quickly if no connection is currently available.
     ///
@@ -635,6 +643,7 @@ extension ChannelPipeline.SynchronousOperations {
     connectionManager: ConnectionManager,
     connectionKeepalive: ClientConnectionKeepalive,
     connectionIdleTimeout: TimeAmount,
+    connectionMaxAge: TimeAmount?,
     httpTargetWindowSize: Int,
     httpMaxFrameSize: Int,
     httpMaxResetStreams: Int,
@@ -672,6 +681,7 @@ extension ChannelPipeline.SynchronousOperations {
         connectionManager: connectionManager,
         multiplexer: h2Multiplexer,
         idleTimeout: connectionIdleTimeout,
+        maxAge: connectionMaxAge,
         keepalive: connectionKeepalive,
         logger: logger
       )

Sources/GRPC/ConnectionManagerChannelProvider.swift

@@ -60,6 +60,8 @@ internal struct DefaultChannelProvider: ConnectionManagerChannelProvider {
   internal var connectionKeepalive: ClientConnectionKeepalive
   @usableFromInline
   internal var connectionIdleTimeout: TimeAmount
+  @usableFromInline
+  internal var connectionMaxAge: TimeAmount?
 
   @usableFromInline
   internal var tlsMode: TLSMode
@@ -100,6 +102,7 @@ internal struct DefaultChannelProvider: ConnectionManagerChannelProvider {
     connectionTarget: ConnectionTarget,
     connectionKeepalive: ClientConnectionKeepalive,
     connectionIdleTimeout: TimeAmount,
+    connectionMaxAge: TimeAmount?,
     tlsMode: TLSMode,
     tlsConfiguration: GRPCTLSConfiguration?,
     httpTargetWindowSize: Int,
@@ -113,6 +116,7 @@ internal struct DefaultChannelProvider: ConnectionManagerChannelProvider {
       connectionTarget: connectionTarget,
       connectionKeepalive: connectionKeepalive,
       connectionIdleTimeout: connectionIdleTimeout,
+      connectionMaxAge: connectionMaxAge,
       tlsMode: tlsMode,
       tlsConfiguration: tlsConfiguration,
       httpTargetWindowSize: httpTargetWindowSize,
@@ -131,6 +135,7 @@ internal struct DefaultChannelProvider: ConnectionManagerChannelProvider {
     connectionTarget: ConnectionTarget,
     connectionKeepalive: ClientConnectionKeepalive,
     connectionIdleTimeout: TimeAmount,
+    connectionMaxAge: TimeAmount?,
     tlsMode: TLSMode,
     tlsConfiguration: GRPCTLSConfiguration?,
     httpTargetWindowSize: Int,
@@ -142,6 +147,7 @@ internal struct DefaultChannelProvider: ConnectionManagerChannelProvider {
     self.connectionTarget = connectionTarget
     self.connectionKeepalive = connectionKeepalive
     self.connectionIdleTimeout = connectionIdleTimeout
+    self.connectionMaxAge = connectionMaxAge
 
     self.tlsMode = tlsMode
     self.tlsConfiguration = tlsConfiguration
@@ -182,6 +188,7 @@ internal struct DefaultChannelProvider: ConnectionManagerChannelProvider {
       connectionTarget: configuration.target,
       connectionKeepalive: configuration.connectionKeepalive,
       connectionIdleTimeout: configuration.connectionIdleTimeout,
+      connectionMaxAge: configuration.connectionMaxAge,
       tlsMode: tlsMode,
       tlsConfiguration: configuration.tlsConfiguration,
       httpTargetWindowSize: configuration.httpTargetWindowSize,
@@ -264,6 +271,7 @@ internal struct DefaultChannelProvider: ConnectionManagerChannelProvider {
             connectionManager: connectionManager,
             connectionKeepalive: self.connectionKeepalive,
             connectionIdleTimeout: self.connectionIdleTimeout,
+            connectionMaxAge: self.connectionMaxAge,
             httpTargetWindowSize: self.httpTargetWindowSize,
             httpMaxFrameSize: self.httpMaxFrameSize,
             httpMaxResetStreams: self.httpMaxResetStreams,

Sources/GRPC/ConnectionPool/GRPCChannelPool.swift

@@ -156,6 +156,14 @@ extension GRPCChannelPool {
     /// If a connection becomes idle, starting a new RPC will automatically create a new connection.
     public var idleTimeout = TimeAmount.minutes(30)
 
+    /// The maximum allowed age of a connection.
+    ///
+    /// If set, no new RPCs will be started on the connection after the connection has been opened
+    /// for this period of time. Existing RPCs will be allowed to continue and the connection will
+    /// close once all RPCs on the connection have finished. If this isn't set then connections have
+    /// no limit on their lifetime.
+    public var maxConnectionAge: TimeAmount? = nil
+
     /// The connection keepalive configuration.
     public var keepalive = ClientConnectionKeepalive()
 

Sources/GRPC/ConnectionPool/PooledChannel.swift

@@ -86,6 +86,7 @@ internal final class PooledChannel: GRPCChannel {
         connectionTarget: configuration.target,
         connectionKeepalive: configuration.keepalive,
         connectionIdleTimeout: configuration.idleTimeout,
+        connectionMaxAge: configuration.maxConnectionAge,
         tlsMode: tlsMode,
         tlsConfiguration: configuration.transportSecurity.tlsConfiguration,
         httpTargetWindowSize: configuration.http2.targetWindowSize,
@@ -100,6 +101,7 @@ internal final class PooledChannel: GRPCChannel {
         connectionTarget: configuration.target,
         connectionKeepalive: configuration.keepalive,
         connectionIdleTimeout: configuration.idleTimeout,
+        connectionMaxAge: configuration.maxConnectionAge,
         tlsMode: tlsMode,
         tlsConfiguration: configuration.transportSecurity.tlsConfiguration,
         httpTargetWindowSize: configuration.http2.targetWindowSize,
@@ -114,6 +116,7 @@ internal final class PooledChannel: GRPCChannel {
       connectionTarget: configuration.target,
       connectionKeepalive: configuration.keepalive,
       connectionIdleTimeout: configuration.idleTimeout,
+      connectionMaxAge: configuration.maxConnectionAge,
       tlsMode: tlsMode,
       tlsConfiguration: configuration.transportSecurity.tlsConfiguration,
       httpTargetWindowSize: configuration.http2.targetWindowSize,

Sources/GRPC/GRPCIdleHandler.swift

@@ -27,11 +27,18 @@ internal final class GRPCIdleHandler: ChannelInboundHandler {
   /// If nil, then we shouldn't schedule idle tasks.
   private let idleTimeout: TimeAmount?
 
+  /// The maximum amount of time the connection is allowed to live before quiescing.
+  private let maxAge: TimeAmount?
+
   /// The ping handler.
   private var pingHandler: PingHandler
 
+  /// The scheduled task which will close the connection gently after the max connection age
+  /// has been reached.
+  private var scheduledMaxAgeClose: Scheduled<Void>?
+
   /// The scheduled task which will close the connection after the keep-alive timeout has expired.
-  private var scheduledClose: Scheduled<Void>?
+  private var scheduledKeepAliveClose: Scheduled<Void>?
 
   /// The scheduled task which will ping.
   private var scheduledPing: RepeatedTask?
@@ -75,6 +82,7 @@ internal final class GRPCIdleHandler: ChannelInboundHandler {
     connectionManager: ConnectionManager,
     multiplexer: HTTP2StreamMultiplexer,
     idleTimeout: TimeAmount,
+    maxAge: TimeAmount?,
     keepalive configuration: ClientConnectionKeepalive,
     logger: Logger
   ) {
@@ -95,6 +103,7 @@ internal final class GRPCIdleHandler: ChannelInboundHandler {
       minimumSentPingIntervalWithoutData: configuration.minimumSentPingIntervalWithoutData
     )
     self.creationTime = .now()
+    self.maxAge = maxAge
   }
 
   init(
@@ -116,6 +125,7 @@ internal final class GRPCIdleHandler: ChannelInboundHandler {
       maximumPingStrikes: configuration.maximumPingStrikes
     )
     self.creationTime = .now()
+    self.maxAge = nil
   }
 
   private func perform(operations: GRPCIdleHandlerStateMachine.Operations) {
@@ -218,8 +228,8 @@ internal final class GRPCIdleHandler: ChannelInboundHandler {
       )
 
     case .cancelScheduledTimeout:
-      self.scheduledClose?.cancel()
-      self.scheduledClose = nil
+      self.scheduledKeepAliveClose?.cancel()
+      self.scheduledKeepAliveClose = nil
 
     case let .schedulePing(delay, timeout):
       self.schedulePing(in: delay, timeout: timeout)
@@ -267,7 +277,7 @@ internal final class GRPCIdleHandler: ChannelInboundHandler {
   }
 
   private func scheduleClose(in timeout: TimeAmount) {
-    self.scheduledClose = self.context?.eventLoop.scheduleTask(in: timeout) {
+    self.scheduledKeepAliveClose = self.context?.eventLoop.scheduleTask(in: timeout) {
       self.stateMachine.logger.debug("keepalive timer expired")
       self.perform(operations: self.stateMachine.shutdownNow())
     }
@@ -334,22 +344,35 @@ internal final class GRPCIdleHandler: ChannelInboundHandler {
       remote: context.remoteAddress
     )
 
+    // If a max age has been set then start a timer. This will only be cancelled when it fires or when
+    // the channel eventually becomes inactive.
+    if let maxAge = self.maxAge {
+      assert(self.scheduledMaxAgeClose == nil)
+      self.scheduledMaxAgeClose = context.eventLoop.scheduleTask(in: maxAge) {
+        let operations = self.stateMachine.reachedMaxAge()
+        self.perform(operations: operations)
+      }
+    }
+
     // No state machine action here.
     switch self.mode {
     case let .client(connectionManager, multiplexer):
       connectionManager.channelActive(channel: context.channel, multiplexer: multiplexer)
     case .server:
       ()
     }
+
     context.fireChannelActive()
   }
 
   func channelInactive(context: ChannelHandlerContext) {
     self.perform(operations: self.stateMachine.channelInactive())
     self.scheduledPing?.cancel()
-    self.scheduledClose?.cancel()
+    self.scheduledKeepAliveClose?.cancel()
+    self.scheduledMaxAgeClose?.cancel()
     self.scheduledPing = nil
-    self.scheduledClose = nil
+    self.scheduledKeepAliveClose = nil
+    self.scheduledMaxAgeClose = nil
     context.fireChannelInactive()
   }
 

Sources/GRPC/GRPCIdleHandlerStateMachine.swift

@@ -465,6 +465,13 @@ struct GRPCIdleHandlerStateMachine {
     return operations
   }
 
+  /// The connection has reached it's max allowable age. Let existing RPCs continue, but don't
+  /// allow any new ones.
+  mutating func reachedMaxAge() -> Operations {
+    // Treat this as if the other side sent us a GOAWAY: gently shutdown the connection.
+    self.receiveGoAway()
+  }
+
   /// We've received a GOAWAY frame from the remote peer. Either the remote peer wants to close the
   /// connection or they're responding to us shutting down the connection.
   mutating func receiveGoAway() -> Operations {

Tests/GRPCTests/ConnectionManagerTests.swift

@@ -165,6 +165,7 @@ extension ConnectionManagerTests {
         connectionManager: manager,
         multiplexer: h2mux,
         idleTimeout: .minutes(5),
+        maxAge: nil,
         keepalive: .init(),
         logger: self.logger
       )
@@ -217,6 +218,7 @@ extension ConnectionManagerTests {
         connectionManager: manager,
         multiplexer: h2mux,
         idleTimeout: .minutes(5),
+        maxAge: nil,
         keepalive: .init(),
         logger: self.logger
       )
@@ -273,6 +275,7 @@ extension ConnectionManagerTests {
           inboundStreamInitializer: nil
         ),
         idleTimeout: .minutes(5),
+        maxAge: nil,
         keepalive: .init(),
         logger: self.logger
       )
@@ -322,6 +325,7 @@ extension ConnectionManagerTests {
           inboundStreamInitializer: nil
         ),
         idleTimeout: .minutes(5),
+        maxAge: nil,
         keepalive: .init(),
         logger: self.logger
       )
@@ -350,6 +354,7 @@ extension ConnectionManagerTests {
           inboundStreamInitializer: nil
         ),
         idleTimeout: .minutes(5),
+        maxAge: nil,
         keepalive: .init(),
         logger: self.logger
       )
@@ -391,6 +396,7 @@ extension ConnectionManagerTests {
         connectionManager: manager,
         multiplexer: h2mux,
         idleTimeout: .minutes(5),
+        maxAge: nil,
         keepalive: .init(),
         logger: self.logger
       )
@@ -464,6 +470,7 @@ extension ConnectionManagerTests {
         connectionManager: manager,
         multiplexer: h2mux,
         idleTimeout: .minutes(5),
+        maxAge: nil,
         keepalive: .init(),
         logger: self.logger
       )
@@ -536,6 +543,7 @@ extension ConnectionManagerTests {
         connectionManager: manager,
         multiplexer: h2mux,
         idleTimeout: .minutes(5),
+        maxAge: nil,
         keepalive: .init(),
         logger: self.logger
       )
@@ -654,6 +662,7 @@ extension ConnectionManagerTests {
         connectionManager: manager,
         multiplexer: h2mux,
         idleTimeout: .minutes(5),
+        maxAge: nil,
         keepalive: .init(),
         logger: self.logger
       )
@@ -730,6 +739,7 @@ extension ConnectionManagerTests {
         connectionManager: manager,
         multiplexer: h2mux,
         idleTimeout: .minutes(5),
+        maxAge: nil,
         keepalive: .init(),
         logger: self.logger
       )
@@ -807,6 +817,7 @@ extension ConnectionManagerTests {
         connectionManager: manager,
         multiplexer: firstH2mux,
         idleTimeout: .minutes(5),
+        maxAge: nil,
         keepalive: .init(),
         logger: self.logger
       )
@@ -855,6 +866,7 @@ extension ConnectionManagerTests {
         connectionManager: manager,
         multiplexer: secondH2mux,
         idleTimeout: .minutes(5),
+        maxAge: nil,
         keepalive: .init(),
         logger: self.logger
       )
@@ -905,6 +917,7 @@ extension ConnectionManagerTests {
         connectionManager: manager,
         multiplexer: h2mux,
         idleTimeout: .minutes(5),
+        maxAge: nil,
         keepalive: .init(),
         logger: self.logger
       )
@@ -1063,6 +1076,7 @@ extension ConnectionManagerTests {
           connectionManager: manager,
           multiplexer: h2mux,
           idleTimeout: .minutes(5),
+          maxAge: nil,
           keepalive: .init(),
           logger: self.logger
         )
@@ -1120,6 +1134,7 @@ extension ConnectionManagerTests {
           connectionManager: manager,
           multiplexer: h2mux,
           idleTimeout: .minutes(5),
+          maxAge: nil,
           keepalive: .init(),
           logger: self.logger
         )
@@ -1201,6 +1216,7 @@ extension ConnectionManagerTests {
           connectionManager: manager,
           multiplexer: multiplexer,
           idleTimeout: .minutes(5),
+          maxAge: nil,
           keepalive: ClientConnectionKeepalive(),
           logger: self.logger
         )
@@ -1314,6 +1330,7 @@ extension ConnectionManagerTests {
         connectionManager: connectionManager,
         multiplexer: multiplexer,
         idleTimeout: .minutes(60),
+        maxAge: nil,
         keepalive: .init(),
         logger: self.clientLogger
       )

Tests/GRPCTests/ConnectionPool/ConnectionPoolTests.swift

@@ -1299,6 +1299,7 @@ extension ChannelController: ConnectionManagerChannelProvider {
       connectionManager: connectionManager,
       multiplexer: multiplexer,
       idleTimeout: .minutes(5),
+      maxAge: nil,
       keepalive: ClientConnectionKeepalive(),
       logger: logger
     )