improve peer test (#288)

Author: xlc

Networking/Sources/Networking/Connection.swift

@@ -28,6 +28,7 @@ enum ConnectionState {
     case connecting(continuations: [CheckedContinuation<Void, Error>])
     case connected(publicKey: Data)
     case closed
+    case closing
     case reconnect(publicKey: Data)
 }
 
@@ -51,7 +52,7 @@ public final class Connection<Handler: StreamHandler>: Sendable, ConnectionInfoP
                 nil
             case let .connected(publicKey):
                 publicKey
-            case .closed:
+            case .closed, .closing:
                 nil
             case let .reconnect(publicKey):
                 publicKey
@@ -99,30 +100,41 @@ public final class Connection<Handler: StreamHandler>: Sendable, ConnectionInfoP
                 for continuation in continuations {
                     continuation.resume(throwing: ConnectionError.closed)
                 }
-                state = .closed
             }
             state = .closed
         }
     }
 
+    func closing() {
+        state.write { state in
+            if case let .connecting(continuations) = state {
+                for continuation in continuations {
+                    continuation.resume(throwing: ConnectionError.closed)
+                }
+            }
+            state = .closing
+        }
+    }
+
     func reconnect(publicKey: Data) {
         state.write { state in
             if case let .connecting(continuations) = state {
                 for continuation in continuations {
                     continuation.resume(throwing: ConnectionError.reconnect)
                 }
-                state = .reconnect(publicKey: publicKey)
             }
             state = .reconnect(publicKey: publicKey)
         }
     }
 
     public var isClosed: Bool {
         state.read {
-            if case .closed = $0 {
-                return true
+            switch $0 {
+            case .closed, .closing:
+                true
+            case .connected, .reconnect, .connecting:
+                false
             }
-            return false
         }
     }
 
@@ -140,11 +152,7 @@ public final class Connection<Handler: StreamHandler>: Sendable, ConnectionInfoP
             switch $0 {
             case .connecting:
                 false
-            case .connected:
-                true
-            case .closed:
-                true
-            case .reconnect:
+            case .connected, .closed, .closing, .reconnect:
                 true
             }
         }
@@ -165,6 +173,7 @@ public final class Connection<Handler: StreamHandler>: Sendable, ConnectionInfoP
     }
 
     public func close(abort: Bool = false) {
+        closing()
         try? connection.shutdown(errorCode: abort ? 1 : 0) // TODO: define some error code
     }
 

Networking/Sources/Networking/Peer.swift

@@ -177,14 +177,17 @@ public final class Peer<Handler: StreamHandler>: Sendable {
     public func broadcast(
         kind: Handler.PresistentHandler.StreamKind, message: Handler.PresistentHandler.Message
     ) {
-        let connections = impl.connections.read { connections in
-            connections.byId.values
-        }
         guard let messageData = try? message.encode() else {
             impl.logger.warning("Failed to encode message: \(message)")
             return
         }
+        let connections = impl.connections.read { connections in
+            connections.byId.values
+        }
         for connection in connections {
+            if connection.isClosed {
+                continue
+            }
             if let stream = try? connection.createPreistentStream(kind: kind) {
                 Task {
                     let res = await Result {
@@ -210,8 +213,9 @@ public final class Peer<Handler: StreamHandler>: Sendable {
     }
 
     // there should be only one connection per peer
+    // exlcude closed connections
     public var peersCount: Int {
-        impl.connections.read { $0.byId.count }
+        impl.connections.read { $0.byId.count { $0.value.isClosed == false } }
     }
 
     public var peersRole: PeerRole {

Networking/Tests/NetworkingTests/PeerTests.swift

@@ -166,8 +166,9 @@ struct PeerTests {
             let con = try peer.connect(to: centerPeer.listenAddress(), role: .builder)
             try await con.ready()
         }
-        // Simulate close connections 1~3s
-        try? await Task.sleep(for: .milliseconds(1000))
+
+        #expect(centerPeer.peersCount == 3)
+
         centerPeer.broadcast(kind: .uniqueA, message: .init(kind: .uniqueA, data: Data("connection rotation strategy".utf8)))
         try? await Task.sleep(for: .milliseconds(100))
         var receivedCount = 0
@@ -546,7 +547,7 @@ struct PeerTests {
         #expect(receivedData == messageData + Data(" response".utf8))
         try? await Task.sleep(for: .milliseconds(100))
         // Simulate abnormal shutdown of connections
-        connection.close(abort: true)
+        try connection.connection.shutdown(errorCode: 1)
         // Wait to simulate downtime & reconnected 3~5s
         try? await Task.sleep(for: .milliseconds(3000))
         peer1.broadcast(
@@ -650,7 +651,7 @@ struct PeerTests {
         #expect(receivedData == messageData + Data(" response".utf8))
         try? await Task.sleep(for: .milliseconds(100))
         // Simulate a peer failure by disconnecting one peer
-        connection.close(abort: false)
+        try connection.connection.shutdown()
         // Wait to simulate downtime
         try? await Task.sleep(for: .milliseconds(200))
         // Reconnect the failing peer