-
Notifications
You must be signed in to change notification settings - Fork 117
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
tart pull: 284% faster pulls with default concurrency setting (#970)
* DiskV2: avoid allocating zero chunk on each zeroSkippingWrite() call * Increase hole granularity size from 64 KiB to 4 MiB * Fetcher: never write to disk, thanks to URLSessionDataDelegate
- Loading branch information
Showing
6 changed files
with
85 additions
and
98 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,80 +1,90 @@ | ||
| import Foundation | ||
| import AsyncAlgorithms | ||
|
|
||
| fileprivate let urlSession = createURLSession() | ||
| fileprivate var urlSessionConfiguration: URLSessionConfiguration { | ||
| let config = URLSessionConfiguration.default | ||
|
|
||
| class DownloadDelegate: NSObject, URLSessionTaskDelegate { | ||
| let progress: Progress | ||
| init(_ progress: Progress) throws { | ||
| self.progress = progress | ||
| } | ||
| // Harbor expects a CSRF token to be present if the HTTP client | ||
| // carries a session cookie between its requests[1] and fails if | ||
| // it was not present[2]. | ||
| // | ||
| // To fix that, we disable the automatic cookies carry in URLSession. | ||
| // | ||
| // [1]: https://github.com/goharbor/harbor/blob/a4c577f9ec4f18396207a5e686433a6ba203d4ef/src/server/middleware/csrf/csrf.go#L78 | ||
| // [2]: https://github.com/cirruslabs/tart/issues/295 | ||
| config.httpShouldSetCookies = false | ||
|
|
||
| func urlSession(_ session: URLSession, didCreateTask task: URLSessionTask) { | ||
| self.progress.addChild(task.progress, withPendingUnitCount: self.progress.totalUnitCount) | ||
| } | ||
| return config | ||
| } | ||
|
|
||
| class Fetcher { | ||
| static func fetch(_ request: URLRequest, viaFile: Bool = false, progress: Progress? = nil) async throws -> (AsyncThrowingChannel<Data, Error>, HTTPURLResponse) { | ||
| let delegate = progress != nil ? try DownloadDelegate(progress!) : nil | ||
| static func fetch(_ request: URLRequest, viaFile: Bool = false, progress: Progress? = nil) async throws -> (AsyncThrowingStream<Data, Error>, HTTPURLResponse) { | ||
| let delegate = Delegate() | ||
| let session = URLSession(configuration: urlSessionConfiguration, delegate: delegate, delegateQueue: nil) | ||
| let task = session.dataTask(with: request) | ||
|
|
||
| let stream = AsyncThrowingStream<Data, Error> { continuation in | ||
| delegate.streamContinuation = continuation | ||
| } | ||
|
|
||
| if viaFile { | ||
| return try await fetchViaFile(request, delegate: delegate) | ||
| let response = try await withCheckedThrowingContinuation { continuation in | ||
| delegate.responseContinuation = continuation | ||
| task.resume() | ||
| } | ||
|
|
||
| return try await fetchViaMemory(request, delegate: delegate) | ||
| return (stream, response as! HTTPURLResponse) | ||
| } | ||
| } | ||
|
|
||
| private static func fetchViaMemory(_ request: URLRequest, delegate: URLSessionTaskDelegate? = nil) async throws -> (AsyncThrowingChannel<Data, Error>, HTTPURLResponse) { | ||
| let dataCh = AsyncThrowingChannel<Data, Error>() | ||
| fileprivate class Delegate: NSObject, URLSessionDataDelegate { | ||
| var responseContinuation: CheckedContinuation<URLResponse, Error>? | ||
| var streamContinuation: AsyncThrowingStream<Data, Error>.Continuation? | ||
|
|
||
| let (data, response) = try await urlSession.data(for: request, delegate: delegate) | ||
| private var buffer: Data = Data() | ||
| private let bufferFlushSize = 16 * 1024 * 1024 | ||
|
|
||
| Task { | ||
| await dataCh.send(data) | ||
| func urlSession( | ||
| _ session: URLSession, | ||
| dataTask: URLSessionDataTask, | ||
| didReceive response: URLResponse, | ||
| completionHandler: @escaping (URLSession.ResponseDisposition) -> Void | ||
| ) { | ||
| // Soft-limit for the maximum buffer capacity | ||
| let capacity = min(response.expectedContentLength, Int64(bufferFlushSize)) | ||
|
|
||
| dataCh.finish() | ||
| } | ||
| // Pre-initialize buffer as we now know the capacity | ||
| buffer = Data(capacity: Int(capacity)) | ||
|
|
||
| return (dataCh, response as! HTTPURLResponse) | ||
| responseContinuation?.resume(returning: response) | ||
| completionHandler(.allow) | ||
| } | ||
|
|
||
| private static func fetchViaFile(_ request: URLRequest, delegate: URLSessionTaskDelegate? = nil) async throws -> (AsyncThrowingChannel<Data, Error>, HTTPURLResponse) { | ||
| let dataCh = AsyncThrowingChannel<Data, Error>() | ||
|
|
||
| let (fileURL, response) = try await urlSession.download(for: request, delegate: delegate) | ||
|
|
||
| // Acquire a handle to the downloaded file and then remove it. | ||
| // | ||
| // This keeps a working reference to that file, yet we don't | ||
| // have to deal with the cleanup any more. | ||
| let mappedFile = try Data(contentsOf: fileURL, options: [.alwaysMapped]) | ||
| try FileManager.default.removeItem(at: fileURL) | ||
| func urlSession( | ||
| _ session: URLSession, | ||
| dataTask: URLSessionDataTask, | ||
| didReceive data: Data | ||
| ) { | ||
| buffer.append(data) | ||
|
|
||
| Task { | ||
| for chunk in (0 ..< mappedFile.count).chunks(ofCount: 64 * 1024 * 1024) { | ||
| await dataCh.send(mappedFile.subdata(in: chunk)) | ||
| } | ||
|
|
||
| dataCh.finish() | ||
| if buffer.count >= bufferFlushSize { | ||
| streamContinuation?.yield(buffer) | ||
| buffer.removeAll(keepingCapacity: true) | ||
| } | ||
|
|
||
| return (dataCh, response as! HTTPURLResponse) | ||
| } | ||
| } | ||
|
|
||
| fileprivate func createURLSession() -> URLSession { | ||
| let config = URLSessionConfiguration.default | ||
|
|
||
| // Harbor expects a CSRF token to be present if the HTTP client | ||
| // carries a session cookie between its requests[1] and fails if | ||
| // it was not present[2]. | ||
| // | ||
| // To fix that, we disable the automatic cookies carry in URLSession. | ||
| // | ||
| // [1]: https://github.com/goharbor/harbor/blob/a4c577f9ec4f18396207a5e686433a6ba203d4ef/src/server/middleware/csrf/csrf.go#L78 | ||
| // [2]: https://github.com/cirruslabs/tart/issues/295 | ||
| config.httpShouldSetCookies = false | ||
| func urlSession( | ||
| _ session: URLSession, | ||
| task: URLSessionTask, | ||
| didCompleteWithError error: Error? | ||
| ) { | ||
| if !buffer.isEmpty { | ||
| streamContinuation?.yield(buffer) | ||
| buffer.removeAll(keepingCapacity: true) | ||
| } | ||
|
|
||
| return URLSession(configuration: config) | ||
| if let error = error { | ||
| streamContinuation?.finish(throwing: error) | ||
| } else { | ||
| streamContinuation?.finish() | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters