fix quicClient

Tun/Punchnet/Actors/SDLContextActor.swift

@@ -181,12 +181,6 @@ actor SDLContextActor {
             try await Task.sleep(for: .seconds(0.3))
             SDLLogger.log("[SDLContext] start quic client: \(self.config.serverHost)")
             
-//            self.quicWorker = Task {
-//                for await message in await quicClient.messageStream {
-//                    await self.handleQUICMessage(message: message)
-//                }
-//            }
-                
             try await withThrowingTaskGroup { group in
                 defer {
                     group.cancelAll()
@@ -196,16 +190,19 @@ actor SDLContextActor {
                     for await message in await quicClient.messageStream {
                         await self.handleQUICMessage(message: message)
                     }
-                    throw SDLQUICClientExit.transportClosed("messageStream finished")
                 }
 
                 group.addTask {
                     let exit = await quicClient.run()
 
                     switch exit {
-                    case .normal, .cancelled:
+                    case .normal:
                         return
-                        
+                    case .cancelled:
+                        if Task.isCancelled {
+                            return
+                        }
+                        throw exit
                     case .transportClosed, .readFailed, .writeFailed:
                         throw exit
                     }
@@ -222,9 +219,11 @@ actor SDLContextActor {
                             throw error
                         }
                     }
-                    throw SDLQUICClientExit.cancelled
                 }
-                try await group.next()
+
+                while let _ = try await group.next() {
+                    ()
+                }
             }
 
         }

Tun/Punchnet/Actors/SDLQuicClient.swift

