vmm: Defer address space allocation

We can ideally defer the address space allocation till we start the
vCPUs for the very first time. Because the VM will not access the memory
until the CPUs start running. Thus there is no need to allocate the
address space eagerly and wait till the time we are going to start the
vCPUs for the first time.

Signed-off-by: Jinank Jain <jinankjain@microsoft.com>

vmm/src/memory_manager.rs

@@ -773,7 +773,7 @@ impl MemoryManager {
         }
     }
 
-    fn allocate_address_space(&mut self) -> Result<(), Error> {
+    pub fn allocate_address_space(&mut self) -> Result<(), Error> {
         let mut list = Vec::new();
 
         for (zone_id, memory_zone) in self.memory_zones.iter() {
@@ -1135,10 +1135,17 @@ impl MemoryManager {
             thp: config.thp,
         };
 
-        memory_manager.allocate_address_space()?;
-
         #[cfg(target_arch = "aarch64")]
-        memory_manager.add_uefi_flash()?;
+        {
+            // For Aarch64 we cannot lazily allocate the address space like we
+            // do for x86, because while restoring a VM from snapshot we would
+            // need the address space to be allocated to properly restore VGIC.
+            // And the restore of VGIC happens before we attempt to run the vCPUs
+            // for the first time, thus we need to allocate the address space
+            // beforehand.
+            memory_manager.allocate_address_space()?;
+            memory_manager.add_uefi_flash()?;
+        }
 
         #[cfg(target_arch = "x86_64")]
         if let Some(sgx_epc_config) = sgx_epc_config {

vmm/src/vm.rs

@@ -2116,6 +2116,18 @@ impl Vm {
             self.vm.tdx_finalize().map_err(Error::FinalizeTdx)?;
         }
 
+        #[cfg(target_arch = "x86_64")]
+        // Note: For x86, always call this function before invoking start boot vcpus.
+        // Otherwise guest would fail to boot because we haven't created the
+        // userspace mappings to update the hypervisor about the memory mappings.
+        // These mappings must be created before we start the vCPU threads for
+        // the very first time.
+        self.memory_manager
+            .lock()
+            .unwrap()
+            .allocate_address_space()
+            .map_err(Error::MemoryManager)?;
+
         self.cpu_manager
             .lock()
             .unwrap()
@@ -2131,6 +2143,18 @@ impl Vm {
     pub fn restore(&mut self) -> Result<()> {
         event!("vm", "restoring");
 
+        #[cfg(target_arch = "x86_64")]
+        // Note: For x86, always call this function before invoking start boot vcpus.
+        // Otherwise guest would fail to boot because we haven't created the
+        // userspace mappings to update the hypervisor about the memory mappings.
+        // These mappings must be created before we start the vCPU threads for
+        // the very first time for the restored VM.
+        self.memory_manager
+            .lock()
+            .unwrap()
+            .allocate_address_space()
+            .map_err(Error::MemoryManager)?;
+
         // Now we can start all vCPUs from here.
         self.cpu_manager
             .lock()