misc: Simplify snapshot/restore by using helper functions

Simplify snapshot & restore code by using generics to specify helper
functions that take / make a Serialize / Deserialize struct

Signed-off-by: Rob Bradford <robert.bradford@intel.com>

arch/src/aarch64/gic/gicv3.rs

@@ -15,8 +15,7 @@ pub mod kvm {
     use std::sync::Arc;
     use std::{boxed::Box, result};
     use vm_migration::{
-        Migratable, MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable,
+        Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable,
-        Transportable,
     };
 
     /// Errors thrown while saving/restoring the GICv3.
@@ -230,42 +229,15 @@ pub mod kvm {
 
         fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
             let gicr_typers = self.gicr_typers.clone();
-            let snapshot = serde_json::to_vec(&self.state(&gicr_typers).unwrap())
+            Snapshot::new_from_state(&self.id(), &self.state(&gicr_typers).unwrap())
-                .map_err(|e| MigratableError::Snapshot(e.into()))?;
-
-            let mut gic_v3_snapshot = Snapshot::new(self.id().as_str());
-            gic_v3_snapshot.add_data_section(SnapshotDataSection {
-                id: format!("{}-section", self.id()),
-                snapshot,
-            });
-
-            Ok(gic_v3_snapshot)
         }
 
         fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-            if let Some(gic_v3_section) = snapshot
-                .snapshot_data
-                .get(&format!("{}-section", self.id()))
-            {
-                let gic_v3_state = match serde_json::from_slice(&gic_v3_section.snapshot) {
-                    Ok(state) => state,
-                    Err(error) => {
-                        return Err(MigratableError::Restore(anyhow!(
-                            "Could not deserialize GICv3 {}",
-                            error
-                        )))
-                    }
-                };
-
             let gicr_typers = self.gicr_typers.clone();
-                return self.set_state(&gicr_typers, &gic_v3_state).map_err(|e| {
+            self.set_state(&gicr_typers, &snapshot.to_state(&self.id())?)
+                .map_err(|e| {
                     MigratableError::Restore(anyhow!("Could not restore GICv3 state {:?}", e))
-                });
+                })
-            }
-
-            Err(MigratableError::Restore(anyhow!(
-                "Could not find GICv3 snapshot section"
-            )))
         }
     }
 

devices/src/ioapic.rs

@@ -20,10 +20,7 @@ use vm_device::interrupt::{
 };
 use vm_device::BusDevice;
 use vm_memory::GuestAddress;
-use vm_migration::{
+use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
-    Migratable, MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable,
-    Transportable,
-};
 use vmm_sys_util::eventfd::EventFd;
 
 #[derive(Serialize, Deserialize)]
@@ -418,38 +415,13 @@ impl Snapshottable for Ioapic {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let snapshot =
+        Snapshot::new_from_state(&self.id, &self.state())
-            serde_json::to_vec(&self.state()).map_err(|e| MigratableError::Snapshot(e.into()))?;
-
-        let mut ioapic_snapshot = Snapshot::new(self.id.as_str());
-        ioapic_snapshot.add_data_section(SnapshotDataSection {
-            id: format!("{}-section", self.id),
-            snapshot,
-        });
-
-        Ok(ioapic_snapshot)
     }
 
     fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        if let Some(ioapic_section) = snapshot.snapshot_data.get(&format!("{}-section", self.id)) {
+        self.set_state(&snapshot.to_state(&self.id)?).map_err(|e| {
-            let ioapic_state = match serde_json::from_slice(&ioapic_section.snapshot) {
+            MigratableError::Restore(anyhow!("Could not restore state for {}: {:?}", self.id, e))
-                Ok(state) => state,
+        })
-                Err(error) => {
-                    return Err(MigratableError::Restore(anyhow!(
-                        "Could not deserialize IOAPIC {}",
-                        error
-                    )))
-                }
-            };
-
-            return self.set_state(&ioapic_state).map_err(|e| {
-                MigratableError::Restore(anyhow!("Could not restore IOAPIC state {:?}", e))
-            });
-        }
-
-        Err(MigratableError::Restore(anyhow!(
-            "Could not find IOAPIC snapshot section"
-        )))
     }
 }
 

devices/src/legacy/gpio_pl061.rs

@@ -8,16 +8,12 @@
 //!
 
 use crate::{read_le_u32, write_le_u32};
-use anyhow::anyhow;
 use std::result;
 use std::sync::{Arc, Barrier};
 use std::{fmt, io};
 use vm_device::interrupt::InterruptSourceGroup;
 use vm_device::BusDevice;
