diff --git a/Sources/SwiftExtensions/AsyncUtils.swift b/Sources/SwiftExtensions/AsyncUtils.swift index 9e57635fd..1d0d4466e 100644 --- a/Sources/SwiftExtensions/AsyncUtils.swift +++ b/Sources/SwiftExtensions/AsyncUtils.swift @@ -176,9 +176,14 @@ package func withTimeout( _ 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] = [] let stream = AsyncThrowingStream { continuation in - let bodyTask = Task { + let bodyTask = Task(priority: priority) { do { let result = try await body() continuation.yield(result) @@ -187,7 +192,7 @@ package func withTimeout( } } - let timeoutTask = Task { + let timeoutTask = Task(priority: priority) { try await Task.sleep(for: duration) continuation.yield(with: .failure(TimeoutError())) bodyTask.cancel() @@ -197,7 +202,7 @@ package func withTimeout( let tasks = mutableTasks - return try await withTaskPriorityChangedHandler { + return try await withTaskPriorityChangedHandler(initialPriority: priority) { for try await value in stream { return value }