Finish remaining features for DiskImageKit

Finish implementing remaining tart commands to
support stacked disk image.

Add integration test and benchmarking.

Author: yzhuang-oai

Sources/tart/Commands/Clone.swift

@@ -31,7 +31,7 @@ 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")
+  @Flag(help: "create a stacked disk using a remote flat macOS OCI image as an immutable base")
   var base: Bool = false
 
   @Option(help: ArgumentHelp("limit automatic pruning to n gigabytes", valueName: "n"))
@@ -56,15 +56,30 @@ struct Clone: AsyncParsableCommand {
       guard remoteName != nil else {
         throw ValidationError("--base requires an OCI source")
       }
+      try DiskImageStack.requireSupport()
     }
 
     if let remoteName, try !ociStorage.hasUsableCachedRecordForClone(remoteName, requireManifest: base) {
       // 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)
+
+      // A stacked clone cannot run without DiskImageKit. Inspect the small
+      // manifest first so macOS 26 does not download all parent layers only
+      // to fail while creating the local writable overlay.
+      if !base {
+        let (manifest, _) = try await registry.pullManifest(reference: remoteName.reference.value)
+        if manifest.layers.contains(where: { $0.mediaType == asifOverlayMediaType }) {
+          try DiskImageStack.requireSupport()
+        }
+      }
+
       try await ociStorage.pull(remoteName, registry: registry, concurrency: concurrency, deduplicate: deduplicate)
     }
 
     let sourceVM = try VMStorageHelper.open(sourceName)
+    if sourceVM.isStackedLocal || sourceVM.isStackedOCIRecord {
+      try DiskImageStack.requireSupport()
+    }
     let tmpVMDir = try VMDirectory.temporary()
 
     // Lock the temporary VM directory to prevent it's garbage collection

Sources/tart/Commands/Import.swift

