Cancel connection backoff on graceful shutdown

Sources/GRPCNIOTransportCore/Client/Connection/LoadBalancers/Subchannel.swift

@@ -250,7 +250,9 @@ extension Subchannel {
       self.event.continuation.yield(.connectivityStateChanged(.shutdown))
       connection.close()
 
-    case .emitShutdownAndFinish:
+    case .emitShutdownAndFinish(let backoffHandle):
+      // Cancel any in-progress backoff (there are no in-flight RPCs to drain when backing off).
+      backoffHandle?.cancel()
       // Connection closed because the load balancer asked it to, so notify the load balancer.
       self.event.continuation.yield(.connectivityStateChanged(.shutdown))
       // At this point there are no more events: close the event streams.
@@ -309,7 +311,11 @@ extension Subchannel {
           .transientFailure(cause: error)
         )
       )
-      group.addTask {
+      // Run the backoff as an individually cancellable task and remember its handle. If the
+      // subchannel is gracefully shut down while backing off, there are no in-flight RPCs to drain,
+      // so shutdown cancels this task rather than waiting for the (potentially long) backoff to
+      // elapse.
+      let handle = group.addCancellableTask {
         do {
           try await Task.sleep(for: duration, tolerance: .zero)
           self.input.continuation.yield(.backedOff)
@@ -319,6 +325,7 @@ extension Subchannel {
           ()
         }
       }
+      self.state.withLock { $0.setBackoffTaskHandle(handle) }
 
     case .finish:
       self.event.continuation.finish()
@@ -438,6 +445,10 @@ extension Subchannel {
       let addresses: [SocketAddress]
       var addressIterator: Array<SocketAddress>.Iterator
       var backoff: Backoff.Iterator
+      /// A handle to the in-progress backoff task, if the subchannel is currently waiting to retry
+      /// a connection. It's `nil` whenever a connection attempt is actually in-flight. Cancelling it
+      /// on graceful shutdown avoids waiting out the backoff when there are no RPCs to drain.
+      var backoffHandle: CancellableTaskHandle? = nil
     }
 
     struct Connected {
@@ -480,6 +491,7 @@ extension Subchannel {
       init(from state: ShuttingDown) {}
       init(from state: GoingAway) {}
       init(from state: NotConnected) {}
+      init(from state: Connecting) {}
     }
 
     mutating func makeConnection(
@@ -520,7 +532,7 @@ extension Subchannel {
 
     enum OnClose {
       case none
-      case emitShutdownAndFinish
+      case emitShutdownAndFinish(cancelling: CancellableTaskHandle?)
       case emitShutdownAndClose(Connection)
       case emitShutdown
     }
@@ -531,12 +543,20 @@ extension Subchannel {
       switch self {
       case .notConnected(let state):
         self = .shutDown(ShutDown(from: state))
-        onShutDown = .emitShutdownAndFinish
+        onShutDown = .emitShutdownAndFinish(cancelling: nil)
 
       case .connecting(let state):
-        // Only emit the shutdown; there's no connection to close yet.
-        self = .shuttingDown(ShuttingDown(from: state))
-        onShutDown = .emitShutdown
+        if let backoffHandle = state.backoffHandle {
+          // The subchannel is waiting to retry a connection. There's no connection to close and no
+          // in-flight RPCs to drain, so cancel the backoff and shut down now rather than waiting for
+          // it to elapse.
+          self = .shutDown(ShutDown(from: state))
+          onShutDown = .emitShutdownAndFinish(cancelling: backoffHandle)
+        } else {
+          // A connection attempt is in-flight; emit the shutdown and wait for it to complete.
+          self = .shuttingDown(ShuttingDown(from: state))
+          onShutDown = .emitShutdown
+        }
 
       case .connected(let state):
         self = .shuttingDown(ShuttingDown(from: state))
@@ -553,6 +573,21 @@ extension Subchannel {
       return onShutDown
     }
 
+    /// Records the handle of the in-progress backoff task so that graceful shutdown can cancel it
+    /// instead of waiting for the backoff to elapse.
+    ///
+    /// This must only be called immediately after ``connectFailed(connector:authority:)`` returns
+    /// `.backoff`, when the subchannel is guaranteed to be `.connecting`.
+    mutating func setBackoffTaskHandle(_ handle: CancellableTaskHandle) {
+      switch self {
+      case .connecting(var state):
+        state.backoffHandle = handle
+        self = .connecting(state)
+      case .notConnected, .connected, .goingAway, .shuttingDown, .shutDown:
+        fatalError("Invalid state")
+      }
+    }
+
     enum OnConnectSucceeded {
       case updateStateToReady
       case finishAndClose(Connection)
@@ -633,7 +668,10 @@ extension Subchannel {
 
     mutating func backedOff() -> OnBackedOff {
       switch self {
-      case .connecting(let state):
+      case .connecting(var state):
+        // The backoff completed; a connection attempt is about to resume, so there's no longer a
+        // backoff task to cancel.
+        state.backoffHandle = nil
         self = .connecting(state)
         return .connect(state.connection)
 

Tests/GRPCNIOTransportCoreTests/Client/Connection/LoadBalancers/SubchannelTests.swift

@@ -279,6 +279,65 @@ final class SubchannelTests: XCTestCase {
     }
   }
 
+  func testShutdownWhileBackingOffCompletesPromptly() async throws {
+    // Regression test: a subchannel sitting in connection backoff has no established
+    // connection and no in-flight RPCs to drain, so graceful shutdown must complete
+    // promptly rather than waiting for the (potentially very long) backoff to elapse.
+    //
+    // Note: we must NOT cancel the task group to end `run()`. Cancellation would tear down
+    // the backoff `Task.sleep` itself, so shutdown would always appear fast and the test
+    // would prove nothing. Instead we let the subchannel finish its own streams in response
+    // to `shutDown()` and measure how long that takes.
+    let backoff = Duration.seconds(30)
+    let subchannel = self.makeSubchannel(
+      address: .unixDomainSocket(path: "not-listening"),
+      connector: .posix(),
+      backoff: .fixed(at: backoff)
+    )
+
+    let clock = ContinuousClock()
+
+    try await withThrowingTaskGroup(of: Void.self) { group in
+      group.addTask {
+        await subchannel.run()
+      }
+
+      var shutdownRequestedAt: ContinuousClock.Instant?
+      for await event in subchannel.events {
+        switch event {
+        case .connectivityStateChanged(.idle):
+          subchannel.connect()
+
+        case .connectivityStateChanged(.transientFailure):
+          // All addresses have been tried and the subchannel is now sleeping in backoff.
+          // Trigger graceful shutdown exactly once.
+          if shutdownRequestedAt == nil {
+            shutdownRequestedAt = clock.now
+            subchannel.shutDown()
+          }
+
+        default:
+          ()
+        }
+      }
+
+      // Exiting this loop means the subchannel finished its own event stream as part of
+      // shutting down (we never cancel the group), so the elapsed time below reflects how
+      // long graceful shutdown actually took. A backing-off subchannel has nothing to
+      // drain, so it should not wait for the backoff `Task.sleep` to complete.
+      let requestedAt = try XCTUnwrap(
+        shutdownRequestedAt,
+        "subchannel never entered the transient failure (backing off) state"
+      )
+      let elapsed = clock.now - requestedAt
+      XCTAssertLessThan(
+        elapsed,
+        backoff / 2,
+        "graceful shutdown waited out the connection backoff (took \(elapsed))"
+      )
+    }
+  }
+
   func testIdleTimeout() async throws {
     let server = TestServer(eventLoopGroup: .singletonMultiThreadedEventLoopGroup)
     let address = try await server.bind()