-use vm_migration::{
+use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
-    Migratable, MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable,
-    Transportable,
-};
 
 const OFS_DATA: u64 = 0x400; // Data Register
 const GPIODIR: u64 = 0x400; // Direction Register
@@ -317,38 +313,12 @@ impl Snapshottable for Gpio {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let snapshot =
+        Snapshot::new_from_state(&self.id, &self.state())
-            serde_json::to_vec(&self.state()).map_err(|e| MigratableError::Snapshot(e.into()))?;
-
-        let mut gpio_snapshot = Snapshot::new(self.id.as_str());
-        gpio_snapshot.add_data_section(SnapshotDataSection {
-            id: format!("{}-section", self.id),
-            snapshot,
-        });
-
-        Ok(gpio_snapshot)
     }
 
     fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        if let Some(gpio_section) = snapshot.snapshot_data.get(&format!("{}-section", self.id)) {
+        self.set_state(&snapshot.to_state(&self.id)?);
-            let gpio_state = match serde_json::from_slice(&gpio_section.snapshot) {
+        Ok(())
-                Ok(state) => state,
-                Err(error) => {
-                    return Err(MigratableError::Restore(anyhow!(
-                        "Could not deserialize GPIO {}",
-                        error
-                    )))
-                }
-            };
-
-            self.set_state(&gpio_state);
-
-            return Ok(());
-        }
-
-        Err(MigratableError::Restore(anyhow!(
-            "Could not find the GPIO snapshot section"
-        )))
     }
 }
 

devices/src/legacy/serial.rs

@@ -5,16 +5,12 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE-BSD-3-Clause file.
 
-use anyhow::anyhow;
 use std::collections::VecDeque;
 use std::sync::{Arc, Barrier};
 use std::{io, result};
 use vm_device::interrupt::InterruptSourceGroup;
 use vm_device::BusDevice;
-use vm_migration::{
+use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
-    Migratable, MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable,
-    Transportable,
-};
 use vmm_sys_util::errno::Result;
 
 const LOOP_SIZE: usize = 0x40;
@@ -292,38 +288,12 @@ impl Snapshottable for Serial {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let snapshot =
+        Snapshot::new_from_state(&self.id, &self.state())
-            serde_json::to_vec(&self.state()).map_err(|e| MigratableError::Snapshot(e.into()))?;
-
-        let mut serial_snapshot = Snapshot::new(self.id.as_str());
-        serial_snapshot.add_data_section(SnapshotDataSection {
-            id: format!("{}-section", self.id),
-            snapshot,
-        });
-
-        Ok(serial_snapshot)
     }
 
     fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        if let Some(serial_section) = snapshot.snapshot_data.get(&format!("{}-section", self.id)) {
+        self.set_state(&snapshot.to_state(&self.id)?);
-            let serial_state = match serde_json::from_slice(&serial_section.snapshot) {
+        Ok(())
-                Ok(state) => state,
-                Err(error) => {
-                    return Err(MigratableError::Restore(anyhow!(
-                        "Could not deserialize SERIAL {}",
-                        error
-                    )))
-                }
-            };
-
-            self.set_state(&serial_state);
-
-            return Ok(());
-        }
-
-        Err(MigratableError::Restore(anyhow!(
-            "Could not find the serial snapshot section"
-        )))
     }
 }
 

devices/src/legacy/uart_pl011.rs

@@ -7,17 +7,13 @@
 //!
 
 use crate::{read_le_u32, write_le_u32};
-use anyhow::anyhow;
 use std::collections::VecDeque;
 use std::fmt;
 use std::sync::{Arc, Barrier};
 use std::{io, result};
 use vm_device::interrupt::InterruptSourceGroup;
 use vm_device::BusDevice;
-use vm_migration::{
+use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
-    Migratable, MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable,
-    Transportable,
-};
 
 /* Registers */
 const UARTDR: u64 = 0;
