Add optional Virtio memory balloon support

Sources/tart/Commands/Run.swift

@@ -282,6 +282,20 @@ struct Run: AsyncParsableCommand {
   @Flag(help: ArgumentHelp("Disable the pointer"))
   var noPointer: Bool = false
 
+  @Option(help: ArgumentHelp(
+    "Ask the guest to reduce its memory footprint to the target size in megabytes using the memory balloon device (e.g. --balloon-target-memory=8192)",
+    discussion: """
+    Requires the VM to have the memory balloon device enabled first:
+
+    tart set <name> --memory-balloon true
+
+    The reclaim is best-effort: the guest OS must support the virtio-balloon device and may release
+    less memory than requested, or none at all. Linux guests generally support it, while macOS guests
+    may show limited or no practical memory reduction. The guest still sees the full configured
+    memory size — this is not dynamic memory expansion nor transparent memory overcommit.
+    """, valueName: "MB"))
+  var balloonTargetMemory: UInt64?
+
   @Flag(help: ArgumentHelp("Disable the keyboard"))
   var noKeyboard: Bool = false
 
@@ -374,6 +388,15 @@ struct Run: AsyncParsableCommand {
       }
     }
 
+    if let balloonTargetMemory = balloonTargetMemory {
+      if suspendable {
+        throw ValidationError("--balloon-target-memory cannot be used with --suspendable")
+      }
+
+      let config = try VMConfig.init(fromURL: vmDir.configURL)
+      try Self.validateBalloonTargetMemory(balloonTargetMemory, vmConfig: config)
+    }
+
     #if arch(arm64) && compiler(>=6.4)
       if provisioningOpts != nil {
         if #unavailable(macOS 27) {
@@ -394,6 +417,31 @@ struct Run: AsyncParsableCommand {
     }
   }
 
