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 {}