82 lines
2.7 KiB
Swift
82 lines
2.7 KiB
Swift
import Foundation
|
|
|
|
package enum BoundedTaskJoinOutcome<Value: Sendable> {
|
|
case value(Value)
|
|
case failure(any Error)
|
|
case timedOut
|
|
}
|
|
|
|
package final class BoundedTaskJoin<Value: Sendable>: @unchecked Sendable {
|
|
private let lock = NSLock()
|
|
private let sourceTask: Task<Value, Error>
|
|
private var outcome: BoundedTaskJoinOutcome<Value>?
|
|
private var continuation: CheckedContinuation<BoundedTaskJoinOutcome<Value>, Never>?
|
|
private var observerTask: Task<Void, Never>?
|
|
private var timeoutTask: Task<Void, Never>?
|
|
|
|
package init(sourceTask: Task<Value, Error>) {
|
|
self.sourceTask = sourceTask
|
|
}
|
|
|
|
package func value(joinGrace: Duration) async -> BoundedTaskJoinOutcome<Value> {
|
|
await withTaskCancellationHandler {
|
|
await withCheckedContinuation { continuation in
|
|
self.lock.lock()
|
|
if let outcome = self.outcome {
|
|
self.lock.unlock()
|
|
continuation.resume(returning: outcome)
|
|
return
|
|
}
|
|
|
|
self.continuation = continuation
|
|
let sourceTask = self.sourceTask
|
|
self.observerTask = Task { [weak self] in
|
|
do {
|
|
let value = try await sourceTask.value
|
|
self?.resolve(.value(value), cancelSource: false)
|
|
} catch {
|
|
self?.resolve(.failure(error), cancelSource: false)
|
|
}
|
|
}
|
|
self.timeoutTask = Task { [weak self] in
|
|
do {
|
|
if joinGrace > .zero {
|
|
try await Task.sleep(for: joinGrace)
|
|
}
|
|
self?.resolve(.timedOut, cancelSource: true)
|
|
} catch {
|
|
// The source completed or the caller canceled the race.
|
|
}
|
|
}
|
|
self.lock.unlock()
|
|
}
|
|
} onCancel: {
|
|
self.resolve(.failure(CancellationError()), cancelSource: true)
|
|
}
|
|
}
|
|
|
|
private func resolve(_ outcome: BoundedTaskJoinOutcome<Value>, cancelSource: Bool) {
|
|
self.lock.lock()
|
|
guard self.outcome == nil else {
|
|
self.lock.unlock()
|
|
return
|
|
}
|
|
|
|
self.outcome = outcome
|
|
let continuation = self.continuation
|
|
self.continuation = nil
|
|
let observerTask = self.observerTask
|
|
let timeoutTask = self.timeoutTask
|
|
self.observerTask = nil
|
|
self.timeoutTask = nil
|
|
self.lock.unlock()
|
|
|
|
if cancelSource {
|
|
self.sourceTask.cancel()
|
|
}
|
|
observerTask?.cancel()
|
|
timeoutTask?.cancel()
|
|
continuation?.resume(returning: outcome)
|
|
}
|
|
}
|