mana: keepalive feature

Cargo.lock

@@ -3865,6 +3865,7 @@ dependencies = [
  "gdma_defs",
  "getrandom 0.3.2",
  "inspect",
+ "mana_save_restore",
  "mesh",
  "net_backend",
  "net_backend_resources",
@@ -3880,6 +3881,13 @@ dependencies = [
  "zerocopy 0.8.24",
 ]
 
+[[package]]
+name = "mana_save_restore"
+version = "0.0.0"
+dependencies = [
+ "mesh",
+]
+
 [[package]]
 name = "managed"
 version = "0.8.0"
@@ -4309,6 +4317,7 @@ dependencies = [
  "futures-concurrency",
  "guestmem",
  "inspect",
+ "mana_save_restore",
  "memory_range",
  "mesh",
  "net_backend_resources",
@@ -4378,6 +4387,7 @@ dependencies = [
  "guestmem",
  "inspect",
  "mana_driver",
+ "mana_save_restore",
  "mesh",
  "net_backend",
  "pal_async",
@@ -4403,6 +4413,7 @@ dependencies = [
  "futures-concurrency",
  "guestmem",
  "inspect",
+ "mana_save_restore",
  "mesh",
  "net_backend",
  "parking_lot",
@@ -4445,6 +4456,7 @@ dependencies = [
  "hvdef",
  "inspect",
  "inspect_counters",
+ "mana_save_restore",
  "mesh",
  "net_backend",
  "net_backend_resources",
@@ -7698,6 +7710,7 @@ dependencies = [
  "pci_core",
  "tracing",
  "user_driver",
+ "vmcore",
 ]
 
 [[package]]

Cargo.toml

@@ -229,6 +229,7 @@ gdma_defs = { path = "vm/devices/net/gdma_defs" }
 gdma_resources = { path = "vm/devices/net/gdma_resources" }
 linux_net_bindings = { path = "vm/devices/net/linux_net_bindings" }
 mana_driver = { path = "vm/devices/net/mana_driver" }
+mana_save_restore = { path = "vm/devices/net/mana_save_restore" }
 vfio_sys = { path = "vm/devices/user_driver/vfio_sys" }
 net_backend = { path = "vm/devices/net/net_backend" }
 net_backend_resources = { path = "vm/devices/net/net_backend_resources" }

openhcl/host_fdt_parser/src/lib.rs

@@ -237,6 +237,8 @@ pub struct ParsedDeviceTree<
     pub nvme_keepalive: bool,
     /// The physical address of the VTL0 alias mapping, if one is configured.
     pub vtl0_alias_map: Option<u64>,
+    /// Indicates that Host supports MANA keep-alive.
+    pub mana_keepalive: bool,
 }
 
 /// The memory allocation mode provided by the host. This determines how OpenHCL
@@ -316,6 +318,7 @@ impl<
             entropy: None,
             device_dma_page_count: None,
             nvme_keepalive: false,
+            mana_keepalive: false,
             vtl0_alias_map: None,
         }
     }
@@ -735,6 +738,7 @@ impl<
             device_dma_page_count: _,
             nvme_keepalive: _,
             vtl0_alias_map: _,
+            mana_keepalive: _,
         } = storage;
 
         *device_tree_size = parser.total_size;

openhcl/lower_vtl_permissions_guard/src/lib.rs

@@ -108,4 +108,8 @@ impl<T: DmaClient> DmaClient for LowerVtlMemorySpawner<T> {
     fn attach_pending_buffers(&self) -> Result<Vec<MemoryBlock>> {
         anyhow::bail!("restore is not supported for LowerVtlMemorySpawner")
     }
+
+    fn get_dma_buffer(&self, _len: usize, _base_pfn: u64) -> Result<MemoryBlock> {
+        anyhow::bail!("get is not supported for LowerVtlMemorySpawner")
+    }
 }

openhcl/openhcl_boot/src/host_params/dt.rs

@@ -528,6 +528,7 @@ impl PartitionInfo {
             entropy,
             vtl0_alias_map: _,
             nvme_keepalive,
+            mana_keepalive,
         } = storage;
 
         assert!(!vtl2_used_ranges.is_empty());
@@ -550,6 +551,7 @@ impl PartitionInfo {
         *gic = parsed.gic.clone();
         *entropy = parsed.entropy.clone();
         *nvme_keepalive = parsed.nvme_keepalive;
+        *mana_keepalive = parsed.mana_keepalive;
 
         Ok(Some(storage))
     }

openhcl/openhcl_boot/src/host_params/mod.rs

@@ -94,6 +94,8 @@ pub struct PartitionInfo {
     pub vtl0_alias_map: Option<u64>,
     /// Host is compatible with DMA preservation / NVMe keep-alive.
     pub nvme_keepalive: bool,
+    /// Host is compatible with MANA keep-alive.
+    pub mana_keepalive: bool,
 }
 
 impl PartitionInfo {
@@ -125,6 +127,7 @@ impl PartitionInfo {
             entropy: None,
             vtl0_alias_map: None,
             nvme_keepalive: false,
+            mana_keepalive: false,
         }
     }
 

openhcl/openhcl_boot/src/main.rs

@@ -259,6 +259,10 @@ fn build_kernel_command_line(
         write!(cmdline, "OPENHCL_NVME_KEEP_ALIVE=1 ")?;
     }
 
+    if partition_info.mana_keepalive && !partition_info.vtl2_pool_memory.is_empty() {
+        write!(cmdline, "OPENHCL_MANA_KEEP_ALIVE=1 ")?;
+    }
+
     if let Some(sidecar) = sidecar {
         write!(cmdline, "{} ", sidecar.kernel_command_line())?;
     }
@@ -956,6 +960,7 @@ mod test {
             entropy: None,
             vtl0_alias_map: None,
             nvme_keepalive: false,
+            mana_keepalive: false,
         }
     }
 