@@ -361,38 +357,12 @@ impl Snapshottable for Pl011 {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let snapshot =
+        Snapshot::new_from_state(&self.id, &self.state())
-            serde_json::to_vec(&self.state()).map_err(|e| MigratableError::Snapshot(e.into()))?;
-
-        let mut pl011_snapshot = Snapshot::new(self.id.as_str());
-        pl011_snapshot.add_data_section(SnapshotDataSection {
-            id: format!("{}-section", self.id),
-            snapshot,
-        });
-
-        Ok(pl011_snapshot)
     }
 
     fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        if let Some(pl011_section) = snapshot.snapshot_data.get(&format!("{}-section", self.id)) {
+        self.set_state(&snapshot.to_state(&self.id)?);
-            let pl011_state = match serde_json::from_slice(&pl011_section.snapshot) {
+        Ok(())
-                Ok(state) => state,
-                Err(error) => {
-                    return Err(MigratableError::Restore(anyhow!(
-                        "Could not deserialize PL011 {}",
-                        error
-                    )))
-                }
-            };
-
-            self.set_state(&pl011_state);
-
-            return Ok(());
-        }
-
-        Err(MigratableError::Restore(anyhow!(
-            "Could not find the PL011 snapshot section"
-        )))
     }
 }
 

pci/src/configuration.rs

@@ -4,11 +4,10 @@
 
 use crate::device::BarReprogrammingParams;
 use crate::{MsixConfig, PciInterruptPin};
-use anyhow::anyhow;
 use byteorder::{ByteOrder, LittleEndian};
 use std::fmt::{self, Display};
 use std::sync::{Arc, Mutex};
-use vm_migration::{MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable};
+use vm_migration::{MigratableError, Pausable, Snapshot, Snapshottable};
 
 // The number of 32bit registers in the config space, 4096 bytes.
 const NUM_CONFIGURATION_REGISTERS: usize = 1024;
@@ -890,43 +889,12 @@ impl Snapshottable for PciConfiguration {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let snapshot =
+        Snapshot::new_from_state(&self.id(), &self.state())
-            serde_json::to_vec(&self.state()).map_err(|e| MigratableError::Snapshot(e.into()))?;
-
-        let mut config_snapshot = Snapshot::new(self.id().as_str());
-        config_snapshot.add_data_section(SnapshotDataSection {
-            id: format!("{}-section", self.id()),
-            snapshot,
-        });
-
-        Ok(config_snapshot)
     }
 
     fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        if let Some(config_section) = snapshot
+        self.set_state(&snapshot.to_state(&self.id())?);
-            .snapshot_data
+        Ok(())
-            .get(&format!("{}-section", self.id()))
-        {
-            let config_state = match serde_json::from_slice(&config_section.snapshot) {
-                Ok(state) => state,
-                Err(error) => {
-                    return Err(MigratableError::Restore(anyhow!(
-                        "Could not deserialize {}: {}",
-                        self.id(),
-                        error
-                    )))
-                }
-            };
-
-            self.set_state(&config_state);
-
-            return Ok(());
-        }
-
-        Err(MigratableError::Restore(anyhow!(
-            "Could not find {} snapshot section",
-            self.id()
-        )))
     }
 }
 

pci/src/msix.rs

@@ -16,7 +16,7 @@ use vm_device::interrupt::{
     InterruptIndex, InterruptSourceConfig, InterruptSourceGroup, MsiIrqSourceConfig,
 };
 use vm_memory::ByteValued;
-use vm_migration::{MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable};
+use vm_migration::{MigratableError, Pausable, Snapshot, Snapshottable};
 
 const MAX_MSIX_VECTORS_PER_DEVICE: u16 = 2048;
 const MSIX_TABLE_ENTRIES_MODULO: u64 = 16;
@@ -432,41 +432,18 @@ impl Snapshottable for MsixConfig {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let snapshot =
+        Snapshot::new_from_state(&self.id(), &self.state())
-            serde_json::to_vec(&self.state()).map_err(|e| MigratableError::Snapshot(e.into()))?;
-
-        let mut msix_snapshot = Snapshot::new(self.id().as_str());
-        msix_snapshot.add_data_section(SnapshotDataSection {
-            id: format!("{}-section", self.id()),
-            snapshot,
-        });
-
-        Ok(msix_snapshot)
     }
 
     fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        if let Some(msix_section) = snapshot
+        self.set_state(&snapshot.to_state(&self.id())?)
-            .snapshot_data
+            .map_err(|e| {
-            .get(&format!("{}-section", self.id()))
+                MigratableError::Restore(anyhow!(
-        {
+                    "Could not restore state for {}: {:?}",
-            let msix_state = match serde_json::from_slice(&msix_section.snapshot) {
+                    self.id(),
-                Ok(state) => state,
+                    e
-                Err(error) => {
+                ))
-                    return Err(MigratableError::Restore(anyhow!(
+            })
-                        "Could not deserialize MSI-X {}",
-                        error
-                    )))
-                }
-            };
-
-            return self.set_state(&msix_state).map_err(|e| {
-                MigratableError::Restore(anyhow!("Could not restore MSI-X state {:?}", e))
-            });
-        }
-
-        Err(MigratableError::Restore(anyhow!(
-            "Could not find MSI-X snapshot section"
-        )))
     }
 }
 

