Support chunked uploads (#159)

* Support chunked uploads

* Rebase fixes

* Bump swift http client

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",

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"),

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 {

Sources/tart/OCI/Registry.swift

@@ -2,6 +2,7 @@ import Foundation
 import NIOCore
 import NIOHTTP1
 import AsyncHTTPClient
+import Algorithms
 import NIOPosix
 
 enum RegistryError: Error {
@@ -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,25 +174,33 @@ 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!)
     }
 

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<OCIManifestLayer>()
 
     // 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,

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()