Complete stacked disk command support

Author: yzhuang-oai

Sources/tart/Commands/Clone.swift

@@ -56,15 +56,28 @@ 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)
+
+      // Fail before pulling disk content when this host cannot create a writable stacked disk.
+      if !stacked {
+        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.isStackedVM || sourceVM.isStackedCachedImage {
+      try DiskImageStack.requireSupport()
+    }
     let tmpVMDir = try VMDirectory.temporary()
 
     // Lock the temporary VM directory to prevent it's garbage collection
@@ -126,7 +139,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
     }
 

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