diff --git a/Package.resolved b/Package.resolved index 2499e91..3c84d2c 100644 --- a/Package.resolved +++ b/Package.resolved @@ -5,8 +5,8 @@ "kind" : "remoteSourceControl", "location" : "https://github.com/swift-server/async-http-client", "state" : { - "revision" : "24425989dadab6d6e4167174791a23d4e2a6d0c3", - "version" : "1.10.0" + "revision" : "df87a860fdc41a595d5ca67f74cde9adbccc099a", + "version" : "1.11.4" } }, { @@ -18,6 +18,15 @@ "revision" : "772883073d044bc754d401cabb6574624eb3778f" } }, + { + "identity" : "swift-algorithms", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-algorithms", + "state" : { + "revision" : "b14b7f4c528c942f121c8b860b9410b2bf57825e", + "version" : "1.0.0" + } + }, { "identity" : "swift-argument-parser", "kind" : "remoteSourceControl", @@ -27,6 +36,15 @@ "version" : "1.1.2" } }, + { + "identity" : "swift-atomics", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-atomics.git", + "state" : { + "revision" : "919eb1d83e02121cdb434c7bfc1f0c66ef17febe", + "version" : "1.0.2" + } + }, { "identity" : "swift-case-paths", "kind" : "remoteSourceControl", @@ -68,8 +86,8 @@ "kind" : "remoteSourceControl", "location" : "https://github.com/apple/swift-nio-http2.git", "state" : { - "revision" : "72bcaf607b40d7c51044f65b0f5ed8581a911832", - "version" : "1.21.0" + "revision" : "108ac15087ea9b79abb6f6742699cf31de0e8772", + "version" : "1.22.0" } }, { @@ -90,6 +108,15 @@ "version" : "1.12.0" } }, + { + "identity" : "swift-numerics", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-numerics", + "state" : { + "revision" : "0a5bc04095a675662cf24757cc0640aa2204253b", + "version" : "1.0.2" + } + }, { "identity" : "swift-parsing", "kind" : "remoteSourceControl", diff --git a/Package.swift b/Package.swift index 4d11d5d..302c721 100644 --- a/Package.swift +++ b/Package.swift @@ -13,10 +13,12 @@ let package = Package( .package(url: "https://github.com/apple/swift-argument-parser", from: "1.1.2"), .package(url: "https://github.com/mhdhejazi/Dynamic", branch: "master"), .package(url: "https://github.com/pointfreeco/swift-parsing", from: "0.9.2"), - .package(url: "https://github.com/swift-server/async-http-client", from: "1.10.0"), + .package(url: "https://github.com/swift-server/async-http-client", from: "1.11.4"), + .package(url: "https://github.com/apple/swift-algorithms", from: "1.0.0"), ], targets: [ .executableTarget(name: "tart", dependencies: [ + .product(name: "Algorithms", package: "swift-algorithms"), .product(name: "ArgumentParser", package: "swift-argument-parser"), .product(name: "AsyncHTTPClient", package: "async-http-client"), .product(name: "Dynamic", package: "Dynamic"), diff --git a/Sources/tart/Commands/Push.swift b/Sources/tart/Commands/Push.swift index c7eb414..9e51ac3 100644 --- a/Sources/tart/Commands/Push.swift +++ b/Sources/tart/Commands/Push.swift @@ -12,6 +12,14 @@ struct Push: AsyncParsableCommand { @Argument(help: "remote VM name(s)") var remoteNames: [String] + @Option(help: ArgumentHelp("chunk size in MB if registry supports chunked uploads", + discussion: """ + By default monolithic method is used for uploading blobs to the registry but some registries support a more efficient chunked method. + For example, AWS Elastic Container Registry supports only chunks larger than 5MB but GitHub Container Registry supports only chunks smaller than 4MB. Google Container Registry on the other hand doesn't support chunked uploads at all. + Please refer to the documentation of your particular registry in order to see if this option is suitable for you and what's the recommended chunk size. + """)) + var chunkSize: Int = 0 + @Flag(help: ArgumentHelp("cache pushed images locally", discussion: "Increases disk usage, but saves time if you're going to pull the pushed images later.")) var populateCache: Bool = false @@ -42,7 +50,11 @@ struct Push: AsyncParsableCommand { defaultLogger.appendNewLine("pushing \(localName) to " + "\(registryIdentifier.host)/\(registryIdentifier.namespace)\(remoteNamesForRegistry.referenceNames())...") - let pushedRemoteName = try await localVMDir.pushToRegistry(registry: registry, references: remoteNamesForRegistry.map{ $0.reference.value }) + let pushedRemoteName = try await localVMDir.pushToRegistry( + registry: registry, + references: remoteNamesForRegistry.map{ $0.reference.value }, + chunkSizeMb: chunkSize + ) // Populate the local cache (if requested) if populateCache { diff --git a/Sources/tart/OCI/Registry.swift b/Sources/tart/OCI/Registry.swift index 0789b30..5b42970 100644 --- a/Sources/tart/OCI/Registry.swift +++ b/Sources/tart/OCI/Registry.swift @@ -2,6 +2,7 @@ import Foundation import NIOCore import NIOHTTP1 import AsyncHTTPClient +import Algorithms import NIOPosix enum RegistryError: Error { @@ -84,7 +85,7 @@ class Registry { deinit { try! httpClient.syncShutdown() } - + let baseURL: URL let namespace: String let credentialsProvider: CredentialsProvider @@ -101,9 +102,9 @@ class Registry { } convenience init( - host: String, - namespace: String, - credentialsProvider: CredentialsProvider = KeychainCredentialsProvider() + host: String, + namespace: String, + credentialsProvider: CredentialsProvider = KeychainCredentialsProvider() ) throws { var baseURLComponents = URLComponents() @@ -162,7 +163,7 @@ class Registry { return URLComponents(url: uploadLocation.absolutize(baseURL), resolvingAgainstBaseURL: true)! } - public func pushBlob(fromData: Data, chunkSize: Int = 5 * 1024 * 1024) async throws -> String { + public func pushBlob(fromData: Data, chunkSizeMb: Int = 0) async throws -> String { // Initiate a blob upload let postResponse = try await endpointRequest(.POST, "\(namespace)/blobs/uploads/", headers: ["Content-Length": "0"]) @@ -173,27 +174,35 @@ class Registry { } // Figure out where to upload the blob - let uploadLocation = try uploadLocationFromResponse(postResponse) - - // Upload the blob - let headers = [ - "Content-Length": "\(fromData.count)", - "Content-Type": "application/octet-stream", - ] + var uploadLocation = try uploadLocationFromResponse(postResponse) let digest = Digest.hash(fromData) - let parameters = [ - "digest": digest, - ] - let putResponse = try await rawRequest(.PUT, uploadLocation, headers: headers, parameters: parameters, - body: fromData) - if putResponse.status != .created { - let body = try await postResponse.body.readTextResponse() - throw RegistryError.UnexpectedHTTPStatusCode(when: "pushing blob (PUT) to \(uploadLocation)", - code: putResponse.status.code, details: body ?? "") + var uploadedBytes = 0 + let chunks = fromData.chunks(ofCount: chunkSizeMb == 0 ? fromData.count : chunkSizeMb * 1_000_000) + for (index, chunk) in chunks.enumerated() { + let lastChunk = index == (chunks.count - 1) + let response = try await rawRequest( + lastChunk ? .PUT : .PATCH, + uploadLocation, + headers: [ + "Content-Type": "application/octet-stream", + "Content-Range": "\(uploadedBytes)-\(uploadedBytes + chunk.count - 1)", + ], + parameters: lastChunk ? ["digest": digest] : [:], + body: chunk + ) + let expectedStatus: HTTPResponseStatus = lastChunk ? .created : .accepted + if response.status != expectedStatus { + let body = try await response.body.readTextResponse() + throw RegistryError.UnexpectedHTTPStatusCode(when: "streaming blob to \(uploadLocation)", + code: response.status.code, details: body ?? "") + } + uploadedBytes += chunk.count + // Update location for the next chunk + uploadLocation = try uploadLocationFromResponse(response) } - + return digest } @@ -248,6 +257,7 @@ class Registry { request.headers.add(name: key, value: value) } if body != nil { + request.headers.add(name: "Content-Length", value: "\(body!.count)") request.body = HTTPClientRequest.Body.bytes(body!) } diff --git a/Sources/tart/VMDirectory+OCI.swift b/Sources/tart/VMDirectory+OCI.swift index cbaa256..1789ac7 100644 --- a/Sources/tart/VMDirectory+OCI.swift +++ b/Sources/tart/VMDirectory+OCI.swift @@ -98,13 +98,14 @@ extension VMDirectory { try nvram.close() } - func pushToRegistry(registry: Registry, references: [String]) async throws -> RemoteName { + func pushToRegistry(registry: Registry, references: [String], chunkSizeMb: Int) async throws -> RemoteName { var layers = Array() // Read VM's config and push it as blob let config = try VMConfig(fromURL: configURL) let configJSON = try JSONEncoder().encode(config) - let configDigest = try await registry.pushBlob(fromData: configJSON) + defaultLogger.appendNewLine("pushing config...") + let configDigest = try await registry.pushBlob(fromData: configJSON, chunkSizeMb: chunkSizeMb) layers.append(OCIManifestLayer(mediaType: Self.configMediaType, size: configJSON.count, digest: configDigest)) // Progress @@ -126,21 +127,21 @@ extension VMDirectory { return data } - while let chunk = try compressingFilter.readData(ofLength: Self.layerLimitBytes) { - let chunkDigest = try await registry.pushBlob(fromData: chunk) - layers.append(OCIManifestLayer(mediaType: Self.diskMediaType, size: chunk.count, digest: chunkDigest)) + while let compressedLayerData = try compressingFilter.readData(ofLength: Self.layerLimitBytes) { + let layerDigest = try await registry.pushBlob(fromData: compressedLayerData, chunkSizeMb: chunkSizeMb) + layers.append(OCIManifestLayer(mediaType: Self.diskMediaType, size: compressedLayerData.count, digest: layerDigest)) } // Read VM's NVRAM and push it as blob defaultLogger.appendNewLine("pushing NVRAM...") let nvram = try FileHandle(forReadingFrom: nvramURL).readToEnd()! - let nvramDigest = try await registry.pushBlob(fromData: nvram) + let nvramDigest = try await registry.pushBlob(fromData: nvram, chunkSizeMb: chunkSizeMb) layers.append(OCIManifestLayer(mediaType: Self.nvramMediaType, size: nvram.count, digest: nvramDigest)) // Craft a stub OCI config for Docker Hub compatibility let ociConfigJSON = try OCIConfig().toJSON() - let ociConfigDigest = try await registry.pushBlob(fromData: ociConfigJSON) + let ociConfigDigest = try await registry.pushBlob(fromData: ociConfigJSON, chunkSizeMb: chunkSizeMb) let manifest = OCIManifest( config: OCIManifestConfig(size: ociConfigJSON.count, digest: ociConfigDigest), layers: layers, diff --git a/Tests/TartTests/RegistryTests.swift b/Tests/TartTests/RegistryTests.swift index 3ae757b..14fd99c 100644 --- a/Tests/TartTests/RegistryTests.swift +++ b/Tests/TartTests/RegistryTests.swift @@ -42,13 +42,13 @@ final class RegistryTests: XCTestCase { XCTAssertEqual(pushedBlob, pulledBlob) } - func testPushPullBlobHuge() async throws { + func testPushPullBlobHugeInChunks() async throws { // Generate a large enough blob let fh = FileHandle(forReadingAtPath: "/dev/urandom")! let largeBlobToPush = try fh.read(upToCount: 768 * 1024 * 1024)! // Push it - let largeBlobDigest = try await registry.pushBlob(fromData: largeBlobToPush) + let largeBlobDigest = try await registry.pushBlob(fromData: largeBlobToPush, chunkSizeMb: 10) // Pull it var pulledLargeBlob = Data()