From d900fdb379d5cac4531a0d275e1875d90beac8d6 Mon Sep 17 00:00:00 2001 From: anlicheng <244108715@qq.com> Date: Tue, 28 Apr 2026 16:23:01 +0800 Subject: [PATCH] fix quicClient --- Tun/Punchnet/Actors/SDLContextActor.swift | 31 ++-- Tun/Punchnet/Actors/SDLQuicClient.swift | 179 ++++++++++++++-------- 2 files changed, 128 insertions(+), 82 deletions(-) diff --git a/Tun/Punchnet/Actors/SDLContextActor.swift b/Tun/Punchnet/Actors/SDLContextActor.swift index eff735a..89c298e 100644 --- a/Tun/Punchnet/Actors/SDLContextActor.swift +++ b/Tun/Punchnet/Actors/SDLContextActor.swift @@ -181,36 +181,33 @@ 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() } - + group.addTask { 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 } } - + group.addTask { for await event in await quicClient.eventStream { switch event { @@ -222,11 +219,13 @@ actor SDLContextActor { throw error } } - throw SDLQUICClientExit.cancelled } - try await group.next() + + while let _ = try await group.next() { + () + } } - + } await self.supervisor.addWorker(name: "udpHole") { diff --git a/Tun/Punchnet/Actors/SDLQuicClient.swift b/Tun/Punchnet/Actors/SDLQuicClient.swift index 5ca0f73..d2881b4 100644 --- a/Tun/Punchnet/Actors/SDLQuicClient.swift +++ b/Tun/Punchnet/Actors/SDLQuicClient.swift @@ -66,6 +66,8 @@ actor SDLQUICClient { public var eventStream: AsyncStream private let eventCont: AsyncStream.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 - SDLLogger.log("[SDLQUICClient] new state: \(state)", for: .debug) - switch 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.stateUpdateHandler = { [weak self] state in + Task { + await self?.handleConnectionStateUpdate(state) } } connection.start(queue: self.queue) } - func waitReady(timeout: Duration = .seconds(5)) async throws { - try await withThrowingTaskGroup(of: Void.self) { group in - group.addTask { - try await self.readyState.waitReady() - } - - group.addTask { - try await Task.sleep(for: timeout) - throw SDLQUICError.timeout - } - - try await group.next() - group.cancelAll() + private func handleConnectionStateUpdate(_ state: NWConnection.State) async { + SDLLogger.log("[SDLQUICClient] new state: \(state)", for: .debug) + switch state { + case .ready: + await self.readyState.markReady() + case .failed(let error): + await self.readyState.markFailed(error) + self.emitEvent(.failed(error)) + case .cancelled: + await self.readyState.markCancelled() + self.emitEvent(.cancelled) + case .setup, .preparing: + 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 - group.addTask { - await self.readLoop() + await withTaskCancellationHandler { + await withTaskGroup(of: SDLQUICClientExit.self) { group in + 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 } - - group.addTask { - await self.heartbeatLoop() + } onCancel: { + Task { + 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) } - deinit { + private func finishStreams() { + guard !self.didFinishStreams else { + return + } + + self.didFinishStreams = true self.messageCont.finish() + self.eventCont.finish() + } + + deinit { + self.connection.cancel() + self.finishStreams() } } @@ -209,26 +224,58 @@ extension SDLQUICClient { } private var state: State = .idle - private var continuations: [CheckedContinuation] = [] + private var continuations: [UUID: CheckedContinuation] = [:] - func waitReady() async throws { - switch state { - case .ready: - return + func waitReady(timeout: Duration) async throws { + let id = UUID() + let timeoutTask = Task { + try? await Task.sleep(for: timeout) + if Task.isCancelled { + return + } + await self.cancelWaiter(id: id, throwing: SDLQUICError.timeout) + } - case .failed(let error): - throw error + defer { + timeoutTask.cancel() + } - case .cancelled: - throw CancellationError() - - case .idle, .connecting: + try await withTaskCancellationHandler { 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) { + 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()) } }