openhcl/openhcl_dma_manager/src/lib.rs

@@ -369,7 +369,7 @@ impl OpenhclDmaManager {
 
         if let Some(private_pool) = &self.private_pool {
             private_pool
-                .validate_restore(false)
+                .validate_restore(true)
                 .context("failed to validate restore for private pool")?
         }
 
@@ -429,6 +429,22 @@ impl DmaClientBacking {
             DmaClientBacking::LockedMemoryLowerVtl(spawner) => spawner.attach_pending_buffers(),
         }
     }
+
+    fn get_dma_buffer(
+        &self,
+        len: usize,
+        base_pfn: u64,
+    ) -> anyhow::Result<user_driver::memory::MemoryBlock> {
+        match self {
+            DmaClientBacking::SharedPool(allocator) => allocator.get_dma_buffer(len, base_pfn),
+            DmaClientBacking::PrivatePool(allocator) => allocator.get_dma_buffer(len, base_pfn),
+            DmaClientBacking::LockedMemory(spawner) => spawner.get_dma_buffer(len, base_pfn),
+            DmaClientBacking::PrivatePoolLowerVtl(spawner) => spawner.get_dma_buffer(len, base_pfn),
+            DmaClientBacking::LockedMemoryLowerVtl(spawner) => {
+                spawner.get_dma_buffer(len, base_pfn)
+            }
+        }
+    }
 }
 
 /// An OpenHCL dma client. This client implements inspect to allow seeing what
@@ -450,4 +466,12 @@ impl DmaClient for OpenhclDmaClient {
     fn attach_pending_buffers(&self) -> anyhow::Result<Vec<user_driver::memory::MemoryBlock>> {
         self.backing.attach_pending_buffers()
     }
+
+    fn get_dma_buffer(
+        &self,
+        len: usize,
+        base_pfn: u64,
+    ) -> anyhow::Result<user_driver::memory::MemoryBlock> {
+        self.backing.get_dma_buffer(len, base_pfn)
+    }
 }

openhcl/underhill_core/src/dispatch/mod.rs

