nvme_driver: test: hardening the save_restore testing for the nvme driver to verify keepalive functionality

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

@@ -138,6 +138,14 @@ impl IoQueue {
             cpu: *cpu,
         })
     }
+
+    #[cfg(test)]
+    pub(crate) async fn verify_restore(&self, saved_state: &IoQueueSavedState, mem: MemoryBlock) {
+        self.queue.verify_restore(&saved_state.queue_data, mem);
+
+        assert_eq!(saved_state.iv, self.iv as u32);
+        assert_eq!(saved_state.cpu, self.cpu);
+    }
 }
 
 #[derive(Debug, Inspect)]
@@ -683,6 +691,49 @@ impl<T: DeviceBacking> NvmeDriver<T> {
     pub fn update_servicing_flags(&mut self, nvme_keepalive: bool) {
         self.nvme_keepalive = nvme_keepalive;
     }
+
+    /// Given an input of the saved state from which the driver was constructed and the underlying
+    /// memory, this validates the current driver.
+    #[cfg(test)]
+    pub(crate) async fn verify_restore(&mut self, saved_state: &NvmeDriverSavedState, mem: MemoryBlock) {
+        if let Some(task) = self.task.as_mut() {
+            task.stop().await;
+            let worker = task.task();
+
+            // Verify Admin Queue
+            match (&saved_state.worker_data.admin, &worker.admin) {
+                (None, None) => (),
+                (Some(admin_saved_state), Some(admin)) => {
+                    // TODO: [expand-verify-restore-functionality] Currrently providing base_pfn value in u64, this might panic
+                    let admin_saved_mem = mem.subblock(admin_saved_state.base_pfn.try_into().unwrap(), admin_saved_state.mem_len);
+                    admin.verify_restore(&admin_saved_state, admin_saved_mem).await;
+                },
+                _ => panic!("admin queue states do not match"),
+            };
+
+            // Verify I/O queues in strict ordering
+            assert_eq!(saved_state.worker_data.io.len(), worker.io.len());
+            for index in 0..saved_state.worker_data.io.len() {
+                let io_saved_state = saved_state.worker_data.io[index].clone();
+                let io_saved_mem = mem.subblock(io_saved_state.queue_data.base_pfn.try_into().unwrap(), io_saved_state.queue_data.mem_len);
+                worker.io[index].verify_restore(&io_saved_state, io_saved_mem);
+            }
+            task.start();
+        } else {
+            panic!("task cannot be None() after restore");
+        }
+
+        assert_eq!(saved_state.device_id, self.device_id);
+
+        if let Some(identify) = &self.identify {
+            assert_eq!(saved_state.identify_ctrl.as_bytes(), identify.as_bytes());
+        } else {
+            panic!("idenitfy value cannot be None after restore");
+        }
+
+        // TODO: [expand-verify-restore-functionality] Namespace save is currently not supported.
+        assert!(self.nvme_keepalive);
+    }
 }
 
 async fn handle_asynchronous_events(

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

@@ -163,6 +163,19 @@ impl PendingCommands {
             next_cid_high_bits: Wrapping(*next_cid_high_bits),
         })
     }
+
+    /// Given the saved state, verifies the state of the PendingCommands to match the saved state
+    #[cfg(test)]
+    pub(crate) fn verify_restore(&self, saved_state: &PendingCommandsSavedState) {
+        // TODO: [expand-verify-restore-functionality] cid_key_bits are currently unused during restore. 
+        assert_eq!(saved_state.commands.len(), self.commands.len());
+
+        for (index, command) in &self.commands {
+            command.verify_restore(&saved_state.commands[index]);
+        }
+
+        assert_eq!(saved_state.next_cid_high_bits, self.next_cid_high_bits.0);
+    }
 }
 
 impl QueuePair {
@@ -324,6 +337,33 @@ impl QueuePair {
             Some(handler_data),
         )
     }
+
+    /// Given the saved state of a queue pair, this verifies the constructed queue pair.
+    /// Input memory block should already be constructed from the offsets.
+    #[cfg(test)]
+    pub(crate) async fn verify_restore(&self, saved_state: &QueuePairSavedState, saved_mem: MemoryBlock) {
+        // Entire memory region is checked below. No need for the the handler to check it again.
+        let _ = self.issuer.send.call(Req::Verify, saved_state.handler_data.clone()).await;
+
+        // `cancel` and `issuers` params are runtime parameters so we don't check underlying values.
+        let mut saved_mem_data: [u8; PAGE_SIZE] = [0; PAGE_SIZE];
+        let mut self_mem_data: [u8; PAGE_SIZE] = [0; PAGE_SIZE];
+
+        assert_eq!(saved_mem.len(), self.mem.len());
+
+        for pfn in 0..(saved_mem.len()/PAGE_SIZE) {
+            saved_mem.read_at(pfn * PAGE_SIZE, &mut saved_mem_data);
+            self.mem.read_at(pfn * PAGE_SIZE, &mut self_mem_data);
+
+            for i in 0..PAGE_SIZE {
+                assert_eq!(saved_mem_data[i], self_mem_data[i]);
+            }
+        }
+
+        assert_eq!(saved_state.qid, self.qid);
+        assert_eq!(saved_state.sq_entries, self.sq_entries);
+        assert_eq!(saved_state.cq_entries, self.cq_entries);
+    }
 }
 
 /// An error issuing an NVMe request.
