Complete stacked disk command support (#1315)

Sources/tart/Commands/Clone.swift

@@ -56,15 +56,37 @@ struct Clone: AsyncParsableCommand {
       guard remoteName != nil else {
         throw ValidationError("--stacked requires a remote image")
       }
+      try DiskImageStack.requireSupport()
     }
 
     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)
+      var resolvedManifest: (manifest: OCIManifest, data: Data)?
+
+      // Fail before pulling disk content when this host cannot create a writable stacked disk.
+      if !stacked {
+        let (manifest, manifestData) = try await registry.pullManifest(reference: remoteName.reference.value)
+        if manifest.layers.contains(where: { $0.mediaType == asifOverlayMediaType }) {
+          try DiskImageStack.requireSupport()
+        }
+        resolvedManifest = (manifest, manifestData)
+      }
+
+      try await ociStorage.pull(
+        remoteName,
+        registry: registry,
+        concurrency: concurrency,
+        deduplicate: deduplicate,
+        requireManifest: stacked,
+        resolvedManifest: resolvedManifest
+      )
     }
 
     let sourceVM = try VMStorageHelper.open(sourceName)
+    if sourceVM.isStackedVM || sourceVM.isStackedCachedImage {
+      try DiskImageStack.requireSupport()
+    }
     let tmpVMDir = try VMDirectory.temporary()
 
     // Lock the temporary VM directory to prevent it's garbage collection
@@ -126,7 +148,7 @@ struct Clone: AsyncParsableCommand {
         }
       }
     }, onCancel: {
-      try? FileManager.default.removeItem(at: tmpVMDir.baseURL)
+      try? tmpVMDir.removeFromDisk()
     })
   }
 }

Sources/tart/VMDirectory+OCI.swift