virtio-devices/src/block.rs

@@ -16,7 +16,6 @@ use super::{
 };
 use crate::seccomp_filters::{get_seccomp_filter, Thread};
 use crate::VirtioInterrupt;
-use anyhow::anyhow;
 use block_util::{
     async_io::AsyncIo, async_io::AsyncIoError, async_io::DiskFile, build_disk_image_id, Request,
     RequestType, VirtioBlockConfig,
@@ -34,10 +33,7 @@ use std::thread;
 use std::{collections::HashMap, convert::TryInto};
 use virtio_bindings::bindings::virtio_blk::*;
 use vm_memory::{ByteValued, Bytes, GuestAddressSpace, GuestMemoryAtomic, GuestMemoryMmap};
-use vm_migration::{
+use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
-    Migratable, MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable,
-    Transportable,
-};
 use vmm_sys_util::eventfd::EventFd;
 
 const SECTOR_SHIFT: u8 = 9;
@@ -678,37 +674,12 @@ impl Snapshottable for Block {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let snapshot =
+        Snapshot::new_from_state(&self.id(), &self.state())
-            serde_json::to_vec(&self.state()).map_err(|e| MigratableError::Snapshot(e.into()))?;
-
-        let mut block_snapshot = Snapshot::new(self.id.as_str());
-        block_snapshot.add_data_section(SnapshotDataSection {
-            id: format!("{}-section", self.id),
-            snapshot,
-        });
-
-        Ok(block_snapshot)
     }
 
     fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        if let Some(block_section) = snapshot.snapshot_data.get(&format!("{}-section", self.id)) {
+        self.set_state(&snapshot.to_state(&self.id)?);
-            let block_state = match serde_json::from_slice(&block_section.snapshot) {
+        Ok(())
-                Ok(state) => state,
-                Err(error) => {
-                    return Err(MigratableError::Restore(anyhow!(
-                        "Could not deserialize BLOCK {}",
-                        error
-                    )))
-                }
-            };
-
-            self.set_state(&block_state);
-            return Ok(());
-        }
-
-        Err(MigratableError::Restore(anyhow!(
-            "Could not find BLOCK snapshot section"
-        )))
     }
 }
 impl Transportable for Block {}

virtio-devices/src/console.rs

@@ -9,7 +9,6 @@ use super::{
 };
 use crate::seccomp_filters::{get_seccomp_filter, Thread};
 use crate::VirtioInterrupt;
-use anyhow::anyhow;
 use libc::EFD_NONBLOCK;
 use seccomp::{SeccompAction, SeccompFilter};
 use serde::ser::{Serialize, SerializeStruct, Serializer};
@@ -24,10 +23,7 @@ use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 use std::sync::{Arc, Barrier, Mutex};
 use std::thread;
 use vm_memory::{ByteValued, Bytes, GuestAddressSpace, GuestMemoryAtomic, GuestMemoryMmap};
-use vm_migration::{
+use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
-    Migratable, MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable,
-    Transportable,
-};
 use vmm_sys_util::eventfd::EventFd;
 
 const QUEUE_SIZE: u16 = 256;
@@ -512,37 +508,12 @@ impl Snapshottable for Console {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let snapshot =
+        Snapshot::new_from_state(&self.id, &self.state())
-            serde_json::to_vec(&self.state()).map_err(|e| MigratableError::Snapshot(e.into()))?;
-
-        let mut console_snapshot = Snapshot::new(self.id.as_str());
-        console_snapshot.add_data_section(SnapshotDataSection {
-            id: format!("{}-section", self.id),
-            snapshot,
-        });
-
-        Ok(console_snapshot)
     }
 
     fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        if let Some(console_section) = snapshot.snapshot_data.get(&format!("{}-section", self.id)) {
+        self.set_state(&snapshot.to_state(&self.id)?);
-            let console_state = match serde_json::from_slice(&console_section.snapshot) {
+        Ok(())
-                Ok(state) => state,
-                Err(error) => {
-                    return Err(MigratableError::Restore(anyhow!(
-                        "Could not deserialize CONSOLE {}",
-                        error
-                    )))
-                }
-            };
-
-            self.set_state(&console_state);
-            return Ok(());
-        }
-
-        Err(MigratableError::Restore(anyhow!(
-            "Could not find CONSOLE snapshot section"
-        )))
     }
 }
 impl Transportable for Console {}