@@ -583,6 +623,8 @@ enum Req {
     Command(Rpc<spec::Command, spec::Completion>),
     Inspect(inspect::Deferred),
     Save(Rpc<(), Result<QueueHandlerSavedState, anyhow::Error>>),
+    #[cfg(test)]
+    Verify(Rpc<QueueHandlerSavedState, ()>),
 }
 
 #[derive(Inspect)]
@@ -668,6 +710,14 @@ impl QueueHandler {
                         // Do not allow any more processing after save completed.
                         break;
                     }
+                    #[cfg(test)]
+                    Req::Verify(verify_state) => {
+                        let saved_state = verify_state.input();
+
+                        self.sq.verify_restore(&saved_state.sq_state);
+                        self.cq.verify_restore(&saved_state.cq_state);
+                        self.commands.verify_restore(&saved_state.pending_cmds);
+                    }
                 },
                 Event::Completion(completion) => {
                     assert_eq!(completion.sqid, self.sq.id());
@@ -718,6 +768,14 @@ impl QueueHandler {
     }
 }
 
+#[cfg(test)]
+impl PendingCommand {
+    /// Verifies the value of the pending command from a saved state
+    pub(crate) fn verify_restore(&self, saved_state: &PendingCommandSavedState) {
+        assert_eq!(saved_state.command, self.command);
+    }
+}
+
 pub(crate) fn admin_cmd(opcode: spec::AdminOpcode) -> spec::Command {
     spec::Command {
         cdw0: spec::Cdw0::new().with_opcode(opcode.0),

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

@@ -106,6 +106,17 @@ impl SubmissionQueue {
             mem,
         })
     }
+
+    /// Given the saved state, checks the state of the submission queue. Does not verify memory.
+    // TODO: Can this be an associated function instead?
+    #[cfg(test)]
+    pub(crate) fn verify_restore(&self, saved_state: &SubmissionQueueSavedState) {
+        assert_eq!(saved_state.sqid, self.sqid);
+        assert_eq!(saved_state.head, self.head);
+        assert_eq!(saved_state.tail, self.tail);
+        assert_eq!(saved_state.committed_tail, self.committed_tail);
+        assert_eq!(saved_state.len, self.len);
+    }
 }
 
 #[derive(Inspect)]
@@ -201,6 +212,16 @@ impl CompletionQueue {
             mem,
         })
     }
