When withTimeout hits a timeout, don’t wait for cooperative cancellation before returning

Previously, when we hit the timeout in `withTimeout`, we just cancelled `body` but relied on cooperative cancellation before the caller could continue. Since the caller is obviously no longer interested in the operation, we can return with `CancellationError` immediately.

Fixes #883
rdar://116705684

Author: ahoppen

Sources/SwiftExtensions/AsyncUtils.swift

@@ -176,18 +176,46 @@ package func withTimeout<T: Sendable>(
   _ duration: Duration,
   _ body: @escaping @Sendable () async throws -> T
 ) async throws -> T {
-  try await withThrowingTaskGroup(of: T.self) { taskGroup in
-    taskGroup.addTask {
-      try await Task.sleep(for: duration)
-      throw TimeoutError()
+  var mutableTasks: [Task<Void, Error>] = []
+  let stream = AsyncThrowingStream<T, Error> { continuation in
+    let bodyTask = Task<Void, Error> {
+      do {
+        let result = try await body()
+        continuation.yield(result)
+      } catch {
+        continuation.yield(with: .failure(error))
+      }
     }
-    taskGroup.addTask {
-      return try await body()
+
+    let timeoutTask = Task {
+      try await Task.sleep(for: duration)
+      bodyTask.cancel()
+      continuation.yield(with: .failure(TimeoutError()))
     }
-    for try await value in taskGroup {
-      taskGroup.cancelAll()
+    mutableTasks = [bodyTask, timeoutTask]
+  }
+
+  let tasks = mutableTasks
+
+  return try await withTaskPriorityChangedHandler {
+    for try await value in stream {
       return value
     }
-    throw CancellationError()
+    // The only reason for the loop above to terminate is if the Task got cancelled or if the continuation finishes
+    // (which it never does).
+    if Task.isCancelled {
+      for task in tasks {
+        task.cancel()
+      }
+      throw CancellationError()
+    } else {
+      preconditionFailure("Continuation never finishes")
+    }
+  } taskPriorityChanged: {
+    for task in tasks {
+      Task(priority: Task.currentPriority) {
+        _ = try? await task.value
+      }
+    }
   }
 }

Tests/SKSupportTests/AsyncUtilsTests.swift

@@ -29,4 +29,31 @@ final class AsyncUtilsTests: XCTestCase {
     }
     try await fulfillmentOfOrThrow([expectation])
   }
+
+  func testWithTimeoutReturnsImmediatelyEvenIfBodyDoesntCooperateInCancellation() async throws {
+    let start = Date()
+    await assertThrowsError(
+      try await withTimeout(.seconds(0.1)) { sleep(10) }
+    ) { error in
+      XCTAssert(error is TimeoutError, "Received unexpected error \(error)")
+    }
+    XCTAssert(Date().timeIntervalSince(start) < 5)
+  }
+
+  func testWithTimeoutEscalatesPriority() async throws {
+    let expectation = self.expectation(description: "Timeout started")
+    let task = Task(priority: .background) {
+      // We don't actually hit the timeout. It's just a large value.
+      try await withTimeout(.seconds(defaultTimeout * 2)) {
+        expectation.fulfill()
+        try await repeatUntilExpectedResult {
+          return Task.currentPriority > .background
+        }
+      }
+    }
+    try await fulfillmentOfOrThrow([expectation])
+    try await Task(priority: .high) {
+      try await task.value
+    }.value
+  }
 }