disk_nvme: add specific test for DMA allocation failure handling

Cargo.lock

@@ -4801,6 +4801,7 @@ dependencies = [
  "nvme_resources",
  "nvme_spec",
  "nvme_test",
+ "page_pool_alloc",
  "pal_async",
  "parking_lot",
  "pci_core",

vm/devices/storage/disk_nvme/nvme_driver/Cargo.toml

@@ -35,6 +35,7 @@ nvme.workspace = true
 pci_core.workspace = true
 scsi_buffers.workspace = true
 test_with_tracing.workspace = true
+page_pool_alloc.workspace = true
 user_driver_emulated_mock.workspace = true
 nvme_test.workspace = true
 nvme_resources.workspace = true

vm/devices/storage/disk_nvme/nvme_driver/src/tests.rs

@@ -24,11 +24,14 @@ use parking_lot::Mutex;
 use pci_core::msi::MsiInterruptSet;
 use scsi_buffers::OwnedRequestBuffers;
 use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
 use test_with_tracing::test;
 use user_driver::DeviceBacking;
 use user_driver::DeviceRegisterIo;
 use user_driver::DmaClient;
 use user_driver::interrupt::DeviceInterrupt;
+use user_driver_emulated_mock::DeviceTestDmaClientCallbacks;
 use user_driver_emulated_mock::DeviceTestMemory;
 use user_driver_emulated_mock::EmulatedDevice;
 use user_driver_emulated_mock::Mapping;
@@ -59,12 +62,54 @@ async fn test_nvme_command_fault(driver: DefaultDriver) {
 
 #[async_test]
 async fn test_nvme_driver_direct_dma(driver: DefaultDriver) {
-    test_nvme_driver(driver, true).await;
+    test_nvme_driver(
+        driver,
+        NvmeTestConfig {
+            allow_dma: true,
+            fail_at_driver_create: false,
+            fail_at_io_issuer: false,
+        },
+    )
+    .await;
 }
 
 #[async_test]
 async fn test_nvme_driver_bounce_buffer(driver: DefaultDriver) {
-    test_nvme_driver(driver, false).await;
+    test_nvme_driver(
+        driver,
+        NvmeTestConfig {
+            allow_dma: false,
+            fail_at_driver_create: false,
+            fail_at_io_issuer: false,
+        },
+    )
+    .await;
+}
+
+#[async_test]
+async fn test_nvme_driver_fails_to_create_dma_alloc_failure(driver: DefaultDriver) {
+    test_nvme_driver(
+        driver,
+        NvmeTestConfig {
+            allow_dma: true,
+            fail_at_driver_create: true,
+            fail_at_io_issuer: false,
+        },
+    )
+    .await;
+}
+
+#[async_test]
+async fn test_nvme_driver_fallback_due_to_dma_alloc_failures(driver: DefaultDriver) {
+    test_nvme_driver(
+        driver,
+        NvmeTestConfig {
+            allow_dma: true,
+            fail_at_driver_create: false,
+            fail_at_io_issuer: true,
+        },
+    )
+    .await;
 }
 
 #[async_test]
@@ -146,16 +191,35 @@ async fn test_nvme_ioqueue_invalid_mqes(driver: DefaultDriver) {
     assert!(driver.is_err());
 }
 
-async fn test_nvme_driver(driver: DefaultDriver, allow_dma: bool) {
+struct NvmeTestConfig {
+    allow_dma: bool,
+    fail_at_driver_create: bool,
+    fail_at_io_issuer: bool,
+}
+
+async fn test_nvme_driver(driver: DefaultDriver, config: NvmeTestConfig) {
     const MSIX_COUNT: u16 = 2;
     const IO_QUEUE_COUNT: u16 = 64;
     const CPU_COUNT: u32 = 64;
 
+    let NvmeTestConfig {
+        allow_dma,
+        fail_at_driver_create,
+        fail_at_io_issuer,
+    } = config;
+
     // Arrange: Create 8MB of space. First 4MB for the device and second 4MB for the payload.
     let pages = 1024; // 4MB
     let device_test_memory = DeviceTestMemory::new(pages * 2, allow_dma, "test_nvme_driver");
     let guest_mem = device_test_memory.guest_memory(); // Access to 0-8MB
-    let dma_client = device_test_memory.dma_client(); // Access 0-4MB
+
+    let fail_alloc = Arc::new(AtomicBool::new(false));
+    let dma_client = Arc::new(user_driver_emulated_mock::DeviceTestDmaClient::new(
+        device_test_memory.dma_client(),
+        NvmeTestDmaClientCallbacks {
+            fail_alloc: fail_alloc.clone(),
+        },
+    )); // Access 0-4MB
     let payload_mem = device_test_memory.payload_mem(); // Access 4-8MB. This will allow dma if the `allow_dma` flag is set.
 
     // Arrange: Create the NVMe controller and driver.
@@ -178,6 +242,14 @@ async fn test_nvme_driver(driver: DefaultDriver, allow_dma: bool) {
         .await
         .unwrap();
     let device = NvmeTestEmulatedDevice::new(nvme, msi_set, dma_client.clone());
+
+    if fail_at_driver_create {
+        fail_alloc.store(true, Ordering::SeqCst);
+        let driver_result = NvmeDriver::new(&driver_source, CPU_COUNT, device, false).await;
+        assert!(driver_result.is_err());
+        return;
+    }
+
     let driver = NvmeDriver::new(&driver_source, CPU_COUNT, device, false)
         .await
         .unwrap();
@@ -198,6 +270,11 @@ async fn test_nvme_driver(driver: DefaultDriver, allow_dma: bool) {
         .await
         .unwrap();
 
+    let fallback_count = driver.fallback_cpu_count();
+    if fail_at_io_issuer {
+        fail_alloc.store(true, Ordering::SeqCst);
+    }
+
     // Act: Read 16384 bytes of data from disk starting at LBA 0.
     namespace
         .read(
@@ -217,6 +294,13 @@ async fn test_nvme_driver(driver: DefaultDriver, allow_dma: bool) {
     assert_eq!(&v[512..1536], &[0xcc; 1024]);
     assert!(v[1536..].iter().all(|&x| x == 0));
 
+    // Assert: there should be another fallback CPU only if we forced allocation failure.
+    if fail_at_io_issuer {
+        assert_eq!(driver.fallback_cpu_count(), fallback_count + 1); // New CPU
+    } else {
+        assert_eq!(driver.fallback_cpu_count(), fallback_count);
+    }
+
     namespace
         .deallocate(
             0,
@@ -236,27 +320,30 @@ async fn test_nvme_driver(driver: DefaultDriver, allow_dma: bool) {
         .await
         .unwrap();
 
-    assert_eq!(driver.fallback_cpu_count(), 0);
-
-    // Test the fallback queue functionality.
-    namespace
-        .read(
-            63,
-            0,
-            32,
-            &payload_mem,
-            buf_range.buffer(&guest_mem).range(),
-        )
-        .await
-        .unwrap();
-
-    assert_eq!(driver.fallback_cpu_count(), 1);
-
-    let mut v = [0; 4096];
-    payload_mem.read_at(0, &mut v).unwrap();
-    assert_eq!(&v[..512], &[0; 512]);
-    assert_eq!(&v[512..1024], &[0xcc; 512]);
-    assert!(v[1024..].iter().all(|&x| x == 0));
+    // Test the fallback queue functionality (because of MSI-X limitations). Only do this
+    // if the test didn't already force a fallback due to DMA client allocation failures.
+    if !fail_at_io_issuer {
+        assert_eq!(driver.fallback_cpu_count(), 0);
+
+        namespace
+            .read(
+                63,
+                0,
+                32,
+                &payload_mem,
+                buf_range.buffer(&guest_mem).range(),
+            )
+            .await
+            .unwrap();
+
+        assert_eq!(driver.fallback_cpu_count(), 1);
+
+        let mut v = [0; 4096];
+        payload_mem.read_at(0, &mut v).unwrap();
+        assert_eq!(&v[..512], &[0; 512]);
+        assert_eq!(&v[512..1024], &[0xcc; 512]);
+        assert!(v[1024..].iter().all(|&x| x == 0));
+    }
 
     driver.shutdown().await;
 }
@@ -498,3 +585,30 @@ impl<T: MmioIntercept + Send> DeviceRegisterIo for NvmeTestMapping<T> {
         self.mapping.write_u64(offset, data);
     }
 }
+
+struct NvmeTestDmaClientCallbacks {
+    fail_alloc: Arc<AtomicBool>,
+}
+
+impl DeviceTestDmaClientCallbacks for NvmeTestDmaClientCallbacks {
+    fn allocate_dma_buffer(
+        &self,
+        inner: &page_pool_alloc::PagePoolAllocator,
+        size: usize,
+    ) -> anyhow::Result<user_driver::memory::MemoryBlock> {
+        match self.fail_alloc.load(Ordering::SeqCst) {
+            true => anyhow::bail!("alloc failed"),
+            false => inner.allocate_dma_buffer(size),
+        }
+    }
+
+    fn attach_pending_buffers(
+        &self,
+        inner: &page_pool_alloc::PagePoolAllocator,
+    ) -> anyhow::Result<Vec<user_driver::memory::MemoryBlock>> {
+        match self.fail_alloc.load(Ordering::SeqCst) {
+            true => anyhow::bail!("alloc failed"),
+            false => inner.attach_pending_buffers(),
+        }
+    }
+}

vm/devices/user_driver_emulated_mock/src/lib.rs

@@ -265,3 +265,95 @@ impl DeviceTestMemory {
         self.allocator.clone()
     }
 }
+
+/// Callbacks for the [`DeviceTestDmaClient`]. Tests supply these to customize the behaviour of the dma client.
+pub trait DeviceTestDmaClientCallbacks: Sync + Send {
+    /// Called when the DMA client needs to allocate a new DMA buffer.
+    fn allocate_dma_buffer(
+        &self,
+        allocator: &PagePoolAllocator,
+        total_size: usize,
+    ) -> anyhow::Result<user_driver::memory::MemoryBlock>;
+
+    /// Called when the DMA client needs to attach pending buffers.
+    fn attach_pending_buffers(
+        &self,
+        inner: &PagePoolAllocator,
+    ) -> anyhow::Result<Vec<user_driver::memory::MemoryBlock>>;
+}
+
+/// A DMA client that uses a [`PagePoolAllocator`] as the backing. It can be customized through the use of
+/// [`DeviceTestDmaClientCallbacks`] to modify its behaviour for testing purposes.
+///
+/// # Example
+/// ```rust
+/// use std::sync::Arc;
+/// use user_driver::DmaClient;
+/// use user_driver_emulated_mock::DeviceTestDmaClient;
+/// use page_pool_alloc::PagePoolAllocator;
+///
+/// struct MyCallbacks;
+/// impl user_driver_emulated_mock::DeviceTestDmaClientCallbacks for MyCallbacks {
+///     fn allocate_dma_buffer(
+///         &self,
+///         allocator: &page_pool_alloc::PagePoolAllocator,
+///         total_size: usize,
+///     ) -> anyhow::Result<user_driver::memory::MemoryBlock> {
+///         // Custom test logic here, for example:
+///         anyhow::bail!("allocation failed for testing");
+///     }
+///
+///     fn attach_pending_buffers(
+///         &self,
+///         allocator: &page_pool_alloc::PagePoolAllocator,
+///     ) -> anyhow::Result<Vec<user_driver::memory::MemoryBlock>> {
+///         // Custom test logic here, for example:
+///         anyhow::bail!("attachment failed for testing");
+///     }
+/// }
+///
+/// // Use the above in a test ...
+/// fn test_dma_client() {
+///     let pages = 1000;
+///     let device_test_memory = user_driver_emulated_mock::DeviceTestMemory::new(
+///         pages,
+///         true,
+///         "test_dma_client",
+///     );
+///     let page_pool_allocator = device_test_memory.dma_client();
+///     let dma_client = DeviceTestDmaClient::new(page_pool_allocator, MyCallbacks);
+///
+///     // Use dma_client in tests...
+///     assert!(dma_client.allocate_dma_buffer(4096).is_err());
+/// }
+/// ```
+#[derive(Inspect)]
+#[inspect(transparent)]
+pub struct DeviceTestDmaClient<C>
+where
+    C: DeviceTestDmaClientCallbacks,
+{
+    inner: Arc<PagePoolAllocator>,
+    #[inspect(skip)]
+    callbacks: C,
+}
+
+impl<C: DeviceTestDmaClientCallbacks> DeviceTestDmaClient<C> {
+    /// Creates a new [`DeviceTestDmaClient`] with the given inner allocator.
+    pub fn new(inner: Arc<PagePoolAllocator>, callbacks: C) -> Self {
+        Self { inner, callbacks }
+    }
+}
+
+impl<C: DeviceTestDmaClientCallbacks> DmaClient for DeviceTestDmaClient<C> {
+    fn allocate_dma_buffer(
+        &self,
+        total_size: usize,
+    ) -> anyhow::Result<user_driver::memory::MemoryBlock> {
+        self.callbacks.allocate_dma_buffer(&self.inner, total_size)
+    }
+
+    fn attach_pending_buffers(&self) -> anyhow::Result<Vec<user_driver::memory::MemoryBlock>> {
+        self.callbacks.attach_pending_buffers(&self.inner)
+    }
+}