From f83cba84afbb2aa8ec5e65da3854696f83ae9971 Mon Sep 17 00:00:00 2001 From: Yibo Zhuang Date: Mon, 17 Aug 2026 14:22:18 -0700 Subject: [PATCH] Support stacked VM archive import and export (#1314) --- Sources/tart/Commands/Import.swift | 11 +- Sources/tart/DiskImageStack.swift | 36 +++-- Sources/tart/VMDirectory+Archive.swift | 153 +++++++++++++++++- .../VMDirectoryDiskImageStackTests.swift | 144 +++++++++++++++++ Tests/TartTests/VMDirectoryLayoutTests.swift | 42 +++-- 5 files changed, 345 insertions(+), 41 deletions(-) diff --git a/Sources/tart/Commands/Import.swift b/Sources/tart/Commands/Import.swift index 7fc955c..a3779bd 100644 --- a/Sources/tart/Commands/Import.swift +++ b/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? tmpVMDir.removeFromDisk() + } // 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.isStackedVM || tmpVMDir.isStackedCachedImage { - 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: { @@ -50,7 +51,7 @@ struct Import: AsyncParsableCommand { try lock.unlock() }, onCancel: { - try? FileManager.default.removeItem(at: tmpVMDir.baseURL) + try? tmpVMDir.removeFromDisk() }) } } diff --git a/Sources/tart/DiskImageStack.swift b/Sources/tart/DiskImageStack.swift index 2a60e76..3b67e8e 100644 --- a/Sources/tart/DiskImageStack.swift +++ b/Sources/tart/DiskImageStack.swift @@ -45,6 +45,22 @@ struct DiskImageStack { 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 block layout 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. @@ -69,21 +85,7 @@ struct DiskImageStack { #if canImport(DiskImageKit) if #available(macOS 27.0, *) { let image = try DiskImage(opening: .open(url: url, mode: .readOnly)) - let matchesFormat = switch expectedFormat { - case .raw: - image.format == .raw - case .asif: - image.format == .asif - } - guard matchesFormat else { - throw DiskImageStackError.invalidDiskImage(url, "base disk format does not match") - } - guard image.layerType == nil, image.parentUUID == nil else { - throw DiskImageStackError.invalidDiskImage(url, "base disk must not be an overlay") - } - if expectedFormat == .asif && image.layerUUID == nil { - throw DiskImageStackError.invalidDiskImage(url, "ASIF base disk is missing a UUID") - } + try validateBase(image, at: url, expectedFormat: expectedFormat) return DiskImageBlockLayout( blockSize: UInt64(image.blockSize.rawValue), @@ -214,7 +216,7 @@ struct DiskImageStack { } let baseImage = try DiskImage(opening: .open(url: baseURL, mode: .readOnly)) - try validateBase(baseImage, at: baseURL, expectedFormat: baseFormat) + try Self.validateBase(baseImage, at: baseURL, expectedFormat: baseFormat) var image = baseImage @@ -239,7 +241,7 @@ struct DiskImageStack { } @available(macOS 27.0, *) - private func validateBase( + private static func validateBase( _ image: DiskImage, at url: URL, expectedFormat: DiskImageFormat diff --git a/Sources/tart/VMDirectory+Archive.swift b/Sources/tart/VMDirectory+Archive.swift index 969a84b..0d4ca12 100644 --- a/Sources/tart/VMDirectory+Archive.swift +++ b/Sources/tart/VMDirectory+Archive.swift @@ -1,3 +1,4 @@ +import Foundation import System import AppleArchive @@ -10,8 +11,14 @@ fileprivate let permissions = FilePermissions(rawValue: 0o644) // [2]: https://developer.apple.com/documentation/compression/algorithm/lzfse extension VMDirectory { func exportToArchive(path: String) throws { - guard !isStackedVM && !isStackedCachedImage else { - throw RuntimeError.ExportFailed("exporting stacked VMs is not supported yet") + let temporaryArchive = try stackedArchiveDirectoryIfNeeded() + let archiveSourceURL = temporaryArchive?.vmDirectory.baseURL ?? baseURL + + defer { + if let temporaryArchive { + try? temporaryArchive.lock.unlock() + try? temporaryArchive.vmDirectory.removeFromDisk() + } } guard let fileStream = ArchiveByteStream.fileStream( @@ -53,7 +60,7 @@ extension VMDirectory { return } - try encodeStream.writeDirectoryContents(archiveFrom: FilePath(baseURL.path), keySet: keySet) + try encodeStream.writeDirectoryContents(archiveFrom: FilePath(archiveSourceURL.path), keySet: keySet) } func importFromArchive(path: String) throws { @@ -96,5 +103,145 @@ extension VMDirectory { } _ = try ArchiveStream.process(readingFrom: decodeStream, writingTo: extractStream) + + if isStackedVM { + try restoreStackedArchive() + } } + + /// Builds a self-contained staging directory for a stacked archive, if this + /// directory currently resolves to a stacked VM or cached image. + private func stackedArchiveDirectoryIfNeeded() throws -> (vmDirectory: VMDirectory, lock: FileLock)? { + guard isStackedVM || isStackedCachedImage else { + return nil + } + try DiskImageStack.requireSupport() + + let contentStore = try ContentStore() + let archiveVMDir = try VMDirectory.temporary() + let archiveVMDirLock = try FileLock(lockURL: archiveVMDir.baseURL) + try archiveVMDirLock.lock() + + do { + let stagedSource: (isStackedVM: Bool, contentDigests: [String])? = try contentStore.withPruneLock { + () -> (isStackedVM: Bool, contentDigests: [String])? in + // OCI tags are mutable symlinks. Resolve one digest record while tag + // replacement and cached-image deletion are blocked, then copy every + // source-owned file before releasing the lock. + let sourceVMDir = VMDirectory(baseURL: baseURL.resolvingSymlinksInPath()) + guard sourceVMDir.isStackedVM || sourceVMDir.isStackedCachedImage else { + throw RuntimeError.ExportFailed("VM changed while preparing export, retry the command") + } + + let sourceIsStackedVM = sourceVMDir.isStackedVM + let sourceVMLock: PIDLock? + if sourceIsStackedVM { + let lock = try sourceVMDir.lock() + guard try lock.trylock() else { + throw RuntimeError.ExportFailed("VM \"\(sourceVMDir.name)\" must be stopped before export") + } + sourceVMLock = lock + + // Holding the PID lock proves that the VM is not running. A saved + // state file is the remaining suspended state that must reject export. + guard !FileManager.default.fileExists(atPath: sourceVMDir.stateURL.path) else { + try? lock.unlock() + throw RuntimeError.ExportFailed("VM \"\(sourceVMDir.name)\" must be stopped before export") + } + } else { + sourceVMLock = nil + } + defer { try? sourceVMLock?.unlock() } + + try FileManager.default.copyItem(at: sourceVMDir.configURL, to: archiveVMDir.configURL) + try FileManager.default.copyItem(at: sourceVMDir.nvramURL, to: archiveVMDir.nvramURL) + try FileManager.default.copyItem(at: sourceVMDir.manifestURL, to: archiveVMDir.manifestURL) + if sourceIsStackedVM { + try FileManager.default.copyItem(at: sourceVMDir.overlayURL, to: archiveVMDir.overlayURL) + } + + return (sourceIsStackedVM, try archiveVMDir.diskContentDigests()) + } + + guard let stagedSource else { + try archiveVMDirLock.unlock() + try archiveVMDir.removeFromDisk() + return nil + } + + if !stagedSource.isStackedVM { + try archiveVMDir.diskImageStack().createWritableOverlay() + } + + // The staged manifest is now an in-progress reference, so immutable + // content remains protected while these potentially large copies run + // without holding the global prune lock. + for contentDigest in stagedSource.contentDigests { + guard let sourceURL = try contentStore.existingContentURL(for: contentDigest) else { + throw RuntimeError.ExportFailed("VM is missing cached disk content \(contentDigest)") + } + + let destinationURL = try contentStore.contentURL( + for: contentDigest, + under: archiveContentStoreURL(in: archiveVMDir) + ) + try FileManager.default.createDirectory( + at: destinationURL.deletingLastPathComponent(), + withIntermediateDirectories: true + ) + try FileManager.default.copyItem(at: sourceURL, to: destinationURL) + } + + return (archiveVMDir, archiveVMDirLock) + } catch { + try? archiveVMDirLock.unlock() + try? archiveVMDir.removeFromDisk() + throw error + } + } + + /// Restores immutable files from an archive into the shared content store, + /// removes the archive-only payload, then validates the resulting stack. + private func restoreStackedArchive() throws { + try DiskImageStack.requireSupport() + + let contentStore = try ContentStore() + // The extracted manifest is already a reference; synchronize publication + // with a concurrent prune before installing its immutable content. + try contentStore.synchronizePublishedReferences() + for contentDigest in try diskContentDigests() { + if try contentStore.existingContentURL(for: contentDigest) != nil { + continue + } + + let archivedContentURL = try contentStore.contentURL( + for: contentDigest, + under: archiveContentStoreURL(in: self) + ) + guard FileManager.default.fileExists(atPath: archivedContentURL.path) else { + throw RuntimeError.ImportFailed("archive is missing disk content \(contentDigest)") + } + + let temporaryURL = try contentStore.temporaryContentURL(for: contentDigest) + do { + try FileManager.default.copyItem(at: archivedContentURL, to: temporaryURL) + _ = try contentStore.install(temporaryURL, contentDigest: contentDigest) + } catch { + try? FileManager.default.removeItem(at: temporaryURL) + throw error + } + } + + try FileManager.default.removeItem(at: archiveContentStoreURL(in: self)) + try? FileManager.default.removeItem(at: stateURL) + + // Opening the attachment validates the reconstructed immutable stack and + // imported writable overlay before the VM enters local storage. + _ = try diskImageStack().makeAttachment() + } + + private func archiveContentStoreURL(in vmDir: VMDirectory) -> URL { + vmDir.baseURL.appendingPathComponent("content", isDirectory: true) + } + } diff --git a/Tests/TartTests/VMDirectoryDiskImageStackTests.swift b/Tests/TartTests/VMDirectoryDiskImageStackTests.swift index 63ae535..42603bd 100644 --- a/Tests/TartTests/VMDirectoryDiskImageStackTests.swift +++ b/Tests/TartTests/VMDirectoryDiskImageStackTests.swift @@ -90,6 +90,128 @@ import XCTest let image = try DiskImage(opening: .open(url: stacked.overlayURL, mode: .readOnly)) XCTAssertEqual(image.blockCount, 1_000_000_000 / 512) XCTAssertEqual(try stacked.diskSizeBytes(), 1_000_000_000) + + let manifest = try OCIManifest(fromJSON: Data(contentsOf: stacked.manifestURL)) + XCTAssertEqual(manifest.diskBlockSize(), 512) + XCTAssertEqual(manifest.diskBlockCount(), 8) + + _ = try stacked.diskImageStack(contentStore: contentStore).makeAttachment() + } + + func testStackedArchiveRoundTripsImmutableContentAndOverlay() throws { + try withTemporaryTartHome { + let source = try flatSource() + let stacked = try temporaryVMDirectory() + try source.cloneAsStackedBase(to: stacked, generateMAC: false) + + let contentDigest = try Digest.hash(source.diskURL) + let contentStore = try ContentStore() + let archivedOverlayDigest = try Digest.hash(stacked.overlayURL) + let archiveURL = try temporaryDirectory().appendingPathComponent("stacked.tvm") + try stacked.exportToArchive(path: archiveURL.path) + + let cachedBaseURL = try XCTUnwrap(try contentStore.existingContentURL(for: contentDigest)) + // Import must repair a corrupt cache entry from the valid archive + // instead of discarding the archive copy as an apparent cache hit. + try Data("corrupt".utf8).write(to: cachedBaseURL) + XCTAssertNil(try contentStore.existingContentURL(for: contentDigest)) + + let imported = try temporaryVMDirectory() + try imported.importFromArchive(path: archiveURL.path) + + XCTAssertTrue(imported.isStackedVM) + XCTAssertEqual(try Digest.hash(imported.overlayURL), archivedOverlayDigest) + XCTAssertNotNil(try contentStore.existingContentURL(for: contentDigest)) + _ = try imported.diskImageStack().makeAttachment() + } + } + + func testStackedArchiveRejectsCorruptImmutableContent() throws { + try withTemporaryTartHome { + let source = try flatSource() + let stacked = try temporaryVMDirectory() + try source.cloneAsStackedBase(to: stacked, generateMAC: false) + + let contentDigest = try Digest.hash(source.diskURL) + let contentStore = try ContentStore() + let cachedBaseURL = try XCTUnwrap(try contentStore.contentURLIfPresent(for: contentDigest)) + try Data("corrupt".utf8).write(to: cachedBaseURL) + + let archiveURL = try temporaryDirectory().appendingPathComponent("stacked.tvm") + XCTAssertThrowsError(try stacked.exportToArchive(path: archiveURL.path)) { error in + guard case RuntimeError.ExportFailed(let message) = error else { + return XCTFail("unexpected error: \(error)") + } + XCTAssertEqual(message, "VM is missing cached disk content \(contentDigest)") + } + } + } + + func testStackedOCIArchiveSurvivesConcurrentRecordDeletion() throws { + try withTemporaryTartHome { + let source = try flatSource() + let stacked = try temporaryVMDirectory() + try source.cloneAsStackedBase(to: stacked, generateMAC: false) + + let manifest = try OCIManifest(fromJSON: Data(contentsOf: stacked.manifestURL)) + let storage = try VMStorageOCI() + let record = try storage.create(RemoteName( + host: "example.com", + namespace: "org/image", + reference: Reference(digest: try manifest.digest()) + )) + try FileManager.default.copyItem(at: stacked.configURL, to: record.configURL) + try FileManager.default.copyItem(at: stacked.nvramURL, to: record.nvramURL) + try FileManager.default.copyItem(at: stacked.manifestURL, to: record.manifestURL) + XCTAssertTrue(record.isStackedCachedImage) + + let archiveURL = try temporaryDirectory().appendingPathComponent("stacked-race.tvm") + let contentStore = try ContentStore() + let lockHeld = DispatchSemaphore(value: 0) + let releaseLock = DispatchSemaphore(value: 0) + let exportStarted = DispatchSemaphore(value: 0) + let exportFinished = DispatchSemaphore(value: 0) + let deletionStarted = DispatchSemaphore(value: 0) + let deletionFinished = DispatchSemaphore(value: 0) + + DispatchQueue.global().async { + try? contentStore.withPruneLock { + lockHeld.signal() + releaseLock.wait() + } + } + XCTAssertEqual(lockHeld.wait(timeout: .now() + 1), .success) + + // Queue export first so it is the next prune-lock waiter, then queue + // deletion behind it. Export must finish staging everything it needs + // before deletion can remove the source cached image. + DispatchQueue.global().async { + exportStarted.signal() + try? record.exportToArchive(path: archiveURL.path) + exportFinished.signal() + } + XCTAssertEqual(exportStarted.wait(timeout: .now() + 1), .success) + Thread.sleep(forTimeInterval: 0.1) + + DispatchQueue.global().async { + deletionStarted.signal() + try? record.delete() + deletionFinished.signal() + } + XCTAssertEqual(deletionStarted.wait(timeout: .now() + 1), .success) + XCTAssertEqual(exportFinished.wait(timeout: .now() + 0.1), .timedOut) + XCTAssertEqual(deletionFinished.wait(timeout: .now() + 0.1), .timedOut) + + releaseLock.signal() + XCTAssertEqual(exportFinished.wait(timeout: .now() + 5), .success) + XCTAssertEqual(deletionFinished.wait(timeout: .now() + 5), .success) + XCTAssertFalse(FileManager.default.fileExists(atPath: record.baseURL.path)) + + let imported = try temporaryVMDirectory() + try imported.importFromArchive(path: archiveURL.path) + XCTAssertTrue(imported.isStackedVM) + _ = try imported.diskImageStack().makeAttachment() + } } func testResolvesPublishedOverlayFromManifestAndCache() throws { @@ -168,6 +290,28 @@ import XCTest return try ContentStore(baseURL: url) } + private func temporaryEntries() throws -> [URL] { + try FileManager.default.contentsOfDirectory( + at: Config().tartTmpDir, + includingPropertiesForKeys: nil + ) + } + + private func withTemporaryTartHome(_ body: () throws -> Void) throws { + let home = try temporaryDirectory() + let previousHome = ProcessInfo.processInfo.environment["TART_HOME"] + setenv("TART_HOME", home.path, 1) + defer { + if let previousHome { + setenv("TART_HOME", previousHome, 1) + } else { + unsetenv("TART_HOME") + } + } + + try body() + } + private func temporaryVMDirectory() throws -> VMDirectory { VMDirectory(baseURL: try temporaryDirectory()) } diff --git a/Tests/TartTests/VMDirectoryLayoutTests.swift b/Tests/TartTests/VMDirectoryLayoutTests.swift index 364a5fd..5f8fcce 100644 --- a/Tests/TartTests/VMDirectoryLayoutTests.swift +++ b/Tests/TartTests/VMDirectoryLayoutTests.swift @@ -75,29 +75,39 @@ final class VMDirectoryLayoutTests: XCTestCase { ) } - func testStackedExportIsRejected() throws { + func testStackedCachedImageAccountingUsesManifestBlockLayout() throws { let vmDir = try temporaryVMDirectory() - try touch(vmDir.configURL) - try touch(vmDir.nvramURL) - try touch(vmDir.manifestURL) - try touch(vmDir.overlayURL) - let archiveURL = vmDir.baseURL.appendingPathComponent("export.tvm") - XCTAssertThrowsError(try vmDir.exportToArchive(path: archiveURL.path)) { error in - guard case RuntimeError.ExportFailed(let message) = error else { - return XCTFail("unexpected error: \(error)") - } - XCTAssertEqual(message, "exporting stacked VMs is not supported yet") + try Data("config".utf8).write(to: vmDir.configURL) + try Data("nvram".utf8).write(to: vmDir.nvramURL) + try stackedManifest(blockSize: 512, blockCount: 8).toJSON().write(to: vmDir.manifestURL) + + XCTAssertEqual( + try vmDir.sizeBytes(), + try vmDir.configURL.sizeBytes() + vmDir.nvramURL.sizeBytes() + ) + XCTAssertEqual( + try vmDir.allocatedSizeBytes(), + try vmDir.configURL.allocatedSizeBytes() + vmDir.nvramURL.allocatedSizeBytes() + ) + XCTAssertEqual(try vmDir.diskSizeBytes(), 4096) + } + + func testStackedArchiveRequiresMacOS27() throws { + if #available(macOS 27.0, *) { + throw XCTSkip("macOS 26 compatibility test") } - XCTAssertFalse(FileManager.default.fileExists(atPath: archiveURL.path)) - try FileManager.default.removeItem(at: vmDir.overlayURL) - XCTAssertTrue(vmDir.isStackedCachedImage) + let vmDir = try temporaryVMDirectory() + try Data("config".utf8).write(to: vmDir.configURL) + try Data("nvram".utf8).write(to: vmDir.nvramURL) + try Data("overlay".utf8).write(to: vmDir.overlayURL) + try stackedManifest(blockSize: 512, blockCount: 8).toJSON().write(to: vmDir.manifestURL) + let archiveURL = try temporaryVMDirectory().baseURL.appendingPathComponent("stacked.tvm") XCTAssertThrowsError(try vmDir.exportToArchive(path: archiveURL.path)) { error in - guard case RuntimeError.ExportFailed(let message) = error else { + guard case DiskImageStackError.unavailable = error else { return XCTFail("unexpected error: \(error)") } - XCTAssertEqual(message, "exporting stacked VMs is not supported yet") } XCTAssertFalse(FileManager.default.fileExists(atPath: archiveURL.path)) }