virtio-devices/src/iommu.rs

@@ -10,7 +10,6 @@ use super::{
 };
 use crate::seccomp_filters::{get_seccomp_filter, Thread};
 use crate::{DmaRemapping, VirtioInterrupt, VirtioInterruptType};
-use anyhow::anyhow;
 use seccomp::{SeccompAction, SeccompFilter};
 use std::collections::BTreeMap;
 use std::fmt::{self, Display};
@@ -27,10 +26,7 @@ use vm_memory::{
     Address, ByteValued, Bytes, GuestAddress, GuestAddressSpace, GuestMemoryAtomic,
     GuestMemoryError, GuestMemoryMmap,
 };
-use vm_migration::{
+use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
-    Migratable, MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable,
-    Transportable,
-};
 use vmm_sys_util::eventfd::EventFd;
 
 /// Queues sizes
@@ -975,37 +971,12 @@ impl Snapshottable for Iommu {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let snapshot =
+        Snapshot::new_from_state(&self.id, &self.state())
-            serde_json::to_vec(&self.state()).map_err(|e| MigratableError::Snapshot(e.into()))?;
-
-        let mut iommu_snapshot = Snapshot::new(self.id.as_str());
-        iommu_snapshot.add_data_section(SnapshotDataSection {
-            id: format!("{}-section", self.id),
-            snapshot,
-        });
-
-        Ok(iommu_snapshot)
     }
 
     fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        if let Some(iommu_section) = snapshot.snapshot_data.get(&format!("{}-section", self.id)) {
+        self.set_state(&snapshot.to_state(&self.id)?);
-            let iommu_state = match serde_json::from_slice(&iommu_section.snapshot) {
+        Ok(())
-                Ok(state) => state,
-                Err(error) => {
-                    return Err(MigratableError::Restore(anyhow!(
-                        "Could not deserialize IOMMU {}",
-                        error
-                    )))
-                }
-            };
-
-            self.set_state(&iommu_state);
-            return Ok(());
-        }
-
-        Err(MigratableError::Restore(anyhow!(
-            "Could not find IOMMU snapshot section"
-        )))
     }
 }
 impl Transportable for Iommu {}

virtio-devices/src/net.rs

@@ -17,7 +17,6 @@ use super::{
 };
 use crate::seccomp_filters::{get_seccomp_filter, Thread};
 use crate::VirtioInterrupt;
-use anyhow::anyhow;
 use net_util::{
     open_tap, MacAddr, NetCounters, NetQueuePair, OpenTapError, RxVirtio, Tap, TapError, TxVirtio,
 };
@@ -34,10 +33,7 @@ use std::{collections::HashMap, convert::TryInto};
 use virtio_bindings::bindings::virtio_net::*;
 use virtio_bindings::bindings::virtio_ring::VIRTIO_RING_F_EVENT_IDX;
 use vm_memory::{ByteValued, GuestAddressSpace, GuestMemoryAtomic, GuestMemoryMmap};
-use vm_migration::{
+use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
-    Migratable, MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable,
-    Transportable,
-};
 use vmm_sys_util::eventfd::EventFd;
 
 // The guest has made a buffer available to receive a frame into.
@@ -685,37 +681,12 @@ impl Snapshottable for Net {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let snapshot =
+        Snapshot::new_from_state(&self.id, &self.state())
-            serde_json::to_vec(&self.state()).map_err(|e| MigratableError::Snapshot(e.into()))?;
-
-        let mut net_snapshot = Snapshot::new(self.id.as_str());
-        net_snapshot.add_data_section(SnapshotDataSection {
-            id: format!("{}-section", self.id),
-            snapshot,
-        });
-
-        Ok(net_snapshot)
     }
 
     fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        if let Some(net_section) = snapshot.snapshot_data.get(&format!("{}-section", self.id)) {
+        self.set_state(&snapshot.to_state(&self.id)?);
-            let net_state = match serde_json::from_slice(&net_section.snapshot) {
+        Ok(())
-                Ok(state) => state,
-                Err(error) => {
-                    return Err(MigratableError::Restore(anyhow!(
-                        "Could not deserialize NET {}",
-                        error
-                    )))
-                }
-            };
-
-            self.set_state(&net_state);
-            return Ok(());
-        }
-
-        Err(MigratableError::Restore(anyhow!(
-            "Could not find NET snapshot section"
-        )))
     }
 }
 impl Transportable for Net {}

virtio-devices/src/pmem.rs

