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

Resolves #5376

Author: osy

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

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
     }
 }

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?
@@ -138,7 +140,6 @@ final class UTMAppleVirtualMachine: UTMVirtualMachine {
     }
 
     private func _start(options: UTMVirtualMachineStartOptions) async throws {
-        try await createAppleVM()
         let boot = await config.system.boot
         try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) -> Void in
             vmQueue.async {
@@ -173,8 +174,14 @@ final class UTMAppleVirtualMachine: UTMVirtualMachine {
         }
         state = .starting
         do {
+            let isSuspended = await registryEntry.isSuspended
             try await beginAccessingResources()
-            try await _start(options: options)
+            try await createAppleVM()
+            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
@@ -194,6 +201,7 @@ final class UTMAppleVirtualMachine: UTMVirtualMachine {
         } catch {
             await stopAccesingResources()
             state = .stopped
+            try? await deleteSnapshot()
             throw error
         }
     }
@@ -328,16 +336,109 @@ 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
+        }
+        await registryEntry.setIsSuspended(false)
+        try FileManager.default.removeItem(at: vmSavedStateURL)
     }
 
+    #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)
+            try await _resume()
+        } catch {
+            state = .stopped
+            throw error
+        }
+        state = .started
+        try await deleteSnapshot(name: name)
+        #else
+        throw UTMAppleVirtualMachineError.operationNotAvailable
+        #endif
     }
 
     private func _resume() async throws {
@@ -388,6 +489,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 +633,7 @@ extension UTMAppleVirtualMachine: VZVirtualMachineDelegate {
     func guestDidStop(_ virtualMachine: VZVirtualMachine) {
         vmQueue.async { [self] in
             apple = nil
+            saveSnapshotError = nil
         }
         sharedDirectoriesChanged = nil
         Task { @MainActor in

Services/UTMQemuVirtualMachine.swift

@@ -465,13 +465,13 @@ extension UTMQemuVirtualMachine {
     }
 
     private func _deleteSnapshot(name: String) async throws {
+        await registryEntry.setIsSuspended(false)
         if let monitor = await monitor { // if QEMU is running
             let result = try await monitor.qemuDeleteSnapshot(name)
             if result.localizedCaseInsensitiveContains("Error") {
                 throw UTMQemuVirtualMachineError.qemuError(result)
             }
         }
-        await registryEntry.setIsSuspended(false)
     }
 
     func deleteSnapshot(name: String? = nil) async throws {

Services/UTMVirtualMachine.swift

@@ -498,6 +498,7 @@ extension UTMVirtualMachine {
         Task {
             do {
                 try await resume()
+                try? await deleteSnapshot(name: nil)
             } catch {
                 delegate?.virtualMachine(self, didErrorWithMessage: error.localizedDescription)
             }