vm(apple): implement snapshot save/restore for macOS 14

Resolves #5376
utmapp · Aug 23, 2023 · 6d4ea1f · 6d4ea1f
1 parent 5967988
commit 6d4ea1f
Show file tree

Hide file tree

Showing 3 changed files with 121 additions and 4 deletions.
diff --git a/Configuration/QEMUConstant.swift b/Configuration/QEMUConstant.swift
@@ -390,6 +390,7 @@ enum QEMUPackageFileName: String {
     case debugLog = "debug.log"
     case efiVariables = "efi_vars.fd"
     case tpmData = "tpmdata"
+    case vmState = "vmstate"
 }
 
 // MARK: Supported features

diff --git a/Configuration/UTMAppleConfigurationBoot.swift b/Configuration/UTMAppleConfigurationBoot.swift
@@ -39,6 +39,7 @@ struct UTMAppleConfigurationBoot: Codable {
     var linuxCommandLine: String?
     var linuxInitialRamdiskURL: URL?
     var efiVariableStorageURL: URL?
+    var vmSavedStateURL: URL?
     var hasUefiBoot: Bool = false
 
     /// IPSW for installing macOS. Not saved.
@@ -78,6 +79,7 @@ struct UTMAppleConfigurationBoot: Codable {
         if let efiVariableStoragePath = try container.decodeIfPresent(String.self, forKey: .efiVariableStoragePath) {
             efiVariableStorageURL = dataURL.appendingPathComponent(efiVariableStoragePath)
         }
+        vmSavedStateURL = dataURL.appendingPathComponent(QEMUPackageFileName.vmState.rawValue)
     }
 
     init(for operatingSystem: OperatingSystem, linuxKernelURL: URL? = nil) throws {
@@ -189,6 +191,9 @@ extension UTMAppleConfigurationBoot {
             self.efiVariableStorageURL = efiVariableStorageURL
             urls.append(efiVariableStorageURL)
         }
+        let vmSavedStateURL = dataURL.appendingPathComponent(QEMUPackageFileName.vmState.rawValue)
+        self.vmSavedStateURL = vmSavedStateURL
+        urls.append(vmSavedStateURL)
         return urls
     }
 }
diff --git a/Services/UTMAppleVirtualMachine.swift b/Services/UTMAppleVirtualMachine.swift
@@ -100,6 +100,8 @@ final class UTMAppleVirtualMachine: UTMVirtualMachine {
     /// This variable MUST be synchronized by `vmQueue`
     private(set) var apple: VZVirtualMachine?
 
+    private var saveSnapshotError: Error?
+
     private var installProgress: Progress?
 
     private var progressObserver: NSKeyValueObservation?
@@ -173,8 +175,13 @@ final class UTMAppleVirtualMachine: UTMVirtualMachine {
         }
         state = .starting
         do {
+            let isSuspended = await registryEntry.isSuspended
             try await beginAccessingResources()
-            try await _start(options: options)
+            if isSuspended && !options.contains(.bootRecovery) {
+                try await restoreSnapshot()
+            } else {
+                try await _start(options: options)
+            }
             if #available(macOS 12, *) {
                 Task { @MainActor in
                     sharedDirectoriesChanged = config.sharedDirectoriesPublisher.sink { [weak self] newShares in
@@ -328,16 +335,108 @@ final class UTMAppleVirtualMachine: UTMVirtualMachine {
         }
     }
 
+    #if arch(arm64)
+    @available(macOS 14, *)
+    private func _saveSnapshot(url: URL) async throws {
+        try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
+            vmQueue.async {
+                guard let apple = self.apple else {
+                    continuation.resume(throwing: UTMAppleVirtualMachineError.operationNotAvailable)
+                    return
+                }
+                apple.saveMachineStateTo(url: url) { error in
+                    if let error = error {
+                        continuation.resume(throwing: error)
+                    } else {
+                        continuation.resume()
+                    }
+                }
+            }
+        }
+    }
+    #endif
+
     func saveSnapshot(name: String? = nil) async throws {
-        // FIXME: implement this
+        guard #available(macOS 14, *) else {
+            return
+        }
+        #if arch(arm64)
+        guard let vmSavedStateURL = await config.system.boot.vmSavedStateURL else {
+            return
+        }
+        if let saveSnapshotError = saveSnapshotError {
+            throw saveSnapshotError
+        }
+        if state == .started {
+            try await pause()
+        }
+        guard state == .paused else {
+            return
+        }
+        state = .saving
+        defer {
+            state = .paused
+        }
+        try await _saveSnapshot(url: vmSavedStateURL)
+        await registryEntry.setIsSuspended(true)
+        #endif
     }
 
     func deleteSnapshot(name: String? = nil) async throws {
-        // FIXME: implement this
+        guard let vmSavedStateURL = await config.system.boot.vmSavedStateURL else {
+            return
+        }
+        try FileManager.default.removeItem(at: vmSavedStateURL)
+        await registryEntry.setIsSuspended(false)
     }
 
+    #if arch(arm64)
+    @available(macOS 14, *)
+    private func _restoreSnapshot(url: URL) async throws {
+        try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
+            vmQueue.async {
+                guard let apple = self.apple else {
+                    continuation.resume(throwing: UTMAppleVirtualMachineError.operationNotAvailable)
+                    return
+                }
+                apple.restoreMachineStateFrom(url: url) { error in
+                    if let error = error {
+                        continuation.resume(throwing: error)
+                    } else {
+                        continuation.resume()
+                    }
+                }
+            }
+        }
+    }
+    #endif
+
     func restoreSnapshot(name: String? = nil) async throws {
-        // FIXME: implement this
+        guard #available(macOS 14, *) else {
+            throw UTMAppleVirtualMachineError.operationNotAvailable
+        }
+        #if arch(arm64)
+        guard let vmSavedStateURL = await config.system.boot.vmSavedStateURL else {
+            throw UTMAppleVirtualMachineError.operationNotAvailable
+        }
+        if state == .started {
+            try await stop(usingMethod: .force)
+        }
+        guard state == .stopped || state == .starting else {
+            throw UTMAppleVirtualMachineError.operationNotAvailable
+        }
+        state = .restoring
+        do {
+            try await _restoreSnapshot(url: vmSavedStateURL)
+        } catch {
+            state = .stopped
+            throw error
+        }
+        state = .started
+        try await deleteSnapshot(name: name)
+        #else
+        throw UTMAppleVirtualMachineError.operationNotAvailable
+        #endif
     }
 
     private func _resume() async throws {
@@ -388,6 +487,17 @@ final class UTMAppleVirtualMachine: UTMVirtualMachine {
         vmQueue.async { [self] in
             apple = VZVirtualMachine(configuration: vzConfig, queue: vmQueue)
             apple!.delegate = self
+            saveSnapshotError = nil
+            #if arch(arm64)
+            if #available(macOS 14, *) {
+                do {
+                    try vzConfig.validateSaveRestoreSupport()
+                } catch {
+                    // save this for later when we want to use snapshots
+                    saveSnapshotError = error
+                }
+            }
+            #endif
         }
     }
 
@@ -521,6 +631,7 @@ extension UTMAppleVirtualMachine: VZVirtualMachineDelegate {
     func guestDidStop(_ virtualMachine: VZVirtualMachine) {
         vmQueue.async { [self] in
             apple = nil
+            saveSnapshotError = nil
         }
         sharedDirectoriesChanged = nil
         Task { @MainActor in