110 lines
2.8 KiB
Swift
110 lines
2.8 KiB
Swift
//
|
|
// AsyncOneShot.swift
|
|
// Tun
|
|
//
|
|
// Created by Codex on 2026/4/28.
|
|
//
|
|
|
|
import Foundation
|
|
|
|
enum AsyncOneShotError: Error {
|
|
case timeout
|
|
}
|
|
|
|
actor AsyncOneShot<Value: Sendable> {
|
|
private struct Waiter {
|
|
let continuation: CheckedContinuation<Value, Error>
|
|
let timeoutTask: Task<Void, Never>?
|
|
}
|
|
|
|
private var result: Result<Value, Error>?
|
|
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<Value, Error>,
|
|
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 {
|
|
self.cancelWaiter(id: id, throwing: timeoutError)
|
|
}
|
|
} catch {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
self.waiters[id] = Waiter(continuation: continuation, timeoutTask: timeoutTask)
|
|
}
|
|
|
|
private func complete(_ result: Result<Value, Error>) {
|
|
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)
|
|
}
|
|
}
|