@@ -14,7 +14,6 @@ use super::{
 };
 use crate::seccomp_filters::{get_seccomp_filter, Thread};
 use crate::{VirtioInterrupt, VirtioInterruptType};
-use anyhow::anyhow;
 use seccomp::{SeccompAction, SeccompFilter};
 use std::fmt::{self, Display};
 use std::fs::File;
@@ -29,10 +28,7 @@ use vm_memory::{
     Address, ByteValued, Bytes, GuestAddress, GuestAddressSpace, GuestMemoryAtomic,
     GuestMemoryError, GuestMemoryMmap, MmapRegion,
 };
-use vm_migration::{
+use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
-    Migratable, MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable,
-    Transportable,
-};
 use vmm_sys_util::eventfd::EventFd;
 
 const QUEUE_SIZE: u16 = 256;
@@ -461,37 +457,12 @@ impl Snapshottable for Pmem {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let snapshot =
+        Snapshot::new_from_state(&self.id, &self.state())
-            serde_json::to_vec(&self.state()).map_err(|e| MigratableError::Snapshot(e.into()))?;
-
-        let mut pmem_snapshot = Snapshot::new(self.id.as_str());
-        pmem_snapshot.add_data_section(SnapshotDataSection {
-            id: format!("{}-section", self.id),
-            snapshot,
-        });
-
-        Ok(pmem_snapshot)
     }
 
     fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        if let Some(pmem_section) = snapshot.snapshot_data.get(&format!("{}-section", self.id)) {
+        self.set_state(&snapshot.to_state(&self.id)?);
-            let pmem_state = match serde_json::from_slice(&pmem_section.snapshot) {
+        Ok(())
-                Ok(state) => state,
-                Err(error) => {
-                    return Err(MigratableError::Restore(anyhow!(
-                        "Could not deserialize PMEM {}",
-                        error
-                    )))
-                }
-            };
-
-            self.set_state(&pmem_state);
-            return Ok(());
-        }
-
-        Err(MigratableError::Restore(anyhow!(
-            "Could not find PMEM snapshot section"
-        )))
     }
 }
 

virtio-devices/src/rng.rs

@@ -10,7 +10,6 @@ use super::{
 };
 use crate::seccomp_filters::{get_seccomp_filter, Thread};
 use crate::{VirtioInterrupt, VirtioInterruptType};
-use anyhow::anyhow;
 use seccomp::{SeccompAction, SeccompFilter};
 use std::fs::File;
 use std::io;
@@ -20,10 +19,7 @@ use std::sync::atomic::AtomicBool;
 use std::sync::{Arc, Barrier};
 use std::thread;
 use vm_memory::{Bytes, GuestAddressSpace, GuestMemoryAtomic, GuestMemoryMmap};
-use vm_migration::{
+use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
-    Migratable, MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable,
-    Transportable,
-};
 use vmm_sys_util::eventfd::EventFd;
 
 const QUEUE_SIZE: u16 = 256;
@@ -303,37 +299,12 @@ impl Snapshottable for Rng {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let snapshot =
+        Snapshot::new_from_state(&self.id, &self.state())
-            serde_json::to_vec(&self.state()).map_err(|e| MigratableError::Snapshot(e.into()))?;
-
-        let mut rng_snapshot = Snapshot::new(self.id.as_str());
-        rng_snapshot.add_data_section(SnapshotDataSection {
-            id: format!("{}-section", self.id),
-            snapshot,
-        });
-
-        Ok(rng_snapshot)
     }
 
     fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        if let Some(rng_section) = snapshot.snapshot_data.get(&format!("{}-section", self.id)) {
+        self.set_state(&snapshot.to_state(&self.id)?);
-            let rng_state = match serde_json::from_slice(&rng_section.snapshot) {
+        Ok(())
-                Ok(state) => state,
-                Err(error) => {
-                    return Err(MigratableError::Restore(anyhow!(
-                        "Could not deserialize RNG {}",
-                        error
-                    )))
-                }
-            };
-
-            self.set_state(&rng_state);
-            return Ok(());
-        }
-
-        Err(MigratableError::Restore(anyhow!(
-            "Could not find RNG snapshot section"
-        )))
     }
 }
 

virtio-devices/src/transport/pci_common_config.rs

@@ -8,12 +8,11 @@
 extern crate byteorder;
 
 use crate::{Queue, VirtioDevice};
-use anyhow::anyhow;
 use byteorder::{ByteOrder, LittleEndian};
 use std::sync::atomic::{AtomicU16, Ordering};
 use std::sync::{Arc, Mutex};
 use vm_memory::GuestAddress;
