From 8e761b054090ed14c64d3c2ba8a62373ef414acf Mon Sep 17 00:00:00 2001 From: anlicheng <244108715@qq.com> Date: Tue, 28 Apr 2026 17:29:01 +0800 Subject: [PATCH] add AsyncOneShot --- Tun/Punchnet/Actors/SDLQuicClient.swift | 124 ++---------------------- Tun/Punchnet/AsyncOneShot.swift | 109 +++++++++++++++++++++ 2 files changed, 115 insertions(+), 118 deletions(-) create mode 100644 Tun/Punchnet/AsyncOneShot.swift diff --git a/Tun/Punchnet/Actors/SDLQuicClient.swift b/Tun/Punchnet/Actors/SDLQuicClient.swift index 6f90313..b5184ee 100644 --- a/Tun/Punchnet/Actors/SDLQuicClient.swift +++ b/Tun/Punchnet/Actors/SDLQuicClient.swift @@ -18,7 +18,6 @@ enum SDLQUICError: Error { case timeout case decodeError(String) case packetTooLarge - case waitReadyAlreadyInProgress } enum SDLQUICEvent: Error { @@ -57,17 +56,7 @@ actor SDLQUICClient { // 最大缓冲区区为2M private let maxBufferSize: Int - private enum ReadyStatus { - case idle - case connecting - case ready - case failed(Error) - case cancelled - } - - private var readyStatus: ReadyStatus = .idle - private var readyContinuation: CheckedContinuation? - private var readyTimeoutTask: Task? + private let readyLatch = AsyncOneShot() // 消息流 public var messageStream: AsyncStream @@ -117,22 +106,20 @@ actor SDLQUICClient { SDLLogger.log("[SDLQUICClient] new state: \(state)", for: .debug) switch state { case .ready: - self.markReady() + await self.readyLatch.succeed(()) case .failed(let error): - self.markFailed(error) + await self.readyLatch.fail(error) self.emitEvent(.failed(error)) case .cancelled: - self.markCancelled() + await self.readyLatch.fail(SDLQUICError.connectionCancelled) self.emitEvent(.cancelled) - case .setup, .preparing: - self.markConnecting() default: () } } func waitReady(timeout: Duration = .seconds(5)) async throws { - try await self.waitReadyUntilStateChanged(timeout: timeout) + try await self.readyLatch.wait(timeout: timeout, timeoutError: SDLQUICError.timeout) } func run() async -> SDLQUICClientExit { @@ -187,7 +174,7 @@ actor SDLQUICClient { func stop() async { self.connection.cancel() - self.markCancelled() + await self.readyLatch.fail(SDLQUICError.connectionCancelled) self.finishStreams() } @@ -211,105 +198,6 @@ actor SDLQUICClient { } -// --MARK: Ready状态机 -extension SDLQUICClient { - - private func waitReadyUntilStateChanged(timeout: Duration) async throws { - try Task.checkCancellation() - - try await withTaskCancellationHandler { - try await withCheckedThrowingContinuation { continuation in - switch self.readyStatus { - case .ready: - continuation.resume() - case .failed(let error): - continuation.resume(throwing: error) - case .cancelled: - continuation.resume(throwing: SDLQUICError.connectionCancelled) - case .idle, .connecting: - guard self.readyContinuation == nil else { - continuation.resume(throwing: SDLQUICError.waitReadyAlreadyInProgress) - return - } - self.readyContinuation = continuation - self.readyTimeoutTask?.cancel() - self.readyTimeoutTask = Task { - do { - try await Task.sleep(for: timeout) - if !Task.isCancelled { - await self.cancelReadyWaiter(throwing: SDLQUICError.timeout) - } - } catch { - return - } - } - } - } - } onCancel: { - Task { - await self.cancelReadyWaiter(throwing: CancellationError()) - } - } - } - - private func cancelReadyWaiter(throwing error: Error) { - guard let continuation = self.readyContinuation else { - return - } - - self.readyContinuation = nil - self.readyTimeoutTask?.cancel() - self.readyTimeoutTask = nil - continuation.resume(throwing: error) - } - - private func markConnecting() { - switch self.readyStatus { - case .idle: - self.readyStatus = .connecting - default: - break - } - } - - private func markReady() { - self.readyStatus = .ready - self.resumeReadyWaiter() - } - - private func markFailed(_ error: Error) { - self.readyStatus = .failed(error) - self.resumeReadyWaiter(throwing: error) - } - - private func markCancelled() { - self.readyStatus = .cancelled - self.resumeReadyWaiter(throwing: SDLQUICError.connectionCancelled) - } - - private func resumeReadyWaiter() { - guard let continuation = self.readyContinuation else { - return - } - - self.readyContinuation = nil - self.readyTimeoutTask?.cancel() - self.readyTimeoutTask = nil - continuation.resume() - } - - private func resumeReadyWaiter(throwing error: Error) { - guard let continuation = self.readyContinuation else { - return - } - - self.readyContinuation = nil - self.readyTimeoutTask?.cancel() - self.readyTimeoutTask = nil - continuation.resume(throwing: error) - } -} - // --MARK: Reader extension SDLQUICClient { private func readLoop() async -> SDLQUICClientExit { diff --git a/Tun/Punchnet/AsyncOneShot.swift b/Tun/Punchnet/AsyncOneShot.swift new file mode 100644 index 0000000..a50d140 --- /dev/null +++ b/Tun/Punchnet/AsyncOneShot.swift @@ -0,0 +1,109 @@ +// +// AsyncOneShot.swift +// Tun +// +// Created by Codex on 2026/4/28. +// + +import Foundation + +enum AsyncOneShotError: Error { + case timeout +} + +actor AsyncOneShot { + private struct Waiter { + let continuation: CheckedContinuation + let timeoutTask: Task? + } + + private var result: Result? + private var waiters: [UUID: Waiter] = [:] + + func wait() async throws -> Value { + try await self.waitInternal(timeout: nil, timeoutError: AsyncOneShotError.timeout) + } + + func wait(timeout: Duration, timeoutError: Error = AsyncOneShotError.timeout) async throws -> Value { + try await self.waitInternal(timeout: timeout, timeoutError: timeoutError) + } + + func succeed(_ value: Value) { + self.complete(.success(value)) + } + + func fail(_ error: Error) { + self.complete(.failure(error)) + } + + private func waitInternal(timeout: Duration?, timeoutError: Error) async throws -> Value { + let id = UUID() + + return try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { continuation in + self.addWaiter( + id: id, + continuation: continuation, + timeout: timeout, + timeoutError: timeoutError + ) + } + } onCancel: { + Task { + await self.cancelWaiter(id: id, throwing: CancellationError()) + } + } + } + + private func addWaiter( + id: UUID, + continuation: CheckedContinuation, + timeout: Duration?, + timeoutError: Error + ) { + if let result { + continuation.resume(with: result) + return + } + + let timeoutTask = timeout.map { timeout in + Task { + do { + try await Task.sleep(for: timeout) + if !Task.isCancelled { + await self.cancelWaiter(id: id, throwing: timeoutError) + } + } catch { + return + } + } + } + + self.waiters[id] = Waiter(continuation: continuation, timeoutTask: timeoutTask) + } + + private func complete(_ result: Result) { + guard self.result == nil else { + return + } + + self.result = result + + let waiters = self.waiters + self.waiters.removeAll() + + for waiter in waiters.values { + waiter.timeoutTask?.cancel() + waiter.continuation.resume(with: result) + } + } + + private func cancelWaiter(id: UUID, throwing error: Error) { + guard let waiter = self.waiters.removeValue(forKey: id) else { + return + } + + waiter.timeoutTask?.cancel() + waiter.continuation.resume(throwing: error) + } +}