[feature/virtio-mem] basic snapshot/restore + metrics

src/firecracker/src/api_server/request/hotplug/memory.rs

@@ -25,8 +25,9 @@ pub(crate) fn parse_get_memory_hotplug() -> Result<ParsedRequest, RequestError>
 
 #[cfg(test)]
 mod tests {
-    use vmm::devices::virtio::mem::{VIRTIO_MEM_DEFAULT_BLOCK_SIZE, VIRTIO_MEM_DEFAULT_SLOT_SIZE};
-    use vmm::utils::bytes_to_mib;
+    use vmm::devices::virtio::mem::{
+        VIRTIO_MEM_DEFAULT_BLOCK_SIZE_MIB, VIRTIO_MEM_DEFAULT_SLOT_SIZE_MIB,
+    };
 
     use super::*;
     use crate::api_server::parsed_request::tests::vmm_action_from_request;
@@ -47,8 +48,8 @@ mod tests {
         }"#;
         let expected_config = MemoryHotplugConfig {
             total_size_mib: 2048,
-            block_size_mib: bytes_to_mib(VIRTIO_MEM_DEFAULT_BLOCK_SIZE),
-            slot_size_mib: bytes_to_mib(VIRTIO_MEM_DEFAULT_SLOT_SIZE),
+            block_size_mib: VIRTIO_MEM_DEFAULT_BLOCK_SIZE_MIB,
+            slot_size_mib: VIRTIO_MEM_DEFAULT_SLOT_SIZE_MIB,
         };
         assert_eq!(
             vmm_action_from_request(parse_put_memory_hotplug(&Body::new(body)).unwrap()),

src/vmm/src/device_manager/pci_mngr.rs

@@ -20,6 +20,8 @@ use crate::devices::virtio::block::device::Block;
 use crate::devices::virtio::block::persist::{BlockConstructorArgs, BlockState};
 use crate::devices::virtio::device::VirtioDevice;
 use crate::devices::virtio::generated::virtio_ids;
+use crate::devices::virtio::mem::VirtioMem;
+use crate::devices::virtio::mem::persist::{VirtioMemConstructorArgs, VirtioMemState};
 use crate::devices::virtio::net::Net;
 use crate::devices::virtio::net::persist::{NetConstructorArgs, NetState};
 use crate::devices::virtio::rng::Entropy;
@@ -240,6 +242,8 @@ pub struct PciDevicesState {
     pub mmds: Option<MmdsState>,
     /// Entropy device state.
     pub entropy_device: Option<VirtioDeviceState<EntropyState>>,
+    /// Memory device state.
+    pub memory_device: Option<VirtioDeviceState<VirtioMemState>>,
 }
 
 pub struct PciDevicesConstructorArgs<'a> {
@@ -388,6 +392,20 @@ impl<'a> Persist<'a> for PciDevices {
                         transport_state,
                     })
                 }
+                virtio_ids::VIRTIO_ID_MEM => {
+                    let mem_dev = locked_virtio_dev
+                        .as_mut_any()
+                        .downcast_mut::<VirtioMem>()
+                        .unwrap();
+                    let device_state = mem_dev.save();
+
+                    state.memory_device = Some(VirtioDeviceState {
+                        device_id: mem_dev.id().to_string(),
+                        pci_device_bdf,
+                        device_state,
+                        transport_state,
+                    })
+                }
                 _ => unreachable!(),
             }
         }
@@ -566,6 +584,29 @@ impl<'a> Persist<'a> for PciDevices {
                 .unwrap()
         }
 
+        if let Some(memory_device) = &state.memory_device {
+            let ctor_args = VirtioMemConstructorArgs::new(Arc::clone(constructor_args.vm));
+
+            let device = Arc::new(Mutex::new(
+                VirtioMem::restore(ctor_args, &memory_device.device_state).unwrap(),
+            ));
+
+            constructor_args
+                .vm_resources
+                .update_from_restored_device(SharedDeviceType::VirtioMem(device.clone()))
+                .unwrap();
+
+            pci_devices
+                .restore_pci_device(
+                    constructor_args.vm,
+                    device,
+                    &memory_device.device_id,
+                    &memory_device.transport_state,
+                    constructor_args.event_manager,
+                )
+                .unwrap()
+        }
+
         Ok(pci_devices)
     }
 }
