Skip to content

Commit

Permalink
Fix a race condition that caused withTimeout to not escalate the pr…
Browse files Browse the repository at this point in the history
…iority of the body

rdar://137678566
  • Loading branch information
ahoppen committed Oct 24, 2024
1 parent 9e9579d commit b601867
Showing 1 changed file with 8 additions and 3 deletions.
11 changes: 8 additions & 3 deletions Sources/SwiftExtensions/AsyncUtils.swift
Original file line number Diff line number Diff line change
Expand Up @@ -176,9 +176,14 @@ package func withTimeout<T: Sendable>(
_ duration: Duration,
_ body: @escaping @Sendable () async throws -> T
) async throws -> T {
// Get the priority with which to launch the body task here so that we can pass the same priority as the initial
// priority to `withTaskPriorityChangedHandler`. Otherwise, we can get into a race condition where bodyTask gets
// launched with a low priority, then the priority gets elevated before we call with `withTaskPriorityChangedHandler`,
// we thus don't receive a `taskPriorityChanged` and hence never increase the priority of `bodyTask`.
let priority = Task.currentPriority
var mutableTasks: [Task<Void, Error>] = []
let stream = AsyncThrowingStream<T, Error> { continuation in
let bodyTask = Task<Void, Error> {
let bodyTask = Task<Void, Error>(priority: priority) {
do {
let result = try await body()
continuation.yield(result)
Expand All @@ -187,7 +192,7 @@ package func withTimeout<T: Sendable>(
}
}

let timeoutTask = Task {
let timeoutTask = Task(priority: priority) {
try await Task.sleep(for: duration)
continuation.yield(with: .failure(TimeoutError()))
bodyTask.cancel()
Expand All @@ -197,7 +202,7 @@ package func withTimeout<T: Sendable>(

let tasks = mutableTasks

return try await withTaskPriorityChangedHandler {
return try await withTaskPriorityChangedHandler(initialPriority: priority) {
for try await value in stream {
return value
}
Expand Down

0 comments on commit b601867

Please sign in to comment.