nvme_driver: add timeout to controller ready/reset wait loops

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

@@ -21,6 +21,7 @@ use crate::queue_pair::QueuePair;
 use crate::queue_pair::admin_cmd;
 use crate::registers::Bar0;
 use crate::registers::DeviceRegisters;
+use crate::registers::ready_timeout;
 use crate::save_restore::NvmeDriverSavedState;
 use anyhow::Context as _;
 use futures::StreamExt;
@@ -31,6 +32,7 @@ use mesh::rpc::Rpc;
 use mesh::rpc::RpcSend;
 use pal_async::task::Spawn;
 use pal_async::task::Task;
+use pal_async::timer::Instant;
 use parking_lot::RwLock;
 use save_restore::NvmeDriverWorkerSavedState;
 use std::collections::HashMap;
@@ -425,6 +427,7 @@ impl<D: DeviceBacking> NvmeDriver<D> {
         );
 
         // Wait for the controller to be ready.
+        let deadline = Instant::now().saturating_add(ready_timeout(worker.registers.cap));
         let mut backoff = Backoff::new(&self.driver);
         loop {
             let csts = worker.registers.bar0.csts();
@@ -448,6 +451,13 @@ impl<D: DeviceBacking> NvmeDriver<D> {
             if csts.rdy() {
                 break;
             }
+            // Give up if the controller never reports ready within CAP.TO.
+            if Instant::now() >= deadline {
+                anyhow::bail!(
+                    "timed out waiting for controller ready, csts: {:#x}",
+                    csts_val
+                );
+            }
             backoff.back_off().await;
         }
         drop(ctrl_enable_span);

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

@@ -6,13 +6,25 @@
 use super::spec;
 use inspect::Inspect;
 use pal_async::driver::Driver;
+use pal_async::timer::Instant;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering::Relaxed;
+use std::time::Duration;
 use tracing::instrument;
 use user_driver::DeviceBacking;
 use user_driver::DeviceRegisterIo;
 use user_driver::backoff::Backoff;
 
+/// Maximum time to wait for CSTS.RDY to change after toggling CC.EN, derived
+/// from CAP.TO (in 500ms units). A CAP.TO of 0 is not useful, so fall back to a
+/// floor of a few seconds.
+pub(crate) fn ready_timeout(cap: spec::Cap) -> Duration {
+    match cap.to() {
+        0 => Duration::from_secs(3),
+        to => Duration::from_millis(500) * to as u32,
+    }
+}
+
 #[derive(Inspect)]
 #[inspect(extra = "Self::inspect_extra")]
 pub(crate) struct DeviceRegisters<T: DeviceBacking> {
@@ -117,6 +129,7 @@ impl<T: DeviceRegisterIo + Inspect> Bar0<T> {
     pub async fn reset(&self, driver: &dyn Driver) -> Result<(), u32> {
         let cc = self.cc().with_en(false);
         self.set_cc(cc);
+        let deadline = Instant::now().saturating_add(ready_timeout(self.cap()));
         let mut backoff = Backoff::new(driver);
         // Loop until either RDY bit is cleared
         // or CSTS read returns -1 which means
@@ -129,6 +142,10 @@ impl<T: DeviceRegisterIo + Inspect> Bar0<T> {
             if u32::from(csts) == !0 {
                 break Err(!0);
             }
+            // Give up if the controller does not clear RDY within CAP.TO.
+            if Instant::now() >= deadline {
+                break Err(u32::from(csts));
+            }
             backoff.back_off().await;
         }
     }

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

@@ -283,6 +283,46 @@ async fn test_nvme_ioqueue_invalid_mqes(driver: DefaultDriver) {
     assert!(driver.is_err());
 }
 
+#[async_test]
+async fn test_nvme_controller_ready_timeout(driver: DefaultDriver) {
+    const MSIX_COUNT: u16 = 2;
+    const IO_QUEUE_COUNT: u16 = 64;
+    const CPU_COUNT: u32 = 64;
+
+    // Memory setup
+    let pages = 1000;
+    let device_test_memory =
+        DeviceTestMemory::new(pages, false, "test_nvme_controller_ready_timeout");
+    let guest_mem = device_test_memory.guest_memory();
+    let dma_client = device_test_memory.dma_client();
+
+    let driver_source = VmTaskDriverSource::new(SingleDriverBackend::new(driver));
+    let msi_conn = MsiConnection::new();
+    let dma_target = DmaTarget::new(AssignedBusRange::new(), 0, guest_mem.clone(), &msi_conn);
+    let nvme = nvme::NvmeController::new(
+        &driver_source,
+        &dma_target,
+        &mut ExternallyManagedMmioIntercepts,
+        NvmeControllerCaps {
+            msix_count: MSIX_COUNT,
+            max_io_queues: IO_QUEUE_COUNT,
+            subsystem_id: Guid::new_random(),
+        },
+    );
+
+    let mut device = NvmeTestEmulatedDevice::new(nvme, msi_conn, dma_client.clone());
+
+    // Report a short CAP.TO and pin CSTS so RDY never sets, forcing the
+    // controller-ready wait loop to hit its deadline.
+    let cap: Cap = Cap::new().with_to(1);
+    device.set_mock_response_u64(Some((0, cap.into())));
+    device.set_mock_response_u32(Some((0x1c, 0)));
+
+    let driver = NvmeDriver::new(&driver_source, CPU_COUNT, device, false).await;
+
+    assert!(driver.is_err());
+}
+
 struct NvmeTestConfig {
     allow_dma: bool,
     fail_at_driver_create: bool,
@@ -538,7 +578,11 @@ impl<T: PciConfigSpace + MmioIntercept + InspectMut, U: DmaClient> NvmeTestEmula
         }
     }
 
-    // TODO: set_mock_response_u32 is intentionally not implemented to avoid dead code.
+    pub fn set_mock_response_u32(&mut self, mapping: Option<(usize, u32)>) {
+        let mut mock_response = self.mocked_response_u32.lock();
+        *mock_response = mapping;
+    }
+
     pub fn set_mock_response_u64(&mut self, mapping: Option<(usize, u64)>) {
         let mut mock_response = self.mocked_response_u64.lock();
         *mock_response = mapping;