From 39e1b84423fff6c6c0425784097ec666a94c9da8 Mon Sep 17 00:00:00 2001 From: Nikolay Edigaryev Date: Tue, 18 Oct 2022 16:54:28 +0400 Subject: [PATCH] Use URLSession.dataTask() with delegate instead of URLSession.bytes() (#282) * Use URLSession.dataTask() with delegate instead of URLSession.bytes() * Use URLSession.shared instead of creating a new one each time --- Sources/tart/Fetcher.swift | 46 +++++++++++++++++++++++++++++++++ Sources/tart/OCI/Registry.swift | 32 +++++++++++------------ Sources/tart/VM.swift | 4 +-- 3 files changed, 63 insertions(+), 19 deletions(-) create mode 100644 Sources/tart/Fetcher.swift diff --git a/Sources/tart/Fetcher.swift b/Sources/tart/Fetcher.swift new file mode 100644 index 0000000..b048861 --- /dev/null +++ b/Sources/tart/Fetcher.swift @@ -0,0 +1,46 @@ +import Foundation +import AsyncAlgorithms + +class Fetcher: NSObject, URLSessionTaskDelegate, URLSessionDelegate, URLSessionDataDelegate { + let responseCh = AsyncThrowingChannel() + let dataCh = AsyncThrowingChannel() + + func fetch(_ request: URLRequest) async throws -> (AsyncThrowingChannel, URLResponse) { + let task = URLSession.shared.dataTask(with: request) + task.delegate = self + task.resume() + + // Wait for the response and only then return + var iter = responseCh.makeAsyncIterator() + let response = try await iter.next()! + + return (dataCh, response) + } + + func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive response: URLResponse) async -> URLSession.ResponseDisposition { + await responseCh.send(response) + + return .allow + } + + func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) { + let sema = DispatchSemaphore(value: 0) + + Task { + await dataCh.send(data) + sema.signal() + } + + sema.wait() + } + + func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) { + if let error = error { + // Premature termination + responseCh.fail(error) + dataCh.fail(error) + } else { + dataCh.finish() + } + } +} diff --git a/Sources/tart/OCI/Registry.swift b/Sources/tart/OCI/Registry.swift index cabf61c..05f00df 100644 --- a/Sources/tart/OCI/Registry.swift +++ b/Sources/tart/OCI/Registry.swift @@ -2,8 +2,6 @@ import Foundation import Algorithms import AsyncAlgorithms -let chunkSizeBytes = 1 * 1024 * 1024 - enum RegistryError: Error { case UnexpectedHTTPStatusCode(when: String, code: Int, details: String = "") case MissingLocationHeader @@ -31,11 +29,11 @@ extension Data { } } -extension URLSession.AsyncBytes { +extension AsyncThrowingChannel { func asData() async throws -> Data { var result = Data() - for try await chunk in chunks(ofCount: chunkSizeBytes) { + for try await chunk in self { result += chunk } @@ -228,14 +226,14 @@ class Registry { } public func pullBlob(_ digest: String, handler: (Data) throws -> Void) async throws { - let (bytes, response) = try await bytesRequest(.GET, endpointURL("\(namespace)/blobs/\(digest)")) + let (channel, response) = try await channelRequest(.GET, endpointURL("\(namespace)/blobs/\(digest)")) if response.statusCode != HTTPCode.Ok.rawValue { - let body = try await bytes.asData().asText() + let body = try await channel.asData().asText() throw RegistryError.UnexpectedHTTPStatusCode(when: "pulling blob", code: response.statusCode, details: body) } - for try await part in bytes.chunks(ofCount: chunkSizeBytes) { + for try await part in channel { try Task.checkCancellation() try handler(Data(part)) @@ -256,20 +254,20 @@ class Registry { body: Data? = nil, doAuth: Bool = true ) async throws -> (Data, HTTPURLResponse) { - let (bytes, response) = try await bytesRequest(method, urlComponents, + let (channel, response) = try await channelRequest(method, urlComponents, headers: headers, parameters: parameters, body: body, doAuth: doAuth) - return (try await bytes.asData(), response) + return (try await channel.asData(), response) } - private func bytesRequest( + private func channelRequest( _ method: HTTPMethod, _ urlComponents: URLComponents, headers: Dictionary = Dictionary(), parameters: Dictionary = Dictionary(), body: Data? = nil, doAuth: Bool = true - ) async throws -> (URLSession.AsyncBytes, HTTPURLResponse) { + ) async throws -> (AsyncThrowingChannel, HTTPURLResponse) { var urlComponents = urlComponents if urlComponents.queryItems == nil && !parameters.isEmpty { @@ -294,14 +292,14 @@ class Registry { currentAuthToken = nil } - var (bytes, response) = try await authAwareRequest(request: request) + var (channel, response) = try await authAwareRequest(request: request) if doAuth && response.statusCode == HTTPCode.Unauthorized.rawValue { try await auth(response: response) - (bytes, response) = try await authAwareRequest(request: request) + (channel, response) = try await authAwareRequest(request: request) } - return (bytes, response) + return (channel, response) } private func auth(response: HTTPURLResponse) async throws { @@ -373,7 +371,7 @@ class Registry { return nil } - private func authAwareRequest(request: URLRequest) async throws -> (URLSession.AsyncBytes, HTTPURLResponse) { + private func authAwareRequest(request: URLRequest) async throws -> (AsyncThrowingChannel, HTTPURLResponse) { var request = request if let token = currentAuthToken { @@ -381,8 +379,8 @@ class Registry { request.addValue(value, forHTTPHeaderField: name) } - let (bytes, response) = try await URLSession.shared.bytes(for: request) + let (channel, response) = try await Fetcher().fetch(request) - return (bytes, response as! HTTPURLResponse) + return (channel, response as! HTTPURLResponse) } } diff --git a/Sources/tart/VM.swift b/Sources/tart/VM.swift index 12e6c15..ee267e5 100644 --- a/Sources/tart/VM.swift +++ b/Sources/tart/VM.swift @@ -63,7 +63,7 @@ class VM: NSObject, VZVirtualMachineDelegate, ObservableObject { static func retrieveIPSW(remoteURL: URL) async throws -> URL { // Check if we already have this IPSW in cache - let (bytes, response) = try await URLSession.shared.bytes(from: remoteURL) + let (channel, response) = try await Fetcher().fetch(URLRequest(url: remoteURL)) if let hash = (response as! HTTPURLResponse).value(forHTTPHeaderField: "x-amz-meta-digest-sha256") { let ipswLocation = try IPSWCache().locationFor(fileName: "sha256:\(hash).ipsw") @@ -90,7 +90,7 @@ class VM: NSObject, VZVirtualMachineDelegate, ObservableObject { let fileHandle = try FileHandle(forWritingTo: temporaryLocation) let digest = Digest() - for try await chunk in bytes.chunks(ofCount: chunkSizeBytes) { + for try await chunk in channel { let chunkAsData = Data(chunk) fileHandle.write(chunkAsData) digest.update(chunkAsData)