[DNM] nvme: fused keepalive (#3811)

openhcl/underhill_core/src/nvme_manager/device.rs

@@ -82,6 +82,7 @@ impl CreateNvmeDriver for VfioNvmeDriverSpawner {
         pci_id: &str,
         vp_count: u32,
         save_restore_supported: bool,
+        fused_keepalive_device: bool,
         mut saved_state: Option<&NvmeDriverSavedState>,
     ) -> Result<Box<dyn NvmeDevice>, NvmeSpawnerError> {
         // Gracefully tear down old state & reset device if a saved state is
@@ -175,6 +176,7 @@ impl CreateNvmeDriver for VfioNvmeDriverSpawner {
                 vfio_device,
                 saved_state,
                 self.is_isolated,
+                fused_keepalive_device,
             )
             .instrument(tracing::info_span!("nvme_driver_restore"))
             .await
@@ -186,6 +188,7 @@ impl CreateNvmeDriver for VfioNvmeDriverSpawner {
                 vp_count,
                 self.nvme_always_flr,
                 self.is_isolated,
+                fused_keepalive_device,
                 dma_clients,
             )
             .await?
@@ -224,6 +227,7 @@ impl VfioNvmeDriverSpawner {
         vp_count: u32,
         nvme_always_flr: bool,
         is_isolated: bool,
+        fused_keepalive_device: bool,
         dma_clients: VfioDmaClients,
     ) -> Result<nvme_driver::NvmeDriver<VfioDevice>, NvmeSpawnerError> {
         let mut last_err = None;
@@ -243,6 +247,7 @@ impl VfioNvmeDriverSpawner {
                 pci_id,
                 vp_count,
                 is_isolated,
+                fused_keepalive_device,
                 dma_clients.clone(),
             )
             .await
@@ -280,6 +285,7 @@ impl VfioNvmeDriverSpawner {
         pci_id: &str,
         vp_count: u32,
         is_isolated: bool,
+        fused_keepalive_device: bool,
         dma_clients: VfioDmaClients,
     ) -> Result<nvme_driver::NvmeDriver<VfioDevice>, NvmeSpawnerError> {
         let device = VfioDevice::new(driver_source, pci_id, dma_clients)
@@ -291,10 +297,20 @@ impl VfioNvmeDriverSpawner {
         // TODO: For now, any isolation means use bounce buffering. This
         // needs to change when we have nvme devices that support DMA to
         // confidential memory.
-        nvme_driver::NvmeDriver::new(driver_source, vp_count, device, is_isolated)
-            .instrument(tracing::info_span!("nvme_driver_new", pci_id))
-            .await
-            .map_err(NvmeSpawnerError::DeviceInitFailed)
+        nvme_driver::NvmeDriver::new(
+            driver_source,
+            vp_count,
+            device,
+            is_isolated,
+            fused_keepalive_device,
+        )
+        .instrument(tracing::info_span!(
+            "nvme_driver_new",
+            pci_id,
+            fused_keepalive_device
+        ))
+        .await
+        .map_err(NvmeSpawnerError::DeviceInitFailed)
     }
 
     fn try_update_reset_method(pci_id: &str, method: PciDeviceResetMethod, label: &str) {
@@ -364,6 +380,7 @@ impl NvmeDriverManager {
         pci_id: &str,
         vp_count: u32,
         save_restore_supported: bool,
+        fused_keepalive_device: bool,
         device: Option<Box<dyn NvmeDevice>>,
         nvme_driver_spawner: Arc<dyn CreateNvmeDriver>,
     ) -> anyhow::Result<Self> {
@@ -375,6 +392,7 @@ impl NvmeDriverManager {
             pci_id: pci_id.into(),
             vp_count,
             save_restore_supported,
+            fused_keepalive_device,
             driver: device,
             nvme_driver_spawner,
         };
@@ -494,6 +512,9 @@ struct NvmeDriverManagerWorker {
     vp_count: u32,
     /// Whether the running environment (specifically the VTL2 memory layout) allows save/restore.
     save_restore_supported: bool,
+    /// WORKAROUND: a subset of devices require "fused keepalive". This flag signals to the NVMe
+    /// driver that this device's admin queues may be unusable after a servicing event
+    fused_keepalive_device: bool,
     #[inspect(skip)]
     nvme_driver_spawner: Arc<dyn CreateNvmeDriver>,
     driver: Option<Box<dyn NvmeDevice>>,
@@ -528,6 +549,7 @@ impl NvmeDriverManagerWorker {
                                     &self.pci_id,
                                     self.vp_count,
                                     self.save_restore_supported,
+                                    self.fused_keepalive_device,
                                     None,
                                 )
                                 .await?;

openhcl/underhill_core/src/nvme_manager/manager.rs

@@ -5,6 +5,7 @@ use crate::nvme_manager::CreateNvmeDriver;
 use crate::nvme_manager::device::NvmeDriverManager;
 use crate::nvme_manager::device::NvmeDriverManagerClient;
 use crate::nvme_manager::device::NvmeDriverShutdownOptions;
+use crate::nvme_manager::is_nvme_fused_keepalive_device;
 use crate::nvme_manager::save_restore::NvmeManagerSavedState;
 use crate::nvme_manager::save_restore::NvmeSavedDiskConfig;
 use crate::servicing::NvmeSavedState;
@@ -14,6 +15,7 @@ use disk_backend::resolve::ResolveDiskParameters;
 use disk_backend::resolve::ResolvedDisk;
 use futures::StreamExt;
 use futures::future::join_all;
+use futures::future::try_join_all;
 use inspect::Inspect;
 use mesh::MeshPayload;
 use mesh::rpc::Rpc;
@@ -338,6 +340,8 @@ impl NvmeManagerWorker {
         // Note: `client` exists outside of the devices write lock. This is safe:
         // the mesh client will fail appropriately if shutdown comes in between inserting
         // this entry and the call to `load_driver()`.
+        let fused_keepalive_device = is_nvme_fused_keepalive_device(&pci_id);
+
         let client = {
             let mut guard = context.devices.write();
 
@@ -360,6 +364,7 @@ impl NvmeManagerWorker {
                             &pci_id,
                             context.vp_count,
                             context.save_restore_supported,
+                            fused_keepalive_device,
                             None, // No device yet,
                             context.nvme_driver_spawner.clone(),
                         )?;
@@ -445,29 +450,39 @@ impl NvmeManagerWorker {
         saved_state: &NvmeManagerSavedState,
         save_restore_supported: bool,
     ) -> anyhow::Result<()> {
-        let mut restored_devices: HashMap<String, NvmeDriverManager> = HashMap::new();
-
-        for disk in &saved_state.nvme_disks {
+        let context = &self.context;
+        let created = try_join_all(saved_state.nvme_disks.iter().map(|disk| {
             let pci_id = disk.pci_id.clone();
-            let nvme_driver = self
-                .context
-                .nvme_driver_spawner
-                .create_driver(
-                    &self.context.driver_source,
-                    &pci_id,
-                    saved_state.cpu_count,
-                    save_restore_supported,
-                    Some(&disk.driver_state),
-                )
-                .await?;
+            async move {
+                let fused_keepalive_device = is_nvme_fused_keepalive_device(&pci_id);
+
+                let nvme_driver = context
+                    .nvme_driver_spawner
+                    .create_driver(
+                        &context.driver_source,
+                        &pci_id,
+                        saved_state.cpu_count,
+                        save_restore_supported,
+                        fused_keepalive_device,
+                        Some(&disk.driver_state),
+                    )
+                    .await?;
+
+                anyhow::Ok((pci_id, fused_keepalive_device, nvme_driver))
+            }
+        }))
+        .await?;
 
+        let mut restored_devices: HashMap<String, NvmeDriverManager> = HashMap::new();
+        for (pci_id, fused_keepalive_device, nvme_driver) in created {
             restored_devices.insert(
-                disk.pci_id.clone(),
+                pci_id.clone(),
                 NvmeDriverManager::new(
                     &self.context.driver_source,
                     &pci_id,
                     self.context.vp_count,
                     true, // save_restore_supported is always `true` when restoring.
+                    fused_keepalive_device,
                     Some(nvme_driver),
                     self.context.nvme_driver_spawner.clone(),
                 )?,
@@ -748,6 +763,7 @@ mod tests {
             pci_id: &str,
             _vp_count: u32,
             _save_restore_supported: bool,
+            _fused_keepalive_device: bool,
             _saved_state: Option<&NvmeDriverSavedState>,
         ) -> Result<Box<dyn NvmeDevice>, NvmeSpawnerError> {
             if self.fail_create.load(Ordering::SeqCst) {
@@ -1068,9 +1084,16 @@ mod tests {
         ));
 
         // Create a driver manager
-        let driver_manager =
-            NvmeDriverManager::new(&driver_source, "0000:00:04.0", 4, false, None, spawner)
-                .unwrap();
+        let driver_manager = NvmeDriverManager::new(
+            &driver_source,
+            "0000:00:04.0",
+            4,
+            false,
+            false,
+            None,
+            spawner,
+        )
+        .unwrap();
 
         let client = driver_manager.client().clone();
 
@@ -1153,7 +1176,8 @@ mod tests {
             &driver_source,
             "0000:00:05.0",
             4,
-            true, // save_restore_supported
+            true,  // save_restore_supported
+            false, // fused_keepalive_device
             None,
             spawner,
         )

openhcl/underhill_core/src/nvme_manager/mod.rs

@@ -67,6 +67,14 @@ pub struct NamespaceError {
     source: NvmeSpawnerError,
 }
 
+/// PCI vendor ID, as it appears in the sysfs `vendor` file (e.g. `0x0100`),
+/// for NVMe devices that require fused keepalive device mode.
+const FUSED_DEVICE_VENDOR_ID: &str = "0x1414";
+
+/// PCI device ID, as it appears in the sysfs `device` file (e.g. `0x0100`),
+/// for NVMe devices that require fused keepalive device mode.
+const FUSED_DEVICE_DEVICE_ID: &str = "0xb111";
+
 #[derive(Debug, Error)]
 pub enum NvmeSpawnerError {
     #[error("failed to initialize vfio device")]
@@ -108,6 +116,48 @@ pub trait CreateNvmeDriver: Inspect + Send + Sync {
         pci_id: &str,
         vp_count: u32,
         save_restore_supported: bool,
+        fused_keepalive_device: bool,
         saved_state: Option<&nvme_driver::save_restore::NvmeDriverSavedState>,
     ) -> Result<Box<dyn NvmeDevice>, NvmeSpawnerError>;
 }
+
+/// Returns whether the given PCI device requires fused keepalive device mode
+pub(crate) fn is_nvme_fused_keepalive_device(pci_id: &str) -> bool {
+    match read_pci_vendor_device_ids(pci_id) {
+        Ok((vendor_id, device_id)) => {
+            vendor_id == FUSED_DEVICE_VENDOR_ID && device_id == FUSED_DEVICE_DEVICE_ID
+        }
+        Err(err) => {
+            tracing::warn!(
+                pci_id = %pci_id,
+                error = err.as_ref() as &dyn std::error::Error,
+                "failed to read PCI vendor/device IDs; treating device as a fused keepalive device"
+            );
+            true
+        }
+    }
+}
+
+/// Reads the sysfs `vendor` and `device` files for the given PCI device,
+/// returning the trimmed contents (e.g. `"0x0100"`).
+///
+/// Callers should invoke this once per device and cache the result, since
+/// the values do not change for the lifetime of the device.
+fn read_pci_vendor_device_ids(pci_id: &str) -> anyhow::Result<(String, String)> {
+    let devpath = std::path::Path::new("/sys/bus/pci/devices").join(pci_id);
+    let vendor = fs_err::read_to_string(devpath.join("vendor"))?
+        .trim_end()
+        .to_owned();
+    let device = fs_err::read_to_string(devpath.join("device"))?
+        .trim_end()
+        .to_owned();
+
+    tracing::info!(
+        pci_id = %pci_id,
+        vendor = %vendor,
+        device = %device,
+        "read PCI vendor/device IDs"
+    );
+
+    Ok((vendor, device))
+}

openhcl/underhill_core/src/nvme_manager/save_restore_helpers.rs

@@ -174,6 +174,7 @@ mod tests {
         IoQueueSavedState {
             cpu,
             iv: qid as u32,
+            unmapped: false,
             queue_data: QueuePairSavedState {
                 mem_len: 0,
                 base_pfn: 0,

vm/devices/storage/disk_nvme/nvme_driver/fuzz/fuzz_nvme_driver.rs

@@ -64,7 +64,15 @@ impl FuzzNvmeDriver {
             .unwrap();
 
         let device = FuzzEmulatedDevice::new(nvme, msi_set, mem.dma_client());
-        let mut nvme_driver = NvmeDriver::new(&driver_source, cpu_count, device, false).await?; // TODO: [use-arbitrary-input]
+        let fused_keepalive_device: bool = arbitrary_data::<bool>()?;
+        let mut nvme_driver = NvmeDriver::new(
+            &driver_source,
+            cpu_count,
+            device,
+            false,
+            fused_keepalive_device,
+        )
+        .await?; // TODO: [use-arbitrary-input]
         let namespace = nvme_driver.namespace(1).await?; // TODO: [use-arbitrary-input]
 
         Ok(Self {