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

@@ -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"

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