Cancellation

swhitty · Sep 8, 2024 · 0799953 · 0799953
1 parent 08027b9
commit 0799953
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 1 deletion.
diff --git a/Sources/Timeout.swift b/Sources/Timeout.swift
@@ -56,7 +56,11 @@ public func withThrowingTimeout<T>(
             throw TimeoutError(timeout: seconds)
         }
 
-        let bodyResult = await bodyTask.result
+        let bodyResult = await withTaskCancellationHandler {
+            await bodyTask.result
+        } onCancel: {
+            bodyTask.cancel()
+        }
         timeoutTask.cancel()
         let timeoutResult = await timeoutTask.result
 

diff --git a/Tests/TimeoutTests.swift b/Tests/TimeoutTests.swift
@@ -90,6 +90,21 @@ struct TimeoutTests {
             }
         }
     }
+
+    @Test
+    func timeout_cancels() async {
+        let task = Task {
+            try await withThrowingTimeout(seconds: 1) {
+                try await Task.sleep(nanoseconds: 1_000_000_000)
+            }
+        }
+
+        task.cancel()
+
+        await #expect(throws: CancellationError.self) {
+            try await task.value
+        }
+    }
 }
 
 public struct NonSendable<T> {

diff --git a/Tests/TimeoutXCTests.swift b/Tests/TimeoutXCTests.swift
@@ -91,6 +91,23 @@ final class TimeoutTests: XCTestCase {
             XCTAssertTrue(error is TimeoutError)
         }
     }
+
+    func timeout_cancels() async {
+        let task = Task {
+            try await withThrowingTimeout(seconds: 1) {
+                try await Task.sleep(nanoseconds: 1_000_000_000)
+            }
+        }
+
+        task.cancel()
+
+        do {
+            _ = try await task.value
+            XCTFail("Expected Error")
+        } catch {
+            XCTAssertTrue(error is CancellationError)
+        }
+    }
 }
 
 public struct NonSendable<T> {