virtio-devices: mem: Add snapshot/restore support

Adding the snapshot/restore support along with migration as well,
allowing a VM with virtio-mem devices attached to be properly
migrated.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>

virtio-devices/src/mem.rs

@@ -33,12 +33,16 @@ use std::result;
 use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 use std::sync::mpsc;
 use std::sync::{Arc, Barrier, Mutex};
+use versionize::{VersionMap, Versionize, VersionizeResult};
+use versionize_derive::Versionize;
 use vm_device::dma_mapping::ExternalDmaMapping;
 use vm_memory::{
     Address, ByteValued, Bytes, GuestAddress, GuestAddressSpace, GuestMemoryAtomic,
     GuestMemoryError, GuestMemoryRegion,
 };
-use vm_migration::{Migratable, MigratableError, Pausable, Snapshottable, Transportable};
+use vm_migration::{
+    Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable, VersionMapped,
+};
 use vmm_sys_util::eventfd::EventFd;
 
 const QUEUE_SIZE: u16 = 128;
@@ -159,8 +163,8 @@ struct VirtioMemResp {
 unsafe impl ByteValued for VirtioMemResp {}
 
 #[repr(C)]
-#[derive(Copy, Clone, Debug, Default)]
-struct VirtioMemConfig {
+#[derive(Copy, Clone, Debug, Default, Versionize)]
+pub struct VirtioMemConfig {
     // Block size and alignment. Cannot change.
     block_size: u64,
     // Valid with VIRTIO_MEM_F_ACPI_PXM. Cannot change.
@@ -320,14 +324,14 @@ impl Request {
 }
 
 pub struct ResizeSender {
-    size: Arc<AtomicU64>,
+    hotplugged_size: Arc<AtomicU64>,
     tx: mpsc::Sender<Result<(), Error>>,
     evt: EventFd,
 }
 
 impl ResizeSender {
     fn size(&self) -> u64 {
-        self.size.load(Ordering::Acquire)
+        self.hotplugged_size.load(Ordering::Acquire)
     }
 
     fn send(&self, r: Result<(), Error>) -> Result<(), mpsc::SendError<Result<(), Error>>> {
@@ -338,7 +342,7 @@ impl ResizeSender {
 impl Clone for ResizeSender {
     fn clone(&self) -> Self {
         ResizeSender {
-            size: self.size.clone(),
+            hotplugged_size: self.hotplugged_size.clone(),
             tx: self.tx.clone(),
             evt: self
                 .evt
@@ -349,18 +353,18 @@ impl Clone for ResizeSender {
 }
 
 pub struct Resize {
-    size: Arc<AtomicU64>,
+    hotplugged_size: Arc<AtomicU64>,
     tx: mpsc::Sender<Result<(), Error>>,
     rx: mpsc::Receiver<Result<(), Error>>,
     evt: EventFd,
 }
 
 impl Resize {
-    pub fn new() -> io::Result<Self> {
+    pub fn new(hotplugged_size: u64) -> io::Result<Self> {
         let (tx, rx) = mpsc::channel();
 
         Ok(Resize {
-            size: Arc::new(AtomicU64::new(0)),
+            hotplugged_size: Arc::new(AtomicU64::new(hotplugged_size)),
             tx,
             rx,
             evt: EventFd::new(EFD_NONBLOCK)?,
@@ -369,25 +373,28 @@ impl Resize {
 
     pub fn new_resize_sender(&self) -> Result<ResizeSender, Error> {
         Ok(ResizeSender {
-            size: self.size.clone(),
+            hotplugged_size: self.hotplugged_size.clone(),
             tx: self.tx.clone(),
             evt: self.evt.try_clone().map_err(Error::EventFdTryCloneFail)?,
         })
     }
 
-    pub fn work(&self, size: u64) -> Result<(), Error> {
-        self.size.store(size, Ordering::Release);
+    pub fn work(&self, desired_size: u64) -> Result<(), Error> {
+        self.hotplugged_size.store(desired_size, Ordering::Release);
         self.evt.write(1).map_err(Error::EventFdWriteFail)?;
         self.rx.recv().map_err(Error::MpscRecvFail)?
     }
 }
 
-struct BlocksState(Vec<bool>);
+#[derive(Clone, Versionize)]
+pub struct BlocksState {
+    bitmap: Vec<bool>,
+}
 
 impl BlocksState {
     fn is_range_state(&self, first_block_index: usize, nb_blocks: u16, plug: bool) -> bool {
         for state in self
-            .0
+            .bitmap
             .iter()
             .skip(first_block_index)
             .take(nb_blocks as usize)
@@ -401,7 +408,7 @@ impl BlocksState {
 
     fn set_range(&mut self, first_block_index: usize, nb_blocks: u16, plug: bool) {
         for state in self
-            .0
+            .bitmap
             .iter_mut()
             .skip(first_block_index)
             .take(nb_blocks as usize)
@@ -411,7 +418,7 @@ impl BlocksState {
     }
 
     fn inner(&self) -> &Vec<bool> {
-        &self.0
+        &self.bitmap
     }
 }
 
@@ -740,6 +747,16 @@ pub enum VirtioMemMappingSource {
     Device(u32),
 }
 
+#[derive(Versionize)]
+pub struct MemState {
+    pub avail_features: u64,
+    pub acked_features: u64,
+    pub config: VirtioMemConfig,
+    pub blocks_state: BlocksState,
+}
+
+impl VersionMapped for MemState {}
+
 pub struct Mem {
     common: VirtioCommon,
     id: String,
@@ -838,11 +855,9 @@ impl Mem {
             seccomp_action,
             hugepages,
             dma_mapping_handlers: Arc::new(Mutex::new(BTreeMap::new())),
-            blocks_state: Arc::new(Mutex::new(BlocksState(vec![
-                false;
-                (config.region_size / config.block_size)
-                    as usize
-            ]))),
+            blocks_state: Arc::new(Mutex::new(BlocksState {
+                bitmap: vec![false; (config.region_size / config.block_size) as usize],
+            })),
             exit_evt,
         })
     }
@@ -899,6 +914,22 @@ impl Mem {
 
         Ok(())
     }
+
+    fn state(&self) -> MemState {
+        MemState {
+            avail_features: self.common.avail_features,
+            acked_features: self.common.acked_features,
+            config: *(self.config.lock().unwrap()),
+            blocks_state: self.blocks_state.lock().unwrap().clone(),
+        }
+    }
+
+    fn set_state(&mut self, state: &MemState) {
+        self.common.avail_features = state.avail_features;
+        self.common.acked_features = state.acked_features;
+        *(self.config.lock().unwrap()) = state.config;
+        *(self.blocks_state.lock().unwrap()) = state.blocks_state.clone();
+    }
 }
 
 impl Drop for Mem {
@@ -1007,6 +1038,15 @@ impl Snapshottable for Mem {
     fn id(&self) -> String {
         self.id.clone()
     }
+
+    fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
+        Snapshot::new_from_versioned_state(&self.id(), &self.state())
+    }
+
+    fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
+        self.set_state(&snapshot.to_versioned_state(&self.id)?);
+        Ok(())
+    }
 }
 impl Transportable for Mem {}
 impl Migratable for Mem {}

vmm/src/device_manager.rs

@@ -128,7 +128,6 @@ const GPIO_DEVICE_NAME_PREFIX: &str = "_gpio";
 const CONSOLE_DEVICE_NAME: &str = "_console";
 const DISK_DEVICE_NAME_PREFIX: &str = "_disk";
 const FS_DEVICE_NAME_PREFIX: &str = "_fs";
-const MEM_DEVICE_NAME_PREFIX: &str = "_mem";
 const BALLOON_DEVICE_NAME: &str = "_balloon";
 const NET_DEVICE_NAME_PREFIX: &str = "_net";
 const PMEM_DEVICE_NAME_PREFIX: &str = "_pmem";
@@ -2670,20 +2669,19 @@ impl DeviceManager {
 
         let mm = self.memory_manager.clone();
         let mm = mm.lock().unwrap();
-        for (_memory_zone_id, memory_zone) in mm.memory_zones().iter() {
+        for (memory_zone_id, memory_zone) in mm.memory_zones().iter() {
             if let Some(virtio_mem_zone) = memory_zone.virtio_mem_zone() {
-                let id = self.next_device_name(MEM_DEVICE_NAME_PREFIX)?;
-                info!("Creating virtio-mem device: id = {}", id);
+                info!("Creating virtio-mem device: id = {}", memory_zone_id);
 
                 #[cfg(all(target_arch = "x86_64", not(feature = "acpi")))]
                 let node_id: Option<u16> = None;
                 #[cfg(any(target_arch = "aarch64", feature = "acpi"))]
-                let node_id = numa_node_id_from_memory_zone_id(&self.numa_nodes, _memory_zone_id)
+                let node_id = numa_node_id_from_memory_zone_id(&self.numa_nodes, memory_zone_id)
                     .map(|i| i as u16);
 
                 let virtio_mem_device = Arc::new(Mutex::new(
                     virtio_devices::Mem::new(
-                        id.clone(),
+                        memory_zone_id.clone(),
                         virtio_mem_zone.region(),
                         virtio_mem_zone
                             .resize_handler()
@@ -2705,16 +2703,16 @@ impl DeviceManager {
                 devices.push((
                     Arc::clone(&virtio_mem_device) as VirtioDeviceArc,
                     false,
-                    id.clone(),
+                    memory_zone_id.clone(),
                 ));
 
                 // Fill the device tree with a new node. In case of restore, we
                 // know there is nothing to do, so we can simply override the
                 // existing entry.
-                self.device_tree
-                    .lock()
-                    .unwrap()
-                    .insert(id.clone(), device_node!(id, virtio_mem_device));
+                self.device_tree.lock().unwrap().insert(
+                    memory_zone_id.clone(),
+                    device_node!(memory_zone_id, virtio_mem_device),
+                );
             }
         }
 

vmm/src/memory_manager.rs

@@ -698,11 +698,12 @@ impl MemoryManager {
 
                         virtio_mem_regions.push(region.clone());
 
+                        let hotplugged_size = zone.hotplugged_size.unwrap_or(0);
                         memory_zone.virtio_mem_zone = Some(VirtioMemZone {
                             region,
-                            resize_handler: virtio_devices::Resize::new()
+                            resize_handler: virtio_devices::Resize::new(hotplugged_size)
                                 .map_err(Error::EventFdFail)?,
-                            hotplugged_size: zone.hotplugged_size.unwrap_or(0),
+                            hotplugged_size,
                             hugepages: zone.hugepages,
                         });