ConnectionPool: fix continuation leak when idle-timer reschedules (#641)

When the idle-timeout timer fires while a lease request is queued
(because a keep-alive is consuming the connection's only stream),
`connectionIdleTimerTriggered` takes the `rescheduleIdleTimer` path.
That path replaced the stale idle timer with a fresh one but dropped
the old timer's `cancellationContinuation` — a
`CheckedContinuation<Void, Never>` registered by
`ConnectionPool.runTimer`'s child task in `withCheckedContinuation`.
Dropped without being resumed, the Swift runtime eventually emits:

    SWIFT TASK CONTINUATION MISUSE: runTimer(_:in:) leaked its
    continuation without resuming it.

Fix:

- `ConnectionState.rescheduleIdleTimer()` and
  `ConnectionGroup.rescheduleIdleTimer(_:)` now also return the old
  idle timer's `TimerCancellationToken`.
- `PoolStateMachine.connectionIdleTimerTriggered(_:)` forwards that
  token via `.makeConnectionsCancelAndScheduleTimers` so
  `ConnectionPool.runConnectionAction` resumes the continuation before
  scheduling the new timer.

A new regression test in `PoolStateMachineTests` drives the state
machine into the exact race and asserts the returned action surfaces
the old token. The test fails on `main` and passes with this fix.

Signed-off-by: Lukas Tetzlaff <lukastetzlaff@gmail.com>
Co-authored-by: Fabian Fett <fabianfett@apple.com>

Author: ltetzlaff

Sources/ConnectionPoolModule/PoolStateMachine+ConnectionGroup.swift

@@ -698,7 +698,7 @@ extension PoolStateMachine {
         }
 
         @inlinable
-        mutating func rescheduleIdleTimer(_ connectionID: Connection.ID) -> ConnectionTimer? {
+        mutating func rescheduleIdleTimer(_ connectionID: Connection.ID) -> (ConnectionTimer, TimerCancellationToken?)? {
             guard let index = self.connections.firstIndex(where: { $0.id == connectionID }) else {
                 // because of a race this connection (connection close runs against trigger of timeout)
                 // was already removed from the state machine.

Sources/ConnectionPoolModule/PoolStateMachine+ConnectionState.swift

@@ -336,16 +336,16 @@ extension PoolStateMachine {
         }
 
         @inlinable
-        mutating func rescheduleIdleTimer() -> ConnectionTimer? {
+        mutating func rescheduleIdleTimer() -> (ConnectionTimer, TimerCancellationToken?)? {
             switch self.state {
             case .backingOff, .starting:
                 preconditionFailure("Invalid state: \(self.state)")
 
-            case .idle(let connection, let maxStreams, let keepAlive, .some):
+            case .idle(let connection, let maxStreams, let keepAlive, .some(let oldIdleTimer)):
                 let idleTimerState = self._nextTimer()
                 let idleTimer = ConnectionTimer(timerID: idleTimerState.timerID, connectionID: self.id, usecase: .idleTimeout)
                 self.state = .idle(connection, maxStreams: maxStreams, keepAlive: keepAlive, idleTimer: idleTimerState)
-                return idleTimer
+                return (idleTimer, oldIdleTimer.cancellationContinuation)
 
             case .idle(_,_,_, .none):
                 return nil

Sources/ConnectionPoolModule/PoolStateMachine.swift

@@ -620,8 +620,26 @@ struct PoolStateMachine<
             // might pickup the lease request (which would cancel the idle timeout) or the lease request might
             // already be handled by another (currently busy connection), in which case the idle timeout can
             // trigger eventually.
-            let timer = self.connections.rescheduleIdleTimer(connectionID)
-            return .init(request: .none, connection: .scheduleTimers(timer.map { [self.mapTimers($0)] } ?? []))
+            guard let (newTimer, oldCancellationToken) = self.connections.rescheduleIdleTimer(connectionID) else {
+                return .none()
+            }
+            let scheduledTimers = Max2Sequence(self.mapTimers(newTimer))
+            // We must propagate the old idle timer's cancellation continuation so that the
+            // `ConnectionPool.runTimer` child task that stored it can be resumed. Dropping it here
+            // produces a "SWIFT TASK CONTINUATION MISUSE: runTimer(_:in:) leaked its continuation
+            // without resuming it" runtime warning (the stored `CheckedContinuation` is never
+            // resumed and eventually deallocated).
+            if let oldCancellationToken {
+                return .init(
+                    request: .none,
+                    connection: .makeConnectionsCancelAndScheduleTimers(
+                        .init(),
+                        .init(element: oldCancellationToken),
+                        scheduledTimers
+                    )
+                )
+            }
+            return .init(request: .none, connection: .scheduleTimers(scheduledTimers))
         }
 
         guard let closeAction = self.connections.closeConnectionIfIdle(connectionID) else {

Tests/ConnectionPoolModuleTests/PoolStateMachineTests.swift

@@ -1398,6 +1398,96 @@ typealias TestPoolStateMachine = PoolStateMachine<
         // Should only lease 1 (the waiting request count)
         #expect(leasedRequests.count == 1)
     }
+
+    /// Regression test: when the idle-timeout timer fires while a lease request is waiting
+    /// in the queue (because a keep-alive is currently consuming the connection's only
+    /// stream), `connectionIdleTimerTriggered` reschedules a fresh idle timer. Prior to the
+    /// fix, the old idle timer's `TimerCancellationToken` was silently dropped, which in
+    /// `ConnectionPool.runTimer` manifests as:
+    ///
+    ///     SWIFT TASK CONTINUATION MISUSE: runTimer(_:in:) leaked its continuation
+    ///     without resuming it.
+    ///
+    /// because the stored `CheckedContinuation<Void, Never>` is never resumed. This test
+    /// drives the state machine into that exact race and asserts the returned action
+    /// surfaces the old cancellation token so the caller can resume it.
+    @available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *)
+    @Test func testIdleTimerReschedulePreservesOldCancellationToken() {
+        var configuration = PoolConfiguration()
+        configuration.minimumConnectionCount = 0
+        configuration.maximumConnectionSoftLimit = 1
+        configuration.maximumConnectionHardLimit = 1
+        configuration.keepAliveDuration = .seconds(2)
+        configuration.idleTimeoutDuration = .seconds(4)
+
+        var stateMachine = TestPoolStateMachine(
+            configuration: configuration,
+            generator: .init(),
+            timerCancellationTokenType: MockTimerCancellationToken.self,
+            clock: MockClock()
+        )
+
+        // Lease a request — triggers creation of a demand connection.
+        let request1 = MockRequest(connectionType: MockConnection.self)
+        let leaseAction1 = stateMachine.leaseConnection(request1)
+        guard case .makeConnection(let makeRequest, _) = leaseAction1.connection else {
+            Issue.record("expected .makeConnection, got \(leaseAction1.connection)")
+            return
+        }
+
+        // Establish the connection — because the request is queued it is leased immediately.
+        let connection = MockConnection(id: makeRequest.connectionID)
+        let establishedAction = stateMachine.connectionEstablished(connection, maxStreams: 1)
+        #expect(establishedAction.request == .leaseConnection(.init(element: request1), connection))
+
+        // Release the connection — parks with both keep-alive and idle timers because this
+        // is a demand connection (index >= minimumConcurrentConnections).
+        let releaseAction = stateMachine.releaseConnection(connection, streams: 1)
+        let keepAliveTimer = TestPoolStateMachine.Timer(
+            .init(timerID: 0, connectionID: connection.id, usecase: .keepAlive),
+            duration: configuration.keepAliveDuration!
+        )
+        let idleTimer = TestPoolStateMachine.Timer(
+            .init(timerID: 1, connectionID: connection.id, usecase: .idleTimeout),
+            duration: configuration.idleTimeoutDuration
+        )
+        #expect(releaseAction.connection == .scheduleTimers([keepAliveTimer, idleTimer]))
+
+        // Register cancellation tokens for both timers (simulates the child task in
+        // `ConnectionPool.runTimer` storing its continuation via `timerScheduled`).
+        let keepAliveToken = MockTimerCancellationToken(keepAliveTimer)
+        let idleToken = MockTimerCancellationToken(idleTimer)
+        #expect(stateMachine.timerScheduled(keepAliveTimer, cancelContinuation: keepAliveToken) == nil)
+        #expect(stateMachine.timerScheduled(idleTimer, cancelContinuation: idleToken) == nil)
+
+        // Keep-alive fires — the connection transitions to keepAlive: .running, idleTimer: .some.
+        let keepAliveFired = stateMachine.timerTriggered(keepAliveTimer)
+        #expect(keepAliveFired.connection == .runKeepAlive(connection, keepAliveToken))
+
+        // Queue a second request. Because the running keep-alive consumes the only stream,
+        // the request stays in the queue.
+        let request2 = MockRequest(connectionType: MockConnection.self)
+        _ = stateMachine.leaseConnection(request2)
+
+        // Idle timer fires while the keep-alive is still running and the queue is non-empty.
+        // This is the `rescheduleIdleTimer` code path.
+        let newIdleTimer = TestPoolStateMachine.Timer(
+            .init(timerID: 2, connectionID: connection.id, usecase: .idleTimeout),
+            duration: configuration.idleTimeoutDuration
+        )
+        let idleFired = stateMachine.timerTriggered(idleTimer)
+
+        // The returned action must carry the old idle timer's cancellation token so the
+        // pool can resume its continuation. Before the fix this was an empty
+        // `.scheduleTimers([newIdleTimer])` and `idleToken` was dropped on the floor.
+        #expect(
+            idleFired.connection == .makeConnectionsCancelAndScheduleTimers(
+                .init(),
+                .init(element: idleToken),
+                .init(newIdleTimer)
+            )
+        )
+    }
 }
 
 struct SomeError: Error {}