Skip to content

Commit 036931d

Browse files
fabianfettfabianfett
authored
Fixes Crash in ConnectionPoolStateMachine (#438)
- Correctly handle Connection closes while running a keep alive (fix: #436) - Add further keep alive tests - Restructure MockClock quite a bit
1 parent c826992 commit 036931d

File tree

8 files changed

+356
-160
lines changed

8 files changed

+356
-160
lines changed

‎Sources/ConnectionPoolModule/PoolStateMachine+ConnectionGroup.swift‎

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -385,7 +385,8 @@ extension PoolStateMachine {
385385
@inlinable
386386
mutating func keepAliveSucceeded(_ connectionID: Connection.ID) -> (Int, AvailableConnectionContext)? {
387387
guard let index = self.connections.firstIndex(where: { $0.id == connectionID }) else {
388-
preconditionFailure("A connection that we don't know was released? Something is very wrong...")
388+
// keepAliveSucceeded can race against, closeIfIdle, shutdowns or connection errors
389+
return nil
389390
}
390391

391392
guard let connectionInfo = self.connections[index].keepAliveSucceeded() else {
@@ -430,15 +431,8 @@ extension PoolStateMachine {
430431

431432
self.stats.idle -= 1
432433
self.stats.closing += 1
433-
434-
// if idleState.runningKeepAlive {
435-
// self.stats.runningKeepAlive -= 1
436-
// if self.keepAliveReducesAvailableStreams {
437-
// self.stats.availableStreams += 1
438-
// }
439-
// }
440-
441-
self.stats.availableStreams -= closeAction.maxStreams
434+
self.stats.runningKeepAlive -= closeAction.runningKeepAlive ? 1 : 0
435+
self.stats.availableStreams -= closeAction.maxStreams - closeAction.usedStreams
442436

443437
return CloseAction(
444438
connection: closeAction.connection!,

‎Sources/ConnectionPoolModule/PoolStateMachine+ConnectionState.swift‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -496,20 +496,25 @@ extension PoolStateMachine {
496496
var usedStreams: UInt16
497497
@usableFromInline
498498
var maxStreams: UInt16
499+
@usableFromInline
500+
var runningKeepAlive: Bool
501+
499502

500503
@inlinable
501504
init(
502505
connection: Connection?,
503506
previousConnectionState: PreviousConnectionState,
504507
cancelTimers: Max2Sequence<TimerCancellationToken>,
505508
usedStreams: UInt16,
506-
maxStreams: UInt16
509+
maxStreams: UInt16,
510+
runningKeepAlive: Bool
507511
) {
508512
self.connection = connection
509513
self.previousConnectionState = previousConnectionState
510514
self.cancelTimers = cancelTimers
511515
self.usedStreams = usedStreams
512516
self.maxStreams = maxStreams
517+
self.runningKeepAlive = runningKeepAlive
513518
}
514519
}
515520

@@ -526,7 +531,8 @@ extension PoolStateMachine {
526531
idleTimerState?.cancellationContinuation
527532
),
528533
usedStreams: keepAlive.usedStreams,
529-
maxStreams: maxStreams
534+
maxStreams: maxStreams,
535+
runningKeepAlive: keepAlive.isRunning
530536
)
531537

532538
case .leased, .closed:
@@ -559,7 +565,8 @@ extension PoolStateMachine {
559565
idleTimerState?.cancellationContinuation
560566
),
561567
usedStreams: keepAlive.usedStreams,
562-
maxStreams: maxStreams
568+
maxStreams: maxStreams,
569+
runningKeepAlive: keepAlive.isRunning
563570
)
564571

565572
case .leased(let connection, usedStreams: let usedStreams, maxStreams: let maxStreams, var keepAlive):
@@ -571,7 +578,8 @@ extension PoolStateMachine {
571578
keepAlive.cancelTimerIfScheduled()
572579
),
573580
usedStreams: keepAlive.usedStreams + usedStreams,
574-
maxStreams: maxStreams
581+
maxStreams: maxStreams,
582+
runningKeepAlive: keepAlive.isRunning
575583
)
576584

577585
case .backingOff(let timer):
@@ -581,7 +589,8 @@ extension PoolStateMachine {
581589
previousConnectionState: .backingOff,
582590
cancelTimers: Max2Sequence(timer.cancellationContinuation),
583591
usedStreams: 0,
584-
maxStreams: 0
592+
maxStreams: 0,
593+
runningKeepAlive: false
585594
)
586595
}
587596
}

‎Tests/ConnectionPoolModuleTests/ConnectionPoolTests.swift‎

Lines changed: 153 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ final class ConnectionPoolTests: XCTestCase {
1616
configuration: config,
1717
idGenerator: ConnectionIDGenerator(),
1818
requestType: ConnectionRequest<MockConnection>.self,
19-
keepAliveBehavior: MockPingPongBehavior(keepAliveFrequency: nil),
19+
keepAliveBehavior: MockPingPongBehavior(keepAliveFrequency: nil, connectionType: MockConnection.self),
2020
observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self),
2121
clock: ContinuousClock()
2222
) {
@@ -74,7 +74,7 @@ final class ConnectionPoolTests: XCTestCase {
7474
configuration: config,
7575
idGenerator: ConnectionIDGenerator(),
7676
requestType: ConnectionRequest<MockConnection>.self,
77-
keepAliveBehavior: MockPingPongBehavior(keepAliveFrequency: nil),
77+
keepAliveBehavior: MockPingPongBehavior(keepAliveFrequency: nil, connectionType: MockConnection.self),
7878
observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self),
7979
clock: clock
8080
) {
@@ -119,7 +119,7 @@ final class ConnectionPoolTests: XCTestCase {
119119
configuration: config,
120120
idGenerator: ConnectionIDGenerator(),
121121
requestType: ConnectionRequest<MockConnection>.self,
122-
keepAliveBehavior: MockPingPongBehavior(keepAliveFrequency: nil),
122+
keepAliveBehavior: MockPingPongBehavior(keepAliveFrequency: nil, connectionType: MockConnection.self),
123123
observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self),
124124
clock: clock
125125
) {
@@ -135,7 +135,7 @@ final class ConnectionPoolTests: XCTestCase {
135135
throw ConnectionCreationError()
136136
}
137137

138-
await clock.timerScheduled()
138+
await clock.nextTimerScheduled()
139139

140140
taskGroup.cancelAll()
141141
}
@@ -156,7 +156,7 @@ final class ConnectionPoolTests: XCTestCase {
156156
configuration: config,
157157
idGenerator: ConnectionIDGenerator(),
158158
requestType: ConnectionRequest<MockConnection>.self,
159-
keepAliveBehavior: MockPingPongBehavior(keepAliveFrequency: nil),
159+
keepAliveBehavior: MockPingPongBehavior(keepAliveFrequency: nil, connectionType: MockConnection.self),
160160
observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self),
161161
clock: ContinuousClock()
162162
) {
@@ -220,6 +220,154 @@ final class ConnectionPoolTests: XCTestCase {
220220
XCTAssert(hasFinished.load(ordering: .relaxed))
221221
XCTAssertEqual(factory.runningConnections.count, 0)
222222
}
223+
224+
func testKeepAliveWorks() async throws {
225+
let clock = MockClock()
226+
let factory = MockConnectionFactory<MockClock>()
227+
let keepAliveDuration = Duration.seconds(30)
228+
let keepAlive = MockPingPongBehavior(keepAliveFrequency: keepAliveDuration, connectionType: MockConnection.self)
229+
230+
var mutableConfig = ConnectionPoolConfiguration()
231+
mutableConfig.minimumConnectionCount = 0
232+
mutableConfig.maximumConnectionSoftLimit = 1
233+
mutableConfig.maximumConnectionHardLimit = 1
234+
let config = mutableConfig
235+
236+
let pool = ConnectionPool(
237+
configuration: config,
238+
idGenerator: ConnectionIDGenerator(),
239+
requestType: ConnectionRequest<MockConnection>.self,
240+
keepAliveBehavior: keepAlive,
241+
observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self),
242+
clock: clock
243+
) {
244+
try await factory.makeConnection(id: $0, for: $1)
245+
}
246+
247+
try await withThrowingTaskGroup(of: Void.self) { taskGroup in
248+
taskGroup.addTask {
249+
await pool.run()
250+
}
251+
252+
async let lease1ConnectionAsync = pool.leaseConnection()
253+
254+
let connection = await factory.nextConnectAttempt { connectionID in
255+
return 1
256+
}
257+
258+
let lease1Connection = try await lease1ConnectionAsync
259+
XCTAssert(connection === lease1Connection)
260+
261+
pool.releaseConnection(lease1Connection)
262+
263+
// keep alive 1
264+
265+
// validate that a keep alive timer and an idle timeout timer is scheduled
266+
var expectedInstants: Set<MockClock.Instant> = [.init(keepAliveDuration), .init(config.idleTimeout)]
267+
let deadline1 = await clock.nextTimerScheduled()
268+
print(deadline1)
269+
XCTAssertNotNil(expectedInstants.remove(deadline1))
270+
let deadline2 = await clock.nextTimerScheduled()
271+
print(deadline2)
272+
XCTAssertNotNil(expectedInstants.remove(deadline2))
273+
XCTAssert(expectedInstants.isEmpty)
274+
275+
// move clock forward to keep alive
276+
let newTime = clock.now.advanced(by: keepAliveDuration)
277+
clock.advance(to: newTime)
278+
print("clock advanced to: \(newTime)")
279+
280+
await keepAlive.nextKeepAlive { keepAliveConnection in
281+
defer { print("keep alive 1 has run") }
282+
XCTAssertTrue(keepAliveConnection === lease1Connection)
283+
return true
284+
}
285+
286+
// keep alive 2
287+
288+
let deadline3 = await clock.nextTimerScheduled()
289+
XCTAssertEqual(deadline3, clock.now.advanced(by: keepAliveDuration))
290+
print(deadline3)
291+
292+
// race keep alive vs timeout
293+
clock.advance(to: clock.now.advanced(by: keepAliveDuration))
294+
295+
taskGroup.cancelAll()
296+
297+
for connection in factory.runningConnections {
298+
connection.closeIfClosing()
299+
}
300+
}
301+
}
302+
303+
func testKeepAliveWorksRacesAgainstShutdown() async throws {
304+
let clock = MockClock()
305+
let factory = MockConnectionFactory<MockClock>()
306+
let keepAliveDuration = Duration.seconds(30)
307+
let keepAlive = MockPingPongBehavior(keepAliveFrequency: keepAliveDuration, connectionType: MockConnection.self)
308+
309+
var mutableConfig = ConnectionPoolConfiguration()
310+
mutableConfig.minimumConnectionCount = 0
311+
mutableConfig.maximumConnectionSoftLimit = 1
312+
mutableConfig.maximumConnectionHardLimit = 1
313+
let config = mutableConfig
314+
315+
let pool = ConnectionPool(
316+
configuration: config,
317+
idGenerator: ConnectionIDGenerator(),
318+
requestType: ConnectionRequest<MockConnection>.self,
319+
keepAliveBehavior: keepAlive,
320+
observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self),
321+
clock: clock
322+
) {
323+
try await factory.makeConnection(id: $0, for: $1)
324+
}
325+
326+
try await withThrowingTaskGroup(of: Void.self) { taskGroup in
327+
taskGroup.addTask {
328+
await pool.run()
329+
}
330+
331+
async let lease1ConnectionAsync = pool.leaseConnection()
332+
333+
let connection = await factory.nextConnectAttempt { connectionID in
334+
return 1
335+
}
336+
337+
let lease1Connection = try await lease1ConnectionAsync
338+
XCTAssert(connection === lease1Connection)
339+
340+
pool.releaseConnection(lease1Connection)
341+
342+
// keep alive 1
343+
344+
// validate that a keep alive timer and an idle timeout timer is scheduled
345+
var expectedInstants: Set<MockClock.Instant> = [.init(keepAliveDuration), .init(config.idleTimeout)]
346+
let deadline1 = await clock.nextTimerScheduled()
347+
print(deadline1)
348+
XCTAssertNotNil(expectedInstants.remove(deadline1))
349+
let deadline2 = await clock.nextTimerScheduled()
350+
print(deadline2)
351+
XCTAssertNotNil(expectedInstants.remove(deadline2))
352+
XCTAssert(expectedInstants.isEmpty)
353+
354+
clock.advance(to: clock.now.advanced(by: keepAliveDuration))
355+
356+
await keepAlive.nextKeepAlive { keepAliveConnection in
357+
defer { print("keep alive 1 has run") }
358+
XCTAssertTrue(keepAliveConnection === lease1Connection)
359+
return true
360+
}
361+
362+
taskGroup.cancelAll()
363+
print("cancelled")
364+
365+
for connection in factory.runningConnections {
366+
connection.closeIfClosing()
367+
}
368+
}
369+
}
370+
223371
}
224372

225373

0 commit comments

Comments
 (0)