Async-await URLSesionTask API wrapper.

AsyncURLSessionTask.swift

// TODO: - Change raw completion handlers once Swift issue is resolved.
// https://github.com/apple/swift/issues/60488
public typealias DataTaskCompletion = (Data?, URLResponse?, Error?) -> Void
public typealias DownloadTaskCompletion = (URL?, URLResponse?, Error?) -> Void

/// async-await URLSessionTask wrapper with `cancel` and `suspend` functionality.
public class AsyncURLSessionTask: Identifiable {

    enum State {
        case ready
        case executing(URLSessionTask)
        case suspended(URLSessionTask)
        case cancelled
    }

    /// Identifier that is used for URLSessionTask `taskDescription` property setup.
    /// Helps to identify a network request among other requests in URLSesion.
    public var identifier: UUID = UUID()

    var state: State = .ready
    private(set) var session: URLSession

    // MARK: - Life cycle

    init(session: URLSession) {
        self.session = session
    }

    // MARK: - Public

    public func cancel() async {
        guard let task = await session.allTasks.first(where: { $0.taskDescription == identifier.uuidString }) else {
            return
        }
        switch state {
        case .ready:
            task.cancel()
            state = .cancelled

        case .executing(let task):
            task.cancel()
            state = .cancelled

        case .suspended(let task):
            task.cancel()
            state = .cancelled

        case .cancelled:
            break
        }
    }

    public func suspend() async {
        guard let task = await session.allTasks.first(where: { $0.taskDescription == identifier.uuidString }) else {
            return
        }
        switch state {
        case .ready:
            task.suspend()
            state = .suspended(task)

        case .executing(let task):
            task.suspend()
            state = .suspended(task)

        case .suspended, .cancelled:
            break
        }
    }

    public func resume() async {
        guard let task = await session.allTasks.first(where: { $0.taskDescription == identifier.uuidString }) else {
            return
        }
        switch state {
        case .ready:
            task.resume()
        case .executing:
            break
        case .suspended(let task):
            task.resume()
        case .cancelled:
            break
        }
    }

    // MARK: Data

    func data(with url: URL) async throws -> (URL, URLResponse) {
        try await download(with: URLRequest(url: url))
    }

    func data(with request: URLRequest) async throws -> (Data, URLResponse) {
        return try await withTaskCancellationHandler {
            try await withCheckedThrowingContinuation { continuation in
                Task {
                    self.dataTask(for: request) { data, response, error in
                        guard let data, let response else {
                            continuation.resume(throwing: error ?? URLError(.badServerResponse))
                            return
                        }

                        continuation.resume(returning: (data, response))
                    }
                }
            }
        } onCancel: {
            Task { await self.cancel() }
        }
    }

    // MARK: Download

    func download(with url: URL) async throws -> (URL, URLResponse) {
        try await download(with: URLRequest(url: url))
    }

    func download(with request: URLRequest) async throws -> (URL, URLResponse) {
        let taskId = self.identifier
        return try await withTaskCancellationHandler {
            try await withCheckedThrowingContinuation { continuation in
                Task {
                    self.downloadTask(for: request) { location, response, error in
                        guard let location, let response else {
                            continuation.resume(throwing: error ?? URLError(.badServerResponse))
                            return
                        }

                        // since continuation can happen later, let’s figure out where to store it ...

                        guard let cachesDirectory = FileManager.default.urls(for: .cachesDirectory,
                                                                             in: .userDomainMask).first else {
                            return
                        }
                        let tempURL = cachesDirectory
                            .appendingPathComponent(taskId.uuidString)
                            .appendingPathExtension(request.url?.pathExtension ?? "")

                        // ... and move it to there

                        do {
                            try FileManager.default.moveItem(at: location, to: tempURL)
                        } catch {
                            continuation.resume(throwing: error)
                            return
                        }

                        continuation.resume(returning: (tempURL, response))
                    }
                }
            }
        } onCancel: {
            Task { await self.cancel() }
        }
    }

    // MARK: Upload

    func upload(to url: URL, data: Data) async throws -> (Data, URLResponse) {
        var request = URLRequest(url: url)
        request.httpBody = data
        return try await upload(with: request, bodyData: data)
    }

    func upload(to url: URL, fileURL: URL) async throws -> (Data, URLResponse) {
        let data = try? Data(contentsOf: fileURL)
        return try await upload(with: URLRequest(url: url), bodyData: data)
    }

    func upload(with request: URLRequest, bodyData: Data?) async throws -> (Data, URLResponse) {
        return try await withTaskCancellationHandler {
            try await withCheckedThrowingContinuation { continuation in
                Task {
                    self.uploadTask(for: request, bodyData: bodyData ?? Data()) { data, response, error in
                        guard let data, let response else {
                            continuation.resume(throwing: error ?? URLError(.badServerResponse))
                            return
                        }

                        continuation.resume(returning: (data, response))
                    }
                }
            }
        } onCancel: {
            Task { await self.cancel() }
        }
    }

    // MARK: - Private

    private func dataTask(for request: URLRequest, completionHandler: @Sendable @escaping (Data?, URLResponse?, Error?) -> Void) {
        if case .cancelled = state {
            completionHandler(nil, nil, CancellationError())
            return
        }

        let task = self.session.dataTask(with: request, completionHandler: completionHandler)
        task.taskDescription = identifier.uuidString

        state = .executing(task)
        task.resume()
    }

    private func downloadTask(for request: URLRequest, completionHandler: @Sendable @escaping (URL?, URLResponse?, Error?) -> Void) {
        if case .cancelled = state {
            completionHandler(nil, nil, CancellationError())
            return
        }

        let task = self.session.downloadTask(with: request, completionHandler: completionHandler)
        task.taskDescription = identifier.uuidString

        state = .executing(task)
        task.resume()
    }

    // https://developer.apple.com/forums/thread/52073
    // While the signature of the method does mark from as a Data?,
    // the task gets automatically cancelled when it is missing/nil.
    private func uploadTask(for request: URLRequest, bodyData: Data, completionHandler: @Sendable @escaping (Data?, URLResponse?, Error?) -> Void) {
        if case .cancelled = state {
            completionHandler(nil, nil, CancellationError())
            return
        }

        let task = self.session.uploadTask(with: request, from: bodyData, completionHandler: completionHandler)
        task.taskDescription = identifier.uuidString

        state = .executing(task)
        task.resume()
    }

}