-use vm_migration::{MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable};
+use vm_migration::{MigratableError, Pausable, Snapshot, Snapshottable};
 
 #[derive(Clone, Serialize, Deserialize)]
 pub struct VirtioPciCommonConfigState {
@@ -291,43 +290,12 @@ impl Snapshottable for VirtioPciCommonConfig {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let snapshot =
+        Snapshot::new_from_state(&self.id(), &self.state())
-            serde_json::to_vec(&self.state()).map_err(|e| MigratableError::Snapshot(e.into()))?;
-
-        let mut config_snapshot = Snapshot::new(self.id().as_str());
-        config_snapshot.add_data_section(SnapshotDataSection {
-            id: format!("{}-section", self.id()),
-            snapshot,
-        });
-
-        Ok(config_snapshot)
     }
 
     fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        if let Some(config_section) = snapshot
+        self.set_state(&snapshot.to_state(&self.id())?);
-            .snapshot_data
+        Ok(())
-            .get(&format!("{}-section", self.id()))
-        {
-            let config_state = match serde_json::from_slice(&config_section.snapshot) {
-                Ok(state) => state,
-                Err(error) => {
-                    return Err(MigratableError::Restore(anyhow!(
-                        "Could not deserialize {}: {}",
-                        self.id(),
-                        error
-                    )))
-                }
-            };
-
-            self.set_state(&config_state);
-
-            return Ok(());
-        }
-
-        Err(MigratableError::Restore(anyhow!(
-            "Could not find {} snapshot section",
-            self.id()
-        )))
     }
 }
 

virtio-devices/src/transport/pci_device.rs

@@ -41,10 +41,7 @@ use vm_memory::{
     Address, ByteValued, GuestAddress, GuestAddressSpace, GuestMemoryAtomic, GuestMemoryMmap,
     GuestUsize, Le32,
 };
-use vm_migration::{
+use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
-    Migratable, MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable,
-    Transportable,
-};
 use vm_virtio::{queue, VirtioIommuRemapping, VIRTIO_MSI_NO_VECTOR};
 use vmm_sys_util::{errno::Result, eventfd::EventFd};
 
@@ -1078,14 +1075,7 @@ impl Snapshottable for VirtioPciDevice {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let snapshot =
+        let mut virtio_pci_dev_snapshot = Snapshot::new_from_state(&self.id, &self.state())?;
-            serde_json::to_vec(&self.state()).map_err(|e| MigratableError::Snapshot(e.into()))?;
-
-        let mut virtio_pci_dev_snapshot = Snapshot::new(self.id.as_str());
-        virtio_pci_dev_snapshot.add_data_section(SnapshotDataSection {
-            id: format!("{}-section", self.id),
-            snapshot,
-        });
 
         // Snapshot PciConfiguration
         virtio_pci_dev_snapshot.add_snapshot(self.configuration.snapshot()?);

virtio-devices/src/vsock/device.rs

@@ -17,7 +17,6 @@ use crate::{
     VirtioCommon, VirtioDevice, VirtioDeviceType, VirtioInterruptType, EPOLL_HELPER_EVENT_LAST,
     VIRTIO_F_IN_ORDER, VIRTIO_F_IOMMU_PLATFORM, VIRTIO_F_VERSION_1,
 };
-use anyhow::anyhow;
 /// This is the `VirtioDevice` implementation for our vsock device. It handles the virtio-level
 /// device logic: feature negotiation, device configuration, and device activation.
 /// The run-time device logic (i.e. event-driven data handling) is implemented by
@@ -47,10 +46,7 @@ use std::sync::atomic::AtomicBool;
 use std::sync::{Arc, Barrier, RwLock};
 use std::thread;
 use vm_memory::{GuestAddressSpace, GuestMemoryAtomic, GuestMemoryMmap};
-use vm_migration::{
+use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
-    Migratable, MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable,
-    Transportable,
-};
 use vmm_sys_util::eventfd::EventFd;
 
 const QUEUE_SIZE: u16 = 256;
@@ -508,37 +504,12 @@ where
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let snapshot =
+        Snapshot::new_from_state(&self.id, &self.state())
-            serde_json::to_vec(&self.state()).map_err(|e| MigratableError::Snapshot(e.into()))?;
-
-        let mut vsock_snapshot = Snapshot::new(self.id.as_str());
-        vsock_snapshot.add_data_section(SnapshotDataSection {
-            id: format!("{}-section", self.id),
-            snapshot,
-        });
-
-        Ok(vsock_snapshot)
     }
 
     fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        if let Some(vsock_section) = snapshot.snapshot_data.get(&format!("{}-section", self.id)) {
+        self.set_state(&snapshot.to_state(&self.id)?);
-            let vsock_state = match serde_json::from_slice(&vsock_section.snapshot) {
+        Ok(())
-                Ok(state) => state,
-                Err(error) => {
-                    return Err(MigratableError::Restore(anyhow!(
-                        "Could not deserialize VSOCK {}",
-                        error
-                    )))
-                }
-            };
-
-            self.set_state(&vsock_state);
-            return Ok(());
-        }
-
-        Err(MigratableError::Restore(anyhow!(
-            "Could not find VSOCK snapshot section"
-        )))
     }
 }
 impl<B> Transportable for Vsock<B> where B: VsockBackend + Sync + 'static {}

