From f345f252cf6b788dc41011e67775e335854899b9 Mon Sep 17 00:00:00 2001
From: Simon Whitty <simon@whileloop.com>
Date: Sun, 8 Sep 2024 18:13:46 +1000
Subject: [PATCH] Cancellation

---
 Sources/Timeout.swift      |  6 +++++-
 Tests/TimeoutTests.swift   | 15 +++++++++++++++
 Tests/TimeoutXCTests.swift | 17 +++++++++++++++++
 3 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/Sources/Timeout.swift b/Sources/Timeout.swift
index 7624b35..b7da13b 100644
--- 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
index 69d4053..9c8c885 100644
--- 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
index 0ae3a0e..5c57d2a 100644
--- a/Tests/TimeoutXCTests.swift
+++ b/Tests/TimeoutXCTests.swift
@@ -91,6 +91,23 @@ final class TimeoutTests: XCTestCase {
             XCTAssertTrue(error is TimeoutError)
         }
     }
+
+    func testTimeout_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> {