From 31ab4218f7ce73a40c3304db3137cb430cd53823 Mon Sep 17 00:00:00 2001 From: Nikolay Edigaryev Date: Wed, 11 Dec 2024 21:48:59 +0400 Subject: [PATCH] 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 --- Package.resolved | 20 +--- Package.swift | 2 - Sources/tart/Fetcher.swift | 140 +++++++++++++----------- Sources/tart/OCI/Layerizer/DiskV2.swift | 31 +++--- Sources/tart/OCI/Registry.swift | 7 +- Sources/tart/VM.swift | 1 - 6 files changed, 94 insertions(+), 107 deletions(-) diff --git a/Package.resolved b/Package.resolved index e8a8793..db15534 100644 --- a/Package.resolved +++ b/Package.resolved @@ -1,5 +1,5 @@ { - "originHash" : "6a15657d8cb1d3e2b447f31aff5b47d6a9655d2262e48ca76476ba525435269b", + "originHash" : "22b3726bc4e4c6e9c04ac97cb08a82967feb39960a93d2909768a16e11576748", "pins" : [ { "identity" : "antlr4", @@ -55,15 +55,6 @@ "version" : "1.5.0" } }, - { - "identity" : "swift-async-algorithms", - "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-async-algorithms", - "state" : { - "branch" : "main", - "revision" : "5c8bd186f48c16af0775972700626f0b74588278" - } - }, { "identity" : "swift-atomics", "kind" : "remoteSourceControl", @@ -73,15 +64,6 @@ "version" : "1.2.0" } }, - { - "identity" : "swift-collections", - "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-collections.git", - "state" : { - "revision" : "9bf03ff58ce34478e66aaee630e491823326fd06", - "version" : "1.1.3" - } - }, { "identity" : "swift-log", "kind" : "remoteSourceControl", diff --git a/Package.swift b/Package.swift index c52611b..ed82911 100644 --- a/Package.swift +++ b/Package.swift @@ -13,7 +13,6 @@ let package = Package( .package(url: "https://github.com/apple/swift-argument-parser", from: "1.3.1"), .package(url: "https://github.com/mhdhejazi/Dynamic", branch: "master"), .package(url: "https://github.com/apple/swift-algorithms", from: "1.2.0"), - .package(url: "https://github.com/apple/swift-async-algorithms", branch: "main"), .package(url: "https://github.com/malcommac/SwiftDate", from: "7.0.0"), .package(url: "https://github.com/antlr/antlr4", exact: "4.13.2"), .package(url: "https://github.com/apple/swift-atomics.git", .upToNextMajor(from: "1.2.0")), @@ -29,7 +28,6 @@ let package = Package( targets: [ .executableTarget(name: "tart", dependencies: [ .product(name: "Algorithms", package: "swift-algorithms"), - .product(name: "AsyncAlgorithms", package: "swift-async-algorithms"), .product(name: "ArgumentParser", package: "swift-argument-parser"), .product(name: "Dynamic", package: "Dynamic"), .product(name: "SwiftDate", package: "SwiftDate"), diff --git a/Sources/tart/Fetcher.swift b/Sources/tart/Fetcher.swift index 77a4368..e56c845 100644 --- a/Sources/tart/Fetcher.swift +++ b/Sources/tart/Fetcher.swift @@ -1,69 +1,6 @@ import Foundation -import AsyncAlgorithms -fileprivate let urlSession = createURLSession() - -class DownloadDelegate: NSObject, URLSessionTaskDelegate { - let progress: Progress - init(_ progress: Progress) throws { - self.progress = progress - } - - func urlSession(_ session: URLSession, didCreateTask task: URLSessionTask) { - self.progress.addChild(task.progress, withPendingUnitCount: self.progress.totalUnitCount) - } -} - -class Fetcher { - static func fetch(_ request: URLRequest, viaFile: Bool = false, progress: Progress? = nil) async throws -> (AsyncThrowingChannel, HTTPURLResponse) { - let delegate = progress != nil ? try DownloadDelegate(progress!) : nil - - if viaFile { - return try await fetchViaFile(request, delegate: delegate) - } - - return try await fetchViaMemory(request, delegate: delegate) - } - - private static func fetchViaMemory(_ request: URLRequest, delegate: URLSessionTaskDelegate? = nil) async throws -> (AsyncThrowingChannel, HTTPURLResponse) { - let dataCh = AsyncThrowingChannel() - - let (data, response) = try await urlSession.data(for: request, delegate: delegate) - - Task { - await dataCh.send(data) - - dataCh.finish() - } - - return (dataCh, response as! HTTPURLResponse) - } - - private static func fetchViaFile(_ request: URLRequest, delegate: URLSessionTaskDelegate? = nil) async throws -> (AsyncThrowingChannel, HTTPURLResponse) { - let dataCh = AsyncThrowingChannel() - - 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) - - Task { - for chunk in (0 ..< mappedFile.count).chunks(ofCount: 64 * 1024 * 1024) { - await dataCh.send(mappedFile.subdata(in: chunk)) - } - - dataCh.finish() - } - - return (dataCh, response as! HTTPURLResponse) - } -} - -fileprivate func createURLSession() -> URLSession { +fileprivate var urlSessionConfiguration: URLSessionConfiguration { let config = URLSessionConfiguration.default // Harbor expects a CSRF token to be present if the HTTP client @@ -76,5 +13,78 @@ fileprivate func createURLSession() -> URLSession { // [2]: https://github.com/cirruslabs/tart/issues/295 config.httpShouldSetCookies = false - return URLSession(configuration: config) + return config +} + +class Fetcher { + static func fetch(_ request: URLRequest, viaFile: Bool = false, progress: Progress? = nil) async throws -> (AsyncThrowingStream, HTTPURLResponse) { + let delegate = Delegate() + let session = URLSession(configuration: urlSessionConfiguration, delegate: delegate, delegateQueue: nil) + let task = session.dataTask(with: request) + + let stream = AsyncThrowingStream { continuation in + delegate.streamContinuation = continuation + } + + let response = try await withCheckedThrowingContinuation { continuation in + delegate.responseContinuation = continuation + task.resume() + } + + return (stream, response as! HTTPURLResponse) + } +} + +fileprivate class Delegate: NSObject, URLSessionDataDelegate { + var responseContinuation: CheckedContinuation? + var streamContinuation: AsyncThrowingStream.Continuation? + + private var buffer: Data = Data() + private let bufferFlushSize = 16 * 1024 * 1024 + + 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)) + + // Pre-initialize buffer as we now know the capacity + buffer = Data(capacity: Int(capacity)) + + responseContinuation?.resume(returning: response) + completionHandler(.allow) + } + + func urlSession( + _ session: URLSession, + dataTask: URLSessionDataTask, + didReceive data: Data + ) { + buffer.append(data) + + if buffer.count >= bufferFlushSize { + streamContinuation?.yield(buffer) + buffer.removeAll(keepingCapacity: true) + } + } + + func urlSession( + _ session: URLSession, + task: URLSessionTask, + didCompleteWithError error: Error? + ) { + if !buffer.isEmpty { + streamContinuation?.yield(buffer) + buffer.removeAll(keepingCapacity: true) + } + + if let error = error { + streamContinuation?.finish(throwing: error) + } else { + streamContinuation?.finish() + } + } } diff --git a/Sources/tart/OCI/Layerizer/DiskV2.swift b/Sources/tart/OCI/Layerizer/DiskV2.swift index a00e016..b926a42 100644 --- a/Sources/tart/OCI/Layerizer/DiskV2.swift +++ b/Sources/tart/OCI/Layerizer/DiskV2.swift @@ -7,6 +7,21 @@ class DiskV2: Disk { private static let bufferSizeBytes = 4 * 1024 * 1024 private static let layerLimitBytes = 512 * 1024 * 1024 + // A zero chunk for faster than byte-by-byte comparisons + // + // Assumes that the other Data(...) is equal in size, but it's fine to get a false-negative + // on the last block since it costs only 4 MiB of excess data per 512 MiB layer. + // + // Some simple benchmarks ("sync && sudo purge" command was used to negate the disk caching effects): + // +--------------------------------------+---------------------------------------------------+ + // | Operation | time(1) result | + // +--------------------------------------+---------------------------------------------------+ + // | Data(...) == zeroChunk | 2.16s user 11.71s system 73% cpu 18.928 total | + // | Data(...).contains(where: {$0 != 0}) | 603.68s user 12.97s system 99% cpu 10:22.85 total | + // +--------------------------------------+---------------------------------------------------+ + private static let holeGranularityBytes = 4 * 1024 * 1024 + private static let zeroChunk = Data(count: holeGranularityBytes) + static func push(diskURL: URL, registry: Registry, chunkSizeMb: Int, concurrency: UInt, progress: Progress) async throws -> [OCIManifestLayer] { var pushedLayers: [(index: Int, pushedLayer: OCIManifestLayer)] = [] @@ -215,22 +230,6 @@ class DiskV2: Disk { } private static func zeroSkippingWrite(_ disk: FileHandle, _ rdisk: FileHandle?, _ fsBlockSize: UInt64, _ offset: UInt64, _ data: Data) throws -> UInt64 { - let holeGranularityBytes = 64 * 1024 - - // A zero chunk for faster than byte-by-byte comparisons - // - // Assumes that the other Data(...) is equal in size, but it's fine to get a false-negative - // on the last block since it costs only 64 KiB of excess data per 500 MB layer. - // - // Some simple benchmarks ("sync && sudo purge" command was used to negate the disk caching effects): - // +--------------------------------------+---------------------------------------------------+ - // | Operation | time(1) result | - // +--------------------------------------+---------------------------------------------------+ - // | Data(...) == zeroChunk | 2.16s user 11.71s system 73% cpu 18.928 total | - // | Data(...).contains(where: {$0 != 0}) | 603.68s user 12.97s system 99% cpu 10:22.85 total | - // +--------------------------------------+---------------------------------------------------+ - let zeroChunk = Data(count: holeGranularityBytes) - var offset = offset for chunk in data.chunks(ofCount: holeGranularityBytes) { diff --git a/Sources/tart/OCI/Registry.swift b/Sources/tart/OCI/Registry.swift index c370b90..a9acbe5 100644 --- a/Sources/tart/OCI/Registry.swift +++ b/Sources/tart/OCI/Registry.swift @@ -1,6 +1,5 @@ import Foundation import Algorithms -import AsyncAlgorithms enum RegistryError: Error { case UnexpectedHTTPStatusCode(when: String, code: Int, details: String = "") @@ -31,7 +30,7 @@ extension Data { } } -extension AsyncThrowingChannel { +extension AsyncThrowingStream { func asData() async throws -> Data { var result = Data() @@ -307,7 +306,7 @@ class Registry { body: Data? = nil, doAuth: Bool = true, viaFile: Bool = false - ) async throws -> (AsyncThrowingChannel, HTTPURLResponse) { + ) async throws -> (AsyncThrowingStream, HTTPURLResponse) { var urlComponents = urlComponents if urlComponents.queryItems == nil && !parameters.isEmpty { @@ -413,7 +412,7 @@ class Registry { return nil } - private func authAwareRequest(request: URLRequest, viaFile: Bool = false, doAuth: Bool) async throws -> (AsyncThrowingChannel, HTTPURLResponse) { + private func authAwareRequest(request: URLRequest, viaFile: Bool = false, doAuth: Bool) async throws -> (AsyncThrowingStream, HTTPURLResponse) { var request = request if doAuth { diff --git a/Sources/tart/VM.swift b/Sources/tart/VM.swift index 433df2c..90ab7c2 100644 --- a/Sources/tart/VM.swift +++ b/Sources/tart/VM.swift @@ -1,6 +1,5 @@ import Foundation import Virtualization -import AsyncAlgorithms import Semaphore struct UnsupportedRestoreImageError: Error {