virtio-devices/src/watchdog.rs

@@ -23,10 +23,7 @@ use std::sync::{Arc, Barrier, Mutex};
 use std::thread;
 use std::time::Instant;
 use vm_memory::{Bytes, GuestAddressSpace, GuestMemoryAtomic, GuestMemoryMmap};
-use vm_migration::{
+use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
-    Migratable, MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable,
-    Transportable,
-};
 use vmm_sys_util::eventfd::EventFd;
 
 const QUEUE_SIZE: u16 = 8;
@@ -395,38 +392,12 @@ impl Snapshottable for Watchdog {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let snapshot =
+        Snapshot::new_from_state(&self.id, &self.state())
-            serde_json::to_vec(&self.state()).map_err(|e| MigratableError::Snapshot(e.into()))?;
-
-        let mut watchdog_snapshot = Snapshot::new(self.id.as_str());
-        watchdog_snapshot.add_data_section(SnapshotDataSection {
-            id: format!("{}-section", self.id),
-            snapshot,
-        });
-
-        Ok(watchdog_snapshot)
     }
 
     fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        if let Some(watchdog_section) = snapshot.snapshot_data.get(&format!("{}-section", self.id))
+        self.set_state(&snapshot.to_state(&self.id)?);
-        {
+        Ok(())
-            let watchdog_state = match serde_json::from_slice(&watchdog_section.snapshot) {
-                Ok(state) => state,
-                Err(error) => {
-                    return Err(MigratableError::Restore(anyhow!(
-                        "Could not deserialize watchdog {}",
-                        error
-                    )))
-                }
-            };
-
-            self.set_state(&watchdog_state);
-            return Ok(());
-        }
-
-        Err(MigratableError::Restore(anyhow!(
-            "Could not find watchdog snapshot section"
-        )))
     }
 }
 

vm-migration/src/lib.rs

@@ -9,6 +9,8 @@ extern crate thiserror;
 extern crate serde_derive;
 extern crate vm_memory;
 
+use anyhow::anyhow;
+use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
 pub mod protocol;
@@ -97,6 +99,23 @@ impl Snapshot {
         }
     }
 
+    /// Create from state that can be serialized
+    pub fn new_from_state<T>(id: &str, state: &T) -> Result<Self, MigratableError>
+    where
+        T: Serialize,
+    {
+        let snapshot = serde_json::to_vec(state)
+            .map_err(|e| MigratableError::Snapshot(anyhow!("Error serialising: {} {}", id, e)))?;
+
+        let mut snapshot_data = Snapshot::new(id);
+        snapshot_data.add_data_section(SnapshotDataSection {
+            id: format!("{}-section", id),
+            snapshot,
+        });
+
+        Ok(snapshot_data)
+    }
+
     /// Add a sub-component's Snapshot to the Snapshot.
     pub fn add_snapshot(&mut self, snapshot: Snapshot) {
         self.snapshots
@@ -107,6 +126,20 @@ impl Snapshot {
     pub fn add_data_section(&mut self, section: SnapshotDataSection) {
         self.snapshot_data.insert(section.id.clone(), section);
     }
+
+    /// Generate the state data from the snapshot
+    pub fn to_state<'a, T>(&'a self, id: &str) -> Result<T, MigratableError>
+    where
+        T: Deserialize<'a>,
+    {
+        let section = self
+            .snapshot_data
+            .get(&format!("{}-section", id))
+            .ok_or_else(|| MigratableError::Restore(anyhow!("Missing section for {}", id)))?;
+
+        serde_json::from_slice(&section.snapshot)
+            .map_err(|e| MigratableError::Restore(anyhow!("Error deserialising: {} {}", id, e)))
+    }
 }
 
 /// A snapshottable component can be snapshotted.