@@ -33,6 +33,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;
@@ -106,18 +107,25 @@ pub trait LoadedVmNetworkSettings: Inspect {
         threadpool: &AffinitizedThreadpool,
         uevent_listener: &UeventListener,
         servicing_netvsp_state: &Option<Vec<crate::emuplat::netvsp::SavedState>>,
+        servicing_mana_state: &Option<ManaSavedState>,
         partition: Arc<UhPartition>,
         state_units: &StateUnits,
         vmbus_server: &Option<VmbusServerHandle>,
         dma_client_spawner: DmaClientSpawner,
         is_isolated: bool,
+        mana_keepalive: bool,
     ) -> anyhow::Result<RuntimeSavedState>;
 
     /// Callback when network is removed externally.
     async fn remove_network(&mut self, instance_id: Guid) -> anyhow::Result<()>;
 
     /// Callback after stopping the VM and all workers, in preparation for a VTL2 reboot.
-    async fn unload_for_servicing(&mut self);
+    async fn unload_for_servicing(&mut self, mana_keepalive: bool);
+
+    async fn save(
+        &mut self,
+        mana_keepalive_flag: bool,
+    ) -> Option<Vec<Result<ManaSavedState, anyhow::Error>>>;
 
     /// Handles packet capture related operations.
     async fn packet_capture(
@@ -181,6 +189,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,
 }
@@ -268,7 +277,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 {
@@ -330,7 +339,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;
                             }
@@ -495,6 +504,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 {
@@ -509,14 +519,14 @@ 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.
             let shutdown_mana = async {
                 if let Some(network_settings) = self.network_settings.as_mut() {
                     network_settings
-                        .unload_for_servicing()
+                        .unload_for_servicing(mana_keepalive)
                         .instrument(tracing::info_span!("shutdown_mana"))
                         .await;
                 }
@@ -662,7 +672,8 @@ 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());
 
@@ -671,31 +682,58 @@ impl LoadedVm {
         // Only save NVMe state when there are NVMe controllers and keep alive
         // was enabled.
         let nvme_state = if let Some(n) = &self.nvme_manager {
-            n.save(vf_keepalive_flag)
+            n.save(nvme_keepalive_flag)
                 .instrument(tracing::info_span!("nvme_manager_save"))
                 .await
                 .map(|s| NvmeSavedState { nvme_state: s })
         } else {
             None
         };
 
-        let units = self.save_units().await.context("state unit save failed")?;
-        let vmgs = self
-            .vmgs_thin_client
-            .save()
-            .await
-            .context("vmgs 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.
-        let dma_manager_state = if vf_keepalive_flag {
+        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
         };
 
+        let units = self.save_units().await.context("state unit save failed")?;
+
+        let mana_state = if let Some(network_settings) = &mut self.network_settings {
+            let results = network_settings.save(mana_keepalive_flag).await;
+            let mut saved_states = Vec::new();
+
+            if let Some(results) = results {
+                for result in results {
+                    match result {
+                        Ok(state) => saved_states.push(state),
+                        Err(e) => tracing::warn!("Error saving MANA device state: {:#}", e),
+                    }
+                }
+            }
+
+            if !saved_states.is_empty() {
+                Some(saved_states)
+            } else {
+                None
+            }
+        } else {
+            None
+        };
+
+        units.iter().for_each(|unit| {
+            tracing::info!(unit = unit.name, "saved state unit");
+        });
+
+        let vmgs = self
+            .vmgs_thin_client
+            .save()
+            .await
+            .context("vmgs save failed")?;
+
         let vmbus_client = if let Some(vmbus_client) = &mut self.vmbus_client {
             vmbus_client.stop().await;
             Some(vmbus_client.save().await)
@@ -713,6 +751,7 @@ impl LoadedVm {
                 vmgs: (vmgs, self.vmgs_disk_metadata.clone()),
                 overlay_shutdown_device: self.shutdown_relay.is_some(),
                 nvme_state,
+                mana_state,
                 dma_manager_state,
                 vmbus_client,
             },
@@ -776,11 +815,13 @@ impl LoadedVm {
                 threadpool,
                 &self.uevent_listener,
                 &None, // VF getting added; no existing state
+                &None, // VF getting added; no existing state
                 self.partition.clone(),
                 &self.state_units,
                 &self.vmbus_server,
                 self.dma_manager.client_spawner(),
                 self.isolation.is_isolated(),
+                self.mana_keep_alive,
             )
             .await?;