@@ -114,7 +114,7 @@ extension VMDirectory {
 
   /// Reconstructs one complete immutable base disk or published ASIF overlay
   /// from its Tart disk chunks, unless the shared content store already has a
-  /// verified copy.
+  /// size-matching copy.
   private func pullDiskFile(
     registry: Registry,
     group: TartDiskFileGroup,
@@ -133,7 +133,10 @@ extension VMDirectory {
     try lock.lock()
     defer { try? lock.unlock() }
 
-    if let existingURL = try contentStore.existingContentURL(for: contentDigest) {
+    if let existingURL = try contentStore.contentURLIfPresent(for: contentDigest),
+       let actualSize = UInt64(exactly: try existingURL.sizeBytes()),
+       let expectedSize = group.uncompressedSize(),
+       actualSize == expectedSize {
       progress.completedUnitCount += group.chunks.reduce(0) { $0 + Int64($1.size) }
       return existingURL
     }
@@ -189,7 +192,6 @@ extension VMDirectory {
     var annotations = diskAnnotations
     annotations[uploadTimeAnnotation] = Date().toISO()
     manifest.annotations = annotations
-
     // Manifest
     for reference in references {
       defaultLogger.appendNewLine("pushing manifest for \(reference)...")
@@ -228,6 +230,15 @@ extension VMDirectory {
     }
 
     let localManifest = try OCIManifest(fromJSON: Data(contentsOf: manifestURL))
+    // pushToRegistry() reads config.json before reaching this point. Closing
+    // that read descriptor can release the caller's fcntl PID lock, so take a
+    // fresh lock before hashing, uploading, and inspecting the writable overlay.
+    let stackedDiskLock = try lock()
+    guard try stackedDiskLock.trylock() else {
+      throw RuntimeError.VMIsRunning(name)
+    }
+    defer { try? stackedDiskLock.unlock() }
+
     let inheritedGroups: [TartDiskFileGroup]
     switch try localManifest.tartDiskRepresentation() {
     case .flat(let base) where base.contentDigest != nil:
@@ -250,26 +261,13 @@ extension VMDirectory {
       ))
     }
 
-    // Keep the snapshot out of startup GC while this potentially long push
-    // hashes, uploads, and inspects it.
-    let frozenOverlayDirectory = try VMDirectory.temporary()
-    let frozenOverlayLock = try FileLock(lockURL: frozenOverlayDirectory.baseURL)
-    try frozenOverlayLock.lock()
-    defer {
-      try? frozenOverlayLock.unlock()
-      try? FileManager.default.removeItem(at: frozenOverlayDirectory.baseURL)
-    }
-
-    let frozenOverlayURL = frozenOverlayDirectory.baseURL.appendingPathComponent("overlay.asif")
-    try FileManager.default.copyItem(at: overlayURL, to: frozenOverlayURL)
-
-    let overlaySize = try FileManager.default.attributesOfItem(atPath: frozenOverlayURL.path)[.size] as! Int64
+    let overlaySize = try FileManager.default.attributesOfItem(atPath: overlayURL.path)[.size] as! Int64
     defaultLogger.appendNewLine("pushing overlay...")
     let progress = Progress(totalUnitCount: overlaySize)
     ProgressObserver(progress).log(defaultLogger)
-    let contentDigest = try Digest.hash(frozenOverlayURL)
+    let contentDigest = try Digest.hash(overlayURL)
     let chunks = try await DiskV2.push(
-      diskURL: frozenOverlayURL,
+      diskURL: overlayURL,
       mediaType: asifOverlayMediaType,
       registry: registry,
       chunkSizeMb: chunkSizeMb,
@@ -278,7 +276,7 @@ extension VMDirectory {
     )
     layers.append(contentsOf: annotatedChunks(chunks, kind: .asifOverlay, contentDigest: contentDigest))
 
-    let blockLayout = try DiskImageStack.diskImageBlockLayout(at: frozenOverlayURL)
+    let blockLayout = try DiskImageStack.diskImageBlockLayout(at: overlayURL)
     let diskSize = blockLayout.blockSize.multipliedReportingOverflow(by: blockLayout.blockCount)
     guard !diskSize.overflow else {
       throw DiskImageStackError.invalidBlockLayout("stacked disk block layout overflows UInt64")
@@ -316,6 +314,8 @@ extension VMDirectory {
       return group.chunks
     }
 
+    // Rebuilding transport blobs republishes this file under the pinned
+    // whole-file digest, so validate the cached bytes at this boundary.
     guard let contentURL = try contentStore.existingContentURL(for: contentDigest) else {
       throw RuntimeError.VMMissingFiles("stacked VM is missing cached disk content \(contentDigest)")
     }

Sources/tart/VMDirectory.swift

@@ -236,17 +236,29 @@ struct VMDirectory: Prunable {
     contentStore: ContentStore? = nil
   ) throws {
     if isStackedVM {
-      guard try state() == .Stopped else {
+      // Resolve the stack before taking the config.json PID lock. Reading
+      // config.json after acquiring an fcntl lock would release that lock
+      // when the read file descriptor is closed.
+      let stack = try diskImageStack(contentStore: contentStore)
+      let lock = try lock()
+      guard try lock.trylock() else {
+        throw RuntimeError.VMConfigurationError("VM \"\(name)\" must be stopped before resizing its disk")
+      }
+      defer { try? lock.unlock() }
+
+      // Holding the PID lock proves that the VM is not running. A saved state
+      // file is the remaining suspended state that must also reject resize.
+      guard !FileManager.default.fileExists(atPath: stateURL.path) else {
         throw RuntimeError.VMConfigurationError("VM \"\(name)\" must be stopped before resizing its disk")
       }
 
-      let stack = try diskImageStack(contentStore: contentStore)
       let desiredSizeBytes = UInt64(sizeGB) * 1000 * 1000 * 1000
       guard desiredSizeBytes.isMultiple(of: stack.blockSize) else {
         throw RuntimeError.InvalidDiskSize("new disk size must align to the stacked disk block size")
       }
 
-      try stack.growWritableOverlay(toBlockCount: desiredSizeBytes / stack.blockSize)
+      let desiredBlockCount = desiredSizeBytes / stack.blockSize
+      try stack.growWritableOverlay(toBlockCount: desiredBlockCount)
       return
     }
 

Sources/tart/VMStorageOCI.swift

@@ -56,11 +56,20 @@ class VMStorageOCI: PrunableStorage {
   /// repairing it. Standalone images keep Tart's existing structural cache-hit
   /// behavior; stacked cached images additionally need every immutable disk file in
   /// the shared content store.
-  func hasCompleteCachedImage(_ name: RemoteName, manifest: OCIManifest) throws -> Bool {
+  func hasCompleteCachedImage(
+    _ name: RemoteName,
+    manifest: OCIManifest,
+    requireManifest: Bool = false
+  ) throws -> Bool {
     guard exists(name) else {
       return false
     }
 
+    let vmDir = VMDirectory(baseURL: vmURL(name))
+    if requireManifest && !FileManager.default.fileExists(atPath: vmDir.manifestURL.path) {
+      return false
+    }
+
     guard let missingGroups = try missingStackedDiskFileGroups(for: manifest) else {
       return true
     }
@@ -71,12 +80,17 @@ class VMStorageOCI: PrunableStorage {
   /// The lock-free pull fast path is only useful for a tag that already
   /// points at this digest. New or retargeted tags validate once after taking
   /// the host lock instead of hashing a large stack twice.
-  func hasCompleteLinkedImage(_ name: RemoteName, digestName: RemoteName, manifest: OCIManifest) throws -> Bool {
+  func hasCompleteLinkedImage(
+    _ name: RemoteName,
+    digestName: RemoteName,
+    manifest: OCIManifest,
+    requireManifest: Bool = false
+  ) throws -> Bool {
     guard exists(name), linked(from: name, to: digestName) else {
       return false
     }
 
-    return try hasCompleteCachedImage(digestName, manifest: manifest)
+    return try hasCompleteCachedImage(digestName, manifest: manifest, requireManifest: requireManifest)
   }
 
   /// Bytes that this pull may need to materialize locally. For stacked images
@@ -361,7 +375,14 @@ class VMStorageOCI: PrunableStorage {
     return result
   }
 
-  func pull(_ name: RemoteName, registry: Registry, concurrency: UInt, deduplicate: Bool) async throws {
+  func pull(
+    _ name: RemoteName,
+    registry: Registry,
+    concurrency: UInt,
+    deduplicate: Bool,
+    requireManifest: Bool = false,
+    resolvedManifest: (manifest: OCIManifest, data: Data)? = nil
+  ) async throws {
     OpenTelemetry.instance.contextProvider.activeSpan?.setAttribute(
       key: "oci.image-name",
       value: .string(name.description)
@@ -369,12 +390,23 @@ class VMStorageOCI: PrunableStorage {
 
     defaultLogger.appendNewLine("pulling manifest...")
 
-    let (manifest, manifestData) = try await registry.pullManifest(reference: name.reference.value)
+    let (manifest, manifestData): (OCIManifest, Data)
+    if let resolvedManifest {
+      manifest = resolvedManifest.manifest
+      manifestData = resolvedManifest.data
+    } else {
+      (manifest, manifestData) = try await registry.pullManifest(reference: name.reference.value)
+    }
 
     let digestName = RemoteName(host: name.host, namespace: name.namespace,
                                 reference: Reference(digest: Digest.hash(manifestData)))
 
-    if try hasCompleteLinkedImage(name, digestName: digestName, manifest: manifest) {
+    if try hasCompleteLinkedImage(
+      name,
+      digestName: digestName,
+      manifest: manifest,
+      requireManifest: requireManifest
+    ) {
       // optimistically check if we need to do anything at all before locking
       defaultLogger.appendNewLine("\(digestName) image is already cached and linked!")
       return
@@ -398,12 +430,21 @@ class VMStorageOCI: PrunableStorage {
       throw CancellationError()
     }
 
-    if try !hasCompleteCachedImage(digestName, manifest: manifest) {
+    let digestVMDir = VMDirectory(baseURL: vmURL(digestName))
+    if requireManifest,
+       !FileManager.default.fileExists(atPath: digestVMDir.manifestURL.path),
+       try hasCompleteCachedImage(digestName, manifest: manifest) {
+      // Old Tart versions cached standalone OCI images without manifest.json.
+      // A stacked clone needs the manifest to describe its immutable base, but
+      // the existing disk remains usable and must not be downloaded again.
+      try manifestData.write(to: digestVMDir.manifestURL, options: .atomic)
+    }
+
+    if try !hasCompleteCachedImage(digestName, manifest: manifest, requireManifest: requireManifest) {
       let span = OTel.shared.tracer.spanBuilder(spanName: "pull").setActive(true).startSpan()
       defer { span.end() }
 
       let tmpVMDir = try VMDirectory.temporaryDeterministic(key: name.description)
-      let digestVMDir = VMDirectory(baseURL: vmURL(digestName))
       let preserveExplicitlyPulledMark = digestVMDir.isExplicitlyPulled()
 
       // Open an existing VM directory corresponding to this name, if any,

Tests/TartTests/VMDirectoryDiskImageStackTests.swift

@@ -79,7 +79,37 @@ import XCTest
       XCTAssertTrue(FileManager.default.fileExists(atPath: fresh.overlayURL.path))
     }
 
-    func testResizeDiskGrowsWritableOverlay() throws {
+    func testStackedRemoteAdditionalDiskRetainsTemporaryVM() throws {
+      try withTemporaryTartHome {
+        let source = try flatSource()
+        let stacked = try temporaryVMDirectory()
+        try source.cloneAsStackedBase(to: stacked, generateMAC: false)
+
+        let storage = try VMStorageOCI()
+        let name = try RemoteName("example.com/org/image:latest")
+        let cachedImage = try storage.create(name)
+        try FileManager.default.copyItem(at: stacked.configURL, to: cachedImage.configURL)
+        try FileManager.default.copyItem(at: stacked.nvramURL, to: cachedImage.nvramURL)
+        try FileManager.default.copyItem(at: stacked.manifestURL, to: cachedImage.manifestURL)
+
+        do {
+          let additionalDisk = try AdditionalDisk(parseFrom: name.description)
+          let entries = try temporaryEntries()
+          XCTAssertEqual(entries.count, 1)
+          XCTAssertTrue(VMDirectory(baseURL: entries[0]).isStackedVM)
+
+          try Config().gc()
+          XCTAssertEqual(try temporaryEntries(), entries)
+
+          withExtendedLifetime(additionalDisk) {}
+        }
+
+        try Config().gc()
+        XCTAssertTrue(try temporaryEntries().isEmpty)
+      }
+    }
+
+    func testResizeDiskGrowsWritableOverlayAndPreservesParentGeometry() throws {
       let contentStore = try temporaryContentStore()
       let source = try flatSource()
       let stacked = try temporaryVMDirectory()

Tests/TartTests/VMStorageOCITests.swift

@@ -35,6 +35,8 @@ final class VMStorageOCITests: XCTestCase {
 
       XCTAssertTrue(try storage.hasUsableCachedImageForClone(name))
       XCTAssertFalse(try storage.hasUsableCachedImageForClone(name, requireManifest: true))
+      XCTAssertTrue(try storage.hasCompleteCachedImage(name, manifest: manifest))
+      XCTAssertFalse(try storage.hasCompleteCachedImage(name, manifest: manifest, requireManifest: true))
     }
   }