mana: save and restore mana devices when keepalive is enabled

openhcl/openhcl_boot/src/main.rs

@@ -355,6 +355,7 @@ fn shim_parameters(shim_params_raw_offset: isize) -> ShimParams {
 #[cfg_attr(not(target_arch = "x86_64"), expect(dead_code))]
 mod x86_boot {
     use crate::PageAlign;
+    use crate::boot_logger::log;
     use crate::memory::AddressSpaceManager;
     use crate::single_threaded::OffStackRef;
     use crate::single_threaded::off_stack;
@@ -405,18 +406,41 @@ mod x86_boot {
         boot_params: &mut boot_params,
         ext: &mut E820Ext,
         address_space: &AddressSpaceManager,
+        initrd: Option<Range<u64>>,
     ) -> Result<bool, BuildE820MapError> {
         boot_params.e820_entries = 0;
         let mut entries = boot_params
             .e820_map
             .iter_mut()
             .chain(ext.entries.iter_mut());
 
+        // Find the largest region and the one containing initrd
+        let mut largest_ram: Option<MemoryRange> = None;
+        let mut initrd_hosting_ram: Option<MemoryRange> = None;
+        for (range, typ) in address_space.vtl2_ranges() {
+            if matches!(typ, MemoryVtlType::VTL2_RAM) {
+                if largest_ram.is_none_or(|r| range.len() > r.len()) {
+                    largest_ram = Some(range);
+                }
+
+                if initrd.clone().is_some()
+                    && range.start() <= initrd.clone().unwrap().start
+                    && range.end() >= initrd.clone().unwrap().end
+                {
+                    initrd_hosting_ram = Some(range);
+                }
+            }
+        }
+
         let mut n = 0;
         for (range, typ) in address_space.vtl2_ranges() {
             match typ {
                 MemoryVtlType::VTL2_RAM => {
-                    add_e820_entry(entries.next(), range, E820_RAM)?;
+                    if Some(range) == largest_ram || Some(range) == initrd_hosting_ram {
+                        add_e820_entry(entries.next(), range, E820_RAM)?;
+                    } else {
+                        add_e820_entry(entries.next(), range, E820_RESERVED)?;
+                    }
                     n += 1;
                 }
                 MemoryVtlType::VTL2_CONFIG
@@ -438,6 +462,43 @@ mod x86_boot {
             }
         }
 
+        for (range, typ) in address_space.vtl2_ranges() {
+            let (kind, label) = match typ {
+                MemoryVtlType::VTL2_RAM => {
+                    if Some(range) == largest_ram {
+                        ("RAM", "usable")
+                    } else if Some(range) == initrd_hosting_ram {
+                        ("RAM", "initrd")
+                    } else {
+                        ("RESERVED", "hidden_ram")
+                    }
+                }
+                MemoryVtlType::VTL2_PERSISTED_STATE_HEADER => {
+                    ("RESERVED", "persisted_state_header")
+                }
+                MemoryVtlType::VTL2_PERSISTED_STATE_PROTOBUF => {
+                    ("RESERVED", "persisted_state_payload")
+                }
+                MemoryVtlType::VTL2_GPA_POOL => ("RESERVED", "keep_alive_gpa_pool"),
+                MemoryVtlType::VTL2_BOOTSHIM_LOG_BUFFER => ("RESERVED", "boot_log_buffer"),
+                MemoryVtlType::VTL2_TDX_PAGE_TABLES => ("RESERVED", "tdx_page_tables"),
+                MemoryVtlType::VTL2_CONFIG => ("RESERVED", "config"),
+                MemoryVtlType::VTL2_SIDECAR_IMAGE => ("RESERVED", "sidecar_image"),
+                MemoryVtlType::VTL2_SIDECAR_NODE => ("RESERVED", "sidecar_node"),
+                MemoryVtlType::VTL2_RESERVED => ("RESERVED", "generic_reserved"),
+                _ => ("UNKNOWN", "unexpected"),
+            };
+
+            log!(
+                "e820 {}: {:#x}-{:#x} len={:#x} {}",
+                kind,
+                range.start(),
+                range.end(),
+                range.len(),
+                label
+            );
+        }
+
         let base = n.min(boot_params.e820_map.len());
         boot_params.e820_entries = base as u8;
 
@@ -474,12 +535,13 @@ mod x86_boot {
         boot_params.hdr.ramdisk_image = (initrd.start as u32).into();
         boot_params.ext_ramdisk_image = (initrd.start >> 32) as u32;
         let initrd_len = initrd.end - initrd.start;
+        log!("initrd start: {:#x} end: {:#x}", initrd.start, initrd.end);
         boot_params.hdr.ramdisk_size = (initrd_len as u32).into();
         boot_params.ext_ramdisk_size = (initrd_len >> 32) as u32;
 
         let e820_ext = OffStackRef::leak(off_stack!(E820Ext, zeroed()));
 
-        let used_ext = build_e820_map(boot_params, e820_ext, address_space)
+        let used_ext = build_e820_map(boot_params, e820_ext, address_space, Some(initrd))
             .expect("building e820 map must succeed");
 
         if used_ext {
@@ -1128,7 +1190,7 @@ mod test {
             None,
         );
 
-        assert!(build_e820_map(&mut boot_params, &mut ext, &address_space).is_ok());
+        assert!(build_e820_map(&mut boot_params, &mut ext, &address_space, None).is_ok());
 
         check_e820(
             &boot_params,
@@ -1159,7 +1221,7 @@ mod test {
             Some(reclaim),
         );
 
-        assert!(build_e820_map(&mut boot_params, &mut ext, &address_space).is_ok());
+        assert!(build_e820_map(&mut boot_params, &mut ext, &address_space, None).is_ok());
 
         check_e820(
             &boot_params,
@@ -1195,7 +1257,7 @@ mod test {
             Some(reclaim),
         );
 
-        assert!(build_e820_map(&mut boot_params, &mut ext, &address_space).is_ok());
+        assert!(build_e820_map(&mut boot_params, &mut ext, &address_space, None).is_ok());
 
         check_e820(
             &boot_params,
@@ -1236,7 +1298,7 @@ mod test {
             Some(reclaim),
         );
 
-        assert!(build_e820_map(&mut boot_params, &mut ext, &address_space).is_ok());
+        assert!(build_e820_map(&mut boot_params, &mut ext, &address_space, None).is_ok());
 
         check_e820(
             &boot_params,
@@ -1322,7 +1384,7 @@ mod test {
         let mut ext = FromZeros::new_zeroed();
         let total_ranges = address_space.vtl2_ranges().count();
 
-        let used_ext = build_e820_map(&mut boot_params, &mut ext, &address_space).unwrap();
+        let used_ext = build_e820_map(&mut boot_params, &mut ext, &address_space, None).unwrap();
 
         // Verify that we used the extension
         assert!(used_ext, "should use extension when there are many ranges");

openhcl/underhill_core/src/dispatch/mod.rs

@@ -34,6 +34,7 @@ use hyperv_ic_resources::shutdown::ShutdownRpc;
 use hyperv_ic_resources::shutdown::ShutdownType;
 use igvm_defs::MemoryMapEntryType;
 use inspect::Inspect;
+use mana_driver::save_restore::ManaSavedState;
 use mesh::CancelContext;
 use mesh::MeshPayload;
 use mesh::error::RemoteError;
@@ -114,6 +115,8 @@ pub trait LoadedVmNetworkSettings: Inspect {
         vmbus_server: &Option<VmbusServerHandle>,
         dma_client_spawner: DmaClientSpawner,
         is_isolated: bool,
+        save_restore_supported: bool,
+        mana_state: Option<&ManaSavedState>,
     ) -> anyhow::Result<RuntimeSavedState>;
 
     /// Callback when network is removed externally.
@@ -127,6 +130,9 @@ pub trait LoadedVmNetworkSettings: Inspect {
         &self,
         mut params: PacketCaptureParams<Socket>,
     ) -> anyhow::Result<PacketCaptureParams<Socket>>;
+
+    /// Save the network state for restoration after servicing.
+    async fn save(&mut self) -> Vec<ManaSavedState>;
 }
 
 /// A VM that has been loaded and can be run.
@@ -187,6 +193,7 @@ pub(crate) struct LoadedVm {
     pub _periodic_telemetry_task: Task<()>,
 
     pub nvme_keep_alive: bool,
+    pub mana_keep_alive: bool,
     pub test_configuration: Option<TestScenarioConfig>,
     pub dma_manager: OpenhclDmaManager,
 }
@@ -298,7 +305,7 @@ impl LoadedVm {
                     WorkerRpc::Restart(rpc) => {
                         let state = async {
                             let running = self.stop().await;
-                            match self.save(None, false).await {
+                            match self.save(None, false, false).await {
                                 Ok(servicing_state) => Some((rpc, servicing_state)),
                                 Err(err) => {
                                     if running {
@@ -363,7 +370,7 @@ impl LoadedVm {
                     UhVmRpc::Save(rpc) => {
                         rpc.handle_failable(async |()| {
                             let running = self.stop().await;
-                            let r = self.save(None, false).await;
+                            let r = self.save(None, false, false).await;
                             if running {
                                 self.start(None).await;
                             }
@@ -565,6 +572,7 @@ impl LoadedVm {
         // NOTE: This is set via the corresponding env arg, as this feature is
         // experimental.
         let nvme_keepalive = self.nvme_keep_alive && capabilities_flags.enable_nvme_keepalive();
+        let mana_keepalive = self.mana_keep_alive && capabilities_flags.enable_mana_keepalive();
 
         // Do everything before the log flush under a span.
         let r = async {
@@ -579,7 +587,7 @@ impl LoadedVm {
                 anyhow::bail!("cannot service underhill while paused");
             }
 
-            let mut state = self.save(Some(deadline), nvme_keepalive).await?;
+            let mut state = self.save(Some(deadline), nvme_keepalive, mana_keepalive).await?;
             state.init_state.correlation_id = Some(correlation_id);
 
             // Unload any network devices.
@@ -741,19 +749,33 @@ impl LoadedVm {
     async fn save(
         &mut self,
         _deadline: Option<std::time::Instant>,
-        vf_keepalive_flag: bool,
+        nvme_keepalive_flag: bool,
+        mana_keepalive_flag: bool,
     ) -> anyhow::Result<ServicingState> {
         assert!(!self.state_units.is_running());
 
         let emuplat = (self.emuplat_servicing.save()).context("emuplat save failed")?;
 
+        // Only save dma manager state if we are expected to keep VF devices
+        // alive across save. Otherwise, don't persist the state at all, as
+        // there should be no live DMA across save.
+        //
+        // This has to happen before saving the network state, otherwise its allocations
+        // are marked as Free and are unable to be restored.
+        let dma_manager_state = if nvme_keepalive_flag || mana_keepalive_flag {
+            use vmcore::save_restore::SaveRestore;
+            Some(self.dma_manager.save().context("dma_manager save failed")?)
+        } else {
+            None
+        };
+
         // Only save NVMe state when there are NVMe controllers and keep alive
         // was enabled.
         let nvme_state = if let Some(n) = &self.nvme_manager {
             // DEVNOTE: A subtlety here is that the act of saving the NVMe state also causes the driver
             // to enter a state where subsequent teardown operations will noop. There is a STRONG
             // correlation between save/restore and keepalive.
-            n.save(vf_keepalive_flag)
+            n.save(nvme_keepalive_flag)
                 .instrument(tracing::info_span!("nvme_manager_save", CVM_ALLOWED))
                 .await
                 .map(|s| NvmeSavedState { nvme_state: s })
@@ -762,6 +784,15 @@ impl LoadedVm {
         };
 
         let units = self.save_units().await.context("state unit save failed")?;
+
+        let mana_state = if let Some(network_settings) = &mut self.network_settings
+            && mana_keepalive_flag
+        {
+            Some(network_settings.save().await)
+        } else {
+            None
+        };
+
         let vmgs = if let Some((vmgs_thin_client, vmgs_disk_metadata, _)) = self.vmgs.as_ref() {
             Some((
                 vmgs_thin_client.save().await.context("vmgs save failed")?,
@@ -771,16 +802,6 @@ impl LoadedVm {
             None
         };
 
-        // Only save dma manager state if we are expected to keep VF devices
-        // alive across save. Otherwise, don't persist the state at all, as
-        // there should be no live DMA across save.
-        let dma_manager_state = if vf_keepalive_flag {
-            use vmcore::save_restore::SaveRestore;
-            Some(self.dma_manager.save().context("dma_manager save failed")?)
-        } else {
-            None
-        };
-
         let vmbus_client = if let Some(vmbus_client) = &mut self.vmbus_client {
             vmbus_client.stop().await;
             Some(vmbus_client.save().await)
@@ -800,6 +821,7 @@ impl LoadedVm {
                 nvme_state,
                 dma_manager_state,
                 vmbus_client,
+                mana_state,
             },
             units,
         };
@@ -866,6 +888,8 @@ impl LoadedVm {
                 &self.vmbus_server,
                 self.dma_manager.client_spawner(),
                 self.isolation.is_isolated(),
+                self.mana_keep_alive,
+                None, // No existing mana state
             )
             .await?;