@@ -66,6 +66,8 @@ actor SDLQUICClient {
     public var eventStream: AsyncStream<SDLQUICEvent>
     private let eventCont: AsyncStream<SDLQUICEvent>.Continuation
     
+    private var didFinishStreams = false
+
     private let connection: NWConnection
     private let queue = DispatchQueue(label: "com.sdl.QUICClient.queue") // 专用队列保证线程安全
     
@@ -92,64 +94,58 @@ actor SDLQUICClient {
     }
 
     func start() {
-        connection.stateUpdateHandler = { state in
+        connection.stateUpdateHandler = { [weak self] state in
-            SDLLogger.log("[SDLQUICClient] new state: \(state)", for: .debug)
+            Task {
-            switch state {
+                await self?.handleConnectionStateUpdate(state)
-            case .ready:
-                Task {
-                    await self.readyState.markReady()
-                }
-            case .failed(let error):
-                Task {
-                    await self.readyState.markFailed(error)
-                }
-                self.eventCont.yield(.failed(error))
-            case .cancelled:
-                Task {
-                    await self.readyState.markCancelled()
-                }
-                self.eventCont.yield(.cancelled)
-            case .setup, .preparing:
-                Task {
-                    await self.readyState.markConnecting()
-                }
-            default:
-                ()
             }
         }
         connection.start(queue: self.queue)
     }
     
-    func waitReady(timeout: Duration = .seconds(5)) async throws {
+    private func handleConnectionStateUpdate(_ state: NWConnection.State) async {
-        try await withThrowingTaskGroup(of: Void.self) { group in
+        SDLLogger.log("[SDLQUICClient] new state: \(state)", for: .debug)
-            group.addTask {
+        switch state {
-                try await self.readyState.waitReady()
+        case .ready:
-            }
+            await self.readyState.markReady()
-            
+        case .failed(let error):
-            group.addTask {
+            await self.readyState.markFailed(error)
-                try await Task.sleep(for: timeout)
+            self.emitEvent(.failed(error))
-                throw SDLQUICError.timeout
+        case .cancelled:
-            }
+            await self.readyState.markCancelled()
-            
+            self.emitEvent(.cancelled)
-            try await group.next()
+        case .setup, .preparing:
-            group.cancelAll()
+            await self.readyState.markConnecting()
+        default:
+            ()
         }
     }
     
+    func waitReady(timeout: Duration = .seconds(5)) async throws {
+        try await self.readyState.waitReady(timeout: timeout)
+    }
+
     func run() async -> SDLQUICClientExit {
-        await withTaskGroup(of: SDLQUICClientExit.self) { group in
+        await withTaskCancellationHandler {
-            group.addTask {
+            await withTaskGroup(of: SDLQUICClientExit.self) { group in
-                await self.readLoop()
+                group.addTask {
+                    await self.readLoop()
+                }
+
+                group.addTask {
+                    await self.heartbeatLoop()
+                }
+
+                let exit = await group.next() ?? .normal
+                group.cancelAll()
+                await self.stop()
+                self.finishStreams()
+
+                return exit
             }
-            
+        } onCancel: {
-            group.addTask {
+            Task {
-                await self.heartbeatLoop()
+                await self.stop()
             }
-            
-            let exit = await group.next() ?? .normal
-            group.cancelAll()
-            
-            return exit
         }
     }
     
@@ -163,7 +159,9 @@ actor SDLQUICClient {
         connection.send(content: packet, completion: .contentProcessed { [weak self] error in
             if let error {
                 SDLLogger.log("[SDLQUICClient] send data get error: \(error)", for: .debug)
-                self?.eventCont.yield(.writeFailed(error))
+                Task {
+                    await self?.emitEvent(.writeFailed(error))
+                }
             }
         })
     }
@@ -183,16 +181,33 @@ actor SDLQUICClient {
         return .cancelled
     }
     
-    func stop() {
+    func stop() async {
         self.connection.cancel()
+        await self.readyState.markCancelled()
+        self.finishStreams()
     }
     
-    func close(_ exit: SDLQUICClientExit = .normal) async {
+    private func emitEvent(_ event: SDLQUICEvent) {
+        guard !self.didFinishStreams else {
+            return
+        }
         
+        self.eventCont.yield(event)
+    }
+    
+    private func finishStreams() {
+        guard !self.didFinishStreams else {
+            return
+        }
+
+        self.didFinishStreams = true
+        self.messageCont.finish()
+        self.eventCont.finish()
     }
 
     deinit {
-        self.messageCont.finish()
+        self.connection.cancel()
+        self.finishStreams()
     }
 }
 
@@ -209,26 +224,58 @@ extension SDLQUICClient {
         }
 
         private var state: State = .idle
-        private var continuations: [CheckedContinuation<Void, Error>] = []
+        private var continuations: [UUID: CheckedContinuation<Void, Error>] = [:]
 
-        func waitReady() async throws {
+        func waitReady(timeout: Duration) async throws {
-            switch state {
+            let id = UUID()
-            case .ready:
+            let timeoutTask = Task {
-                return
+                try? await Task.sleep(for: timeout)
+                if Task.isCancelled {
+                    return
+                }
+                await self.cancelWaiter(id: id, throwing: SDLQUICError.timeout)
+            }
 
-            case .failed(let error):
+            defer {
-                throw error
+                timeoutTask.cancel()
+            }
 
-            case .cancelled:
+            try await withTaskCancellationHandler {
-                throw CancellationError()
-
-            case .idle, .connecting:
                 try await withCheckedThrowingContinuation { continuation in
-                    continuations.append(continuation)
+                    self.addWaiter(id: id, continuation: continuation)
+                }
+            } onCancel: {
+                timeoutTask.cancel()
+                Task {
+                    await self.cancelWaiter(id: id, throwing: CancellationError())
                 }
             }
         }
 
+        private func addWaiter(id: UUID, continuation: CheckedContinuation<Void, Error>) {
+            switch state {
+            case .ready:
+                continuation.resume()
+
+            case .failed(let error):
+                continuation.resume(throwing: error)
+
+            case .cancelled:
+                continuation.resume(throwing: CancellationError())
+
+            case .idle, .connecting:
+                continuations[id] = continuation
+            }
+        }
+
+        private func cancelWaiter(id: UUID, throwing error: Error) {
+            guard let continuation = continuations.removeValue(forKey: id) else {
+                return
+            }
+
+            continuation.resume(throwing: error)
+        }
+
         func markConnecting() {
             switch state {
             case .idle:
@@ -244,7 +291,7 @@ extension SDLQUICClient {
             let list = continuations
             continuations.removeAll()
 
-            for continuation in list {
+            for continuation in list.values {
                 continuation.resume()
             }
         }
@@ -255,7 +302,7 @@ extension SDLQUICClient {
             let list = continuations
             continuations.removeAll()
 
-            for continuation in list {
+            for continuation in list.values {
                 continuation.resume(throwing: error)
             }
         }
@@ -266,7 +313,7 @@ extension SDLQUICClient {
             let list = continuations
             continuations.removeAll()
 
-            for continuation in list {
+            for continuation in list.values {
                 continuation.resume(throwing: CancellationError())
             }
         }