Add private pool and keepalive test to VMM, device tree updates (#740)

Start bringing up missing coverage for private pool and NVMe keepalive
to VMM test suite.
This is not complete end-to-end test yet, but brings necessary
infrastructure changes.

Update device tree properties to sync with Windows host changes.

---------

Co-authored-by: Daniel Prilik <[email protected]>

openvmm/hvlite_core/src/worker/vm_loaders/igvm.rs

@@ -467,6 +467,7 @@ fn build_device_tree(
     let p_memory_allocation_mode = root.add_string("memory-allocation-mode")?;
     let p_memory_size = root.add_string("memory-size")?;
     let p_mmio_size = root.add_string("mmio-size")?;
+    let p_vf_keep_alive_devs = root.add_string("device-types")?;
     let mut openhcl = root.start_node("openhcl")?;
 
     let memory_allocation_mode = match vtl2_base_address {
@@ -494,6 +495,12 @@ fn build_device_tree(
             .end_node()?;
     }
 
+    // Indicate that NVMe keep-alive feature is supported by this VMM.
+    openhcl = openhcl
+        .start_node("keep-alive")?
+        .add_str(p_vf_keep_alive_devs, "nvme")?
+        .end_node()?;
+
     root = openhcl.end_node()?;
 
     let bytes_used = root

openvmm/hvlite_helpers/src/underhill.rs

@@ -5,13 +5,15 @@
 
 use anyhow::Context;
 use get_resources::ged::GuestEmulationRequest;
+use get_resources::ged::GuestServicingFlags;
 use hvlite_defs::rpc::VmRpc;
 use mesh::rpc::RpcSend;
 
 /// Replace the running version of Underhill.
 pub async fn service_underhill(
     vm_send: &mesh::Sender<VmRpc>,
     send: &mesh::Sender<GuestEmulationRequest>,
+    flags: GuestServicingFlags,
     file: std::fs::File,
 ) -> anyhow::Result<()> {
     // Stage the IGVM file in the VM worker.
@@ -27,7 +29,7 @@ pub async fn service_underhill(
     // blocked while waiting for the guest.
     tracing::debug!("waiting for guest to send saved state");
     let r = send
-        .call_failable(GuestEmulationRequest::SaveGuestVtl2State, ())
+        .call_failable(GuestEmulationRequest::SaveGuestVtl2State, flags)
         .await
         .context("failed to save VTL2 state");
 

openvmm/openvmm_entry/src/lib.rs

@@ -46,6 +46,7 @@ use futures::StreamExt;
 use futures_concurrency::stream::Merge;
 use gdma_resources::GdmaDeviceHandle;
 use gdma_resources::VportDefinition;
+use get_resources::ged::GuestServicingFlags;
 use guid::Guid;
 use hvlite_defs::config::Config;
 use hvlite_defs::config::DeviceVtl;
@@ -2692,6 +2693,7 @@ async fn run_control(driver: &DefaultDriver, mesh: &VmmMesh, opt: Options) -> an
                         hvlite_helpers::underhill::service_underhill(
                             &vm_rpc,
                             resources.ged_rpc.as_ref().context("no GED")?,
+                            GuestServicingFlags::default(),
                             file.into(),
                         )
                         .await?;

petri/src/lib.rs

@@ -3,9 +3,8 @@
 
 //! A Rust-based testing framework for VMMs.
 //!
-//! At this time - `petri` only supports testing OpenVMM and OpenHCL based VMs.
-//! In the future, we expect `petri` to also support other VMM backends
-//! (notably, Hyper-V).
+//! At this time - `petri` supports testing OpenVMM, OpenHCL,
+//! and Hyper-V based VMs.
 
 #![forbid(unsafe_code)]
 #![warn(missing_docs)]

petri/src/vm/mod.rs

@@ -294,6 +294,13 @@ pub enum IsolationType {
     Tdx,
 }
 
+/// Flags controlling servicing behavior.
+#[derive(Default)]
+pub struct OpenHclServicingFlags {
+    /// Preserve DMA memory for NVMe devices if supported.
+    pub enable_nvme_keepalive: bool,
+}
+
 /// Generates a name for the petri test based on the thread name
 pub fn get_test_name() -> anyhow::Result<String> {
     // Use the current thread name for the test name, both cargo-test and

petri/src/vm/openvmm/runtime.rs

@@ -6,6 +6,7 @@
 use super::PetriVmResourcesOpenVmm;
 use crate::openhcl_diag::OpenHclDiagHandler;
 use crate::worker::Worker;
+use crate::OpenHclServicingFlags;
 use crate::PetriVm;
 use crate::ShutdownKind;
 use anyhow::Context;
@@ -172,7 +173,11 @@ impl PetriVmOpenVmm {
     );
     petri_vm_fn!(
         /// Restarts OpenHCL.
-        pub async fn restart_openhcl(&mut self, new_openhcl: ArtifactHandle<impl petri_artifacts_common::tags::IsOpenhclIgvm>) -> anyhow::Result<()>
+        pub async fn restart_openhcl(
+            &mut self,
+            new_openhcl: ArtifactHandle<impl petri_artifacts_common::tags::IsOpenhclIgvm>,
+            flags: OpenHclServicingFlags
+        ) -> anyhow::Result<()>
     );
     petri_vm_fn!(
         /// Resets the hardware state of the VM, simulating a power cycle.
@@ -345,6 +350,7 @@ impl PetriVmInner {
     async fn restart_openhcl(
         &self,
         new_openhcl: ArtifactHandle<impl petri_artifacts_common::tags::IsOpenhclIgvm>,
+        flags: OpenHclServicingFlags,
     ) -> anyhow::Result<()> {
         let ged_send = self
             .resources
@@ -355,7 +361,7 @@ impl PetriVmInner {
         let igvm_path = self.resources.resolver.resolve(new_openhcl);
         let igvm_file = fs_err::File::open(igvm_path).context("failed to open igvm file")?;
         self.worker
-            .restart_openhcl(ged_send, igvm_file.into())
+            .restart_openhcl(ged_send, flags, igvm_file.into())
             .await
     }
 

petri/src/worker.rs

@@ -1,6 +1,8 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+use crate::OpenHclServicingFlags;
+use get_resources::ged::GuestServicingFlags;
 use hvlite_defs::config::Config;
 use hvlite_defs::rpc::PulseSaveRestoreError;
 use hvlite_defs::rpc::VmRpc;
@@ -59,9 +61,18 @@ impl Worker {
     pub(crate) async fn restart_openhcl(
         &self,
         send: &mesh::Sender<get_resources::ged::GuestEmulationRequest>,
+        flags: OpenHclServicingFlags,
         file: std::fs::File,
     ) -> anyhow::Result<()> {
-        hvlite_helpers::underhill::service_underhill(&self.rpc, send, file).await
+        hvlite_helpers::underhill::service_underhill(
+            &self.rpc,
+            send,
+            GuestServicingFlags {
+                nvme_keepalive: flags.enable_nvme_keepalive,
+            },
+            file,
+        )
+        .await
     }
 
     pub(crate) async fn inspect_all(&self) -> String {

vm/devices/get/get_resources/src/lib.rs

@@ -151,6 +151,13 @@ pub mod ged {
         const ID: &'static str = "ged";
     }
 
+    /// Define servicing behavior.
+    #[derive(MeshPayload, Default)]
+    pub struct GuestServicingFlags {
+        /// Retain memory for DMA-attached devices.
+        pub nvme_keepalive: bool,
+    }
+
     /// Actions a client can request that the Guest Emulation
     /// Device perform.
     #[derive(MeshPayload)]
@@ -160,7 +167,7 @@ pub mod ged {
         /// Wait for VTL2 to start VTL0.
         WaitForVtl0Start(Rpc<(), Result<(), Vtl0StartError>>),
         /// Save VTL2 state.
-        SaveGuestVtl2State(Rpc<(), Result<(), SaveRestoreError>>),
+        SaveGuestVtl2State(Rpc<GuestServicingFlags, Result<(), SaveRestoreError>>),
         /// Update the VTL2 settings.
         ModifyVtl2Settings(Rpc<Vec<u8>, Result<(), ModifyVtl2SettingsError>>),
     }

vm/devices/get/guest_emulation_device/src/lib.rs

@@ -37,6 +37,7 @@ use get_protocol::VmgsIoStatus;
 use get_protocol::MAX_PAYLOAD_SIZE;
 use get_resources::ged::FirmwareEvent;
 use get_resources::ged::GuestEmulationRequest;
+use get_resources::ged::GuestServicingFlags;
 use get_resources::ged::ModifyVtl2SettingsError;
 use get_resources::ged::SaveRestoreError;
 use get_resources::ged::Vtl0StartError;
@@ -316,7 +317,7 @@ pub struct GedChannel<T: RingMem = GpadlRingMem> {
 }
 
 struct InProgressSave {
-    rpc: Rpc<(), Result<(), SaveRestoreError>>,
+    rpc: Rpc<GuestServicingFlags, Result<(), SaveRestoreError>>,
     buffer: Vec<u8>,
 }
 
@@ -535,7 +536,8 @@ impl<T: RingMem + Unpin> GedChannel<T> {
                             get_protocol::GuestNotifications::SAVE_GUEST_VTL2_STATE,
                         ),
                         correlation_id: Guid::ZERO,
-                        capabilities_flags: SaveGuestVtl2StateFlags::new(),
+                        capabilities_flags: SaveGuestVtl2StateFlags::new()
+                            .with_enable_nvme_keepalive(rpc.input().nvme_keepalive),
                         timeout_hint_secs: 60,
                     };
 

vm/devices/get/guest_emulation_device/src/test_utilities.rs

@@ -12,6 +12,7 @@ use get_protocol::HostRequests;
 use get_protocol::SecureBootTemplateType;
 use get_protocol::UefiConsoleMode;
 use get_resources::ged::GuestEmulationRequest;
+use get_resources::ged::GuestServicingFlags;
 use guestmem::GuestMemory;
 use mesh::rpc::RpcSend;
 use pal_async::task::Spawn;
@@ -334,7 +335,10 @@ enum TestTask {
 impl TestGedClient {
     pub async fn test_save_guest_vtl2_state(&mut self) {
         self.sender
-            .call_failable(GuestEmulationRequest::SaveGuestVtl2State, ())
+            .call_failable(
+                GuestEmulationRequest::SaveGuestVtl2State,
+                GuestServicingFlags::default(),
+            )
             .await
             .expect("no failure");
     }

vmm_tests/vmm_tests/tests/tests/x86_64/openhcl_servicing.rs

@@ -5,13 +5,15 @@
 
 use petri::openvmm::PetriVmConfigOpenVmm;
 use petri::ArtifactHandle;
+use petri::OpenHclServicingFlags;
 use petri_artifacts_vmm_test::artifacts::openhcl_igvm::LATEST_LINUX_DIRECT_TEST_X64;
 use vmm_core_defs::HaltReason;
 use vmm_test_macros::openvmm_test;
 
 async fn openhcl_servicing_core(
     config: PetriVmConfigOpenVmm,
     new_openhcl: ArtifactHandle<impl petri_artifacts_common::tags::IsOpenhclIgvm>,
+    flags: OpenHclServicingFlags,
 ) -> anyhow::Result<()> {
     let (mut vm, agent) = config.run().await?;
 
@@ -20,7 +22,7 @@ async fn openhcl_servicing_core(
     // Test that inspect serialization works with the old version.
     vm.test_inspect_openhcl().await?;
 
-    vm.restart_openhcl(new_openhcl).await?;
+    vm.restart_openhcl(new_openhcl, flags).await?;
 
     agent.ping().await?;
 
@@ -36,7 +38,26 @@ async fn openhcl_servicing_core(
 /// Test servicing an OpenHCL VM from the current version to itself.
 #[openvmm_test(openhcl_linux_direct_x64)]
 async fn openhcl_servicing(config: PetriVmConfigOpenVmm) -> Result<(), anyhow::Error> {
-    openhcl_servicing_core(config, LATEST_LINUX_DIRECT_TEST_X64).await
+    openhcl_servicing_core(
+        config,
+        LATEST_LINUX_DIRECT_TEST_X64,
+        OpenHclServicingFlags::default(),
+    )
+    .await
+}
+
+/// Test servicing an OpenHCL VM from the current version to itself
+/// with VF keepalive support.
+#[openvmm_test(openhcl_linux_direct_x64)]
+async fn openhcl_servicing_keepalive(config: PetriVmConfigOpenVmm) -> Result<(), anyhow::Error> {
+    openhcl_servicing_core(
+        config,
+        LATEST_LINUX_DIRECT_TEST_X64,
+        OpenHclServicingFlags {
+            enable_nvme_keepalive: true,
+        },
+    )
+    .await
 }
 
 // TODO: add tests with guest workloads while doing servicing.

vmm_tests/vmm_tests/tests/tests/x86_64/openhcl_uefi.rs

@@ -41,6 +41,16 @@ async fn nvme_relay_shared_pool(config: PetriVmConfigOpenVmm) -> Result<(), anyh
     nvme_relay_test_core(config, "OPENHCL_ENABLE_SHARED_VISIBILITY_POOL=1").await
 }
 
+/// Test an OpenHCL uefi VM with a NVME disk assigned to VTL2 that boots
+/// linux, with vmbus relay. This should expose a disk to VTL0 via vmbus.
+///
+/// Use the private pool override to test the private pool dma path.
+#[openvmm_test(openhcl_uefi_x64[nvme](vhd(ubuntu_2204_server_x64)))]
+async fn nvme_relay_private_pool(config: PetriVmConfigOpenVmm) -> Result<(), anyhow::Error> {
+    // Number of pages to reserve as a private pool.
+    nvme_relay_test_core(config, "OPENHCL_ENABLE_VTL2_GPA_POOL=1024").await
+}
+
 /// Boot the UEFI firmware, with a VTL2 range automatically configured by
 /// hvlite.
 #[openvmm_test(openhcl_uefi_x64(none))]