[DRAFT] nvme_driver: fix overflow computing queue size and don't wait for bad devices forever

openhcl/underhill_core/src/nvme_manager.rs

@@ -303,7 +303,7 @@ impl NvmeManagerWorker {
                 .map_err(InnerError::Vfio)?;
 
                 let driver =
-                    nvme_driver::NvmeDriver::new(&self.driver_source, self.vp_count, device)
+                    nvme_driver::NvmeDriver::new(&self.driver_source, self.vp_count, device, None)
                         .instrument(tracing::info_span!(
                             "nvme_driver_init",
                             pci_id = entry.key()

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

@@ -13,11 +13,13 @@ use nvme::NvmeController;
 use nvme::NvmeControllerCaps;
 use nvme_driver::Namespace;
 use nvme_driver::NvmeDriver;
+use nvme_driver::NVME_TIMEOUT_FUZZER;
 use nvme_spec::nvm::DsmRange;
 use pal_async::DefaultDriver;
 use pci_core::msi::MsiInterruptSet;
 use scsi_buffers::OwnedRequestBuffers;
 use std::convert::TryFrom;
+use std::time::Duration;
 use user_driver::emulated::DeviceSharedMemory;
 use vmcore::vm_task::SingleDriverBackend;
 use vmcore::vm_task::VmTaskDriverSource;
@@ -67,7 +69,8 @@ impl FuzzNvmeDriver {
             .unwrap();
 
         let device = FuzzEmulatedDevice::new(nvme, msi_set, mem);
-        let nvme_driver = NvmeDriver::new(&driver_source, cpu_count, device).await?; // TODO: [use-arbitrary-input]
+        let timeout = Duration::from_millis(NVME_TIMEOUT_FUZZER);
+        let nvme_driver = NvmeDriver::new(&driver_source, cpu_count, device, Some(timeout)).await?; // TODO: [use-arbitrary-input]
         let namespace = nvme_driver.namespace(1).await?; // TODO: [use-arbitrary-input]
 
         Ok(Self {

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

@@ -15,6 +15,8 @@ use crate::NamespaceError;
 use crate::NvmeDriverSavedState;
 use crate::RequestError;
 use crate::NVME_PAGE_SHIFT;
+use crate::NVME_TIMEOUT_INCREMENT;
+use crate::NVME_TIMEOUT_MINIMUM;
 use anyhow::Context as _;
 use futures::future::join_all;
 use futures::StreamExt;
@@ -27,6 +29,7 @@ use pal_async::task::Task;
 use save_restore::NvmeDriverWorkerSavedState;
 use std::sync::Arc;
 use std::sync::OnceLock;
+use std::time::Duration;
 use task_control::AsyncRun;
 use task_control::InspectTask;
 use task_control::TaskControl;
@@ -71,6 +74,8 @@ pub struct NvmeDriver<T: DeviceBacking> {
     namespaces: Vec<Arc<Namespace>>,
     /// Keeps the controller connected (CC.EN==1) while servicing.
     nvme_keepalive: bool,
+    #[inspect(skip)]
+    reset_timeout: Duration,
 }
 
 #[derive(Inspect)]
@@ -165,14 +170,16 @@ enum NvmeWorkerRequest {
 }
 
 impl<T: DeviceBacking> NvmeDriver<T> {
-    /// Initializes the driver.
+    /// Initializes the driver. Only provide reset timeout for testing.
+    /// If no value provided, timeout from device capabilites will be used.
     pub async fn new(
         driver_source: &VmTaskDriverSource,
         cpu_count: u32,
         device: T,
+        reset_timeout: Option<Duration>,
     ) -> anyhow::Result<Self> {
         let pci_id = device.id().to_owned();
-        let mut this = Self::new_disabled(driver_source, cpu_count, device)
+        let mut this = Self::new_disabled(driver_source, cpu_count, device, reset_timeout)
             .instrument(tracing::info_span!("nvme_new_disabled", pci_id))
             .await?;
         match this
@@ -198,6 +205,7 @@ impl<T: DeviceBacking> NvmeDriver<T> {
         driver_source: &VmTaskDriverSource,
         cpu_count: u32,
         mut device: T,
+        reset_timeout: Option<Duration>
     ) -> anyhow::Result<Self> {
         let driver = driver_source.simple();
         let bar0 = Bar0(
@@ -206,10 +214,15 @@ impl<T: DeviceBacking> NvmeDriver<T> {
                 .context("failed to map device registers")?,
         );
 
+        let timeout = match reset_timeout {
+            Some(duration) => duration,
+            None => Duration::from_millis(NVME_TIMEOUT_MINIMUM.max(bar0.cap().to() as u64 * NVME_TIMEOUT_INCREMENT)),
+        };
+
         let cc = bar0.cc();
         if cc.en() || bar0.csts().rdy() {
             if !bar0
-                .reset(&driver)
+                .reset(&driver, timeout.clone())
                 .instrument(tracing::info_span!(
                     "nvme_already_enabled",
                     pci_id = device.id().to_owned()
@@ -254,6 +267,7 @@ impl<T: DeviceBacking> NvmeDriver<T> {
             rescan_event: Default::default(),
             namespaces: vec![],
             nvme_keepalive: false,
+            reset_timeout: timeout,
         })
     }
 
@@ -319,7 +333,7 @@ impl<T: DeviceBacking> NvmeDriver<T> {
                 anyhow::bail!("device is gone");
             }
             if csts.cfs() {
-                worker.registers.bar0.reset(&self.driver).await;
+                worker.registers.bar0.reset(&self.driver, self.reset_timeout.clone()).await;
                 anyhow::bail!("device had fatal error");
             }
             if csts.rdy() {
@@ -399,8 +413,8 @@ impl<T: DeviceBacking> NvmeDriver<T> {
         let max_io_queues = allocated_io_queue_count.min(requested_io_queue_count);
 
         let qsize = {
-            let io_cqsize = QueuePair::MAX_CQ_ENTRIES.min(worker.registers.cap.mqes_z() + 1);
-            let io_sqsize = QueuePair::MAX_SQ_ENTRIES.min(worker.registers.cap.mqes_z() + 1);
+            let io_cqsize = QueuePair::MAX_CQ_ENTRIES.min(worker.registers.cap.mqes_z().saturating_add(1));
+            let io_sqsize = QueuePair::MAX_SQ_ENTRIES.min(worker.registers.cap.mqes_z().saturating_add(1));
 
             // Some hardware (such as ASAP) require that the sq and cq have the same size.
             io_cqsize.min(io_sqsize)
@@ -454,6 +468,7 @@ impl<T: DeviceBacking> NvmeDriver<T> {
 
     fn reset(&mut self) -> impl 'static + Send + std::future::Future<Output = ()> {
         let driver = self.driver.clone();
+        let timeout = self.reset_timeout.clone();
         let mut task = std::mem::take(&mut self.task).unwrap();
         async move {
             task.stop().await;
@@ -468,7 +483,7 @@ impl<T: DeviceBacking> NvmeDriver<T> {
             if let Some(admin) = worker.admin {
                 _admin_responses = admin.shutdown().await;
             }
-            worker.registers.bar0.reset(&driver).await;
+            worker.registers.bar0.reset(&driver, timeout).await;
         }
     }
 
@@ -549,6 +564,8 @@ impl<T: DeviceBacking> NvmeDriver<T> {
         if !bar0.csts().rdy() {
             anyhow::bail!("device is gone");
         }
+
+        let reset_timeout = Duration::from_millis(NVME_TIMEOUT_MINIMUM.max(bar0.cap().to() as u64 * NVME_TIMEOUT_INCREMENT));
 
         let registers = Arc::new(DeviceRegisters::new(bar0));
 
@@ -557,7 +574,7 @@ impl<T: DeviceBacking> NvmeDriver<T> {
             per_cpu: (0..cpu_count).map(|_| OnceLock::new()).collect(),
             send,
         });
-
+ 
         let mut this = Self {
             device_id: device.id().to_owned(),
             task: Some(TaskControl::new(DriverWorkerTask {
@@ -579,6 +596,7 @@ impl<T: DeviceBacking> NvmeDriver<T> {
             rescan_event: Default::default(),
             namespaces: vec![],
             nvme_keepalive: true,
+            reset_timeout,
         };
 
         let task = &mut this.task.as_mut().unwrap();

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

@@ -25,3 +25,6 @@ pub use self::queue_pair::RequestError;
 use nvme_spec as spec;
 
 const NVME_PAGE_SHIFT: u8 = 12;
+const NVME_TIMEOUT_INCREMENT: u64 = 500;
+const NVME_TIMEOUT_MINIMUM: u64 = 10000;
+pub const NVME_TIMEOUT_FUZZER: u64 = 10;

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

@@ -8,6 +8,8 @@ use inspect::Inspect;
 use pal_async::driver::Driver;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering::Relaxed;
+use std::time::Duration;
+use std::time::Instant;
 use tracing::instrument;
 use user_driver::backoff::Backoff;
 use user_driver::DeviceBacking;
@@ -104,10 +106,11 @@ impl<T: DeviceRegisterIo + Inspect> Bar0<T> {
     reg32!(aqa, set_aqa, AQA, spec::Aqa);
 
     #[instrument(skip_all)]
-    pub async fn reset(&self, driver: &dyn Driver) -> bool {
+    pub async fn reset(&self, driver: &dyn Driver, reset_timeout: Duration) -> bool {
         let cc = self.cc().with_en(false);
         self.set_cc(cc);
         let mut backoff = Backoff::new(driver);
+        let start = Instant::now();
         loop {
             let csts = self.csts();
             if !csts.rdy() {
@@ -116,6 +119,9 @@ impl<T: DeviceRegisterIo + Inspect> Bar0<T> {
             if u32::from(csts) == !0 {
                 break false;
             }
+            if start.elapsed() >= reset_timeout {
+                break false;
+            }
             backoff.back_off().await;
         }
     }

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

@@ -74,7 +74,7 @@ async fn test_nvme_driver(driver: DefaultDriver, allow_dma: bool) {
 
     let device = EmulatedDevice::new(nvme, msi_set, mem);
 
-    let driver = NvmeDriver::new(&driver_source, CPU_COUNT, device)
+    let driver = NvmeDriver::new(&driver_source, CPU_COUNT, device, None)
         .await
         .unwrap();
 
@@ -183,7 +183,7 @@ async fn test_nvme_save_restore_inner(driver: DefaultDriver) {
         .unwrap();
 
     let device = EmulatedDevice::new(nvme_ctrl, msi_x, mem);
-    let mut nvme_driver = NvmeDriver::new(&driver_source, CPU_COUNT, device)
+    let mut nvme_driver = NvmeDriver::new(&driver_source, CPU_COUNT, device, None)
         .await
         .unwrap();
     let _ns1 = nvme_driver.namespace(1).await.unwrap();