Add OCI transport and base clone with DiskImageKit (#1304)

* Add OCI transport and base clone with DiskImageKit

* Address stacked OCI pull review feedback

* Stream file digest hashing

* Lock frozen overlays during push

Author: yzhuang-oai

Sources/tart/Commands/Clone.swift

@@ -31,6 +31,9 @@ struct Clone: AsyncParsableCommand {
   @Flag(help: .hidden)
   var deduplicate: Bool = false
 
+  @Flag(help: "create a stacked disk that uses the source image as an immutable base")
+  var stacked: Bool = false
+
   @Option(help: ArgumentHelp("limit automatic pruning to n gigabytes", valueName: "n"))
   var pruneLimit: UInt = 100
 
@@ -47,8 +50,15 @@ struct Clone: AsyncParsableCommand {
   func run() async throws {
     let ociStorage = try VMStorageOCI()
     let localStorage = try VMStorageLocal()
+    let remoteName = try? RemoteName(sourceName)
 
-    if let remoteName = try? RemoteName(sourceName), !ociStorage.exists(remoteName) {
+    if stacked {
+      guard remoteName != nil else {
+        throw ValidationError("--stacked requires a remote image")
+      }
+    }
+
+    if let remoteName, try !ociStorage.hasUsableCachedImageForClone(remoteName, requireManifest: stacked) {
       // Pull the VM in case it's OCI-based and doesn't exist locally yet
       let registry = try Registry(host: remoteName.host, namespace: remoteName.namespace, insecure: insecure)
       try await ociStorage.pull(remoteName, registry: registry, concurrency: concurrency, deduplicate: deduplicate)
@@ -66,9 +76,28 @@ struct Clone: AsyncParsableCommand {
       let lock = try FileLock(lockURL: Config().tartHomeDir)
       try lock.lock()
 
+      let sourceState = try sourceVM.state()
       let generateMAC = try localStorage.hasVMsWithMACAddress(macAddress: sourceVM.macAddress())
-        && sourceVM.state() != .Suspended
-      try sourceVM.clone(to: tmpVMDir, generateMAC: generateMAC)
+        && sourceState != .Suspended
+
+      if stacked {
+        guard sourceVM.isStandalone else {
+          throw ValidationError("--stacked cannot use an image that already has a stacked disk")
+        }
+        guard try VMConfig(fromURL: sourceVM.configURL).os == .darwin else {
+          throw ValidationError("--stacked currently supports only macOS images")
+        }
+        try sourceVM.cloneAsStackedBase(to: tmpVMDir, generateMAC: generateMAC)
+      } else if sourceVM.isStackedCachedImage {
+        try sourceVM.cloneStacked(to: tmpVMDir, copyWritableOverlay: false, generateMAC: generateMAC)
+      } else if sourceVM.isStackedVM {
+        guard sourceState == .Stopped else {
+          throw RuntimeError.VMConfigurationError("VM \"\(sourceName)\" must be stopped before cloning")
+        }
+        try sourceVM.cloneStacked(to: tmpVMDir, copyWritableOverlay: true, generateMAC: generateMAC)
+      } else {
+        try sourceVM.clone(to: tmpVMDir, generateMAC: generateMAC)
+      }
 
       try localStorage.move(newName, from: tmpVMDir)
 
@@ -78,11 +107,23 @@ struct Clone: AsyncParsableCommand {
       // is not actually claiming new space until the VM is started and it writes something to disk.
       //
       // So, once we clone the VM let's try to claim the rest of space for the VM to run without errors.
-      let unallocatedBytes = try sourceVM.sizeBytes() - sourceVM.allocatedSizeBytes()
-      // Avoid reclaiming an excessive amount of disk space.
-      let reclaimBytes = min(unallocatedBytes, Int(pruneLimit) * 1024 * 1024 * 1024)
-      if reclaimBytes > 0 {
-        try Prune.reclaimIfNeeded(UInt64(reclaimBytes), sourceVM)
+      if sourceVM.isStandalone {
+        let unallocatedBytes = try sourceVM.sizeBytes() - sourceVM.allocatedSizeBytes()
+        // Avoid reclaiming an excessive amount of disk space.
+        let reclaimBytes = min(unallocatedBytes, Int(pruneLimit) * 1024 * 1024 * 1024)
+        if reclaimBytes > 0 {
+          try Prune.reclaimIfNeeded(UInt64(reclaimBytes), sourceVM)
+        }
+      } else if sourceVM.isStackedVM || sourceVM.isStackedCachedImage {
+        let clonedVM = try localStorage.open(newName)
+        // A stacked clone owns only its writable overlay locally, but that
+        // overlay may grow to the full guest-visible disk block layout at
+        // runtime. Reclaim against the clone so it is not pruned itself.
+        let unallocatedBytes = try clonedVM.diskSizeBytes() - clonedVM.allocatedSizeBytes()
+        let reclaimBytes = min(unallocatedBytes, Int(pruneLimit) * 1024 * 1024 * 1024)
+        if reclaimBytes > 0 {
+          try Prune.reclaimIfNeeded(UInt64(reclaimBytes), clonedVM)
+        }
       }
     }, onCancel: {
       try? FileManager.default.removeItem(at: tmpVMDir.baseURL)

Sources/tart/Commands/Get.swift

@@ -31,7 +31,7 @@ struct Get: AsyncParsableCommand {
       OS: vmConfig.os,
       CPU: vmConfig.cpuCount,
       Memory: memorySizeInMb,
-      Disk: HumanReadableByteCount(try vmDir.sizeBytes()) { $0 / 1000 / 1000 / 1000 },
+      Disk: HumanReadableByteCount(try vmDir.diskSizeBytes()) { $0 / 1000 / 1000 / 1000 },
       DiskFormat: vmConfig.diskFormat.rawValue,
       Size: HumanReadableByteCount(try vmDir.allocatedSizeBytes()) {
         String(format: "%.3f", Float($0) / 1000 / 1000 / 1000)
@@ -40,7 +40,6 @@ struct Get: AsyncParsableCommand {
       Running: try vmDir.running(),
       State: try vmDir.state().rawValue
     )
-
     print(format.renderSingle(info))
   }
 }

Sources/tart/Commands/Import.swift

@@ -31,6 +31,11 @@ struct Import: AsyncParsableCommand {
     print("importing...")
     try tmpVMDir.importFromArchive(path: path)
 
+    if tmpVMDir.isStackedVM || tmpVMDir.isStackedCachedImage {
+      try? FileManager.default.removeItem(at: tmpVMDir.baseURL)
+      throw RuntimeError.ImportFailed("importing stacked VMs is not supported yet")
+    }
+
     try await withTaskCancellationHandler(operation: {
       // Acquire a global lock
       let lock = try FileLock(lockURL: Config().tartHomeDir)

Sources/tart/Commands/List.swift

@@ -42,7 +42,7 @@ struct List: AsyncParsableCommand {
         try VMInfo(
           Source: "local",
           Name: name,
-          Disk: HumanReadableByteCount(try vmDir.sizeBytes()) { $0 / 1000 / 1000 / 1000 },
+          Disk: HumanReadableByteCount(try vmDir.diskSizeBytes()) { $0 / 1000 / 1000 / 1000 },
           Size: HumanReadableByteCount(try vmDir.allocatedSizeBytes()) { $0 / 1000 / 1000 / 1000 },
           Accessed: formatAccessDate(try vmDir.accessDate()),
           Running: vmDir.running(),
@@ -56,7 +56,7 @@ struct List: AsyncParsableCommand {
         try VMInfo(
           Source: "OCI",
           Name: name,
-          Disk: HumanReadableByteCount(try vmDir.sizeBytes()) { $0 / 1000 / 1000 / 1000 },
+          Disk: HumanReadableByteCount(try vmDir.diskSizeBytes()) { $0 / 1000 / 1000 / 1000 },
           Size: HumanReadableByteCount(try vmDir.allocatedSizeBytes()) { $0 / 1000 / 1000 / 1000 },
           Accessed: formatAccessDate(try vmDir.accessDate()),
           Running: vmDir.running(),

Sources/tart/Commands/Push.swift

@@ -69,7 +69,7 @@ struct Push: AsyncParsableCommand {
       let references = remoteNamesForRegistry.map{ $0.reference.value }
 
       let pushedRemoteName: RemoteName
-      // If we're pushing a local OCI VM, check if points to an already existing registry manifest
+      // If we're pushing a cached remote image, check if it points to an existing registry manifest
       // and if so, only upload manifests (without config, disk and NVRAM) to the user-specified references
       if let remoteName = try? RemoteName(localName) {
         pushedRemoteName = try await lightweightPushToRegistry(
@@ -78,17 +78,18 @@ struct Push: AsyncParsableCommand {
           references: references
         )
       } else {
-        pushedRemoteName = try await localVMDir.pushToRegistry(
+        let pushedImage = try await localVMDir.pushToRegistry(
           registry: registry,
           references: references,
           chunkSizeMb: chunkSize,
           concurrency: concurrency,
           labels: parseLabels()
         )
+        pushedRemoteName = pushedImage.name
+
         // Populate the local cache (if requested)
         if populateCache {
-          let expectedPushedVMDir = try ociStorage.create(pushedRemoteName)
-          try localVMDir.clone(to: expectedPushedVMDir, generateMAC: false)
+          try ociStorage.populate(pushedImage.name, from: localVMDir, manifest: pushedImage.manifest)
         }
       }
 
@@ -102,7 +103,7 @@ struct Push: AsyncParsableCommand {
   }
 
   func lightweightPushToRegistry(registry: Registry, remoteName: RemoteName, references: [String]) async throws -> RemoteName {
-    // Is the local OCI VM already present in the registry?
+    // Is the cached remote image already present in the registry?
     let digest = try VMStorageOCI().digest(remoteName)
 
     let (remoteManifest, _) = try await registry.pullManifest(reference: digest)

Sources/tart/Commands/Run.swift

@@ -455,10 +455,15 @@ struct Run: AsyncParsableCommand {
       let provisioning = try provisioningOpts.map { try GuestProvisioningOptions($0) }
     #endif
 
+    // Keep these values alive while the VM runs. Some additional disks own a
+    // lock that protects their temporary backing files from Config.gc().
+    let additionalDisks = try additionalDisks()
+    defer { withExtendedLifetime(additionalDisks) {} }
+
     vm = try VM(
       vmDir: vmDir,
       network: userSpecifiedNetwork(vmDir: vmDir) ?? NetworkShared(),
-      additionalStorageDevices: try additionalDiskAttachments(),
+      additionalStorageDevices: additionalDisks.map(\.configuration),
       directorySharingDevices: directoryShares() + rosettaDirectoryShare(),
       serialPorts: serialPorts,
       suspendable: suspendable,
@@ -727,9 +732,9 @@ struct Run: AsyncParsableCommand {
     }
   }
 
-  func additionalDiskAttachments() throws -> [VZStorageDeviceConfiguration] {
+  func additionalDisks() throws -> [AdditionalDisk] {
     try disk.map {
-      try AdditionalDisk(parseFrom: $0).configuration
+      try AdditionalDisk(parseFrom: $0)
     }
   }
 
@@ -952,14 +957,32 @@ struct VMView: NSViewRepresentable {
 
 struct AdditionalDisk {
   let configuration: VZStorageDeviceConfiguration
+  // Retained for as long as the additional disk is attached, so Config.gc()
+  // cannot remove a temporary backing file or stacked-disk directory.
+  private let temporaryDiskLock: FileLock?
 
   init(parseFrom: String) throws {
     let (diskPath, readOnly, syncModeRaw, cachingModeRaw) = Self.parseOptions(parseFrom)
 
-    self.configuration = try Self.craft(diskPath, readOnly: readOnly, syncModeRaw: syncModeRaw, cachingModeRaw: cachingModeRaw)
+    self = try Self.craft(
+      diskPath,
+      readOnly: readOnly,
+      syncModeRaw: syncModeRaw,
+      cachingModeRaw: cachingModeRaw
+    )
+  }
+
+  private init(configuration: VZStorageDeviceConfiguration, temporaryDiskLock: FileLock? = nil) {
+    self.configuration = configuration
+    self.temporaryDiskLock = temporaryDiskLock
   }
 
-  static func craft(_ diskPath: String, readOnly diskReadOnly: Bool, syncModeRaw: String, cachingModeRaw: String) throws -> VZStorageDeviceConfiguration {
+  private static func craft(
+    _ diskPath: String,
+    readOnly diskReadOnly: Bool,
+    syncModeRaw: String,
+    cachingModeRaw: String
+  ) throws -> AdditionalDisk {
     let diskURL = URL(string: diskPath)
 
     if (["nbd", "nbds", "nbd+unix", "nbds+unix"].contains(diskURL?.scheme)) {
@@ -974,7 +997,7 @@ struct AdditionalDisk {
         synchronizationMode: try VZDiskSynchronizationMode(syncModeRaw)
       )
 
-      return VZVirtioBlockDeviceConfiguration(attachment: nbdAttachment)
+      return AdditionalDisk(configuration: VZVirtioBlockDeviceConfiguration(attachment: nbdAttachment))
     }
 
     // Expand the tilde (~) since at this point we're dealing with a local path,
@@ -1005,13 +1028,33 @@ struct AdditionalDisk {
       let blockAttachment = try VZDiskBlockDeviceStorageDeviceAttachment(fileHandle: FileHandle(fileDescriptor: fd, closeOnDealloc: true),
                                                                          readOnly: diskReadOnly, synchronizationMode: try VZDiskSynchronizationMode(syncModeRaw))
 
-      return VZVirtioBlockDeviceConfiguration(attachment: blockAttachment)
+      return AdditionalDisk(configuration: VZVirtioBlockDeviceConfiguration(attachment: blockAttachment))
     }
 
     // Support remote VM names in --disk command-line argument
     if let remoteName = try? RemoteName(diskPath) {
       let vmDir = try VMStorageOCI().open(remoteName)
 
+      if vmDir.isStackedCachedImage {
+        // A cached stacked image has no writable top overlay. Create one in a
+        // disposable directory for this additional-disk attachment.
+        let temporaryVMDir = try VMDirectory.temporary()
+        try FileManager.default.copyItem(at: vmDir.configURL, to: temporaryVMDir.configURL)
+        try FileManager.default.copyItem(at: vmDir.nvramURL, to: temporaryVMDir.nvramURL)
+        try FileManager.default.copyItem(at: vmDir.manifestURL, to: temporaryVMDir.manifestURL)
+        let lock = try FileLock(lockURL: temporaryVMDir.baseURL)
+        try lock.lock()
+        let stack = try temporaryVMDir.diskImageStack()
+        try stack.createWritableOverlay()
+        let attachment = try stack.makeAttachment(
+          readOnly: diskReadOnly,
+          cachingMode: try VZDiskImageCachingMode(cachingModeRaw) ?? .automatic,
+          synchronizationMode: try VZDiskImageSynchronizationMode(syncModeRaw)
+        )
+
+        return AdditionalDisk(configuration: VZVirtioBlockDeviceConfiguration(attachment: attachment), temporaryDiskLock: lock)
+      }
+
       // Unfortunately, VZDiskImageStorageDeviceAttachment does not support
       // FileHandle, so we can't easily clone the disk, open it and unlink(2)
       // to simplify the garbage collection, so use an intermediate directory.
@@ -1024,7 +1067,7 @@ struct AdditionalDisk {
 
       let diskImageAttachment = try VZDiskImageStorageDeviceAttachment(url: clonedDiskURL, readOnly: diskReadOnly)
 
-      return VZVirtioBlockDeviceConfiguration(attachment: diskImageAttachment)
+      return AdditionalDisk(configuration: VZVirtioBlockDeviceConfiguration(attachment: diskImageAttachment), temporaryDiskLock: lock)
     }
 
     // Error out if the disk is locked by the host (e.g. it was mounted in Finder),
@@ -1040,7 +1083,7 @@ struct AdditionalDisk {
       synchronizationMode: try VZDiskImageSynchronizationMode(syncModeRaw)
     )
 
-    return VZVirtioBlockDeviceConfiguration(attachment: diskImageAttachment)
+    return AdditionalDisk(configuration: VZVirtioBlockDeviceConfiguration(attachment: diskImageAttachment))
   }
 
   static func parseOptions(_ parseFrom: String) -> (String, Bool, String, String) {

Sources/tart/Commands/Set.swift

@@ -39,6 +39,14 @@ struct Set: AsyncParsableCommand {
 
   func run() async throws {
     let vmDir = try VMStorageLocal().open(name)
+
+    // Replacing disk.img would leave a stacked VM with both disk.img and
+    // overlay.asif, which is not a supported local layout. Reject before
+    // saving any other requested configuration changes.
+    if disk != nil, vmDir.isStackedVM {
+      throw ValidationError("--disk is not supported for VMs with a stacked disk")
+    }
+
     var vmConfig = try VMConfig(fromURL: vmDir.configURL)
 
     if let cpu = cpu {

Sources/tart/ContentStore.swift

@@ -39,15 +39,48 @@ struct ContentStore {
     return targetURL.deletingLastPathComponent().appendingPathComponent(".\(UUID().uuidString).tmp")
   }
 
-  /// Returns a validated cache hit. Corrupt files are treated as misses so a
-  /// later pull can safely rebuild them.
-  func existingContentURL(for contentDigest: String) throws -> URL? {
+  /// Returns a stable staging path so an interrupted registry pull can resume
+  /// reconstructing this content entry on a later attempt.
+  func resumableContentURL(for contentDigest: String) throws -> URL {
+    let targetURL = try contentURL(for: contentDigest)
+
+    return targetURL.deletingLastPathComponent().appendingPathComponent(".\(targetURL.lastPathComponent).partial")
+  }
+
+  /// Returns a stable lock file for serializing reconstruction of one content
+  /// entry. The file is intentionally retained; flock state lives on the file
+  /// descriptor and disappears when the owning process exits.
+  func lockURL(for contentDigest: String) throws -> URL {
+    let targetURL = try contentURL(for: contentDigest)
+
+    let lockURL = targetURL.deletingLastPathComponent().appendingPathComponent(".\(targetURL.lastPathComponent).lock")
+    if !FileManager.default.fileExists(atPath: lockURL.path) {
+      _ = FileManager.default.createFile(atPath: lockURL.path, contents: nil)
+    }
+
+    return lockURL
+  }
+
+  /// Returns a digest-addressed entry without rereading it. Pull verifies
+  /// content hashes before accepting a cache hit; clone only needs a cheap
+  /// structural check, like Tart's existing disk.img path.
+  func contentURLIfPresent(for contentDigest: String) throws -> URL? {
     let url = try contentURL(for: contentDigest)
 
     guard FileManager.default.fileExists(atPath: url.path) else {
       return nil
     }
 
+    return url
+  }
+
+  /// Returns a validated cache hit. Corrupt files are treated as misses so a
+  /// later pull can safely rebuild them.
+  func existingContentURL(for contentDigest: String) throws -> URL? {
+    guard let url = try contentURLIfPresent(for: contentDigest) else {
+      return nil
+    }
+
     guard try Digest.hash(url) == contentDigest else {
       return nil
     }
@@ -57,22 +90,29 @@ struct ContentStore {
 
   /// Move a fully reconstructed temporary file into the cache after verifying
   /// its semantic identity. The caller should create the temporary file with
-  /// `temporaryContentURL(for:)` so rename stays on the same filesystem.
+  /// temporaryContentURL(for:) or resumableContentURL(for:) so rename stays on
+  /// the same filesystem.
   func install(_ temporaryURL: URL, contentDigest: String) throws -> URL {
     let actualDigest = try Digest.hash(temporaryURL)
     guard actualDigest == contentDigest else {
       throw ContentStoreError.contentDigestMismatch(expected: contentDigest, actual: actualDigest)
     }
 
     let targetURL = try contentURL(for: contentDigest)
+    let lock = try FileLock(lockURL: baseURL)
+    try lock.lock()
+    defer { try? lock.unlock() }
 
     if let existingURL = try existingContentURL(for: contentDigest) {
       try? FileManager.default.removeItem(at: temporaryURL)
       return existingURL
     }
 
-    try? FileManager.default.removeItem(at: targetURL)
-    try FileManager.default.moveItem(at: temporaryURL, to: targetURL)
+    if FileManager.default.fileExists(atPath: targetURL.path) {
+      _ = try FileManager.default.replaceItemAt(targetURL, withItemAt: temporaryURL)
+    } else {
+      try FileManager.default.moveItem(at: temporaryURL, to: targetURL)
+    }
 
     return targetURL
   }