tart pull: introduce --populate-cache flag (#103)

* tart pull: introduce --populate-cache flag

* VMStorageOCI: introduce cache() method

* Review comments (#107)

* Rename SetCommand back to Set

Co-authored-by: Fedor Korotkov <fedor.korotkov@gmail.com>

Sources/tart/Commands/Push.swift

@@ -12,6 +12,10 @@ struct Push: AsyncParsableCommand {
   @Argument(help: "remote VM name(s)")
   var remoteNames: [String]
 
+  @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
+
   func run() async throws {
     do {
       let localVMDir = try VMStorageLocal().open(localName)
@@ -35,12 +39,20 @@ struct Push: AsyncParsableCommand {
       for (registryIdentifier, remoteNamesForRegistry) in registryGroups {
         let registry = try Registry(host: registryIdentifier.host, namespace: registryIdentifier.namespace)
 
-        let listOfTagsAndDigests = "{" + remoteNamesForRegistry.map{$0.fullyQualifiedReference }
-                .joined(separator: ",") + "}"
         defaultLogger.appendNewLine("pushing \(localName) to "
-          + "\(registryIdentifier.host)/\(registryIdentifier.namespace)\(listOfTagsAndDigests)...")
+          + "\(registryIdentifier.host)/\(registryIdentifier.namespace)\(remoteNamesForRegistry.referenceNames())...")
 
-        try await localVMDir.pushToRegistry(registry: registry, references: remoteNamesForRegistry.map{ $0.reference })
+        let pushedRemoteName = try await localVMDir.pushToRegistry(registry: registry, references: remoteNamesForRegistry.map{ $0.reference.value })
+
+        // Populate the local cache (if requested)
+        if populateCache {
+          let ociStorage = VMStorageOCI()
+          let expectedPushedVMDir = try ociStorage.create(pushedRemoteName)
+          try localVMDir.clone(to: expectedPushedVMDir, generateMAC: false)
+          for remoteName in remoteNamesForRegistry {
+            try ociStorage.link(from: remoteName, to: pushedRemoteName)
+          }
+        }
       }
 
       Foundation.exit(0)
@@ -51,3 +63,15 @@ struct Push: AsyncParsableCommand {
     }
   }
 }
+
+extension Collection where Element == RemoteName {
+  func referenceNames() -> String {
+    let references = self.map{ $0.reference.fullyQualified }
+
+    switch count {
+    case 0: return "∅"
+    case 1: return references.first!
+    default: return "{" + references.joined(separator: ",") + "}"
+    }
+  }
+}

Sources/tart/Commands/Set.swift

@@ -2,7 +2,7 @@ import ArgumentParser
 import Foundation
 
 struct Set: AsyncParsableCommand {
-  static var configuration = CommandConfiguration(abstract: "Modify VM's configuration")
+  static var configuration = CommandConfiguration(commandName: "set", abstract: "Modify VM's configuration")
 
   @Argument(help: "VM name")
   var name: String

Sources/tart/OCI/Manifest.swift

@@ -8,6 +8,10 @@ struct OCIManifest: Codable, Equatable {
   var mediaType: String = ociManifestMediaType
   var config: OCIManifestConfig
   var layers: [OCIManifestLayer] = Array()
+
+  func digest() throws -> String {
+    try Digest.hash(JSONEncoder().encode(self))
+  }
 }
 
 struct OCIManifestConfig: Codable, Equatable {

Sources/tart/OCI/RemoteName.swift

@@ -1,31 +1,57 @@
 import Foundation
 import Parsing
 
-struct Tail {
-  enum TailType {
+struct Reference: Comparable, Hashable, CustomStringConvertible {
+  enum ReferenceType: Comparable {
     case Tag
     case Digest
   }
 
-  var type: TailType
-  var value: String
-}
+  let type: ReferenceType
+  let value: String
 
-struct RemoteName: Comparable, CustomStringConvertible {
-  var host: String
-  var namespace: String
-  var reference: String = "latest"
-  var fullyQualifiedReference: String {
+  var fullyQualified: String {
     get {
-      if reference.starts(with: "sha256:") {
-        return "@" + reference
+      switch type {
+      case .Tag:
+        return ":" + value
+      case .Digest:
+        return "@" + value
       }
-
-      return ":" + reference
     }
   }
 
-  init(host: String, namespace: String, reference: String) {
+  init(tag: String) {
+    type = .Tag
+    value = tag
+  }
+
+  init(digest: String) {
+    type = .Digest
+    value = digest
+  }
+
+  static func <(lhs: Reference, rhs: Reference) -> Bool {
+    if lhs.type != rhs.type {
+      return lhs.type < rhs.type
+    } else {
+      return lhs.value < rhs.value
+    }
+  }
+
+  var description: String {
+    get {
+      fullyQualified
+    }
+  }
+}
+
+struct RemoteName: Comparable, Hashable, CustomStringConvertible {
+  var host: String
+  var namespace: String
+  var reference: Reference
+
+  init(host: String, namespace: String, reference: Reference) {
     self.host = host
     self.namespace = namespace
     self.reference = reference
@@ -58,13 +84,13 @@ struct RemoteName: Comparable, CustomStringConvertible {
           Parse {
             ":"
             csNormal.map {
-              Tail(type: .Tag, value: String($0))
+              Reference(tag: String($0))
             }
           }
           Parse {
             "@sha256:"
             csHex.map {
-              Tail(type: .Digest, value: "sha256:" + String($0))
+              Reference(digest: "sha256:" + String($0))
             }
           }
         }
@@ -76,9 +102,7 @@ struct RemoteName: Comparable, CustomStringConvertible {
 
     host = String(result.0)
     namespace = String(result.1)
-    if let tail = result.2 {
-      reference = tail.value
-    }
+    reference = result.2 ?? Reference(tag: "latest")
   }
 
   static func <(lhs: RemoteName, rhs: RemoteName) -> Bool {
@@ -92,7 +116,7 @@ struct RemoteName: Comparable, CustomStringConvertible {
   }
 
   var description: String {
-    "\(host)/\(namespace)\(fullyQualifiedReference)"
+    "\(host)/\(namespace)\(reference.fullyQualified)"
   }
 }
 

Sources/tart/VMDirectory+OCI.swift

@@ -98,7 +98,7 @@ extension VMDirectory {
     try nvram.close()
   }
 
-  func pushToRegistry(registry: Registry, references: [String]) async throws {
+  func pushToRegistry(registry: Registry, references: [String]) async throws -> RemoteName {
     var layers = Array<OCIManifestLayer>()
 
     // Read VM's config and push it as blob
@@ -155,6 +155,9 @@ extension VMDirectory {
 
       _ = try await registry.pushManifest(reference: reference, manifest: manifest)
     }
+
+    let pushedReference = Reference(digest: try manifest.digest())
+    return RemoteName(host: registry.baseURL.host!, namespace: registry.namespace, reference: pushedReference)
   }
 }
 

Sources/tart/VMStorageOCI.swift

@@ -64,27 +64,39 @@ class VMStorageOCI {
   func pull(_ name: RemoteName, registry: Registry) async throws {
     defaultLogger.appendNewLine("pulling manifest...")
 
-    let (manifest, manifestData) = try await registry.pullManifest(reference: name.reference)
+    let (manifest, _) = try await registry.pullManifest(reference: name.reference.value)
+
+    if let cacheToVMDir = try cache(name: name, digest: manifest.digest()) {
+      try await cacheToVMDir.pullFromRegistry(registry: registry, manifest: manifest)
+    }
+  }
+
+  func cache(name: RemoteName, digest: String) throws -> VMDirectory? {
+    var result: VMDirectory? = nil
 
-    // Create directory for manifest's digest
     var digestName = name
-    digestName.reference = Digest.hash(manifestData)
+    digestName.reference = Reference(digest: digest)
+
     if !exists(digestName) {
-      let vmDir = try create(digestName)
-      try await vmDir.pullFromRegistry(registry: registry, manifest: manifest)
+      result = try create(digestName)
     } else {
-      defaultLogger.appendNewLine("\(digestName.reference) image is already cached! creating a symlink...")   
+      defaultLogger.appendNewLine("\(digestName) image is already cached! creating a symlink...")
     }
 
-    // Create directory for reference if it's different
-    if digestName != name {
+    if name != digestName {
       // Overwrite the old symbolic link
-      if FileManager.default.fileExists(atPath: vmURL(name).path) {
-        try FileManager.default.removeItem(at: vmURL(name))
-      }
-
-      try FileManager.default.createSymbolicLink(at: vmURL(name), withDestinationURL: vmURL(digestName))
+      try link(from: digestName, to: name)
     }
+
+    return result
+  }
+
+  func link(from: RemoteName, to: RemoteName) throws {
+    if FileManager.default.fileExists(atPath: vmURL(to).path) {
+      try FileManager.default.removeItem(at: vmURL(to))
+    }
+
+    try FileManager.default.createSymbolicLink(at: vmURL(to), withDestinationURL: vmURL(from))
   }
 }
 
@@ -92,7 +104,7 @@ extension URL {
   func appendingRemoteName(_ name: RemoteName) -> URL {
     var result: URL = self
 
-    for pathComponent in (name.host + "/" + name.namespace + "/" + name.reference).split(separator: "/") {
+    for pathComponent in (name.host + "/" + name.namespace + "/" + name.reference.value).split(separator: "/") {
       result = result.appendingPathComponent(String(pathComponent))
     }
 

Tests/TartTests/RemoteNameTests.swift

@@ -3,13 +3,13 @@ import XCTest
 
 final class RemoteNameTests: XCTestCase {
   func testTag() throws {
-    let expectedRemoteName = RemoteName(host: "ghcr.io", namespace: "a/b", reference: "latest")
+    let expectedRemoteName = RemoteName(host: "ghcr.io", namespace: "a/b", reference: Reference(tag: "latest"))
 
     XCTAssertEqual(expectedRemoteName, try RemoteName("ghcr.io/a/b:latest"))
   }
   
   func testComplexTag() throws {
-    let expectedRemoteName = RemoteName(host: "ghcr.io", namespace: "a/b", reference: "1.2.3-RC-1")
+    let expectedRemoteName = RemoteName(host: "ghcr.io", namespace: "a/b", reference: Reference(tag: "1.2.3-RC-1"))
 
     XCTAssertEqual(expectedRemoteName, try RemoteName("ghcr.io/a/b:1.2.3-RC-1"))
   }
@@ -18,7 +18,7 @@ final class RemoteNameTests: XCTestCase {
     let expectedRemoteName = RemoteName(
       host: "ghcr.io",
       namespace: "a/b",
-      reference: "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+      reference: Reference(digest: "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855")
     )
 
     XCTAssertEqual(expectedRemoteName,