Fix race conditions and delayed connections

.github/workflows/ci.yml

@@ -4,11 +4,11 @@ on: push
 
 jobs:
   test:
-    runs-on: macos-14
+    runs-on: macos-15
 
     steps:
       - uses: actions/checkout@v4
-      - name: Select Xcode 15
-        run: sudo xcode-select -s /Applications/Xcode_15.4.app
+      - name: Select Xcode 26
+        run: sudo xcode-select -s /Applications/Xcode_26.3.app
       - name: Test
         run: swift test

Package.swift

@@ -4,7 +4,7 @@ import PackageDescription
 let package = Package(
     name: "Phoenix",
     platforms: [
-        .macOS(.v12), .iOS(.v15), .tvOS(.v15), .watchOS(.v8),
+        .macOS(.v13), .iOS(.v16), .tvOS(.v16), .watchOS(.v9),
     ],
     products: [
         .library(name: "Phoenix", targets: ["Phoenix"]),

Sources/Phoenix/PhoenixSocket.swift

@@ -101,17 +101,32 @@ final actor PhoenixSocket {
     }
 
     func connect() async {
-        guard case .closed = _connectionState.value else { return }
+        switch _connectionState.value {
+        case .closed:
+            break
+        case .waitingToReconnect, .preparingToReconnect:
+            tasks.cancel(forKey: "reconnect")
+            _connectionState.value = .closed(connectionAttempts: 0)
+        case .connecting, .open, .closing:
+            return
+        }
         shouldReconnect = true
         await doConnect()
     }
 
     func disconnect(timeout: TimeInterval? = nil) async {
-        guard let ws = webSocket else { return }
-        await doCloseFromClient(
-            id: ws.id,
-            timeout: timeout?.nanoseconds ?? self.timeout
-        )
+        if let ws = webSocket {
+            await doCloseFromClient(
+                id: ws.id,
+                timeout: timeout?.nanoseconds ?? self.timeout
+            )
+        } else if !_connectionState.value.isClosed {
+            shouldReconnect = false
+            pushes.pause()
+            removeAll()
+            tasks.cancelAll()
+            _connectionState.value = .closed(connectionAttempts: 0)
+        }
     }
 }
 
@@ -182,6 +197,9 @@ extension PhoenixSocket {
     }
 
     private func flush() {
+        tasks.cancel(forKey: "flush")
+        pushes.cancelAwaitingContinuation()
+
         let task = Task { [weak self] in
             guard let self, !Task.isCancelled else { return }
 
@@ -229,6 +247,8 @@ extension PhoenixSocket {
     }
 
     private func listen() {
+        tasks.cancel(forKey: "listen")
+
         let task = Task { [weak self] in
             guard !Task.isCancelled,
                   let ws = await self?.webSocket,
@@ -276,6 +296,8 @@ extension PhoenixSocket {
 
 extension PhoenixSocket {
     private func scheduleHeartbeat() {
+        tasks.cancel(forKey: "heartbeat")
+
         let interval = heartbeatInterval
         let task = Task { [weak self] in
             try await Task.sleep(nanoseconds: interval)
@@ -410,6 +432,7 @@ extension PhoenixSocket {
             os_log("connect", log: .phoenix, type: .debug)
 
             let ws = try await doMakeWebSocket()
+            try Task.checkCancellation()
             _connectionState.value = .connecting(ws)
 
             try await ws.open()
@@ -423,6 +446,7 @@ extension PhoenixSocket {
             scheduleHeartbeat()
 
         } catch {
+            guard !Task.isCancelled else { return }
             _connectionState.value = .closed(connectionAttempts: attempts + 1)
             await doConnect()
         }
@@ -469,7 +493,6 @@ extension PhoenixSocket {
         switch _connectionState.value {
         case let .connecting(ws) where ws.id == id,
              let .open(ws) where ws.id == id:
-
             os_log(
                 "close: %@",
                 log: .phoenix,
@@ -485,7 +508,7 @@ extension PhoenixSocket {
             let connectTask = Task.detached { [weak self] in
                 await self?.doConnect()
             }
-            tasks.add(connectTask)
+            tasks.insert(connectTask, forKey: "reconnect")
 
         default:
             break

Sources/Phoenix/PushBuffer.swift

@@ -123,6 +123,15 @@ Sendable {
         state.access { $0.putBack(push) }
     }
 
+    func cancelAwaitingContinuation() {
+        let cont = state.access { state -> AwaitingPushContinuation? in
+            let cont = state.awaitingPushContinuation
+            state.awaitingPushContinuation = nil
+            return cont
+        }
+        cont?.resume(throwing: CancellationError())
+    }
+
     /// Cancels all in-flight and buffered pushes and invalidates the
     /// buffer with the specified error or `CancellationError`. Any
     /// subsequent calls to `append()`, `appendAndWait()`, or `next()`
@@ -344,6 +353,7 @@ private extension PushBuffer {
         mutating func setTimeout(_ date: Date, makeTask: () -> Task<Void, Error>) {
             if let timeout {
                 guard date < timeout.date else { return }
+                timeout.cancel()
                 self.timeout = Timeout(date: date, task: makeTask())
             } else {
                 timeout = Timeout(date: date, task: makeTask())

Tests/PhoenixTests/PhoenixSocketTests.swift

@@ -586,6 +586,82 @@ final class PhoenixSocketTests: XCTestCase {
         }
     }
 
+    func testDisconnectStopsReconnectionLoop() async throws {
+        let openCount = Locked(0)
+        let shouldFailOpen = Locked(false)
+
+        let socket = PhoenixSocket(
+            url: url,
+            timeout: 0.001,
+            heartbeatInterval: 0.001,
+            makeWebSocket: { [weak self] _, _, _, onOpen, onClose in
+                guard let self else { throw CancellationError() }
+                return fake(onOpen: onOpen, onClose: onClose, open: {
+                    openCount.access { $0 += 1 }
+                    if shouldFailOpen.access({ $0 }) {
+                        throw URLError(.notConnectedToInternet)
+                    }
+                    onOpen()
+                })
+            }
+        )
+
+        await socket.connect()
+        await AssertTrue(socket.connectionState.isOpen)
+
+        shouldFailOpen.access { $0 = true }
+
+        await AssertTrueEventually(openCount.access { $0 } >= 4)
+
+        await socket.disconnect()
+        XCTAssertTrue(socket.connectionState.isClosed)
+
+        try await Task.sleep(nanoseconds: NSEC_PER_MSEC * 10)
+        let countAfterDisconnect = openCount.access { $0 }
+        try await Task.sleep(nanoseconds: NSEC_PER_MSEC * 50)
+        XCTAssertEqual(countAfterDisconnect, openCount.access { $0 })
+    }
+
+    func testConnectAfterDisconnectDuringReconnectionHasNoBackoff() async throws {
+        let openCount = Locked(0)
+        let shouldFailOpen = Locked(false)
+
+        let socket = PhoenixSocket(
+            url: url,
+            timeout: 0.01,
+            heartbeatInterval: 0.01,
+            makeWebSocket: { [weak self] _, _, _, onOpen, onClose in
+                guard let self else { throw CancellationError() }
+                return fake(onOpen: onOpen, onClose: onClose, open: {
+                    openCount.access { $0 += 1 }
+                    if shouldFailOpen.access({ $0 }) {
+                        throw URLError(.notConnectedToInternet)
+                    }
+                    onOpen()
+                })
+            }
+        )
+
+        await socket.connect()
+        await AssertTrue(socket.connectionState.isOpen)
+
+        shouldFailOpen.access { $0 = true }
+
+        await AssertTrueEventually(openCount.access { $0 } >= 4)
+
+        await socket.disconnect()
+        XCTAssertTrue(socket.connectionState.isClosed)
+
+        shouldFailOpen.access { $0 = false }
+
+        let start = ContinuousClock.now
+        await socket.connect()
+        let elapsed = ContinuousClock.now - start
+
+        await AssertTrue(socket.connectionState.isOpen)
+        XCTAssertLessThan(elapsed, .milliseconds(500))
+    }
+
     func testTriggersChannelErrorIfJoining() async throws {
         let didErrorWhileJoining = Locked(false)
 

Tests/PhoenixTests/PushBufferTests.swift

@@ -223,8 +223,8 @@ final class PushBufferTests: XCTestCase {
             }
         }
 
-        let _ = await timeoutTask(push: push1).value
-        let _ = await timeoutTask(push: push2).value
+        _ = await timeoutTask(push: push1).value
+        _ = await timeoutTask(push: push2).value
 
         #if compiler(>=5.8)
             await fulfillment(of: [ex], timeout: 2)
@@ -414,7 +414,7 @@ final class PushBufferTests: XCTestCase {
 
             group.addTask {
                 self.prepareToSend(join)
-                let _ = try await buffer.appendAndWait(join)
+                _ = try await buffer.appendAndWait(join)
                 XCTAssertFalse(didProcessJoin.access { didProcess in
                     let old = didProcess
                     didProcess = true
@@ -929,6 +929,69 @@ final class PushBufferTests: XCTestCase {
         XCTAssertEqual(1, result.access { $0.processCount })
     }
 
+    func testCancelAwaitingContinuationCancelsWaitingIterator() async throws {
+        let buffer = PushBuffer()
+        buffer.resume()
+
+        let didCancel = Locked(false)
+
+        let iteratorTask = Task {
+            do {
+                for try await _ in buffer {
+                    XCTFail("Should not have produced push")
+                }
+            } catch {
+                XCTAssertTrue(error is CancellationError)
+                didCancel.access { $0 = true }
+            }
+        }
+
+        try await Task.sleep(nanoseconds: NSEC_PER_MSEC * 20)
+        buffer.cancelAwaitingContinuation()
+        await iteratorTask.value
+
+        XCTAssertTrue(didCancel.access { $0 })
+    }
+
+    func testCancelAwaitingContinuationAllowsNewIterator() async throws {
+        let buffer = PushBuffer()
+        buffer.resume()
+
+        let firstIteratorCancelled = Locked(false)
+
+        let firstTask = Task {
+            do {
+                for try await _ in buffer {
+                    XCTFail("Should not have produced push")
+                }
+            } catch {
+                firstIteratorCancelled.access { $0 = true }
+            }
+        }
+
+        try await Task.sleep(nanoseconds: NSEC_PER_MSEC * 20)
+        buffer.cancelAwaitingContinuation()
+        await firstTask.value
+        XCTAssertTrue(firstIteratorCancelled.access { $0 })
+
+        let push = makePush(1)
+        let didReceivePush = Locked(false)
+
+        Task {
+            try await Task.sleep(nanoseconds: NSEC_PER_MSEC * 10)
+            try await buffer.append(push) as Void
+        }
+
+        for try await p in buffer {
+            XCTAssertEqual(push, p)
+            didReceivePush.access { $0 = true }
+            buffer.didSend(p)
+            break
+        }
+
+        XCTAssertTrue(didReceivePush.access { $0 })
+    }
+
     func testSlowPushesDoNotDelayOtherPushes() async throws {
         let pushes = makePushes(5)
         let receivedMessages = Locked<[Message]>([])