From 6de90bdec655f3977bf19941416c77069590fc3a Mon Sep 17 00:00:00 2001
From: Rob Bradford <rbradford@meta.com>
Date: Wed, 8 Jul 2026 06:12:15 -0700
Subject: [PATCH] vmm: Error out on migration & snapshot if on-demand restoring

If there is an active on-demand restoration then reject any requests to
migrate or snapshot this VM as the memory will not be available for the
snapshot.

Signed-off-by: Rob Bradford <rbradford@meta.com>
---
 vmm/src/lib.rs | 11 ++++++++++-
 vmm/src/vm.rs  |  7 +++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/vmm/src/lib.rs b/vmm/src/lib.rs
index b9cbd9ca8..38434f208 100644
--- a/vmm/src/lib.rs
+++ b/vmm/src/lib.rs
@@ -2318,6 +2318,9 @@ impl RequestHandler for Vmm {
     fn vm_snapshot(&mut self, destination_url: &str) -> result::Result<(), VmError> {
         match self.vm {
             VmOwnership::Owned(ref mut vm) => {
+                if vm.restoring() {
+                    return Err(VmError::VmRestoring);
+                }
                 // Drain console_info so that FDs are not reused
                 let _ = self.console_info.take();
                 vm.snapshot()
@@ -3091,7 +3094,13 @@ impl RequestHandler for Vmm {
         send_data_migration: VmSendMigrationData,
     ) -> result::Result<(), MigratableError> {
         match self.vm {
-            VmOwnership::Owned(_) => (),
+            VmOwnership::Owned(ref vm) => {
+                if vm.restoring() {
+                    return Err(MigratableError::MigrateSend(anyhow!(
+                        "Cannot migrate while on-demand memory restore is in progress"
+                    )));
+                }
+            }
             VmOwnership::Migration { .. } => {
                 return Err(MigratableError::MigrateSend(anyhow!(
                     "There is already an ongoing migration"
diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index 843603e79..effcd988d 100644
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -201,6 +201,9 @@ pub enum Error {
     #[error("VM is currently migrating and can't be modified")]
     VmMigrating,
 
+    #[error("VM on-demand memory restore is still in progress")]
+    VmRestoring,
+
     #[error("Cannot clone EventFd")]
     EventFdClone(#[source] io::Error),
 
@@ -3041,6 +3044,10 @@ impl Vm {
         self.memory_manager.lock().unwrap().guest_memory()
     }
 
+    pub fn restoring(&self) -> bool {
+        self.memory_manager.lock().unwrap().restoring()
+    }
+
     pub fn device_tree(&self) -> Arc<Mutex<DeviceTree>> {
         self.device_manager.lock().unwrap().device_tree()
     }