diff --git a/Tun/Super/SDLSuperClient.swift b/Tun/Super/SDLSuperClient.swift index 1c5dace..58ab4bf 100644 --- a/Tun/Super/SDLSuperClient.swift +++ b/Tun/Super/SDLSuperClient.swift @@ -10,10 +10,10 @@ import Network final class SDLSuperClient: @unchecked Sendable { private let queue = DispatchQueue(label: "com.sdl.SuperClient.queue") // 专用队列保证线程安全 - // 数据流 public let messageStream: AsyncThrowingStream - private let messageCont: AsyncThrowingStream.Continuation + private let messageContinuation: AsyncThrowingStream.Continuation + private let readySignal = AsyncOneShot() private let stateLock = NSLock() private var isStarted = false private var isStopped = false @@ -25,9 +25,9 @@ final class SDLSuperClient: @unchecked Sendable { init(serverEndpoint: SDLConfiguration.ResolvedServerEndpoint, port: UInt16, maxBufferSize: Int = 2 * 1024 * 1024) { self.maxBufferSize = maxBufferSize - let pairs = AsyncThrowingStream.makeStream(of: SDLSuperMessage.self) + let pairs = AsyncThrowingStream.makeStream(of: SDLSuperMessage.self, bufferingPolicy: .bufferingNewest(1024)) self.messageStream = pairs.stream - self.messageCont = pairs.continuation + self.messageContinuation = pairs.continuation let options = NWProtocolTLS.Options() serverEndpoint.host.withCString { @@ -61,18 +61,30 @@ final class SDLSuperClient: @unchecked Sendable { guard self.markStarted() else { return } - - let stateStream = self.makeStateStream() + defer { self.stop() } - - try await withTaskCancellationHandler { - self.connection.start(queue: self.queue) - try await self.waitUntilReady(stateStream) - try await self.readLoop() - } onCancel: { - self.connection.cancel() + + do { + try await withTaskCancellationHandler { + self.connection.stateUpdateHandler = { [weak self] state in + self?.handleConnectionStateUpdate(state) + } + self.connection.start(queue: self.queue) + try await self.readySignal.wait() + try await self.readLoop() + self.finishMessageStream() + } onCancel: { + self.connection.stateUpdateHandler = nil + self.connection.cancel() + } + } catch is CancellationError { + self.finishMessageStream() + throw CancellationError() + } catch { + self.finishMessageStream(throwing: error) + throw error } } @@ -88,36 +100,28 @@ final class SDLSuperClient: @unchecked Sendable { preconditionFailure("invalid super server IP: \(ip)") } - private func makeStateStream() -> AsyncThrowingStream { - let connection = self.connection - return AsyncThrowingStream(bufferingPolicy: .bufferingNewest(16)) { continuation in - connection.stateUpdateHandler = { [weak self] state in - SDLLogger.log("[SDLSuperClient] new state: \(state)", category: .super) - switch state { - case .failed(let error): - let wrappedError = SDLSuperError.connectionFailed(error) - self?.finishMessageStream(throwing: wrappedError) - continuation.finish(throwing: wrappedError) - case .cancelled: - self?.finishMessageStream(throwing: SDLSuperError.connectionCancelled) - continuation.finish(throwing: SDLSuperError.connectionCancelled) - default: - continuation.yield(state) - } + private func handleConnectionStateUpdate(_ state: NWConnection.State) { + SDLLogger.log("[SDLSuperClient] new state: \(state)", category: .super) + switch state { + case .ready: + Task { + await self.readySignal.succeed(()) } - } - } - - private func waitUntilReady(_ stateStream: AsyncThrowingStream) async throws { - for try await state in stateStream { - try Task.checkCancellation() - - if case .ready = state { - return + case .failed(let error): + let wrappedError = SDLSuperError.connectionFailed(error) + Task { + await self.readySignal.fail(wrappedError) } + self.finishMessageStream(throwing: wrappedError) + case .cancelled: + let error = SDLSuperError.connectionCancelled + Task { + await self.readySignal.fail(error) + } + self.finishMessageStream(throwing: error) + default: + break } - - throw SDLSuperError.connectionCancelled } private func readLoop() async throws { @@ -129,7 +133,7 @@ final class SDLSuperClient: @unchecked Sendable { for frame in frames { try Task.checkCancellation() if let message = SDLSuperCodec.decode(frame: frame) { - self.messageCont.yield(message) + self.messageContinuation.yield(message) } else { throw SDLSuperError.decodeError("invalid message") } @@ -151,10 +155,11 @@ final class SDLSuperClient: @unchecked Sendable { if let error { SDLLogger.log("[SDLSuperClient] send data get error: \(error)", category: .super) self?.finishMessageStream(throwing: SDLSuperError.writeFailed(error)) + self?.connection.cancel() } }) } - + private static func readOnce(connection: NWConnection) async throws -> Data { guard connection.state == .ready else { throw SDLSuperError.connectionCancelled @@ -191,6 +196,9 @@ final class SDLSuperClient: @unchecked Sendable { let connection = self.connection connection.stateUpdateHandler = nil connection.cancel() + Task { + await self.readySignal.fail(SDLSuperError.connectionCancelled) + } self.finishMessageStream() SDLLogger.log("[SDLSuperClient] stopped", category: .super) @@ -237,9 +245,9 @@ final class SDLSuperClient: @unchecked Sendable { } if let error { - self.messageCont.finish(throwing: error) + self.messageContinuation.finish(throwing: error) } else { - self.messageCont.finish() + self.messageContinuation.finish() } } diff --git a/Tun/Super/SDLSuperSession.swift b/Tun/Super/SDLSuperSession.swift index 034333b..0607173 100644 --- a/Tun/Super/SDLSuperSession.swift +++ b/Tun/Super/SDLSuperSession.swift @@ -20,20 +20,8 @@ final class SDLSuperSession: @unchecked Sendable { self.onMessage = onMessage self.client = SDLSuperClient(serverEndpoint: serverEndpoint, port: port) } - + func run() async throws { - try await self.runLoops() - } - - func stop() async { - self.client.stop() - } - - func send(type: SDLPacketType, data: Data) async { - self.client.send(type: type, data: data) - } - - private func runLoops() async throws { SDLLogger.log("[SDLSuperSession] start super client: \(self.serverEndpoint.ip)", category: .super) try await withThrowingTaskGroup(of: Void.self) { group in @@ -56,6 +44,14 @@ final class SDLSuperSession: @unchecked Sendable { _ = try await group.next() } } + + func stop() async { + self.client.stop() + } + + func send(type: SDLPacketType, data: Data) async { + self.client.send(type: type, data: data) + } private func readLoop() async throws { for try await message in self.client.messageStream {