vmm: Allow preserving the source VM after snapshot

Extend the migration protocol with a `preserve_source` option so that a
source VM can be preserved. This benefits the snapshot case where the
offload daemon can now snapshot a VM without tearing it down.

Signed-off-by: Sebastien Boeuf <sboeuf@meta.com>
Assisted-by: Claude:claude-opus-4-8

vmm/src/api/mod.rs

@@ -590,6 +590,10 @@ pub struct VmSendMigrationData {
     /// Send memory across socket without copying
     #[serde(default)]
     pub local: bool,
+    /// Keep the source VM alive in a paused state once the migration is
+    /// complete.
+    #[serde(default)]
+    pub preserve_source: bool,
     /// The maximum downtime the migration aims for.
     ///
     /// Usually, on the order of a few hundred milliseconds.
@@ -621,7 +625,7 @@ pub struct VmSendMigrationData {
 
 impl VmSendMigrationData {
     pub const SYNTAX: &'static str = "VM send migration parameters \
-        \"destination_url=<url>[,local=on|off,\
+        \"destination_url=<url>[,local=on|off,preserve_source=on|off,\
         downtime_ms=<milliseconds>,timeout_s=<seconds>,\
         timeout_strategy=cancel|ignore,connections=<amount>,\
         tls_dir=<path>,memory_mode=precopy|postcopy]\"";
@@ -649,6 +653,7 @@ impl VmSendMigrationData {
         parser
             .add("destination_url")
             .add("local")
+            .add("preserve_source")
             .add("downtime_ms")
             .add("timeout_s")
             .add("timeout_strategy")
@@ -669,6 +674,11 @@ impl VmSendMigrationData {
             .map_err(VmSendMigrationConfigError::ParseError)?
             .unwrap_or(Toggle(false))
             .0;
+        let preserve_source = parser
+            .convert::<Toggle>("preserve_source")
+            .map_err(VmSendMigrationConfigError::ParseError)?
+            .unwrap_or(Toggle(false))
+            .0;
         let downtime_ms = match parser
             .convert::<u64>("downtime_ms")
             .map_err(VmSendMigrationConfigError::ParseError)?
@@ -718,6 +728,7 @@ impl VmSendMigrationData {
         let data = Self {
             destination_url,
             local,
+            preserve_source,
             downtime_ms,
             timeout_s,
             timeout_strategy,
@@ -786,6 +797,12 @@ impl VmSendMigrationData {
             }
         }
 
+        if self.preserve_source && !self.local {
+            return Err(VmSendMigrationConfigError::ValidationError(
+                "preserve_source option is only supported with the local option.".to_string(),
+            ));
+        }
+
         if let Some(tls_dir) = &self.tls_dir {
             validate_tls_dir(tls_dir, TlsEndpoint::Client).map_err(|e| {
                 VmSendMigrationConfigError::ValidationError(format!(
@@ -2408,6 +2425,7 @@ mod unit_tests {
             VmSendMigrationData {
                 destination_url: "tcp:192.168.1.1:8080".to_string(),
                 local: false,
+                preserve_source: false,
                 downtime_ms: NonZeroU64::new(150).unwrap(),
                 timeout_s: VmSendMigrationData::default_timeout_s(),
                 timeout_strategy: Default::default(),
@@ -2429,6 +2447,7 @@ mod unit_tests {
             VmSendMigrationData {
                 destination_url: "tcp:192.168.1.1:8080".to_string(),
                 local: false,
+                preserve_source: false,
                 downtime_ms: NonZeroU64::new(150).unwrap(),
                 timeout_s: NonZeroU64::new(900).unwrap(),
                 timeout_strategy: TimeoutStrategy::Ignore,
@@ -2453,5 +2472,23 @@ mod unit_tests {
             "destination_url=tcp:192.168.1.1:8080,memory_mode=postcopy,connections=4",
         )
         .unwrap_err();
+
+        // preserve_source is accepted together with local (offload snapshot).
+        let data = VmSendMigrationData::parse(
+            "destination_url=unix:/tmp/sock,local=on,preserve_source=on",
+        )
+        .unwrap();
+        assert!(data.preserve_source);
+        assert!(data.local);
+
+        // preserve_source defaults to false when unspecified.
+        let data = VmSendMigrationData::parse("destination_url=unix:/tmp/sock,local=on").unwrap();
+        assert!(!data.preserve_source);
+
+        // preserve_source without local must be rejected.
+        VmSendMigrationData::parse("destination_url=tcp:192.168.1.1:8080,preserve_source=on")
+            .unwrap_err();
+        VmSendMigrationData::parse("destination_url=unix:/tmp/sock,preserve_source=on")
+            .unwrap_err();
     }
 }

vmm/src/lib.rs

@@ -1670,8 +1670,11 @@ impl Vmm {
         };
         transport::send_config(&mut socket, &vm_migration_config)?;
 
-        // Let every Migratable object know about the migration being started.
-        vm.start_migration()?;
+        // Let every Migratable object know about the migration being started
+        // unless the source VM must be preserved.
+        if !send_data_migration.preserve_source {
+            vm.start_migration()?;
+        }
 
         if send_data_migration.local
             || matches!(send_data_migration.memory_mode, MigrationMode::Postcopy)
@@ -1716,8 +1719,11 @@ impl Vmm {
 
         // We release the locks early to enable locking them on the destination host.
         // The VM is already stopped.
-        vm.release_disk_locks()
-            .map_err(|e| MigratableError::UnlockError(anyhow!("{e}")))?;
+        // Keep the locks held if the source VM must be preserved.
+        if !send_data_migration.preserve_source {
+            vm.release_disk_locks()
+                .map_err(|e| MigratableError::UnlockError(anyhow!("{e}")))?;
+        }
 
         // For postcopy, serve faults before sending State so the destination
         // can fault pages in during restore.
@@ -1809,7 +1815,12 @@ impl Vmm {
         }
 
         // Let every Migratable object know about the migration being complete
-        vm.complete_migration()
+        // unless the source VM must be preserved.
+        if send_data_migration.preserve_source {
+            Ok(())
+        } else {
+            vm.complete_migration()
+        }
     }
 
     /// Serve `Command::PageFault` requests from local guest memory on the fault
@@ -2031,6 +2042,7 @@ impl Vmm {
             vm,
             migration_result: migration_res,
             initial_vm_state,
+            preserve_source,
         } = migration_worker_handle.join();
 
         let mut try_resume_vm_after_failed_migration = |mut vm: Vm| {
@@ -2056,6 +2068,10 @@ impl Vmm {
         };
 
         match migration_res {
+            Ok(()) if preserve_source => {
+                // Give the source VM back to the VMM.
+                self.vm = VmOwnership::Owned(vm);
+            }
             Ok(()) => {
                 self.vm = VmOwnership::None;
                 let mut vm = vm;

vmm/src/migration/worker.rs

@@ -127,6 +127,7 @@ impl MigrationWorker {
             vm,
             migration_result,
             initial_vm_state: self.initial_vm_state,
+            preserve_source: self.config.preserve_source,
         }
     }
 
@@ -180,11 +181,13 @@ impl MigrationWorker {
 pub struct MigrationWorkerResult {
     /// The VM that was migrated.
     ///
-    /// If `migration_result` is `Ok`, the VM is paused and can be deleted.
-    /// If `migration_result` is `Err`, the VM can be resumed and given back to
-    /// the VMM.
+    /// If `migration_result` is `Ok`, the VM is paused and can be deleted
+    /// unless `preserve_source` is true, which means the VM is given back
+    /// to the VMM in a paused state.
+    /// If `Err`, the VM can be resumed and given back to the VMM.
     pub vm: Vm,
     /// The result of [`Vmm::send_migration`].
     pub migration_result: Result<(), MigratableError>,
     pub initial_vm_state: VmState,
+    pub preserve_source: bool,
 }