+
+    /// Given the saved state, checks the state of the completion queue. Does not verify memory.
+    #[cfg(test)]
+    pub(crate) fn verify_restore(&self, saved_state: &CompletionQueueSavedState) {
+        assert_eq!(saved_state.cqid, self.cqid);
+        assert_eq!(saved_state.head, self.head);
+        assert_eq!(saved_state.committed_head, self.committed_head);
+        assert_eq!(saved_state.len, self.len);
+        assert_eq!(saved_state.phase, self.phase);
+    }
 }
 
 fn advance(n: u32, l: u32) -> u32 {

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

@@ -20,6 +20,7 @@ use std::sync::Arc;
 use test_with_tracing::test;
 use user_driver::emulated::DeviceSharedMemory;
 use user_driver::emulated::EmulatedDevice;
+use user_driver::emulated::EmulatedDmaAllocator;
 use user_driver::emulated::Mapping;
 use user_driver::interrupt::DeviceInterrupt;
 use user_driver::DeviceBacking;
@@ -233,6 +234,7 @@ async fn test_nvme_driver(driver: DefaultDriver, allow_dma: bool) {
 }
 
 async fn test_nvme_save_restore_inner(driver: DefaultDriver) {
+    // ===== SHARED RESOURCES =====
     const MSIX_COUNT: u16 = 2;
     const IO_QUEUE_COUNT: u16 = 64;
     const CPU_COUNT: u32 = 64;
@@ -261,12 +263,16 @@ async fn test_nvme_save_restore_inner(driver: DefaultDriver) {
         .await
         .unwrap();
 
-    let device = EmulatedDevice::new(nvme_ctrl, msi_x, mem);
+    // ===== FIRST DRIVER INIT =====
+    let device = EmulatedDevice::new(nvme_ctrl, msi_x, mem.clone());
     let mut nvme_driver = NvmeDriver::new(&driver_source, CPU_COUNT, device)
         .await
         .unwrap();
     let _ns1 = nvme_driver.namespace(1).await.unwrap();
     let saved_state = nvme_driver.save().await.unwrap();
+
+    // Tear down the original driver to kill the underlying tasks.
+    nvme_driver.shutdown().await;
     // As of today we do not save namespace data to avoid possible conflict
     // when namespace has changed during servicing.
     // TODO: Review and re-enable in future.
@@ -299,11 +305,20 @@ async fn test_nvme_save_restore_inner(driver: DefaultDriver) {
     // Wait for CSTS.RDY to set.
     backoff.back_off().await;
 
-    let _new_device = EmulatedDevice::new(new_nvme_ctrl, new_msi_x, new_emu_mem);
-    // TODO: Memory restore is disabled for emulated DMA, uncomment once fixed.
-    // let _new_nvme_driver = NvmeDriver::restore(&driver_source, CPU_COUNT, new_device, &saved_state)
-    //     .await
-    //     .unwrap();
+    // ====== SECOND DRIVER INIT =====
+    let mem_new = DeviceSharedMemory::new(base_len, payload_len);
+    let new_device = EmulatedDevice::new(new_nvme_ctrl, new_msi_x, mem_new.clone());
+    let mut new_nvme_driver = NvmeDriver::restore(&driver_source, CPU_COUNT, new_device, &saved_state)
+        .await
+        .unwrap();
+
+
+    // ===== VERIFY RESTORE =====
+    let host_allocator = EmulatedDmaAllocator::new(mem_new.clone());
+    let verify_mem = DmaClient::attach_dma_buffer(&host_allocator, base_len, 0).unwrap();
+
+    // Verify restore functions will panic if verification failed.
+    new_nvme_driver.verify_restore(&saved_state, verify_mem).await;
 }
 
 #[derive(Inspect)]

vm/devices/storage/nvme_spec/src/lib.rs

@@ -127,7 +127,7 @@ pub struct Aqa {
 }
 
 #[repr(C)]
-#[derive(Copy, Clone, Debug, IntoBytes, Immutable, KnownLayout, FromBytes, Inspect)]
+#[derive(Copy, Clone, Debug, IntoBytes, Immutable, KnownLayout, FromBytes, Inspect, PartialEq)]
 pub struct Command {
     pub cdw0: Cdw0,
     pub nsid: u32,
@@ -146,7 +146,7 @@ pub struct Command {
 
 #[derive(Inspect)]
 #[bitfield(u32)]
-#[derive(IntoBytes, Immutable, KnownLayout, FromBytes)]
+#[derive(IntoBytes, Immutable, KnownLayout, FromBytes, PartialEq)]
 pub struct Cdw0 {
     pub opcode: u8,
     #[bits(2)]

vm/devices/user_driver/src/emulated.rs

@@ -224,6 +224,21 @@ impl DeviceSharedMemory {
             state: self.state.clone(),
         })
     }
+
+    // TODO: [nvme-keepalive-testing] 
+    // This is only a stop-gap until we can swap out the back end of nvme tests to use real memory
+    pub fn alloc_specific(&self, len: usize, base_pfn: u64) -> Option<DmaBuffer>{
+        assert!(len % PAGE_SIZE == 0);
+        let count = len / PAGE_SIZE;
+        let start_page = base_pfn as usize;
+
+        let pages = (start_page..start_page + count).map(|p| p as u64).collect();
+        Some(DmaBuffer {
+            mem: self.mem.clone(),
+            pfns: pages,
+            state: self.state.clone(),
+        })
+    }
 }
 
 pub struct DmaBuffer {
@@ -270,15 +285,24 @@ pub struct EmulatedDmaAllocator {
     shared_mem: DeviceSharedMemory,
 }
 
+impl EmulatedDmaAllocator {
+    pub fn new(shared_mem: DeviceSharedMemory) -> Self {
+        EmulatedDmaAllocator {
+            shared_mem,
+        }
+    }
+}
+
 impl DmaClient for EmulatedDmaAllocator {
     fn allocate_dma_buffer(&self, len: usize) -> anyhow::Result<MemoryBlock> {
         let memory = MemoryBlock::new(self.shared_mem.alloc(len).context("out of memory")?);
         memory.as_slice().atomic_fill(0);
         Ok(memory)
     }
 
-    fn attach_dma_buffer(&self, _len: usize, _base_pfn: u64) -> anyhow::Result<MemoryBlock> {
-        anyhow::bail!("restore is not supported for emulated DMA")
+    fn attach_dma_buffer(&self, len: usize, base_pfn: u64) -> anyhow::Result<MemoryBlock> {
+        let memory = MemoryBlock::new( self.shared_mem.alloc_specific(len, base_pfn).context("could not alloc specific. out of memory")?);
+        Ok(memory)
     }
 }