+  static func validateBalloonTargetMemory(_ targetMemoryMB: UInt64, vmConfig: VMConfig) throws {
+    if !vmConfig.memoryBalloon {
+      throw ValidationError("--balloon-target-memory requires the VM to have the memory balloon device enabled,"
+        + " enable it via \"tart set <name> --memory-balloon true\"")
+    }
+
+    let (targetMemoryBytes, overflown) = targetMemoryMB.multipliedReportingOverflow(by: 1024 * 1024)
+    if overflown || targetMemoryBytes > vmConfig.memorySize {
+      throw ValidationError("--balloon-target-memory (\(targetMemoryMB) MB) cannot exceed the VM's"
+        + " configured memory size of \(vmConfig.memorySize / 1024 / 1024) MB")
+    }
+
+    var minimumAllowedMemorySize = VZVirtualMachineConfiguration.minimumAllowedMemorySize
+    if vmConfig.os == .darwin {
+      // macOS guests additionally have a minimum supported memory size
+      // dictated by the restore image they were created from, similarly
+      // to how "tart set --memory" restricts the configured memory size
+      minimumAllowedMemorySize = max(minimumAllowedMemorySize, vmConfig.memorySizeMin)
+    }
+    if targetMemoryBytes < minimumAllowedMemorySize {
+      throw ValidationError("--balloon-target-memory (\(targetMemoryMB) MB) is too small,"
+        + " it should be at least \(minimumAllowedMemorySize / 1024 / 1024) MB")
+    }
+  }
+
   @MainActor
   func runOnMainThread() throws {
     let localStorage = try VMStorageLocal()
@@ -546,6 +594,26 @@ struct Run: AsyncParsableCommand {
           throw error
         }
 
+        if let balloonTargetMemory = balloonTargetMemory {
+          try vm!.setBalloonTargetMemory(balloonTargetMemory * 1024 * 1024)
+          print("asked the guest to reduce its memory footprint to \(balloonTargetMemory) MB"
+            + " (best-effort, requires guest OS support for the virtio-balloon device)")
+
+          // Keep re-applying the balloon target, since a target set before
+          // the guest's virtio-balloon driver has probed the device is lost
+          // when the guest resets the device while booting (the same applies
+          // to guest reboots)
+          Task {
+            while !Task.isCancelled {
+              try? await Task.sleep(nanoseconds: 15_000_000_000)
+
+              if vm!.virtualMachine.state == VZVirtualMachine.State.running {
+                try? vm!.setBalloonTargetMemory(balloonTargetMemory * 1024 * 1024)
+              }
+            }
+          }
+        }
+
         if let vncImpl = vncImpl {
           let vncURL = try await vncImpl.waitForURL(netBridged: !netBridged.isEmpty)
 

Sources/tart/Commands/Set.swift

@@ -14,6 +14,17 @@ struct Set: AsyncParsableCommand {
   @Option(help: "VM memory size in megabytes")
   var memory: UInt64?
 
+  @Option(help: ArgumentHelp("Attach a virtio memory balloon device to the VM (true or false)", discussion: """
+  When enabled, the host can ask the guest to release some of its memory on a best-effort basis
+  (see "tart run --help" on the --balloon-target-memory option for more details).
+
+  Note that this is not dynamic memory expansion nor transparent memory overcommit: the guest
+  always sees the configured memory size, and the reclaim only works when the guest OS supports
+  the virtio-balloon device. Linux guests generally support it, while macOS guests may show
+  limited or no practical memory reduction.
+  """, valueName: "true|false"))
+  var memoryBalloon: Bool?
+
   @Option(help: "VM display resolution in a format of WIDTHxHEIGHT[pt|px]. For example, 1200x800, 1200x800pt or 1920x1080px. Units are treated as hints and default to \"pt\" (points) for macOS VMs and \"px\" (pixels) for Linux VMs when not specified.")
   var display: VMDisplayConfig?
 
@@ -49,6 +60,10 @@ struct Set: AsyncParsableCommand {
       try vmConfig.setMemory(memorySize: memory * 1024 * 1024)
     }
 
+    if let memoryBalloon = memoryBalloon {
+      vmConfig.memoryBalloon = memoryBalloon
+    }
+
     if let display = display {
       if (display.width > 0) {
         vmConfig.display.width = display.width

Sources/tart/VM.swift

@@ -254,6 +254,16 @@ class VM: NSObject, VZVirtualMachineDelegate, ObservableObject {
     }
   }
 
+  @MainActor
+  func setBalloonTargetMemory(_ targetMemoryBytes: UInt64) throws {
+    guard let balloonDevice = virtualMachine.memoryBalloonDevices.first as? VZVirtioTraditionalMemoryBalloonDevice else {
+      throw RuntimeError.VMConfigurationError("VM has no memory balloon device configured,"
+        + " enable it via \"tart set \(name) --memory-balloon true\"")
+    }
+
+    balloonDevice.targetVirtualMachineMemorySize = targetMemoryBytes
+  }
+
   @MainActor
   func connect(toPort: UInt32) async throws -> VZVirtioSocketConnection {
     guard let socketDevice = virtualMachine.socketDevices.first else {
@@ -328,6 +338,48 @@ class VM: NSObject, VZVirtualMachineDelegate, ObservableObject {
     noTrackpad: Bool = false,
     noPointer: Bool = false,
     noKeyboard: Bool = false
+  ) throws -> VZVirtualMachineConfiguration {
+    let configuration = try buildConfiguration(diskURL: diskURL,
+                                               nvramURL: nvramURL, vmConfig: vmConfig,
+                                               network: network, additionalStorageDevices: additionalStorageDevices,
+                                               directorySharingDevices: directorySharingDevices,
+                                               serialPorts: serialPorts,
+                                               suspendable: suspendable,
+                                               nested: nested,
+                                               audio: audio,
+                                               clipboard: clipboard,
+                                               sync: sync,
+                                               caching: caching,
+                                               noTrackpad: noTrackpad,
+                                               noPointer: noPointer,
+                                               noKeyboard: noKeyboard
+    )
+
+    try configuration.validate()
+
+    return configuration
+  }
+
+  // Builds the virtual machine configuration without validating it, since
+  // validation requires the "com.apple.security.virtualization" entitlement
+  // that unit tests don't have
+  static func buildConfiguration(
+    diskURL: URL,
+    nvramURL: URL,
+    vmConfig: VMConfig,
+    network: Network = NetworkShared(),
+    additionalStorageDevices: [VZStorageDeviceConfiguration],
+    directorySharingDevices: [VZDirectorySharingDeviceConfiguration],
+    serialPorts: [VZSerialPortConfiguration],
+    suspendable: Bool = false,
+    nested: Bool = false,
+    audio: Bool = true,
+    clipboard: Bool = true,
+    sync: VZDiskImageSynchronizationMode = .full,
+    caching: VZDiskImageCachingMode? = nil,
+    noTrackpad: Bool = false,
+    noPointer: Bool = false,
+    noKeyboard: Bool = false
   ) throws -> VZVirtualMachineConfiguration {
     let configuration = VZVirtualMachineConfiguration()
 
@@ -423,6 +475,17 @@ class VM: NSObject, VZVirtualMachineDelegate, ObservableObject {
       configuration.entropyDevices = [VZVirtioEntropyDeviceConfiguration()]
     }
 
+    // Memory balloon
+    //
+    // Allows the host to reclaim memory from the guest on a best-effort
+    // basis, provided that the guest OS supports the virtio-balloon device.
+    //
+    // Skipped for suspendable VMs (similarly to the entropy device above)
+    // to not interfere with the save/restore support.
+    if vmConfig.memoryBalloon && !suspendable {
+      configuration.memoryBalloonDevices = [VZVirtioTraditionalMemoryBalloonDeviceConfiguration()]
+    }
+
     // Directory sharing devices
     configuration.directorySharingDevices = directorySharingDevices
 
@@ -444,8 +507,6 @@ class VM: NSObject, VZVirtualMachineDelegate, ObservableObject {
     // Socket device
     configuration.socketDevices = [VZVirtioSocketDeviceConfiguration()]
 
-    try configuration.validate()
-
     return configuration
   }
 

Sources/tart/VMConfig.swift

@@ -26,6 +26,7 @@ enum CodingKeys: String, CodingKey {
   case display
   case displayRefit
   case diskFormat
+  case memoryBalloon
 
   // macOS-specific keys
   case ecid
@@ -66,6 +67,7 @@ struct VMConfig: Codable {
   var display: VMDisplayConfig = VMDisplayConfig()
   var displayRefit: Bool?
   var diskFormat: DiskImageFormat = .raw
+  var memoryBalloon: Bool = false
 
   init(
     platform: Platform,
@@ -140,6 +142,7 @@ struct VMConfig: Codable {
     displayRefit = try container.decodeIfPresent(Bool.self, forKey: .displayRefit)
     let diskFormatString = try container.decodeIfPresent(String.self, forKey: .diskFormat) ?? "raw"
     diskFormat = DiskImageFormat(rawValue: diskFormatString) ?? .raw
+    memoryBalloon = try container.decodeIfPresent(Bool.self, forKey: .memoryBalloon) ?? false
   }
 
   func encode(to encoder: Encoder) throws {
@@ -159,6 +162,12 @@ struct VMConfig: Codable {
       try container.encode(displayRefit, forKey: .displayRefit)
     }
     try container.encode(diskFormat.rawValue, forKey: .diskFormat)
+    // Only write the key when the balloon is enabled, so that configurations
+    // of VMs that don't use this feature remain byte-identical to those
+    // produced by older Tart versions
+    if memoryBalloon {
+      try container.encode(memoryBalloon, forKey: .memoryBalloon)
+    }
   }
 
   mutating func setCPU(cpuCount: Int) throws {

Tests/TartTests/MemoryBalloonTests.swift

@@ -0,0 +1,143 @@
+import Virtualization
+import XCTest
+
+@testable import tart
+
+final class MemoryBalloonTests: XCTestCase {
+  // Configurations created by older Tart versions don't have
+  // the "memoryBalloon" key and should default to a disabled
+  // memory balloon device
+  func testDisabledByDefaultWhenDecodingLegacyConfig() throws {
+    let legacyConfigJSON = """
+    {
+      "version": 1,
+      "os": "linux",
+      "arch": "arm64",
+      "cpuCountMin": 1,
+      "cpuCount": 1,
+      "memorySizeMin": 536870912,
+      "memorySize": 536870912,
+      "macAddress": "5a:00:00:00:00:01"
+    }
+    """
+
+    let vmConfig = try VMConfig(fromJSON: legacyConfigJSON.data(using: .utf8)!)
+
+    XCTAssertFalse(vmConfig.memoryBalloon)
+  }
+
+  func testDisabledByDefaultWhenCreatingNewConfig() throws {
+    let vmConfig = VMConfig(platform: Linux(), cpuCountMin: 1, memorySizeMin: 512 * 1024 * 1024)
+
+    XCTAssertFalse(vmConfig.memoryBalloon)
+
+    // The key shouldn't even be present in the resulting JSON to keep
+    // the configurations of VMs that don't use this feature identical
+    // to those produced by older Tart versions
+    let encodedConfig = String(data: try vmConfig.toJSON(), encoding: .utf8)!
+    XCTAssertFalse(encodedConfig.contains("memoryBalloon"))
+  }
+
+  func testPersistsWhenEnabled() throws {
+    var vmConfig = VMConfig(platform: Linux(), cpuCountMin: 1, memorySizeMin: 512 * 1024 * 1024)
+    vmConfig.memoryBalloon = true
+
+    let roundtrippedVMConfig = try VMConfig(fromJSON: try vmConfig.toJSON())
+
+    XCTAssertTrue(roundtrippedVMConfig.memoryBalloon)
+  }
+
+  func testMemoryBalloonSetArgumentParsing() throws {
+    XCTAssertEqual(try tart.Set.parse(["vm", "--memory-balloon", "true"]).memoryBalloon, true)
+    XCTAssertEqual(try tart.Set.parse(["vm", "--memory-balloon", "false"]).memoryBalloon, false)
+    XCTAssertNil(try tart.Set.parse(["vm"]).memoryBalloon)
+    XCTAssertThrowsError(try tart.Set.parse(["vm", "--memory-balloon", "yes"]))
+  }
+
+  func testBalloonTargetMemoryValidation() throws {
+    var vmConfig = VMConfig(platform: Linux(), cpuCountMin: 1, memorySizeMin: 4096 * 1024 * 1024)
+
+    // Balloon device is not enabled
+    XCTAssertThrowsError(try Run.validateBalloonTargetMemory(2048, vmConfig: vmConfig))
+
+    vmConfig.memoryBalloon = true
+
+    // Target exceeds the configured memory size
+    XCTAssertThrowsError(try Run.validateBalloonTargetMemory(8192, vmConfig: vmConfig))
+
+    // Target multiplication by 1 MB overflows UInt64
+    XCTAssertThrowsError(try Run.validateBalloonTargetMemory(UInt64.max, vmConfig: vmConfig))
+
+    // Target is too small to be safe
+    XCTAssertThrowsError(try Run.validateBalloonTargetMemory(1, vmConfig: vmConfig))
+
+    // Sane target
+    XCTAssertNoThrow(try Run.validateBalloonTargetMemory(2048, vmConfig: vmConfig))
+
+    // Target that is exactly the configured memory size (fully deflated balloon)
+    XCTAssertNoThrow(try Run.validateBalloonTargetMemory(4096, vmConfig: vmConfig))
+  }
+
+  func testBalloonTargetMemoryValidationRespectsDarwinMinimum() throws {
+    // A macOS guest whose restore image requires 4096 MB of memory at minimum
+    var vmConfig = VMConfig(platform: Linux(), cpuCountMin: 1, memorySizeMin: 4096 * 1024 * 1024)
+    vmConfig.os = .darwin
+    vmConfig.memoryBalloon = true
+    try vmConfig.setMemory(memorySize: 8192 * 1024 * 1024)
+
+    // Target below the restore image's minimum supported memory size
+    XCTAssertThrowsError(try Run.validateBalloonTargetMemory(2048, vmConfig: vmConfig))
+
+    // Target at the restore image's minimum supported memory size
+    XCTAssertNoThrow(try Run.validateBalloonTargetMemory(4096, vmConfig: vmConfig))
+
+    // The same minimum doesn't apply to Linux guests, similarly
+    // to how "tart set --memory" doesn't restrict them
+    vmConfig.os = .linux
+    XCTAssertNoThrow(try Run.validateBalloonTargetMemory(2048, vmConfig: vmConfig))
+  }
+
+  func testBalloonDeviceOnlyConfiguredWhenEnabled() throws {
+    // Disabled by default
+    XCTAssertEqual(try craftConfiguration().memoryBalloonDevices.count, 0)
+
+    // Configured when enabled
+    let memoryBalloonDevices = try craftConfiguration(memoryBalloon: true).memoryBalloonDevices
+    XCTAssertEqual(memoryBalloonDevices.count, 1)
+    XCTAssertTrue(memoryBalloonDevices.first is VZVirtioTraditionalMemoryBalloonDeviceConfiguration)
+
+    // Not configured for suspendable VMs, even when enabled
+    XCTAssertEqual(try craftConfiguration(memoryBalloon: true, suspendable: true).memoryBalloonDevices.count, 0)
+  }
+
+  private func craftConfiguration(memoryBalloon: Bool = false, suspendable: Bool = false) throws -> VZVirtualMachineConfiguration {
+    let tmpDir = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
+    try FileManager.default.createDirectory(at: tmpDir, withIntermediateDirectories: true)
+    defer { try? FileManager.default.removeItem(at: tmpDir) }
+
+    let nvramURL = tmpDir.appendingPathComponent("nvram.bin")
+    _ = try VZEFIVariableStore(creatingVariableStoreAt: nvramURL)
+
+    let diskURL = tmpDir.appendingPathComponent("disk.img")
+    FileManager.default.createFile(atPath: diskURL.path, contents: nil)
+    let diskFileHandle = try FileHandle(forWritingTo: diskURL)
+    try diskFileHandle.truncate(atOffset: 512 * 1024 * 1024)
+    try diskFileHandle.close()
+
+    var vmConfig = VMConfig(platform: Linux(), cpuCountMin: 1, memorySizeMin: 1024 * 1024 * 1024)
+    vmConfig.memoryBalloon = memoryBalloon
+
+    // Note: VM.buildConfiguration() is used here instead of VM.craftConfiguration(),
+    // because the latter additionally validates the configuration, which requires
+    // the "com.apple.security.virtualization" entitlement that tests don't have
+    return try VM.buildConfiguration(
+      diskURL: diskURL,
+      nvramURL: nvramURL,
+      vmConfig: vmConfig,
+      additionalStorageDevices: [],
+      directorySharingDevices: [],
+      serialPorts: [],
+      suspendable: suspendable
+    )
+  }
+}