How to use withTaskCancellationHandler properly?

Using Swift
12

Thanks for calling those points out @ole. Here's an updated (and slightly more generalized) version which should address them:

public extension URLSession {
    func data(for request: URLRequest) async throws -> (Data, URLResponse) {
        let dataTaskHolder = CancellableHolder<URLSessionDataTask>()

        return try await withTaskCancellationHandler(
            handler: {
                dataTaskHolder.cancel()
            },
            operation: {
                try await withCheckedThrowingContinuation { continuation in
                    dataTaskHolder.value = self.dataTask(with: request) { data, response, error in
                        guard let data = data, let response = response else {
                            let error = error ?? URLError(.badServerResponse)
                            return continuation.resume(throwing: error)
                        }

                        continuation.resume(returning: (data, response))
                    }

                    dataTaskHolder.value?.resume()
                }
            }
        )
    }
}

private class CancellableHolder<T: Cancellable>: @unchecked Sendable {
    private var lock = NSRecursiveLock()
    private var innerCancellable: T?

    private func synced<Result>(_ action: () throws -> Result) rethrows -> Result {
        lock.lock()
        defer { lock.unlock() }
        return try action()
    }

    var value: T? {
        get { synced { innerCancellable } }
        set { synced { innerCancellable = newValue } }
    }

    func cancel() {
        synced { innerCancellable?.cancel() }
    }
}

protocol Cancellable {
    func cancel()
}

extension URLSessionDataTask: Cancellable {}

What do you think?

3 Likes