nvme: reset shadow doorbell state

vm/devices/storage/nvme/src/queue.rs

@@ -20,6 +20,7 @@ use vmcore::interrupt::Interrupt;
 
 pub struct DoorbellMemory {
     mem: GuestMemory,
+    private_mem: GuestMemory,
     offset: u64,
     event_idx_offset: Option<u64>,
     wakers: Vec<Option<Waker>>,
@@ -29,14 +30,33 @@ pub struct InvalidDoorbell;
 
 impl DoorbellMemory {
     pub fn new(num_qids: u16) -> Self {
+        let private_mem = GuestMemory::allocate((num_qids as usize) << DOORBELL_STRIDE_BITS);
         Self {
-            mem: GuestMemory::allocate((num_qids as usize) << DOORBELL_STRIDE_BITS),
+            mem: private_mem.clone(),
+            private_mem,
             offset: 0,
             event_idx_offset: None,
             wakers: (0..num_qids).map(|_| None).collect(),
         }
     }
 
+    pub fn reset(&mut self) {
+        let Self {
+            mem,
+            private_mem,
+            offset,
+            event_idx_offset,
+            wakers,
+        } = self;
+        private_mem
+            .fill_at(0, 0, wakers.len() << DOORBELL_STRIDE_BITS)
+            .expect("private doorbell memory must be writable");
+        *mem = private_mem.clone();
+        *offset = 0;
+        *event_idx_offset = None;
+        wakers.fill(None);
+    }
+
     /// Update the memory used to store the doorbell values. This is used to
     /// support shadow doorbells, where the values are directly in guest memory.
     pub fn replace_mem(

vm/devices/storage/nvme/src/tests/shadow_doorbell_tests.rs

@@ -14,6 +14,7 @@ use pal_async::DefaultDriver;
 use pal_async::async_test;
 use pci_core::test_helpers::TestPciInterruptController;
 use user_driver::backoff::Backoff;
+use vmcore::device_state::ChangeDeviceState;
 use zerocopy::FromZeros;
 use zerocopy::IntoBytes;
 
@@ -157,6 +158,29 @@ async fn test_setup_shadow_doorbells(driver: DefaultDriver) {
     setup_shadow_doorbells(driver.clone(), &cq_buf, &sq_buf, &gm, &int_controller, None).await;
 }
 
+#[async_test]
+async fn test_reset_shadow_doorbells(driver: DefaultDriver) {
+    let cq_buf = PrpRange::new(vec![CQ_BASE], 0, PAGE_SIZE64).unwrap();
+    let sq_buf = PrpRange::new(vec![SQ_BASE], 0, PAGE_SIZE64).unwrap();
+    let gm = test_memory();
+    let int_controller = TestPciInterruptController::new();
+
+    let mut nvmec =
+        setup_shadow_doorbells(driver, &cq_buf, &sq_buf, &gm, &int_controller, None).await;
+
+    ChangeDeviceState::reset(&mut nvmec).await;
+
+    let shadow_value = 0x1234;
+    gm.write_plain::<u32>(DOORBELL_BUFFER_BASE, &shadow_value)
+        .unwrap();
+    nvmec.write_bar0(0x1000, 0x5678_u32.as_bytes()).unwrap();
+
+    assert_eq!(
+        gm.read_plain::<u32>(DOORBELL_BUFFER_BASE).unwrap(),
+        shadow_value
+    );
+}
+
 #[async_test]
 async fn test_setup_sq_ring_with_shadow(driver: DefaultDriver) {
     let cq_buf = PrpRange::new(vec![CQ_BASE], 0, PAGE_SIZE64).unwrap();

vm/devices/storage/nvme/src/workers/coordinator.rs

@@ -174,6 +174,7 @@ impl NvmeWorkers {
                 }
             }
         }
+        self.doorbells.write().reset();
     }
 }
 

vm/devices/storage/nvme_test/src/queue.rs

@@ -20,6 +20,7 @@ use vmcore::interrupt::Interrupt;
 
 pub struct DoorbellMemory {
     mem: GuestMemory,
+    private_mem: GuestMemory,
     offset: u64,
     event_idx_offset: Option<u64>,
     wakers: Vec<Option<Waker>>,
@@ -29,14 +30,33 @@ pub struct InvalidDoorbell;
 
 impl DoorbellMemory {
     pub fn new(num_qids: u16) -> Self {
+        let private_mem = GuestMemory::allocate((num_qids as usize) << DOORBELL_STRIDE_BITS);
         Self {
-            mem: GuestMemory::allocate((num_qids as usize) << DOORBELL_STRIDE_BITS),
+            mem: private_mem.clone(),
+            private_mem,
             offset: 0,
             event_idx_offset: None,
             wakers: (0..num_qids).map(|_| None).collect(),
         }
     }
 
+    pub fn reset(&mut self) {
+        let Self {
+            mem,
+            private_mem,
+            offset,
+            event_idx_offset,
+            wakers,
+        } = self;
+        private_mem
+            .fill_at(0, 0, wakers.len() << DOORBELL_STRIDE_BITS)
+            .expect("private doorbell memory must be writable");
+        *mem = private_mem.clone();
+        *offset = 0;
+        *event_idx_offset = None;
+        wakers.fill(None);
+    }
+
     /// Update the memory used to store the doorbell values. This is used to
     /// support shadow doorbells, where the values are directly in guest memory.
     pub fn replace_mem(

vm/devices/storage/nvme_test/src/tests/shadow_doorbell_tests.rs

@@ -16,6 +16,7 @@ use pal_async::DefaultDriver;
 use pal_async::async_test;
 use pci_core::test_helpers::TestPciInterruptController;
 use user_driver::backoff::Backoff;
+use vmcore::device_state::ChangeDeviceState;
 use zerocopy::FromZeros;
 use zerocopy::IntoBytes;
 
@@ -162,6 +163,29 @@ async fn test_setup_shadow_doorbells(driver: DefaultDriver) {
     setup_shadow_doorbells(driver.clone(), &cq_buf, &sq_buf, &gm, &int_controller, None).await;
 }
 
+#[async_test]
+async fn test_reset_shadow_doorbells(driver: DefaultDriver) {
+    let cq_buf = PrpRange::new(vec![CQ_BASE], 0, PAGE_SIZE64).unwrap();
+    let sq_buf = PrpRange::new(vec![SQ_BASE], 0, PAGE_SIZE64).unwrap();
+    let gm = test_memory();
+    let int_controller = TestPciInterruptController::new();
+
+    let mut nvmec =
+        setup_shadow_doorbells(driver, &cq_buf, &sq_buf, &gm, &int_controller, None).await;
+
+    ChangeDeviceState::reset(&mut nvmec).await;
+
+    let shadow_value = 0x1234;
+    gm.write_plain::<u32>(DOORBELL_BUFFER_BASE, &shadow_value)
+        .unwrap();
+    nvmec.write_bar0(0x1000, 0x5678_u32.as_bytes()).unwrap();
+
+    assert_eq!(
+        gm.read_plain::<u32>(DOORBELL_BUFFER_BASE).unwrap(),
+        shadow_value
+    );
+}
+
 #[async_test]
 async fn test_setup_sq_ring_with_shadow(driver: DefaultDriver) {
     let cq_buf = PrpRange::new(vec![CQ_BASE], 0, PAGE_SIZE64).unwrap();

vm/devices/storage/nvme_test/src/workers/coordinator.rs

@@ -177,6 +177,7 @@ impl NvmeWorkers {
                 }
             }
         }
+        self.doorbells.write().reset();
     }
 }