@@ -21,6 +21,9 @@ struct Import: AsyncParsableCommand {
 
     // Create a temporary VM directory to which we will load the export file
     let tmpVMDir = try VMDirectory.temporary()
+    defer {
+      try? FileManager.default.removeItem(at: tmpVMDir.baseURL)
+    }
 
     // Lock the temporary VM directory to prevent it's garbage collection
     // while we're running
@@ -30,10 +33,8 @@ struct Import: AsyncParsableCommand {
     // Populate the temporary VM directory with the export file contents
     print("importing...")
     try tmpVMDir.importFromArchive(path: path)
-
-    if tmpVMDir.isStackedLocal || tmpVMDir.isStackedOCIRecord {
-      try? FileManager.default.removeItem(at: tmpVMDir.baseURL)
-      throw RuntimeError.ImportFailed("importing stacked VMs is not supported yet")
+    guard tmpVMDir.initialized else {
+      throw RuntimeError.ImportFailed("archive does not contain a runnable VM")
     }
 
     try await withTaskCancellationHandler(operation: {

Sources/tart/Commands/Prune.swift

@@ -81,27 +81,36 @@ struct Prune: AsyncParsableCommand {
   }
 
   static func pruneSpaceBudget(prunableStorages: [PrunableStorage], spaceBudgetBytes: UInt64) throws {
-    let prunables: [Prunable] = try prunableStorages
-      .flatMap { try $0.prunables() }
-      .sorted { try $0.accessDate() > $1.accessDate() }
-
-    var spaceBudgetBytes = spaceBudgetBytes
-    var prunablesToDelete: [Prunable] = []
-
-    for prunable in prunables {
-      let prunableSizeBytes = UInt64(try prunable.allocatedSizeBytes())
+    while true {
+      let prunables: [Prunable] = try prunableStorages
+        .flatMap { try $0.prunables() }
+        .sorted { try $0.accessDate() > $1.accessDate() }
+
+      var remainingBudgetBytes = spaceBudgetBytes
+      var prunablesToDelete: [Prunable] = []
+
+      for prunable in prunables {
+        let prunableSizeBytes = UInt64(try prunable.allocatedSizeBytes())
+
+        if prunableSizeBytes <= remainingBudgetBytes {
+          // Don't mark for deletion as
+          // there's a budget available
+          remainingBudgetBytes -= prunableSizeBytes
+        } else {
+          // Mark for deletion
+          prunablesToDelete.append(prunable)
+        }
+      }
 
-      if prunableSizeBytes <= spaceBudgetBytes {
-        // Don't mark for deletion as
-        // there's a budget available
-        spaceBudgetBytes -= prunableSizeBytes
-      } else {
-        // Mark for deletion
-        prunablesToDelete.append(prunable)
+      guard !prunablesToDelete.isEmpty else {
+        return
       }
-    }
 
-    try prunablesToDelete.forEach { try $0.delete() }
+      try prunablesToDelete.forEach { try $0.delete() }
+      // Deleting one stacked OCI record can change which remaining record
+      // owns shared immutable content. Rebuild the candidates until the
+      // retained set actually fits the requested budget.
+    }
   }
 
   static func reclaimIfNeeded(_ requiredBytes: UInt64, _ initiator: Prunable? = nil) throws {

Sources/tart/Commands/Run.swift

@@ -1018,7 +1018,9 @@ struct AdditionalDisk {
         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)
+        try ContentStore().withPruneLock {
+          try FileManager.default.copyItem(at: vmDir.manifestURL, to: temporaryVMDir.manifestURL)
+        }
         let stack = try temporaryVMDir.diskImageStack()
         try stack.createWritableOverlay()
         let attachment = try stack.makeAttachment(

Sources/tart/Commands/Set.swift

@@ -41,6 +41,10 @@ struct Set: AsyncParsableCommand {
     let vmDir = try VMStorageLocal().open(name)
     var vmConfig = try VMConfig(fromURL: vmDir.configURL)
 
+    if disk != nil && vmDir.isStackedLocal {
+      throw ValidationError("--disk is not supported for images with a read-only base")
+    }
+
     if let cpu = cpu {
       try vmConfig.setCPU(cpuCount: Int(cpu))
     }

Sources/tart/ContentStore.swift

@@ -16,14 +16,29 @@ struct ContentStore {
   private static let digestPrefix = "sha256:"
 
   let baseURL: URL
+  private let pruneLockURL: URL
 
   init() throws {
     try self.init(baseURL: Config().tartCacheDir.appendingPathComponent("content", isDirectory: true))
   }
 
   init(baseURL: URL) throws {
     self.baseURL = baseURL
+    pruneLockURL = baseURL.appendingPathComponent(".gc.lock")
     try FileManager.default.createDirectory(at: baseURL, withIntermediateDirectories: true)
+    if !FileManager.default.fileExists(atPath: pruneLockURL.path) {
+      _ = FileManager.default.createFile(atPath: pruneLockURL.path, contents: Data())
+    }
+  }
+
+  /// Serializes reference publication with the final reference check and
+  /// deletion of immutable cache entries across Tart processes.
+  func withPruneLock<T>(_ body: () throws -> T) throws -> T {
+    let lock = try FileLock(lockURL: pruneLockURL)
+    try lock.lock()
+    defer { try? lock.unlock() }
+
+    return try body()
   }
 
   func contentURL(for contentDigest: String) throws -> URL {
@@ -65,16 +80,18 @@ struct ContentStore {
     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.
+  /// Returns an immutable digest-addressed entry without rereading it. Files
+  /// are verified when installed and when deciding whether a pull is a cache
+  /// hit; normal clone/run/push paths trust the store like Tart's disk.img.
   func contentURLIfPresent(for contentDigest: String) throws -> URL? {
     let url = try contentURL(for: contentDigest)
 
     guard FileManager.default.fileExists(atPath: url.path) else {
       return nil
     }
 
+    try url.updateAccessDate()
+
     return url
   }
 
@@ -92,6 +109,66 @@ struct ContentStore {
     return url
   }
 
+  /// Returns immutable content files that no retained OCI record or local VM
+  /// references. Callers may prune these like other cache entries.
+  func prunables(excluding referencedContentDigests: Swift.Set<String>) throws -> [URL] {
+    let sha256URL = baseURL.appendingPathComponent("sha256", isDirectory: true)
+    guard let enumerator = FileManager.default.enumerator(
+      at: sha256URL,
+      includingPropertiesForKeys: [.isRegularFileKey],
+      options: [.skipsSubdirectoryDescendants]
+    ) else {
+      return []
+    }
+
+    return try enumerator.compactMap { element in
+      guard let url = element as? URL,
+            try url.resourceValues(forKeys: [.isRegularFileKey]).isRegularFile == true else {
+        return nil
+      }
+
+      let contentDigest = "sha256:\(url.lastPathComponent)"
+      guard (try? validatedDigestHex(contentDigest)) != nil,
+            !referencedContentDigests.contains(contentDigest) else {
+        return nil
+      }
+
+      return url
+    }
+  }
+
+  /// Reconstructs one immutable disk file from transport chunks unless a
+  /// validated copy is already present under its whole-file digest.
+  func pullContent(
+    registry: Registry,
+    chunks: [OCIManifestLayer],
+    contentDigest: String,
+    concurrency: UInt,
+    progress: Progress
+  ) async throws -> URL {
+    if let existingURL = try existingContentURL(for: contentDigest) {
+      progress.completedUnitCount += chunks.reduce(0) { $0 + Int64($1.size) }
+      return existingURL
+    }
+
+    let temporaryURL = try temporaryContentURL(for: contentDigest)
+
+    do {
+      try await DiskV2.pull(
+        registry: registry,
+        diskLayers: chunks,
+        diskURL: temporaryURL,
+        concurrency: concurrency,
+        progress: progress
+      )
+
+      return try install(temporaryURL, contentDigest: contentDigest)
+    } catch {
+      try? FileManager.default.removeItem(at: temporaryURL)
+      throw error
+    }
+  }
+
   /// Move a fully reconstructed temporary file into the cache after verifying
   /// its semantic identity. The caller should create the temporary file with
   /// temporaryContentURL(for:) or partialContentURL(for:) so rename stays on

Sources/tart/DiskImageStack.swift

@@ -48,6 +48,22 @@ struct DiskImageStack: DiskAttachmentSource {
   let blockSize: UInt64
   let blockCount: UInt64
 
+  static var isSupported: Bool {
+    #if canImport(DiskImageKit)
+      if #available(macOS 27.0, *) {
+        return true
+      }
+    #endif
+
+    return false
+  }
+
+  static func requireSupport() throws {
+    guard isSupported else {
+      throw DiskImageStackError.unavailable
+    }
+  }
+
   /// Reads a disk image's current geometry without resolving or validating a
   /// whole stack. This is used for the VM's private writable overlay, whose
   /// size may be newer than the pinned immutable parent manifest.

Sources/tart/OCI/Digest.swift

@@ -7,6 +7,8 @@ enum DigestError: Error {
 }
 
 class Digest {
+  private static let fileBufferSize = 4 * 1024 * 1024
+
   var hash: SHA256 = SHA256()
 
   func update(_ data: Data) {
@@ -22,7 +24,23 @@ class Digest {
   }
 
   static func hash(_ url: URL) throws -> String {
-    hash(try Data(contentsOf: url))
+    let file = try FileHandle(forReadingFrom: url)
+    defer { try? file.close() }
+
+    let digest = Digest()
+    while try autoreleasepool(invoking: {
+      guard let data = try file.read(upToCount: fileBufferSize), !data.isEmpty else {
+        return false
+      }
+
+      digest.update(data)
+      return true
+    }) {
+      // Keep the pool scoped to one read so large disk hashing does not
+      // retain Foundation's temporary buffers until the command exits.
+    }
+
+    return digest.finalize()
   }
 
   static func hash(_ url: URL, offset: UInt64, size: UInt64) throws -> String {
@@ -36,20 +54,31 @@ class Digest {
       throw DigestError.InvalidOffset
     }
 
-    if (offset + size) > fileSize {
+    if size > fileSize - offset {
       throw DigestError.InvalidSize
     }
 
-    // Read a chunk of size ``size`` at offset ``offset``
-    // and calculate it's digest
+    // Read the requested range incrementally and calculate its digest.
     let fh = try FileHandle(forReadingFrom: url)
-    defer { try! fh.close() }
+    defer { try? fh.close() }
 
     try fh.seek(toOffset: offset)
 
-    let data = try fh.read(upToCount: Int(size))!
+    let digest = Digest()
+    var remaining = size
+    while remaining > 0 {
+      try autoreleasepool {
+        let count = Int(min(UInt64(fileBufferSize), remaining))
+        guard let data = try fh.read(upToCount: count), !data.isEmpty else {
+          throw DigestError.InvalidSize
+        }
+
+        digest.update(data)
+        remaining -= UInt64(data.count)
+      }
+    }
 
-    return hash(data)
+    return digest.finalize()
   }
 }
 

Sources/tart/OCI/Manifest.swift

@@ -25,6 +25,10 @@ let uncompressedContentDigestAnnotation = "org.cirruslabs.tart.uncompressed-cont
 let diskFileContentDigestAnnotation = "org.cirruslabs.tart.disk-file-content-digest"
 let diskFileChunkCountAnnotation = "org.cirruslabs.tart.disk-file-chunk-count"
 
+// Keep stacked manifests comfortably within the limit accepted by common OCI
+// registries. Flat manifests predate this feature and keep their behavior.
+let maxStackedManifestSizeBytes = 4 * 1024 * 1024
+
 /// The OCI-layer descriptors whose Tart disk chunks reconstruct one complete
 /// base disk or ASIF overlay.
 struct TartDiskFileGroup: Equatable {
@@ -186,6 +190,22 @@ struct OCIManifest: Codable, Equatable {
   func diskBlockCount() -> UInt64? {
     annotations?[diskBlockCountAnnotation].flatMap(UInt64.init)
   }
+
+  /// Reject an impractically deep stack before submitting its manifest to a
+  /// registry. The limit applies only to the new representation so existing
+  /// flat images remain backwards-compatible.
+  func validateStackedSizeForPush(maxBytes: Int = maxStackedManifestSizeBytes) throws {
+    guard case .stacked = try tartDiskRepresentation() else {
+      return
+    }
+
+    let size = try toJSON().count
+    guard size <= maxBytes else {
+      throw OCIManifestValidationError.invalidLayout(
+        "stacked manifest is \(size) bytes, exceeding the \(maxBytes)-byte limit"
+      )
+    }
+  }
 }
 
 struct OCIConfig: Codable {