Support Task cancellation (#844)

Author: crazytonyli

Package.swift

@@ -11,8 +11,8 @@ let package = Package(
     targets: [
         .binaryTarget(
             name: "WordPressKit",
-            url: "https://github.com/user-attachments/files/20895757/WordPressKit.zip",
-            checksum: "b08eaf182f0399303aadccb1a6dad6cad294a9c8d123d920889b15950c85e08f"
+            url: "https://github.com/user-attachments/files/21518814/WordPressKit.zip",
+            checksum: "a43e82909d851e78dff7fa64edba12635e2e96206c1fcdca075eae02dd4c157e"
         ),
     ]
 )

Podfile.lock

@@ -13,4 +13,4 @@ SPEC CHECKSUMS:
 
 PODFILE CHECKSUM: c0da9313733b88a1d938ba6a329dd46b895c7dea
 
-COCOAPODS: 1.15.2
+COCOAPODS: 1.16.2

Sources/CoreAPI/HTTPClient.swift

@@ -70,40 +70,46 @@ extension URLSession {
             assert(parentProgress.cancellationHandler == nil, "The progress instance's cancellationHandler property must be nil")
         }
 
-        return await withCheckedContinuation { continuation in
-            let completion: @Sendable (Data?, URLResponse?, Error?) -> Void = { data, response, error in
-                let result: WordPressAPIResult<HTTPAPIResponse<Data>, E> = Self.parseResponse(
-                    data: data,
-                    response: response,
-                    error: error,
-                    acceptableStatusCodes: acceptableStatusCodes
-                )
-
-                continuation.resume(returning: result)
-            }
+        let taskHolder = TaskHolder()
+        return await withTaskCancellationHandler {
+            await withCheckedContinuation { continuation in
+                let completion: @Sendable (Data?, URLResponse?, Error?) -> Void = { data, response, error in
+                    let result: WordPressAPIResult<HTTPAPIResponse<Data>, E> = Self.parseResponse(
+                        data: data,
+                        response: response,
+                        error: error,
+                        acceptableStatusCodes: acceptableStatusCodes
+                    )
+
+                    continuation.resume(returning: result)
+                }
 
-            let task: URLSessionTask
+                let task: URLSessionTask
 
-            do {
-                task = try self.task(for: builder, completion: completion)
-            } catch {
-                continuation.resume(returning: .failure(.requestEncodingFailure(underlyingError: error)))
-                return
-            }
+                do {
+                    task = try self.task(for: builder, completion: completion)
+                } catch {
+                    continuation.resume(returning: .failure(.requestEncodingFailure(underlyingError: error)))
+                    return
+                }
 
-            task.resume()
-            taskCreated?(task.taskIdentifier)
+                task.resume()
+                taskCreated?(task.taskIdentifier)
+                Task { await taskHolder.assign(task) }
 
-            if let parentProgress, parentProgress.totalUnitCount > parentProgress.completedUnitCount {
-                let pending = parentProgress.totalUnitCount - parentProgress.completedUnitCount
-                // The Jetpack/WordPress app requires task progress updates to be delievered on the main queue.
-                let progressUpdator = parentProgress.update(totalUnit: pending, with: task.progress, queue: .main)
+                if let parentProgress, parentProgress.totalUnitCount > parentProgress.completedUnitCount {
+                    let pending = parentProgress.totalUnitCount - parentProgress.completedUnitCount
+                    // The Jetpack/WordPress app requires task progress updates to be delievered on the main queue.
+                    let progressUpdator = parentProgress.update(totalUnit: pending, with: task.progress, queue: .main)
 
-                parentProgress.cancellationHandler = { [weak task] in
-                    task?.cancel()
-                    progressUpdator.cancel()
+                    parentProgress.cancellationHandler = { [weak task] in
+                        task?.cancel()
+                        progressUpdator.cancel()
+                    }
                 }
             }
+        } onCancel: {
+            Task { await taskHolder.cancel() }
         }
     }
 
@@ -334,3 +340,15 @@ extension URLSession {
         self.taskData.count
     }
 }
+
+private actor TaskHolder {
+    weak var task: URLSessionTask?
+
+    func assign(_ task: URLSessionTask) {
+        self.task = task
+    }
+
+    func cancel() {
+        task?.cancel()
+    }
+}

Tests/WordPressKitTests/Tests/Utilities/URLSessionHelperTests.swift

@@ -172,6 +172,31 @@ class URLSessionHelperTests: XCTestCase {
         }
     }
 
+    func testTaskCancellation() async throws {
+        // Give a slow HTTP request that takes 0.5 second to complete
+        stub(condition: isPath("/hello")) { _ in
+            let response = HTTPStubsResponse(data: "success".data(using: .utf8)!, statusCode: 200, headers: nil)
+            response.responseTime = 0.5
+            return response
+        }
+
+        let task = Task {
+            await session.perform(request: .init(url: URL(string: "https://wordpress.org/hello")!), errorType: TestError.self)
+        }
+
+        // and cancelling it (in 0.1 second) before it completes
+        try await Task.sleep(nanoseconds: 100_000_000)
+        task.cancel()
+
+        // The result should be an cancellation result
+        let result = await task.value
+        if case let .failure(.connection(urlError)) = result, urlError.code == .cancelled {
+            // Do nothing
+        } else {
+            XCTFail("Unexpected result: \(result)")
+        }
+    }
+
     func testEncodingError() async {
         let underlyingError = NSError(domain: "test", code: 123)
         let builder = HTTPRequestBuilder(url: URL(string: "https://wordpress.org")!)