@@ -583,6 +624,7 @@ mod tests {
     use crate::snapshot::Snapshot;
     use crate::vmm_config::balloon::BalloonDeviceConfig;
     use crate::vmm_config::entropy::EntropyDeviceConfig;
+    use crate::vmm_config::memory_hotplug::MemoryHotplugConfig;
     use crate::vmm_config::net::NetworkInterfaceConfig;
     use crate::vmm_config::vsock::VsockDeviceConfig;
 
@@ -645,6 +687,18 @@ mod tests {
             let entropy_config = EntropyDeviceConfig::default();
             insert_entropy_device(&mut vmm, &mut cmdline, &mut event_manager, entropy_config);
 
+            let memory_hotplug_config = MemoryHotplugConfig {
+                total_size_mib: 1024,
+                block_size_mib: 2,
+                slot_size_mib: 128,
+            };
+            insert_virtio_mem_device(
+                &mut vmm,
+                &mut cmdline,
+                &mut event_manager,
+                memory_hotplug_config,
+            );
+
             Snapshot::new(vmm.device_manager.save())
                 .save(&mut buf.as_mut_slice())
                 .unwrap();
@@ -674,7 +728,6 @@ mod tests {
         let _restored_dev_manager =
             PciDevices::restore(restore_args, &device_manager_state.pci_state).unwrap();
 
-        // TODO(virtio-mem): add memory-hotplug device when snapshot-restore is implemented
         let expected_vm_resources = format!(
             r#"{{
   "balloon": {{
@@ -734,7 +787,11 @@ mod tests {
   "entropy": {{
     "rate_limiter": null
   }},
-  "memory-hotplug": null
+  "memory-hotplug": {{
+    "total_size_mib": 1024,
+    "block_size_mib": 2,
+    "slot_size_mib": 128
+  }}
 }}"#,
             _block_files.last().unwrap().as_path().to_str().unwrap(),
             tmp_sock_file.as_path().to_str().unwrap()

src/vmm/src/device_manager/persist.rs

@@ -25,6 +25,10 @@ use crate::devices::virtio::block::device::Block;
 use crate::devices::virtio::block::persist::{BlockConstructorArgs, BlockState};
 use crate::devices::virtio::device::VirtioDevice;
 use crate::devices::virtio::generated::virtio_ids;
+use crate::devices::virtio::mem::VirtioMem;
+use crate::devices::virtio::mem::persist::{
+    VirtioMemConstructorArgs, VirtioMemPersistError, VirtioMemState,
+};
 use crate::devices::virtio::net::Net;
 use crate::devices::virtio::net::persist::{
     NetConstructorArgs, NetPersistError as NetError, NetState,
@@ -72,6 +76,8 @@ pub enum DevicePersistError {
     MmdsConfig(#[from] MmdsConfigError),
     /// Entropy: {0}
     Entropy(#[from] EntropyError),
+    /// virtio-mem: {0}
+    VirtioMem(#[from] VirtioMemPersistError),
     /// Resource misconfiguration: {0}. Is the snapshot file corrupted?
     ResourcesError(#[from] ResourcesError),
     /// Could not activate device: {0}
@@ -125,6 +131,8 @@ pub struct DeviceStates {
     pub mmds: Option<MmdsState>,
     /// Entropy device state.
     pub entropy_device: Option<VirtioDeviceState<EntropyState>>,
+    /// Memory device state.
+    pub memory_device: Option<VirtioDeviceState<VirtioMemState>>,
 }
 
 /// A type used to extract the concrete `Arc<Mutex<T>>` for each of the device
@@ -136,6 +144,7 @@ pub enum SharedDeviceType {
     Balloon(Arc<Mutex<Balloon>>),
     Vsock(Arc<Mutex<Vsock<VsockUnixBackend>>>),
     Entropy(Arc<Mutex<Entropy>>),
+    VirtioMem(Arc<Mutex<VirtioMem>>),
 }
 
 pub struct MMIODevManagerConstructorArgs<'a> {
@@ -335,6 +344,20 @@ impl<'a> Persist<'a> for MMIODeviceManager {
                         device_info,
                     });
                 }
+                virtio_ids::VIRTIO_ID_MEM => {
+                    let mem = locked_device
+                        .as_mut_any()
+                        .downcast_mut::<VirtioMem>()
+                        .unwrap();
+                    let device_state = mem.save();
+
+                    states.memory_device = Some(VirtioDeviceState {
+                        device_id,
+                        device_state,
+                        transport_state,
+                        device_info,
+                    });
+                }
                 _ => unreachable!(),
             };
 
@@ -549,6 +572,30 @@ impl<'a> Persist<'a> for MMIODeviceManager {
             )?;
         }
 
+        if let Some(memory_state) = &state.memory_device {
+            let ctor_args = VirtioMemConstructorArgs::new(Arc::clone(vm));
+
+            let device = Arc::new(Mutex::new(VirtioMem::restore(
+                ctor_args,
+                &memory_state.device_state,
+            )?));
+
+            constructor_args
+                .vm_resources
+                .update_from_restored_device(SharedDeviceType::VirtioMem(device.clone()))?;
+
+            restore_helper(
+                device.clone(),
+                memory_state.device_state.virtio_state.activated,
+                false,
+                device,
+                &memory_state.device_id,
+                &memory_state.transport_state,
+                &memory_state.device_info,
+                constructor_args.event_manager,
+            )?;
+        }
+
         Ok(dev_manager)
     }
 }
@@ -565,6 +612,7 @@ mod tests {
     use crate::snapshot::Snapshot;
     use crate::vmm_config::balloon::BalloonDeviceConfig;
     use crate::vmm_config::entropy::EntropyDeviceConfig;
+    use crate::vmm_config::memory_hotplug::MemoryHotplugConfig;
     use crate::vmm_config::net::NetworkInterfaceConfig;
     use crate::vmm_config::vsock::VsockDeviceConfig;
 
@@ -582,6 +630,7 @@ mod tests {
                 && self.net_devices == other.net_devices
                 && self.vsock_device == other.vsock_device
                 && self.entropy_device == other.entropy_device
+                && self.memory_device == other.memory_device
         }
     }
 
@@ -666,6 +715,18 @@ mod tests {
             let entropy_config = EntropyDeviceConfig::default();
             insert_entropy_device(&mut vmm, &mut cmdline, &mut event_manager, entropy_config);
 
+            let memory_hotplug_config = MemoryHotplugConfig {
+                total_size_mib: 1024,
+                block_size_mib: 2,
+                slot_size_mib: 128,
+            };
+            insert_virtio_mem_device(
+                &mut vmm,
+                &mut cmdline,
+                &mut event_manager,
+                memory_hotplug_config,
+            );
+
             Snapshot::new(vmm.device_manager.save())
                 .save(&mut buf.as_mut_slice())
                 .unwrap();
@@ -691,7 +752,6 @@ mod tests {
         let _restored_dev_manager =
             MMIODeviceManager::restore(restore_args, &device_manager_state.mmio_state).unwrap();
 
-        // TODO(virtio-mem): add memory-hotplug device when snapshot-restore is implemented
         let expected_vm_resources = format!(
             r#"{{
   "balloon": {{
@@ -751,7 +811,11 @@ mod tests {
   "entropy": {{
     "rate_limiter": null
   }},
-  "memory-hotplug": null
+  "memory-hotplug": {{
+    "total_size_mib": 1024,
+    "block_size_mib": 2,
+    "slot_size_mib": 128
+  }}
 }}"#,
             _block_files.last().unwrap().as_path().to_str().unwrap(),
             tmp_sock_file.as_path().to_str().unwrap()

src/vmm/src/devices/virtio/mem/device.rs

@@ -23,6 +23,7 @@ use crate::devices::virtio::generated::virtio_mem::{
     VIRTIO_MEM_F_UNPLUGGED_INACCESSIBLE, virtio_mem_config,
 };
 use crate::devices::virtio::iov_deque::IovDequeError;
+use crate::devices::virtio::mem::metrics::METRICS;
 use crate::devices::virtio::mem::{VIRTIO_MEM_DEV_ID, VIRTIO_MEM_GUEST_ADDRESS};
 use crate::devices::virtio::queue::{FIRECRACKER_MAX_QUEUE_SIZE, InvalidAvailIdx, Queue};
 use crate::devices::virtio::transport::{VirtioInterrupt, VirtioInterruptType};
@@ -125,12 +126,12 @@ impl VirtioMem {
 
     /// Gets the total hotpluggable size.
     pub fn total_size_mib(&self) -> usize {
-        bytes_to_mib(self.config.region_size.try_into().unwrap())
+        bytes_to_mib(u64_to_usize(self.config.region_size))
     }
 
     /// Gets the block size.
     pub fn block_size_mib(&self) -> usize {
-        bytes_to_mib(self.config.block_size.try_into().unwrap())
+        bytes_to_mib(u64_to_usize(self.config.block_size))
     }
 
     /// Gets the block size.
@@ -140,12 +141,12 @@ impl VirtioMem {
 
     /// Gets the total size of the plugged memory blocks.
     pub fn plugged_size_mib(&self) -> usize {
-        bytes_to_mib(self.config.plugged_size.try_into().unwrap())
+        bytes_to_mib(u64_to_usize(self.config.plugged_size))
     }
 
     /// Gets the requested size
     pub fn requested_size_mib(&self) -> usize {
-        bytes_to_mib(self.config.requested_size.try_into().unwrap())
+        bytes_to_mib(u64_to_usize(self.config.requested_size))
     }
 
     pub fn status(&self) -> VirtioMemStatus {
@@ -170,12 +171,15 @@ impl VirtioMem {
     }
 
     pub(crate) fn process_mem_queue_event(&mut self) {
+        METRICS.queue_event_count.inc();
         if let Err(err) = self.queue_events[MEM_QUEUE].read() {
+            METRICS.queue_event_fails.inc();
             error!("Failed to read mem queue event: {err}");
             return;
         }
 
         if let Err(err) = self.process_mem_queue() {
+            METRICS.queue_event_fails.inc();
             error!("virtio-mem: Failed to process queue: {err}");
         }
     }
@@ -233,7 +237,7 @@ impl VirtioDevice for VirtioMem {
     }
 
     fn read_config(&self, offset: u64, data: &mut [u8]) {
-        let offset: usize = u64_to_usize(offset);
+        let offset = u64_to_usize(offset);
         self.config
             .as_slice()
             .get(offset..offset + data.len())
@@ -263,7 +267,7 @@ impl VirtioDevice for VirtioMem {
             error!(
                 "virtio-mem: VIRTIO_MEM_F_UNPLUGGED_INACCESSIBLE feature not acknowledged by guest"
             );
-            // TODO(virtio-mem): activation failed metric
+            METRICS.activate_fails.inc();
             return Err(ActivateError::RequiredFeatureNotAcked(
                 "VIRTIO_MEM_F_UNPLUGGED_INACCESSIBLE",
             ));
@@ -276,7 +280,7 @@ impl VirtioDevice for VirtioMem {
 
         self.device_state = DeviceState::Activated(ActiveState { mem, interrupt });
         if self.activate_event.write(1).is_err() {
-            // TODO(virtio-mem): activation failed metric
+            METRICS.activate_fails.inc();
             self.device_state = DeviceState::Inactive;
             return Err(ActivateError::EventFd);
         }