vmm: virtio-devices: Restore every VirtioDevice upon creation

Following the new design proposal to improve the restore codepath when
migrating a VM, all virtio devices are supplied with an optional state
they can use to restore from. The restore() implementation every device
was providing has been removed in order to prevent from going through
the restoration twice.

Here is the list of devices now following the new restore design:

- Block (virtio-block)
- Net (virtio-net)
- Rng (virtio-rng)
- Fs (vhost-user-fs)
- Blk (vhost-user-block)
- Net (vhost-user-net)
- Pmem (virtio-pmem)
- Vsock (virtio-vsock)
- Mem (virtio-mem)
- Balloon (virtio-balloon)
- Watchdog (virtio-watchdog)
- Vdpa (vDPA)
- Console (virtio-console)
- Iommu (virtio-iommu)

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

fuzz/fuzz_targets/balloon.rs

@@ -49,6 +49,7 @@ fuzz_target!(|bytes| {
         true,
         SeccompAction::Allow,
         EventFd::new(EFD_NONBLOCK).unwrap(),
+        None,
     )
     .unwrap();
 

fuzz/fuzz_targets/block.rs

@@ -60,6 +60,7 @@ fuzz_target!(|bytes| {
         SeccompAction::Allow,
         None,
         EventFd::new(EFD_NONBLOCK).unwrap(),
+        None,
     )
     .unwrap();
 

fuzz/fuzz_targets/iommu.rs

@@ -64,6 +64,7 @@ fuzz_target!(|bytes| {
         SeccompAction::Allow,
         EventFd::new(EFD_NONBLOCK).unwrap(),
         ((MEM_SIZE - IOVA_SPACE_SIZE) as u64, (MEM_SIZE - 1) as u64),
+        None,
     )
     .unwrap();
 

fuzz/fuzz_targets/mem.rs

@@ -152,6 +152,7 @@ fn create_dummy_virtio_mem(bytes: &[u8; VIRTIO_MEM_DATA_SIZE]) -> (Mem, Arc<Gues
             false,
             EventFd::new(EFD_NONBLOCK).unwrap(),
             blocks_state.clone(),
+            None,
         )
         .unwrap(),
         region,

fuzz/fuzz_targets/pmem.rs

@@ -127,6 +127,7 @@ fn create_dummy_pmem() -> Pmem {
         false,
         SeccompAction::Allow,
         EventFd::new(EFD_NONBLOCK).unwrap(),
+        None,
     )
     .unwrap()
 }

fuzz/fuzz_targets/rng.rs

@@ -63,6 +63,7 @@ fuzz_target!(|bytes| {
         false,
         SeccompAction::Allow,
         EventFd::new(EFD_NONBLOCK).unwrap(),
+        None,
     )
     .unwrap();
 

fuzz/fuzz_targets/watchdog.rs

@@ -38,6 +38,7 @@ fuzz_target!(|bytes| {
         EventFd::new(EFD_NONBLOCK).unwrap(),
         SeccompAction::Allow,
         EventFd::new(EFD_NONBLOCK).unwrap(),
+        None,
     )
     .unwrap();
 

virtio-devices/src/balloon.rs

@@ -360,26 +360,39 @@ impl Balloon {
         free_page_reporting: bool,
         seccomp_action: SeccompAction,
         exit_evt: EventFd,
+        state: Option<BalloonState>,
     ) -> io::Result<Self> {
         let mut queue_sizes = vec![QUEUE_SIZE; MIN_NUM_QUEUES];
-        let mut avail_features = 1u64 << VIRTIO_F_VERSION_1;
+
-        if deflate_on_oom {
+        let (avail_features, acked_features, config) = if let Some(state) = state {
-            avail_features |= 1u64 << VIRTIO_BALLOON_F_DEFLATE_ON_OOM;
+            info!("Restoring virtio-balloon {}", id);
-        }
+            (state.avail_features, state.acked_features, state.config)
+        } else {
+            let mut avail_features = 1u64 << VIRTIO_F_VERSION_1;
+            if deflate_on_oom {
+                avail_features |= 1u64 << VIRTIO_BALLOON_F_DEFLATE_ON_OOM;
+            }
+            if free_page_reporting {
+                avail_features |= 1u64 << VIRTIO_BALLOON_F_REPORTING;
+            }
+
+            let config = VirtioBalloonConfig {
+                num_pages: (size >> VIRTIO_BALLOON_PFN_SHIFT) as u32,
+                ..Default::default()
+            };
+
+            (avail_features, 0, config)
+        };
+
         if free_page_reporting {
-            avail_features |= 1u64 << VIRTIO_BALLOON_F_REPORTING;
             queue_sizes.push(REPORTING_QUEUE_SIZE);
         }
 
-        let config = VirtioBalloonConfig {
-            num_pages: (size >> VIRTIO_BALLOON_PFN_SHIFT) as u32,
-            ..Default::default()
-        };
-
         Ok(Balloon {
             common: VirtioCommon {
                 device_type: VirtioDeviceType::Balloon as u32,
                 avail_features,
+                acked_features,
                 paused_sync: Some(Arc::new(Barrier::new(2))),
                 queue_sizes,
                 min_queues: MIN_NUM_QUEUES as u16,
@@ -418,12 +431,6 @@ impl Balloon {
         }
     }
 
-    fn set_state(&mut self, state: &BalloonState) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        self.config = state.config;
-    }
-
     #[cfg(fuzzing)]
     pub fn wait_for_epoll_threads(&mut self) {
         self.common.wait_for_epoll_threads();
@@ -573,11 +580,6 @@ impl Snapshottable for Balloon {
     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 Balloon {}
 impl Migratable for Balloon {}

virtio-devices/src/block.rs

@@ -416,72 +416,86 @@ impl Block {
         seccomp_action: SeccompAction,
         rate_limiter_config: Option<RateLimiterConfig>,
         exit_evt: EventFd,
+        state: Option<BlockState>,
     ) -> io::Result<Self> {
-        let disk_size = disk_image.size().map_err(|e| {
+        let (disk_nsectors, avail_features, acked_features, config) = if let Some(state) = state {
-            io::Error::new(
+            info!("Restoring virtio-block {}", id);
-                io::ErrorKind::Other,
+            (
-                format!("Failed getting disk size: {}", e),
+                state.disk_nsectors,
+                state.avail_features,
+                state.acked_features,
+                state.config,
             )
-        })?;
-        if disk_size % SECTOR_SIZE != 0 {
-            warn!(
-                "Disk size {} is not a multiple of sector size {}; \
-                 the remainder will not be visible to the guest.",
-                disk_size, SECTOR_SIZE
-            );
-        }
-
-        let mut avail_features = (1u64 << VIRTIO_F_VERSION_1)
-            | (1u64 << VIRTIO_BLK_F_FLUSH)
-            | (1u64 << VIRTIO_BLK_F_CONFIG_WCE)
-            | (1u64 << VIRTIO_BLK_F_BLK_SIZE)
-            | (1u64 << VIRTIO_BLK_F_TOPOLOGY);
-
-        if iommu {
-            avail_features |= 1u64 << VIRTIO_F_IOMMU_PLATFORM;
-        }
-
-        if is_disk_read_only {
-            avail_features |= 1u64 << VIRTIO_BLK_F_RO;
-        }
-
-        let topology = disk_image.topology();
-        info!("Disk topology: {:?}", topology);
-
-        let logical_block_size = if topology.logical_block_size > 512 {
-            topology.logical_block_size
         } else {
-            512
+            let disk_size = disk_image.size().map_err(|e| {
+                io::Error::new(
+                    io::ErrorKind::Other,
+                    format!("Failed getting disk size: {}", e),
+                )
+            })?;
+            if disk_size % SECTOR_SIZE != 0 {
+                warn!(
+                    "Disk size {} is not a multiple of sector size {}; \
+                 the remainder will not be visible to the guest.",
+                    disk_size, SECTOR_SIZE
+                );
+            }
+
+            let mut avail_features = (1u64 << VIRTIO_F_VERSION_1)
+                | (1u64 << VIRTIO_BLK_F_FLUSH)
+                | (1u64 << VIRTIO_BLK_F_CONFIG_WCE)
+                | (1u64 << VIRTIO_BLK_F_BLK_SIZE)
+                | (1u64 << VIRTIO_BLK_F_TOPOLOGY);
+
+            if iommu {
+                avail_features |= 1u64 << VIRTIO_F_IOMMU_PLATFORM;
+            }
+
+            if is_disk_read_only {
+                avail_features |= 1u64 << VIRTIO_BLK_F_RO;
+            }
+
+            let topology = disk_image.topology();
+            info!("Disk topology: {:?}", topology);
+
+            let logical_block_size = if topology.logical_block_size > 512 {
+                topology.logical_block_size
+            } else {
+                512
+            };
+
+            // Calculate the exponent that maps physical block to logical block
+            let mut physical_block_exp = 0;
+            let mut size = logical_block_size;
+            while size < topology.physical_block_size {
+                physical_block_exp += 1;
+                size <<= 1;
+            }
+
+            let disk_nsectors = disk_size / SECTOR_SIZE;
+            let mut config = VirtioBlockConfig {
+                capacity: disk_nsectors,
+                writeback: 1,
+                blk_size: topology.logical_block_size as u32,
+                physical_block_exp,
+                min_io_size: (topology.minimum_io_size / logical_block_size) as u16,
+                opt_io_size: (topology.optimal_io_size / logical_block_size) as u32,
+                ..Default::default()
+            };
+
+            if num_queues > 1 {
+                avail_features |= 1u64 << VIRTIO_BLK_F_MQ;
+                config.num_queues = num_queues as u16;
+            }
+
+            (disk_nsectors, avail_features, 0, config)
         };
 
-        // Calculate the exponent that maps physical block to logical block
-        let mut physical_block_exp = 0;
-        let mut size = logical_block_size;
-        while size < topology.physical_block_size {
-            physical_block_exp += 1;
-            size <<= 1;
-        }
-
-        let disk_nsectors = disk_size / SECTOR_SIZE;
-        let mut config = VirtioBlockConfig {
-            capacity: disk_nsectors,
-            writeback: 1,
-            blk_size: topology.logical_block_size as u32,
-            physical_block_exp,
-            min_io_size: (topology.minimum_io_size / logical_block_size) as u16,
-            opt_io_size: (topology.optimal_io_size / logical_block_size) as u32,
-            ..Default::default()
-        };
-
-        if num_queues > 1 {
-            avail_features |= 1u64 << VIRTIO_BLK_F_MQ;
-            config.num_queues = num_queues as u16;
-        }
-
         Ok(Block {
             common: VirtioCommon {
                 device_type: VirtioDeviceType::Block as u32,
                 avail_features,
+                acked_features,
                 paused_sync: Some(Arc::new(Barrier::new(num_queues + 1))),
                 queue_sizes: vec![queue_size; num_queues],
                 min_queues: 1,
@@ -510,14 +524,6 @@ impl Block {
         }
     }
 
-    fn set_state(&mut self, state: &BlockState) {
-        self.disk_path = state.disk_path.clone().into();
-        self.disk_nsectors = state.disk_nsectors;
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        self.config = state.config;
-    }
-
     fn update_writeback(&mut self) {
         // Use writeback from config if VIRTIO_BLK_F_CONFIG_WCE
         let writeback = if self.common.feature_acked(VIRTIO_BLK_F_CONFIG_WCE.into()) {
@@ -710,11 +716,6 @@ impl Snapshottable for Block {
     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 Block {}
 impl Migratable for Block {}

virtio-devices/src/console.rs

@@ -601,7 +601,7 @@ fn get_win_size(tty: &dyn AsRawFd) -> (u16, u16) {
 impl VersionMapped for ConsoleState {}
 
 impl Console {
-    /// Create a new virtio console device that gets random data from /dev/urandom.
+    /// Create a new virtio console device
     pub fn new(
         id: String,
         endpoint: Endpoint,
@@ -609,20 +609,37 @@ impl Console {
         iommu: bool,
         seccomp_action: SeccompAction,
         exit_evt: EventFd,
+        state: Option<ConsoleState>,
     ) -> io::Result<(Console, Arc<ConsoleResizer>)> {
-        let mut avail_features = 1u64 << VIRTIO_F_VERSION_1 | 1u64 << VIRTIO_CONSOLE_F_SIZE;
+        let (avail_features, acked_features, config, in_buffer) = if let Some(state) = state {
+            info!("Restoring virtio-console {}", id);
+            (
+                state.avail_features,
+                state.acked_features,
+                state.config,
+                state.in_buffer.into(),
+            )
+        } else {
+            let mut avail_features = 1u64 << VIRTIO_F_VERSION_1 | 1u64 << VIRTIO_CONSOLE_F_SIZE;
+            if iommu {
+                avail_features |= 1u64 << VIRTIO_F_IOMMU_PLATFORM;
+            }
 
-        if iommu {
+            (
-            avail_features |= 1u64 << VIRTIO_F_IOMMU_PLATFORM;
+                avail_features,
-        }
+                0,
+                VirtioConsoleConfig::default(),
+                VecDeque::new(),
+            )
+        };
 
         let config_evt = EventFd::new(EFD_NONBLOCK).unwrap();
-        let console_config = Arc::new(Mutex::new(VirtioConsoleConfig::default()));
+        let console_config = Arc::new(Mutex::new(config));
         let resizer = Arc::new(ConsoleResizer {
             config_evt,
             config: console_config.clone(),
             tty: endpoint.out_file().as_ref().map(|t| t.try_clone().unwrap()),
-            acked_features: AtomicU64::new(0),
+            acked_features: AtomicU64::new(acked_features),
         });
 
         resizer.update_console_size();
@@ -633,6 +650,7 @@ impl Console {
                     device_type: VirtioDeviceType::Console as u32,
                     queue_sizes: QUEUE_SIZES.to_vec(),
                     avail_features,
+                    acked_features,
                     paused_sync: Some(Arc::new(Barrier::new(2))),
                     min_queues: NUM_QUEUES as u16,
                     ..Default::default()
@@ -643,7 +661,7 @@ impl Console {
                 resize_pipe,
                 endpoint,
                 seccomp_action,
-                in_buffer: Arc::new(Mutex::new(VecDeque::new())),
+                in_buffer: Arc::new(Mutex::new(in_buffer)),
                 exit_evt,
             },
             resizer,
@@ -658,13 +676,6 @@ impl Console {
             in_buffer: self.in_buffer.lock().unwrap().clone().into(),
         }
     }
-
-    fn set_state(&mut self, state: &ConsoleState) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        *(self.config.lock().unwrap()) = state.config;
-        *(self.in_buffer.lock().unwrap()) = state.in_buffer.clone().into();
-    }
 }
 
 impl Drop for Console {
@@ -786,11 +797,6 @@ impl Snapshottable for Console {
     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 Console {}
 impl Migratable for Console {}

virtio-devices/src/iommu.rs

@@ -850,7 +850,7 @@ type EndpointsState = Vec<(u32, u32)>;
 type DomainsState = Vec<(u32, (Vec<(u64, Mapping)>, bool))>;
 
 #[derive(Versionize)]
-struct IommuState {
+pub struct IommuState {
     avail_features: u64,
     acked_features: u64,
     endpoints: EndpointsState,
@@ -865,7 +865,37 @@ impl Iommu {
         seccomp_action: SeccompAction,
         exit_evt: EventFd,
         msi_iova_space: (u64, u64),
+        state: Option<IommuState>,
     ) -> io::Result<(Self, Arc<IommuMapping>)> {
+        let (avail_features, acked_features, endpoints, domains) = if let Some(state) = state {
+            info!("Restoring virtio-iommu {}", id);
+            (
+                state.avail_features,
+                state.acked_features,
+                state.endpoints.into_iter().collect(),
+                state
+                    .domains
+                    .into_iter()
+                    .map(|(k, v)| {
+                        (
+                            k,
+                            Domain {
+                                mappings: v.0.into_iter().collect(),
+                                bypass: v.1,
+                            },
+                        )
+                    })
+                    .collect(),
+            )
+        } else {
+            let avail_features = 1u64 << VIRTIO_F_VERSION_1
+                | 1u64 << VIRTIO_IOMMU_F_MAP_UNMAP
+                | 1u64 << VIRTIO_IOMMU_F_PROBE
+                | 1u64 << VIRTIO_IOMMU_F_BYPASS_CONFIG;
+
+            (avail_features, 0, BTreeMap::new(), BTreeMap::new())
+        };
+
         let config = VirtioIommuConfig {
             page_size_mask: VIRTIO_IOMMU_PAGE_SIZE_MASK,
             probe_size: PROBE_PROP_SIZE,
@@ -873,8 +903,8 @@ impl Iommu {
         };
 
         let mapping = Arc::new(IommuMapping {
-            endpoints: Arc::new(RwLock::new(BTreeMap::new())),
+            endpoints: Arc::new(RwLock::new(endpoints)),
-            domains: Arc::new(RwLock::new(BTreeMap::new())),
+            domains: Arc::new(RwLock::new(domains)),
             bypass: AtomicBool::new(true),
         });
 
@@ -884,10 +914,8 @@ impl Iommu {
                 common: VirtioCommon {
                     device_type: VirtioDeviceType::Iommu as u32,
                     queue_sizes: QUEUE_SIZES.to_vec(),
-                    avail_features: 1u64 << VIRTIO_F_VERSION_1
+                    avail_features,
-                        | 1u64 << VIRTIO_IOMMU_F_MAP_UNMAP
+                    acked_features,
-                        | 1u64 << VIRTIO_IOMMU_F_PROBE
-                        | 1u64 << VIRTIO_IOMMU_F_BYPASS_CONFIG,
                     paused_sync: Some(Arc::new(Barrier::new(2))),
                     min_queues: NUM_QUEUES as u16,
                     ..Default::default()
@@ -927,26 +955,6 @@ impl Iommu {
         }
     }
 
-    fn set_state(&mut self, state: &IommuState) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        *(self.mapping.endpoints.write().unwrap()) = state.endpoints.clone().into_iter().collect();
-        *(self.mapping.domains.write().unwrap()) = state
-            .domains
-            .clone()
-            .into_iter()
-            .map(|(k, v)| {
-                (
-                    k,
-                    Domain {
-                        mappings: v.0.into_iter().collect(),
-                        bypass: v.1,
-                    },
-                )
-            })
-            .collect();
-    }
-
     fn update_bypass(&mut self) {
         // Use bypass from config if VIRTIO_IOMMU_F_BYPASS_CONFIG has been negotiated
         if !self
@@ -1088,11 +1096,6 @@ impl Snapshottable for Iommu {
     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 Iommu {}
 impl Migratable for Iommu {}

virtio-devices/src/mem.rs

@@ -714,6 +714,7 @@ impl Mem {
         hugepages: bool,
         exit_evt: EventFd,
         blocks_state: Arc<Mutex<BlocksState>>,
+        state: Option<MemState>,
     ) -> io::Result<Mem> {
         let region_len = region.len();
 
@@ -727,43 +728,51 @@ impl Mem {
             ));
         }
 
-        let mut avail_features = 1u64 << VIRTIO_F_VERSION_1;
+        let (avail_features, acked_features, config) = if let Some(state) = state {
+            info!("Restoring virtio-mem {}", id);
+            *(blocks_state.lock().unwrap()) = state.blocks_state.clone();
+            (state.avail_features, state.acked_features, state.config)
+        } else {
+            let mut avail_features = 1u64 << VIRTIO_F_VERSION_1;
 
-        let mut config = VirtioMemConfig {
+            let mut config = VirtioMemConfig {
-            block_size: VIRTIO_MEM_DEFAULT_BLOCK_SIZE,
+                block_size: VIRTIO_MEM_DEFAULT_BLOCK_SIZE,
-            addr: region.start_addr().raw_value(),
+                addr: region.start_addr().raw_value(),
-            region_size: region.len(),
+                region_size: region.len(),
-            usable_region_size: region.len(),
+                usable_region_size: region.len(),
-            plugged_size: 0,
+                plugged_size: 0,
-            requested_size: 0,
+                requested_size: 0,
-            ..Default::default()
+                ..Default::default()
-        };
+            };
 
-        if initial_size != 0 {
+            if initial_size != 0 {
-            config.resize(initial_size).map_err(|e| {
+                config.resize(initial_size).map_err(|e| {
+                    io::Error::new(
+                        io::ErrorKind::Other,
+                        format!(
+                            "Failed to resize virtio-mem configuration to {}: {:?}",
+                            initial_size, e
+                        ),
+                    )
+                })?;
+            }
+
+            if let Some(node_id) = numa_node_id {
+                avail_features |= 1u64 << VIRTIO_MEM_F_ACPI_PXM;
+                config.node_id = node_id;
+            }
+
+            // Make sure the virtio-mem configuration complies with the
+            // specification.
+            config.validate().map_err(|e| {
                 io::Error::new(
                     io::ErrorKind::Other,
-                    format!(
+                    format!("Invalid virtio-mem configuration: {:?}", e),
-                        "Failed to resize virtio-mem configuration to {}: {:?}",
-                        initial_size, e
-                    ),
                 )
             })?;
-        }
 
-        if let Some(node_id) = numa_node_id {
+            (avail_features, 0, config)
-            avail_features |= 1u64 << VIRTIO_MEM_F_ACPI_PXM;
+        };
-            config.node_id = node_id;
-        }
-
-        // Make sure the virtio-mem configuration complies with the
-        // specification.
-        config.validate().map_err(|e| {
-            io::Error::new(
-                io::ErrorKind::Other,
-                format!("Invalid virtio-mem configuration: {:?}", e),
-            )
-        })?;
 
         let host_fd = region
             .file_offset()
@@ -773,6 +782,7 @@ impl Mem {
             common: VirtioCommon {
                 device_type: VirtioDeviceType::Mem as u32,
                 avail_features,
+                acked_features,
                 paused_sync: Some(Arc::new(Barrier::new(2))),
                 queue_sizes: QUEUE_SIZES.to_vec(),
                 min_queues: 1,
@@ -870,13 +880,6 @@ impl Mem {
         }
     }
 
-    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();
-    }
-
     #[cfg(fuzzing)]
     pub fn wait_for_epoll_threads(&mut self) {
         self.common.wait_for_epoll_threads();
@@ -999,11 +1002,6 @@ impl Snapshottable for Mem {
     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 {}

virtio-devices/src/net.rs

@@ -445,45 +445,70 @@ impl Net {
         seccomp_action: SeccompAction,
         rate_limiter_config: Option<RateLimiterConfig>,
         exit_evt: EventFd,
+        state: Option<NetState>,
     ) -> Result<Self> {
         assert!(!taps.is_empty());
 
         let mtu = taps[0].mtu().map_err(Error::TapError)? as u16;
 
-        let mut avail_features = 1 << VIRTIO_NET_F_CSUM
+        let (avail_features, acked_features, config, queue_sizes) = if let Some(state) = state {
-            | 1 << VIRTIO_NET_F_CTRL_GUEST_OFFLOADS
+            info!("Restoring virtio-net {}", id);
-            | 1 << VIRTIO_NET_F_GUEST_CSUM
+            (
-            | 1 << VIRTIO_NET_F_GUEST_ECN
+                state.avail_features,
-            | 1 << VIRTIO_NET_F_GUEST_TSO4
+                state.acked_features,
-            | 1 << VIRTIO_NET_F_GUEST_TSO6
+                state.config,
-            | 1 << VIRTIO_NET_F_GUEST_UFO
+                state.queue_size,
-            | 1 << VIRTIO_NET_F_HOST_ECN
+            )
-            | 1 << VIRTIO_NET_F_HOST_TSO4
-            | 1 << VIRTIO_NET_F_HOST_TSO6
-            | 1 << VIRTIO_NET_F_HOST_UFO
-            | 1 << VIRTIO_NET_F_MTU
-            | 1 << VIRTIO_RING_F_EVENT_IDX
-            | 1 << VIRTIO_F_VERSION_1;
-
-        if iommu {
-            avail_features |= 1u64 << VIRTIO_F_IOMMU_PLATFORM;
-        }
-
-        avail_features |= 1 << VIRTIO_NET_F_CTRL_VQ;
-        let queue_num = num_queues + 1;
-
-        let mut config = VirtioNetConfig::default();
-        if let Some(mac) = guest_mac {
-            build_net_config_space(&mut config, mac, num_queues, Some(mtu), &mut avail_features);
         } else {
-            build_net_config_space_with_mq(&mut config, num_queues, Some(mtu), &mut avail_features);
+            let mut avail_features = 1 << VIRTIO_NET_F_CSUM
-        }
+                | 1 << VIRTIO_NET_F_CTRL_GUEST_OFFLOADS
+                | 1 << VIRTIO_NET_F_GUEST_CSUM
+                | 1 << VIRTIO_NET_F_GUEST_ECN
+                | 1 << VIRTIO_NET_F_GUEST_TSO4
+                | 1 << VIRTIO_NET_F_GUEST_TSO6
+                | 1 << VIRTIO_NET_F_GUEST_UFO
+                | 1 << VIRTIO_NET_F_HOST_ECN
+                | 1 << VIRTIO_NET_F_HOST_TSO4
+                | 1 << VIRTIO_NET_F_HOST_TSO6
+                | 1 << VIRTIO_NET_F_HOST_UFO
+                | 1 << VIRTIO_NET_F_MTU
+                | 1 << VIRTIO_RING_F_EVENT_IDX
+                | 1 << VIRTIO_F_VERSION_1;
+
+            if iommu {
+                avail_features |= 1u64 << VIRTIO_F_IOMMU_PLATFORM;
+            }
+
+            avail_features |= 1 << VIRTIO_NET_F_CTRL_VQ;
+            let queue_num = num_queues + 1;
+
+            let mut config = VirtioNetConfig::default();
+            if let Some(mac) = guest_mac {
+                build_net_config_space(
+                    &mut config,
+                    mac,
+                    num_queues,
+                    Some(mtu),
+                    &mut avail_features,
+                );
+            } else {
+                build_net_config_space_with_mq(
+                    &mut config,
+                    num_queues,
+                    Some(mtu),
+                    &mut avail_features,
+                );
+            }
+
+            (avail_features, 0, config, vec![queue_size; queue_num])
+        };
 
         Ok(Net {
             common: VirtioCommon {
                 device_type: VirtioDeviceType::Net as u32,
                 avail_features,
-                queue_sizes: vec![queue_size; queue_num],
+                acked_features,
+                queue_sizes,
                 paused_sync: Some(Arc::new(Barrier::new((num_queues / 2) + 1))),
                 min_queues: 2,
                 ..Default::default()
@@ -516,6 +541,7 @@ impl Net {
         seccomp_action: SeccompAction,
         rate_limiter_config: Option<RateLimiterConfig>,
         exit_evt: EventFd,
+        state: Option<NetState>,
     ) -> Result<Self> {
         let taps = open_tap(
             if_name,
@@ -538,6 +564,7 @@ impl Net {
             seccomp_action,
             rate_limiter_config,
             exit_evt,
+            state,
         )
     }
 
@@ -552,6 +579,7 @@ impl Net {
         seccomp_action: SeccompAction,
         rate_limiter_config: Option<RateLimiterConfig>,
         exit_evt: EventFd,
+        state: Option<NetState>,
     ) -> Result<Self> {
         let mut taps: Vec<Tap> = Vec::new();
         let num_queue_pairs = fds.len();
@@ -583,6 +611,7 @@ impl Net {
             seccomp_action,
             rate_limiter_config,
             exit_evt,
+            state,
         )
     }
 
@@ -594,13 +623,6 @@ impl Net {
             queue_size: self.common.queue_sizes.clone(),
         }
     }
-
-    fn set_state(&mut self, state: &NetState) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        self.config = state.config;
-        self.common.queue_sizes = state.queue_size.clone();
-    }
 }
 
 impl Drop for Net {
@@ -820,11 +842,6 @@ impl Snapshottable for Net {
     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 Net {}
 impl Migratable for Net {}

virtio-devices/src/pmem.rs

@@ -299,24 +299,32 @@ impl Pmem {
         iommu: bool,
         seccomp_action: SeccompAction,
         exit_evt: EventFd,
+        state: Option<PmemState>,
     ) -> io::Result<Pmem> {
-        let config = VirtioPmemConfig {
+        let (avail_features, acked_features, config) = if let Some(state) = state {
-            start: addr.raw_value().to_le(),
+            info!("Restoring virtio-pmem {}", id);
-            size: (_region.size() as u64).to_le(),
+            (state.avail_features, state.acked_features, state.config)
+        } else {
+            let config = VirtioPmemConfig {
+                start: addr.raw_value().to_le(),
+                size: (_region.size() as u64).to_le(),
+            };
+
+            let mut avail_features = 1u64 << VIRTIO_F_VERSION_1;
+
+            if iommu {
+                avail_features |= 1u64 << VIRTIO_F_IOMMU_PLATFORM;
+            }
+            (avail_features, 0, config)
         };
 
-        let mut avail_features = 1u64 << VIRTIO_F_VERSION_1;
-
-        if iommu {
-            avail_features |= 1u64 << VIRTIO_F_IOMMU_PLATFORM;
-        }
-
         Ok(Pmem {
             common: VirtioCommon {
                 device_type: VirtioDeviceType::Pmem as u32,
                 queue_sizes: QUEUE_SIZES.to_vec(),
                 paused_sync: Some(Arc::new(Barrier::new(2))),
                 avail_features,
+                acked_features,
                 min_queues: 1,
                 ..Default::default()
             },
@@ -338,12 +346,6 @@ impl Pmem {
         }
     }
 
-    fn set_state(&mut self, state: &PmemState) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        self.config = state.config;
-    }
-
     #[cfg(fuzzing)]
     pub fn wait_for_epoll_threads(&mut self) {
         self.common.wait_for_epoll_threads();
@@ -461,11 +463,6 @@ impl Snapshottable for Pmem {
     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 Pmem {}

virtio-devices/src/rng.rs

@@ -174,13 +174,22 @@ impl Rng {
         iommu: bool,
         seccomp_action: SeccompAction,
         exit_evt: EventFd,
+        state: Option<RngState>,
     ) -> io::Result<Rng> {
         let random_file = File::open(path)?;
-        let mut avail_features = 1u64 << VIRTIO_F_VERSION_1;
 
-        if iommu {
+        let (avail_features, acked_features) = if let Some(state) = state {
-            avail_features |= 1u64 << VIRTIO_F_IOMMU_PLATFORM;
+            info!("Restoring virtio-rng {}", id);
-        }
+            (state.avail_features, state.acked_features)
+        } else {
+            let mut avail_features = 1u64 << VIRTIO_F_VERSION_1;
+
+            if iommu {
+                avail_features |= 1u64 << VIRTIO_F_IOMMU_PLATFORM;
+            }
+
+            (avail_features, 0)
+        };
 
         Ok(Rng {
             common: VirtioCommon {
@@ -188,6 +197,7 @@ impl Rng {
                 queue_sizes: QUEUE_SIZES.to_vec(),
                 paused_sync: Some(Arc::new(Barrier::new(2))),
                 avail_features,
+                acked_features,
                 min_queues: 1,
                 ..Default::default()
             },
@@ -205,11 +215,6 @@ impl Rng {
         }
     }
 
-    fn set_state(&mut self, state: &RngState) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-    }
-
     #[cfg(fuzzing)]
     pub fn wait_for_epoll_threads(&mut self) {
         self.common.wait_for_epoll_threads();
@@ -319,11 +324,6 @@ impl Snapshottable for Rng {
     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 Rng {}

virtio-devices/src/vdpa.rs

@@ -64,6 +64,8 @@ pub enum Error {
     ResetOwner(vhost::Error),
     #[error("Failed to set backend specific features: {0}")]
     SetBackendFeatures(vhost::Error),
+    #[error("Failed to set backend configuration: {0}")]
+    SetConfig(vhost::Error),
     #[error("Failed to set eventfd notifying about a configuration change: {0}")]
     SetConfigCall(vhost::Error),
     #[error("Failed to set virtio features: {0}")]
@@ -107,14 +109,11 @@ impl VersionMapped for VdpaState {}
 pub struct Vdpa {
     common: VirtioCommon,
     id: String,
-    mem: GuestMemoryAtomic<GuestMemoryMmap>,
-    device_path: String,
     vhost: Option<VhostKernVdpa<GuestMemoryAtomic<GuestMemoryMmap>>>,
     iova_range: VhostVdpaIovaRange,
     enabled_queues: BTreeMap<usize, bool>,
     backend_features: u64,
     migrating: bool,
-    buffered_maps: Vec<(u64, u64, u64, bool)>,
 }
 
 impl Vdpa {
@@ -123,61 +122,77 @@ impl Vdpa {
         device_path: &str,
         mem: GuestMemoryAtomic<GuestMemoryMmap>,
         num_queues: u16,
-        restoring: bool,
+        state: Option<VdpaState>,
     ) -> Result<Self> {
-        if restoring {
+        let mut vhost = VhostKernVdpa::new(device_path, mem).map_err(Error::CreateVhostVdpa)?;
-            return Ok(Vdpa {
-                common: VirtioCommon {
-                    queue_sizes: vec![1024; num_queues as usize],
-                    min_queues: num_queues,
-                    ..Default::default()
-                },
-                id,
-                mem,
-                device_path: device_path.to_string(),
-                vhost: None,
-                iova_range: VhostVdpaIovaRange { first: 0, last: 0 },
-                enabled_queues: BTreeMap::new(),
-                backend_features: 0,
-                migrating: false,
-                buffered_maps: Vec::new(),
-            });
-        }
-
-        let mut vhost =
-            VhostKernVdpa::new(device_path, mem.clone()).map_err(Error::CreateVhostVdpa)?;
         vhost.set_owner().map_err(Error::SetOwner)?;
-        let device_type = vhost.get_device_id().map_err(Error::GetDeviceId)?;
-        let queue_size = vhost.get_vring_num().map_err(Error::GetVringNum)?;
-        let avail_features = vhost.get_features().map_err(Error::GetFeatures)?;
-        let backend_features = vhost
-            .get_backend_features()
-            .map_err(Error::GetBackendFeatures)?;
-        vhost.set_backend_features_acked(backend_features);
 
-        let iova_range = vhost.get_iova_range().map_err(Error::GetIovaRange)?;
+        let (
+            device_type,
+            avail_features,
+            acked_features,
+            queue_sizes,
+            iova_range,
+            backend_features,
+        ) = if let Some(state) = state {
+            info!("Restoring vDPA {}", id);
 
-        if avail_features & (1u64 << VIRTIO_F_IOMMU_PLATFORM) == 0 {
+            vhost.set_backend_features_acked(state.backend_features);
-            return Err(Error::MissingAccessPlatformVirtioFeature);
+            vhost
-        }
+                .set_config(0, state.config.as_slice())
+                .map_err(Error::SetConfig)?;
+
+            (
+                state.device_type,
+                state.avail_features,
+                state.acked_features,
+                state.queue_sizes,
+                VhostVdpaIovaRange {
+                    first: state.iova_range_first,
+                    last: state.iova_range_last,
+                },
+                state.backend_features,
+            )
+        } else {
+            let device_type = vhost.get_device_id().map_err(Error::GetDeviceId)?;
+            let queue_size = vhost.get_vring_num().map_err(Error::GetVringNum)?;
+            let avail_features = vhost.get_features().map_err(Error::GetFeatures)?;
+            let backend_features = vhost
+                .get_backend_features()
+                .map_err(Error::GetBackendFeatures)?;
+            vhost.set_backend_features_acked(backend_features);
+
+            let iova_range = vhost.get_iova_range().map_err(Error::GetIovaRange)?;
+
+            if avail_features & (1u64 << VIRTIO_F_IOMMU_PLATFORM) == 0 {
+                return Err(Error::MissingAccessPlatformVirtioFeature);
+            }
+
+            (
+                device_type,
+                avail_features,
+                0,
+                vec![queue_size; num_queues as usize],
+                iova_range,
+                backend_features,
+            )
+        };
 
         Ok(Vdpa {
             common: VirtioCommon {
                 device_type,
-                queue_sizes: vec![queue_size; num_queues as usize],
+                queue_sizes,
                 avail_features,
+                acked_features,
                 min_queues: num_queues,
                 ..Default::default()
             },
             id,
-            mem,
-            device_path: device_path.to_string(),
             vhost: Some(vhost),
             iova_range,
             enabled_queues: BTreeMap::new(),
             backend_features,
             migrating: false,
-            buffered_maps: Vec::new(),
         })
     }
 
@@ -318,12 +333,6 @@ impl Vdpa {
         host_vaddr: *const u8,
         readonly: bool,
     ) -> Result<()> {
-        if self.vhost.is_none() {
-            self.buffered_maps
-                .push((iova, size, host_vaddr as u64, readonly));
-            return Ok(());
-        }
-
         let iova_last = iova + size - 1;
         if iova < self.iova_range.first || iova_last > self.iova_range.last {
             return Err(Error::InvalidIovaRange(iova, iova_last));
@@ -373,33 +382,6 @@ impl Vdpa {
             backend_features: self.backend_features,
         })
     }
-
-    fn set_state(&mut self, state: &VdpaState) -> Result<()> {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        self.common.device_type = state.device_type;
-        self.common.queue_sizes = state.queue_sizes.clone();
-        self.iova_range = VhostVdpaIovaRange {
-            first: state.iova_range_first,
-            last: state.iova_range_last,
-        };
-        self.backend_features = state.backend_features;
-
-        let mut vhost = VhostKernVdpa::new(self.device_path.as_str(), self.mem.clone())
-            .map_err(Error::CreateVhostVdpa)?;
-        vhost.set_owner().map_err(Error::SetOwner)?;
-        vhost.set_backend_features_acked(self.backend_features);
-        self.vhost = Some(vhost);
-
-        self.write_config(0, state.config.as_slice());
-
-        let maps: Vec<(u64, u64, u64, bool)> = self.buffered_maps.drain(..).collect();
-        for (iova, size, host_vaddr, readonly) in maps {
-            self.dma_map(iova, size, host_vaddr as *const u8, readonly)?;
-        }
-
-        Ok(())
-    }
 }
 
 impl VirtioDevice for Vdpa {
@@ -514,14 +496,6 @@ impl Snapshottable for Vdpa {
 
         Ok(snapshot)
     }
-
-    fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        self.set_state(&snapshot.to_versioned_state(&self.id)?)
-            .map_err(|e| {
-                MigratableError::Restore(anyhow!("Error restoring vDPA device: {:?}", e))
-            })?;
-        Ok(())
-    }
 }
 
 impl Transportable for Vdpa {}

virtio-devices/src/vhost_user/blk.rs

@@ -69,112 +69,113 @@ impl Blk {
     pub fn new(
         id: String,
         vu_cfg: VhostUserConfig,
-        restoring: bool,
         seccomp_action: SeccompAction,
         exit_evt: EventFd,
         iommu: bool,
+        state: Option<State>,
     ) -> Result<Blk> {
         let num_queues = vu_cfg.num_queues;
 
-        if restoring {
-            // We need 'queue_sizes' to report a number of queues that will be
-            // enough to handle all the potential queues. VirtioPciDevice::new()
-            // will create the actual queues based on this information.
-            return Ok(Blk {
-                common: VirtioCommon {
-                    device_type: VirtioDeviceType::Block as u32,
-                    queue_sizes: vec![vu_cfg.queue_size; num_queues],
-                    paused_sync: Some(Arc::new(Barrier::new(2))),
-                    min_queues: DEFAULT_QUEUE_NUMBER as u16,
-                    ..Default::default()
-                },
-                vu_common: VhostUserCommon {
-                    socket_path: vu_cfg.socket,
-                    vu_num_queues: num_queues,
-                    ..Default::default()
-                },
-                id,
-                config: VirtioBlockConfig::default(),
-                guest_memory: None,
-                epoll_thread: None,
-                seccomp_action,
-                exit_evt,
-                iommu,
-            });
-        }
-
         let mut vu =
             VhostUserHandle::connect_vhost_user(false, &vu_cfg.socket, num_queues as u64, false)?;
 
-        // Filling device and vring features VMM supports.
+        let (avail_features, acked_features, acked_protocol_features, vu_num_queues, config) =
-        let mut avail_features = 1 << VIRTIO_BLK_F_SIZE_MAX
+            if let Some(state) = state {
-            | 1 << VIRTIO_BLK_F_SEG_MAX
+                info!("Restoring vhost-user-block {}", id);
-            | 1 << VIRTIO_BLK_F_GEOMETRY
-            | 1 << VIRTIO_BLK_F_RO
-            | 1 << VIRTIO_BLK_F_BLK_SIZE
-            | 1 << VIRTIO_BLK_F_FLUSH
-            | 1 << VIRTIO_BLK_F_TOPOLOGY
-            | 1 << VIRTIO_BLK_F_CONFIG_WCE
-            | 1 << VIRTIO_BLK_F_DISCARD
-            | 1 << VIRTIO_BLK_F_WRITE_ZEROES
-            | DEFAULT_VIRTIO_FEATURES;
 
-        if num_queues > 1 {
+                vu.set_protocol_features_vhost_user(
-            avail_features |= 1 << VIRTIO_BLK_F_MQ;
+                    state.acked_features,
-        }
+                    state.acked_protocol_features,
+                )?;
 
-        let avail_protocol_features = VhostUserProtocolFeatures::CONFIG
+                (
-            | VhostUserProtocolFeatures::MQ
+                    state.avail_features,
-            | VhostUserProtocolFeatures::CONFIGURE_MEM_SLOTS
+                    state.acked_features,
-            | VhostUserProtocolFeatures::REPLY_ACK
+                    state.acked_protocol_features,
-            | VhostUserProtocolFeatures::INFLIGHT_SHMFD
+                    state.vu_num_queues,
-            | VhostUserProtocolFeatures::LOG_SHMFD;
+                    state.config,
-
+                )
-        let (acked_features, acked_protocol_features) =
-            vu.negotiate_features_vhost_user(avail_features, avail_protocol_features)?;
-
-        let backend_num_queues =
-            if acked_protocol_features & VhostUserProtocolFeatures::MQ.bits() != 0 {
-                vu.socket_handle()
-                    .get_queue_num()
-                    .map_err(Error::VhostUserGetQueueMaxNum)? as usize
             } else {
-                DEFAULT_QUEUE_NUMBER
+                // Filling device and vring features VMM supports.
-            };
+                let mut avail_features = 1 << VIRTIO_BLK_F_SIZE_MAX
+                    | 1 << VIRTIO_BLK_F_SEG_MAX
+                    | 1 << VIRTIO_BLK_F_GEOMETRY
+                    | 1 << VIRTIO_BLK_F_RO
+                    | 1 << VIRTIO_BLK_F_BLK_SIZE
+                    | 1 << VIRTIO_BLK_F_FLUSH
+                    | 1 << VIRTIO_BLK_F_TOPOLOGY
+                    | 1 << VIRTIO_BLK_F_CONFIG_WCE
+                    | 1 << VIRTIO_BLK_F_DISCARD
+                    | 1 << VIRTIO_BLK_F_WRITE_ZEROES
+                    | DEFAULT_VIRTIO_FEATURES;
 
-        if num_queues > backend_num_queues {
+                if num_queues > 1 {
-            error!("vhost-user-blk requested too many queues ({}) since the backend only supports {}\n",
+                    avail_features |= 1 << VIRTIO_BLK_F_MQ;
+                }
+
+                let avail_protocol_features = VhostUserProtocolFeatures::CONFIG
+                    | VhostUserProtocolFeatures::MQ
+                    | VhostUserProtocolFeatures::CONFIGURE_MEM_SLOTS
+                    | VhostUserProtocolFeatures::REPLY_ACK
+                    | VhostUserProtocolFeatures::INFLIGHT_SHMFD
+                    | VhostUserProtocolFeatures::LOG_SHMFD;
+
+                let (acked_features, acked_protocol_features) =
+                    vu.negotiate_features_vhost_user(avail_features, avail_protocol_features)?;
+
+                let backend_num_queues =
+                    if acked_protocol_features & VhostUserProtocolFeatures::MQ.bits() != 0 {
+                        vu.socket_handle()
+                            .get_queue_num()
+                            .map_err(Error::VhostUserGetQueueMaxNum)?
+                            as usize
+                    } else {
+                        DEFAULT_QUEUE_NUMBER
+                    };
+
+                if num_queues > backend_num_queues {
+                    error!("vhost-user-blk requested too many queues ({}) since the backend only supports {}\n",
                 num_queues, backend_num_queues);
-            return Err(Error::BadQueueNum);
+                    return Err(Error::BadQueueNum);
-        }
+                }
 
-        let config_len = mem::size_of::<VirtioBlockConfig>();
+                let config_len = mem::size_of::<VirtioBlockConfig>();
-        let config_space: Vec<u8> = vec![0u8; config_len as usize];
+                let config_space: Vec<u8> = vec![0u8; config_len as usize];
-        let (_, config_space) = vu
+                let (_, config_space) = vu
-            .socket_handle()
+                    .socket_handle()
-            .get_config(
+                    .get_config(
-                VHOST_USER_CONFIG_OFFSET,
+                        VHOST_USER_CONFIG_OFFSET,
-                config_len as u32,
+                        config_len as u32,
-                VhostUserConfigFlags::WRITABLE,
+                        VhostUserConfigFlags::WRITABLE,
-                config_space.as_slice(),
+                        config_space.as_slice(),
-            )
+                    )
-            .map_err(Error::VhostUserGetConfig)?;
+                    .map_err(Error::VhostUserGetConfig)?;
-        let mut config = VirtioBlockConfig::default();
+                let mut config = VirtioBlockConfig::default();
-        if let Some(backend_config) = VirtioBlockConfig::from_slice(config_space.as_slice()) {
+                if let Some(backend_config) = VirtioBlockConfig::from_slice(config_space.as_slice())
-            config = *backend_config;
+                {
-            config.num_queues = num_queues as u16;
+                    config = *backend_config;
-        }
+                    config.num_queues = num_queues as u16;
+                }
+
+                (
+                    acked_features,
+                    // If part of the available features that have been acked,
+                    // the PROTOCOL_FEATURES bit must be already set through
+                    // the VIRTIO acked features as we know the guest would
+                    // never ack it, thus the feature would be lost.
+                    acked_features & VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits(),
+                    acked_protocol_features,
+                    num_queues,
+                    config,
+                )
+            };
 
         Ok(Blk {
             common: VirtioCommon {
                 device_type: VirtioDeviceType::Block as u32,
                 queue_sizes: vec![vu_cfg.queue_size; num_queues],
-                avail_features: acked_features,
+                avail_features,
-                // If part of the available features that have been acked, the
+                acked_features,
-                // PROTOCOL_FEATURES bit must be already set through the VIRTIO
-                // acked features as we know the guest would never ack it, thus
-                // the feature would be lost.
-                acked_features: acked_features & VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits(),
                 paused_sync: Some(Arc::new(Barrier::new(2))),
                 min_queues: DEFAULT_QUEUE_NUMBER as u16,
                 ..Default::default()
@@ -183,7 +184,7 @@ impl Blk {
                 vu: Some(Arc::new(Mutex::new(vu))),
                 acked_protocol_features,
                 socket_path: vu_cfg.socket,
-                vu_num_queues: num_queues,
+                vu_num_queues,
                 ..Default::default()
             },
             id,
@@ -205,24 +206,6 @@ impl Blk {
             vu_num_queues: self.vu_common.vu_num_queues,
         }
     }
-
-    fn set_state(&mut self, state: &State) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        self.config = state.config;
-        self.vu_common.acked_protocol_features = state.acked_protocol_features;
-        self.vu_common.vu_num_queues = state.vu_num_queues;
-
-        if let Err(e) = self
-            .vu_common
-            .restore_backend_connection(self.common.acked_features)
-        {
-            error!(
-                "Failed restoring connection with vhost-user backend: {:?}",
-                e
-            );
-        }
-    }
 }
 
 impl Drop for Blk {
@@ -392,11 +375,6 @@ impl Snapshottable for Blk {
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
         self.vu_common.snapshot(&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 Blk {}
 

virtio-devices/src/vhost_user/fs.rs

@@ -315,102 +315,107 @@ impl Fs {
         queue_size: u16,
         cache: Option<(VirtioSharedMemoryList, MmapRegion)>,
         seccomp_action: SeccompAction,
-        restoring: bool,
         exit_evt: EventFd,
         iommu: bool,
+        state: Option<State>,
     ) -> Result<Fs> {
         let mut slave_req_support = false;
 
         // Calculate the actual number of queues needed.
         let num_queues = NUM_QUEUE_OFFSET + req_num_queues;
 
-        if restoring {
-            // We need 'queue_sizes' to report a number of queues that will be
-            // enough to handle all the potential queues. VirtioPciDevice::new()
-            // will create the actual queues based on this information.
-            return Ok(Fs {
-                common: VirtioCommon {
-                    device_type: VirtioDeviceType::Fs as u32,
-                    queue_sizes: vec![queue_size; num_queues],
-                    paused_sync: Some(Arc::new(Barrier::new(2))),
-                    min_queues: 1,
-                    ..Default::default()
-                },
-                vu_common: VhostUserCommon {
-                    socket_path: path.to_string(),
-                    vu_num_queues: num_queues,
-                    ..Default::default()
-                },
-                id,
-                config: VirtioFsConfig::default(),
-                cache,
-                slave_req_support,
-                seccomp_action,
-                guest_memory: None,
-                epoll_thread: None,
-                exit_evt,
-                iommu,
-            });
-        }
-
         // Connect to the vhost-user socket.
         let mut vu = VhostUserHandle::connect_vhost_user(false, path, num_queues as u64, false)?;
 
-        // Filling device and vring features VMM supports.
+        let (
-        let avail_features = DEFAULT_VIRTIO_FEATURES;
+            avail_features,
+            acked_features,
+            acked_protocol_features,
+            vu_num_queues,
+            config,
+            slave_req_support,
+        ) = if let Some(state) = state {
+            info!("Restoring vhost-user-fs {}", id);
 
-        let mut avail_protocol_features = VhostUserProtocolFeatures::MQ
+            vu.set_protocol_features_vhost_user(
-            | VhostUserProtocolFeatures::CONFIGURE_MEM_SLOTS
+                state.acked_features,
-            | VhostUserProtocolFeatures::REPLY_ACK
+                state.acked_protocol_features,
-            | VhostUserProtocolFeatures::INFLIGHT_SHMFD
+            )?;
-            | VhostUserProtocolFeatures::LOG_SHMFD;
-        let slave_protocol_features =
-            VhostUserProtocolFeatures::SLAVE_REQ | VhostUserProtocolFeatures::SLAVE_SEND_FD;
-        if cache.is_some() {
-            avail_protocol_features |= slave_protocol_features;
-        }
 
-        let (acked_features, acked_protocol_features) =
+            (
-            vu.negotiate_features_vhost_user(avail_features, avail_protocol_features)?;
+                state.avail_features,
+                state.acked_features,
+                state.acked_protocol_features,
+                state.vu_num_queues,
+                state.config,
+                state.slave_req_support,
+            )
+        } else {
+            // Filling device and vring features VMM supports.
+            let avail_features = DEFAULT_VIRTIO_FEATURES;
 
-        let backend_num_queues =
+            let mut avail_protocol_features = VhostUserProtocolFeatures::MQ
-            if acked_protocol_features & VhostUserProtocolFeatures::MQ.bits() != 0 {
+                | VhostUserProtocolFeatures::CONFIGURE_MEM_SLOTS
-                vu.socket_handle()
+                | VhostUserProtocolFeatures::REPLY_ACK
-                    .get_queue_num()
+                | VhostUserProtocolFeatures::INFLIGHT_SHMFD
-                    .map_err(Error::VhostUserGetQueueMaxNum)? as usize
+                | VhostUserProtocolFeatures::LOG_SHMFD;
-            } else {
+            let slave_protocol_features =
-                DEFAULT_QUEUE_NUMBER
+                VhostUserProtocolFeatures::SLAVE_REQ | VhostUserProtocolFeatures::SLAVE_SEND_FD;
-            };
+            if cache.is_some() {
+                avail_protocol_features |= slave_protocol_features;
+            }
 
-        if num_queues > backend_num_queues {
+            let (acked_features, acked_protocol_features) =
-            error!(
+                vu.negotiate_features_vhost_user(avail_features, avail_protocol_features)?;
+
+            let backend_num_queues =
+                if acked_protocol_features & VhostUserProtocolFeatures::MQ.bits() != 0 {
+                    vu.socket_handle()
+                        .get_queue_num()
+                        .map_err(Error::VhostUserGetQueueMaxNum)? as usize
+                } else {
+                    DEFAULT_QUEUE_NUMBER
+                };
+
+            if num_queues > backend_num_queues {
+                error!(
                 "vhost-user-fs requested too many queues ({}) since the backend only supports {}\n",
                 num_queues, backend_num_queues
             );
-            return Err(Error::BadQueueNum);
+                return Err(Error::BadQueueNum);
-        }
+            }
 
-        if acked_protocol_features & slave_protocol_features.bits()
+            if acked_protocol_features & slave_protocol_features.bits()
-            == slave_protocol_features.bits()
+                == slave_protocol_features.bits()
-        {
+            {
-            slave_req_support = true;
+                slave_req_support = true;
-        }
+            }
 
-        // Create virtio-fs device configuration.
+            // Create virtio-fs device configuration.
-        let mut config = VirtioFsConfig::default();
+            let mut config = VirtioFsConfig::default();
-        let tag_bytes_vec = tag.to_string().into_bytes();
+            let tag_bytes_vec = tag.to_string().into_bytes();
-        config.tag[..tag_bytes_vec.len()].copy_from_slice(tag_bytes_vec.as_slice());
+            config.tag[..tag_bytes_vec.len()].copy_from_slice(tag_bytes_vec.as_slice());
-        config.num_request_queues = req_num_queues as u32;
+            config.num_request_queues = req_num_queues as u32;
 
-        Ok(Fs {
+            (
-            common: VirtioCommon {
+                acked_features,
-                device_type: VirtioDeviceType::Fs as u32,
-                avail_features: acked_features,
                 // If part of the available features that have been acked, the
                 // PROTOCOL_FEATURES bit must be already set through the VIRTIO
                 // acked features as we know the guest would never ack it, thus
                 // the feature would be lost.
-                acked_features: acked_features & VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits(),
+                acked_features & VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits(),
+                acked_protocol_features,
+                num_queues,
+                config,
+                slave_req_support,
+            )
+        };
+
+        Ok(Fs {
+            common: VirtioCommon {
+                device_type: VirtioDeviceType::Fs as u32,
+                avail_features,
+                acked_features,
                 queue_sizes: vec![queue_size; num_queues],
                 paused_sync: Some(Arc::new(Barrier::new(2))),
                 min_queues: 1,
@@ -420,7 +425,7 @@ impl Fs {
                 vu: Some(Arc::new(Mutex::new(vu))),
                 acked_protocol_features,
                 socket_path: path.to_string(),
-                vu_num_queues: num_queues,
+                vu_num_queues,
                 ..Default::default()
             },
             id,
@@ -445,25 +450,6 @@ impl Fs {
             slave_req_support: self.slave_req_support,
         }
     }
-
-    fn set_state(&mut self, state: &State) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        self.config = state.config;
-        self.vu_common.acked_protocol_features = state.acked_protocol_features;
-        self.vu_common.vu_num_queues = state.vu_num_queues;
-        self.slave_req_support = state.slave_req_support;
-
-        if let Err(e) = self
-            .vu_common
-            .restore_backend_connection(self.common.acked_features)
-        {
-            error!(
-                "Failed restoring connection with vhost-user backend: {:?}",
-                e
-            );
-        }
-    }
 }
 
 impl Drop for Fs {
@@ -663,11 +649,6 @@ impl Snapshottable for Fs {
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
         self.vu_common.snapshot(&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 Fs {}
 

virtio-devices/src/vhost_user/net.rs

@@ -74,115 +74,123 @@ impl Net {
         vu_cfg: VhostUserConfig,
         server: bool,
         seccomp_action: SeccompAction,
-        restoring: bool,
         exit_evt: EventFd,
         iommu: bool,
+        state: Option<State>,
     ) -> Result<Net> {
         let mut num_queues = vu_cfg.num_queues;
 
-        if restoring {
-            // We need 'queue_sizes' to report a number of queues that will be
-            // enough to handle all the potential queues. Including the control
-            // queue (with +1) will guarantee that. VirtioPciDevice::new() will
-            // create the actual queues based on this information.
-            return Ok(Net {
-                common: VirtioCommon {
-                    device_type: VirtioDeviceType::Net as u32,
-                    queue_sizes: vec![vu_cfg.queue_size; num_queues + 1],
-                    paused_sync: Some(Arc::new(Barrier::new(2))),
-                    min_queues: DEFAULT_QUEUE_NUMBER as u16,
-                    ..Default::default()
-                },
-                vu_common: VhostUserCommon {
-                    socket_path: vu_cfg.socket,
-                    vu_num_queues: num_queues,
-                    server,
-                    ..Default::default()
-                },
-                id,
-                config: VirtioNetConfig::default(),
-                guest_memory: None,
-                ctrl_queue_epoll_thread: None,
-                epoll_thread: None,
-                seccomp_action,
-                exit_evt,
-                iommu,
-            });
-        }
-
-        // Filling device and vring features VMM supports.
-        let mut avail_features = 1 << VIRTIO_NET_F_CSUM
-            | 1 << VIRTIO_NET_F_GUEST_CSUM
-            | 1 << VIRTIO_NET_F_GUEST_TSO4
-            | 1 << VIRTIO_NET_F_GUEST_TSO6
-            | 1 << VIRTIO_NET_F_GUEST_ECN
-            | 1 << VIRTIO_NET_F_GUEST_UFO
-            | 1 << VIRTIO_NET_F_HOST_TSO4
-            | 1 << VIRTIO_NET_F_HOST_TSO6
-            | 1 << VIRTIO_NET_F_HOST_ECN
-            | 1 << VIRTIO_NET_F_HOST_UFO
-            | 1 << VIRTIO_NET_F_MRG_RXBUF
-            | 1 << VIRTIO_NET_F_CTRL_VQ
-            | 1 << VIRTIO_F_RING_EVENT_IDX
-            | 1 << VIRTIO_F_VERSION_1
-            | VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits();
-
-        if mtu.is_some() {
-            avail_features |= 1u64 << VIRTIO_NET_F_MTU;
-        }
-
-        let mut config = VirtioNetConfig::default();
-        build_net_config_space(&mut config, mac_addr, num_queues, mtu, &mut avail_features);
-
         let mut vu =
             VhostUserHandle::connect_vhost_user(server, &vu_cfg.socket, num_queues as u64, false)?;
 
-        let avail_protocol_features = VhostUserProtocolFeatures::MQ
+        let (avail_features, acked_features, acked_protocol_features, vu_num_queues, config) =
-            | VhostUserProtocolFeatures::CONFIGURE_MEM_SLOTS
+            if let Some(state) = state {
-            | VhostUserProtocolFeatures::REPLY_ACK
+                info!("Restoring vhost-user-net {}", id);
-            | VhostUserProtocolFeatures::INFLIGHT_SHMFD
-            | VhostUserProtocolFeatures::LOG_SHMFD;
 
-        let (mut acked_features, acked_protocol_features) =
+                // The backend acknowledged features must not contain
-            vu.negotiate_features_vhost_user(avail_features, avail_protocol_features)?;
+                // VIRTIO_NET_F_MAC since we don't expect the backend
+                // to handle it.
+                let backend_acked_features = state.acked_features & !(1 << VIRTIO_NET_F_MAC);
 
-        let backend_num_queues =
+                vu.set_protocol_features_vhost_user(
-            if acked_protocol_features & VhostUserProtocolFeatures::MQ.bits() != 0 {
+                    backend_acked_features,
-                vu.socket_handle()
+                    state.acked_protocol_features,
-                    .get_queue_num()
+                )?;
-                    .map_err(Error::VhostUserGetQueueMaxNum)? as usize
+
+                // If the control queue feature has been negotiated, let's
+                // increase the number of queues.
+                if state.acked_features & (1 << VIRTIO_NET_F_CTRL_VQ) != 0 {
+                    num_queues += 1;
+                }
+
+                (
+                    state.avail_features,
+                    state.acked_features,
+                    state.acked_protocol_features,
+                    state.vu_num_queues,
+                    state.config,
+                )
             } else {
-                DEFAULT_QUEUE_NUMBER
+                // Filling device and vring features VMM supports.
-            };
+                let mut avail_features = 1 << VIRTIO_NET_F_CSUM
+                    | 1 << VIRTIO_NET_F_GUEST_CSUM
+                    | 1 << VIRTIO_NET_F_GUEST_TSO4
+                    | 1 << VIRTIO_NET_F_GUEST_TSO6
+                    | 1 << VIRTIO_NET_F_GUEST_ECN
+                    | 1 << VIRTIO_NET_F_GUEST_UFO
+                    | 1 << VIRTIO_NET_F_HOST_TSO4
+                    | 1 << VIRTIO_NET_F_HOST_TSO6
+                    | 1 << VIRTIO_NET_F_HOST_ECN
+                    | 1 << VIRTIO_NET_F_HOST_UFO
+                    | 1 << VIRTIO_NET_F_MRG_RXBUF
+                    | 1 << VIRTIO_NET_F_CTRL_VQ
+                    | 1 << VIRTIO_F_RING_EVENT_IDX
+                    | 1 << VIRTIO_F_VERSION_1
+                    | VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits();
 
-        if num_queues > backend_num_queues {
+                if mtu.is_some() {
-            error!("vhost-user-net requested too many queues ({}) since the backend only supports {}\n",
+                    avail_features |= 1u64 << VIRTIO_NET_F_MTU;
+                }
+
+                let mut config = VirtioNetConfig::default();
+                build_net_config_space(&mut config, mac_addr, num_queues, mtu, &mut avail_features);
+
+                let avail_protocol_features = VhostUserProtocolFeatures::MQ
+                    | VhostUserProtocolFeatures::CONFIGURE_MEM_SLOTS
+                    | VhostUserProtocolFeatures::REPLY_ACK
+                    | VhostUserProtocolFeatures::INFLIGHT_SHMFD
+                    | VhostUserProtocolFeatures::LOG_SHMFD;
+
+                let (mut acked_features, acked_protocol_features) =
+                    vu.negotiate_features_vhost_user(avail_features, avail_protocol_features)?;
+
+                let backend_num_queues =
+                    if acked_protocol_features & VhostUserProtocolFeatures::MQ.bits() != 0 {
+                        vu.socket_handle()
+                            .get_queue_num()
+                            .map_err(Error::VhostUserGetQueueMaxNum)?
+                            as usize
+                    } else {
+                        DEFAULT_QUEUE_NUMBER
+                    };
+
+                if num_queues > backend_num_queues {
+                    error!("vhost-user-net requested too many queues ({}) since the backend only supports {}\n",
                 num_queues, backend_num_queues);
-            return Err(Error::BadQueueNum);
+                    return Err(Error::BadQueueNum);
-        }
+                }
 
-        // If the control queue feature has been negotiated, let's increase
+                // If the control queue feature has been negotiated, let's increase
-        // the number of queues.
+                // the number of queues.
-        let vu_num_queues = num_queues;
+                let vu_num_queues = num_queues;
-        if acked_features & (1 << VIRTIO_NET_F_CTRL_VQ) != 0 {
+                if acked_features & (1 << VIRTIO_NET_F_CTRL_VQ) != 0 {
-            num_queues += 1;
+                    num_queues += 1;
-        }
+                }
 
-        // Make sure the virtio feature to set the MAC address is exposed to
+                // Make sure the virtio feature to set the MAC address is exposed to
-        // the guest, even if it hasn't been negotiated with the backend.
+                // the guest, even if it hasn't been negotiated with the backend.
-        acked_features |= 1 << VIRTIO_NET_F_MAC;
+                acked_features |= 1 << VIRTIO_NET_F_MAC;
+
+                (
+                    acked_features,
+                    // If part of the available features that have been acked,
+                    // the PROTOCOL_FEATURES bit must be already set through
+                    // the VIRTIO acked features as we know the guest would
+                    // never ack it, thus the feature would be lost.
+                    acked_features & VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits(),
+                    acked_protocol_features,
+                    vu_num_queues,
+                    config,
+                )
+            };
 
         Ok(Net {
             id,
             common: VirtioCommon {
                 device_type: VirtioDeviceType::Net as u32,
                 queue_sizes: vec![vu_cfg.queue_size; num_queues],
-                avail_features: acked_features,
+                avail_features,
-                // If part of the available features that have been acked, the
+                acked_features,
-                // PROTOCOL_FEATURES bit must be already set through the VIRTIO
-                // acked features as we know the guest would never ack it, thus
-                // the feature would be lost.
-                acked_features: acked_features & VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits(),
                 paused_sync: Some(Arc::new(Barrier::new(2))),
                 min_queues: DEFAULT_QUEUE_NUMBER as u16,
                 ..Default::default()
@@ -214,28 +222,6 @@ impl Net {
             vu_num_queues: self.vu_common.vu_num_queues,
         }
     }
-
-    fn set_state(&mut self, state: &State) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        self.config = state.config;
-        self.vu_common.acked_protocol_features = state.acked_protocol_features;
-        self.vu_common.vu_num_queues = state.vu_num_queues;
-
-        // The backend acknowledged features must not contain VIRTIO_NET_F_MAC
-        // since we don't expect the backend to handle it.
-        let backend_acked_features = self.common.acked_features & !(1 << VIRTIO_NET_F_MAC);
-
-        if let Err(e) = self
-            .vu_common
-            .restore_backend_connection(backend_acked_features)
-        {
-            error!(
-                "Failed restoring connection with vhost-user backend: {:?}",
-                e
-            );
-        }
-    }
 }
 
 impl Drop for Net {
@@ -425,11 +411,6 @@ impl Snapshottable for Net {
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
         self.vu_common.snapshot(&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 Net {}
 

virtio-devices/src/vsock/device.rs

@@ -337,6 +337,7 @@ where
 {
     /// Create a new virtio-vsock device with the given VM CID and vsock
     /// backend.
+    #[allow(clippy::too_many_arguments)]
     pub fn new(
         id: String,
         cid: u64,
@@ -345,17 +346,25 @@ where
         iommu: bool,
         seccomp_action: SeccompAction,
         exit_evt: EventFd,
+        state: Option<VsockState>,
     ) -> io::Result<Vsock<B>> {
-        let mut avail_features = 1u64 << VIRTIO_F_VERSION_1 | 1u64 << VIRTIO_F_IN_ORDER;
+        let (avail_features, acked_features) = if let Some(state) = state {
+            info!("Restoring virtio-vsock {}", id);
+            (state.avail_features, state.acked_features)
+        } else {
+            let mut avail_features = 1u64 << VIRTIO_F_VERSION_1 | 1u64 << VIRTIO_F_IN_ORDER;
 
-        if iommu {
+            if iommu {
-            avail_features |= 1u64 << VIRTIO_F_IOMMU_PLATFORM;
+                avail_features |= 1u64 << VIRTIO_F_IOMMU_PLATFORM;
-        }
+            }
+            (avail_features, 0)
+        };
 
         Ok(Vsock {
             common: VirtioCommon {
                 device_type: VirtioDeviceType::Vsock as u32,
                 avail_features,
+                acked_features,
                 paused_sync: Some(Arc::new(Barrier::new(2))),
                 queue_sizes: QUEUE_SIZES.to_vec(),
                 min_queues: NUM_QUEUES as u16,
@@ -376,11 +385,6 @@ where
             acked_features: self.common.acked_features,
         }
     }
-
-    fn set_state(&mut self, state: &VsockState) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-    }
 }
 
 impl<B> Drop for Vsock<B>
@@ -515,11 +519,6 @@ where
     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<B> Transportable for Vsock<B> where B: VsockBackend + Sync + 'static {}
 impl<B> Migratable for Vsock<B> where B: VsockBackend + Sync + 'static {}

virtio-devices/src/vsock/mod.rs

@@ -276,6 +276,7 @@ mod tests {
                     false,
                     seccompiler::SeccompAction::Trap,
                     EventFd::new(EFD_NONBLOCK).unwrap(),
+                    None,
                 )
                 .unwrap(),
             }

virtio-devices/src/watchdog.rs

@@ -213,26 +213,44 @@ impl Watchdog {
         reset_evt: EventFd,
         seccomp_action: SeccompAction,
         exit_evt: EventFd,
+        state: Option<WatchdogState>,
     ) -> io::Result<Watchdog> {
-        let avail_features = 1u64 << VIRTIO_F_VERSION_1;
+        let mut last_ping_time = None;
+        let (avail_features, acked_features) = if let Some(state) = state {
+            info!("Restoring virtio-watchdog {}", id);
+
+            // When restoring enable the watchdog if it was previously enabled.
+            // We reset the timer to ensure that we don't unnecessarily reboot
+            // due to the offline time.
+            if state.enabled {
+                last_ping_time = Some(Instant::now());
+            }
+
+            (state.avail_features, state.acked_features)
+        } else {
+            (1u64 << VIRTIO_F_VERSION_1, 0)
+        };
+
         let timer_fd = timerfd_create().map_err(|e| {
             error!("Failed to create timer fd {}", e);
             e
         })?;
         let timer = unsafe { File::from_raw_fd(timer_fd) };
+
         Ok(Watchdog {
             common: VirtioCommon {
                 device_type: VirtioDeviceType::Watchdog as u32,
                 queue_sizes: QUEUE_SIZES.to_vec(),
                 paused_sync: Some(Arc::new(Barrier::new(2))),
                 avail_features,
+                acked_features,
                 min_queues: 1,
                 ..Default::default()
             },
             id,
             seccomp_action,
             reset_evt,
-            last_ping_time: Arc::new(Mutex::new(None)),
+            last_ping_time: Arc::new(Mutex::new(last_ping_time)),
             timer,
             exit_evt,
         })
@@ -246,16 +264,6 @@ impl Watchdog {
         }
     }
 
-    fn set_state(&mut self, state: &WatchdogState) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        // When restoring enable the watchdog if it was previously enabled. We reset the timer
-        // to ensure that we don't unnecessarily reboot due to the offline time.
-        if state.enabled {
-            self.last_ping_time.lock().unwrap().replace(Instant::now());
-        }
-    }
-
     #[cfg(fuzzing)]
     pub fn wait_for_epoll_threads(&mut self) {
         self.common.wait_for_epoll_threads();
@@ -409,11 +417,6 @@ impl Snapshottable for Watchdog {
     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 Watchdog {}

vm-migration/src/lib.rs

@@ -236,6 +236,23 @@ impl Snapshot {
     }
 }
 
+pub fn snapshot_from_id(snapshot: Option<&Snapshot>, id: &str) -> Option<Snapshot> {
+    snapshot.and_then(|s| s.snapshots.get(id).map(|s| *s.clone()))
+}
+
+pub fn versioned_state_from_id<T>(
+    snapshot: Option<&Snapshot>,
+    id: &str,
+) -> Result<Option<T>, MigratableError>
+where
+    T: Versionize + VersionMapped,
+{
+    snapshot
+        .and_then(|s| s.snapshots.get(id).map(|s| *s.clone()))
+        .map(|s| s.to_versioned_state(id))
+        .transpose()
+}
+
 /// A snapshottable component can be snapshotted.
 pub trait Snapshottable: Pausable {
     /// The snapshottable component id.

vmm/src/device_manager.rs

@@ -96,8 +96,8 @@ use vm_memory::{Address, GuestAddress, GuestUsize, MmapRegion};
 #[cfg(target_arch = "x86_64")]
 use vm_memory::{GuestAddressSpace, GuestMemory};
 use vm_migration::{
-    protocol::MemoryRangeTable, Migratable, MigratableError, Pausable, Snapshot,
+    protocol::MemoryRangeTable, snapshot_from_id, versioned_state_from_id, Migratable,
-    SnapshotDataSection, Snapshottable, Transportable,
+    MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable, Transportable,
 };
 use vm_virtio::AccessPlatform;
 use vm_virtio::VirtioDeviceType;
@@ -466,6 +466,9 @@ pub enum DeviceManagerError {
 
     /// Error activating virtio device
     VirtioActivate(ActivateError),
+
+    /// Failed retrieving device state from snapshot
+    RestoreGetState(MigratableError),
 }
 pub type DeviceManagerResult<T> = result::Result<T, DeviceManagerError>;
 
@@ -940,6 +943,8 @@ pub struct DeviceManager {
 
     // Addresses for ACPI platform devices e.g. ACPI PM timer, sleep/reset registers
     acpi_platform_addresses: AcpiPlatformAddresses,
+
+    snapshot: Option<Snapshot>,
 }
 
 impl DeviceManager {
@@ -958,6 +963,7 @@ impl DeviceManager {
         restoring: bool,
         boot_id_list: BTreeSet<String>,
         timestamp: Instant,
+        snapshot: Option<Snapshot>,
     ) -> DeviceManagerResult<Arc<Mutex<Self>>> {
         trace_scoped!("DeviceManager::new");
 
@@ -1085,6 +1091,7 @@ impl DeviceManager {
             timestamp,
             pending_activations: Arc::new(Mutex::new(Vec::default())),
             acpi_platform_addresses: AcpiPlatformAddresses::default(),
+            snapshot,
         };
 
         let device_manager = Arc::new(Mutex::new(device_manager));
@@ -1238,6 +1245,8 @@ impl DeviceManager {
                     .try_clone()
                     .map_err(DeviceManagerError::EventFd)?,
                 self.get_msi_iova_space(),
+                versioned_state_from_id(self.snapshot.as_ref(), iommu_id.as_str())
+                    .map_err(DeviceManagerError::RestoreGetState)?,
             )
             .map_err(DeviceManagerError::CreateVirtioIommu)?;
             let device = Arc::new(Mutex::new(device));
@@ -1886,6 +1895,8 @@ impl DeviceManager {
             self.exit_evt
                 .try_clone()
                 .map_err(DeviceManagerError::EventFd)?,
+            versioned_state_from_id(self.snapshot.as_ref(), id.as_str())
+                .map_err(DeviceManagerError::RestoreGetState)?,
         )
         .map_err(DeviceManagerError::CreateVirtioConsole)?;
         let virtio_console_device = Arc::new(Mutex::new(virtio_console_device));
@@ -2031,6 +2042,8 @@ impl DeviceManager {
 
         info!("Creating virtio-block device: {:?}", disk_cfg);
 
+        let snapshot = snapshot_from_id(self.snapshot.as_ref(), id.as_str());
+
         let (virtio_device, migratable_device) = if disk_cfg.vhost_user {
             let socket = disk_cfg.vhost_socket.as_ref().unwrap().clone();
             let vu_cfg = VhostUserConfig {
@@ -2042,12 +2055,15 @@ impl DeviceManager {
                 match virtio_devices::vhost_user::Blk::new(
                     id.clone(),
                     vu_cfg,
-                    self.restoring,
                     self.seccomp_action.clone(),
                     self.exit_evt
                         .try_clone()
                         .map_err(DeviceManagerError::EventFd)?,
                     self.force_iommu,
+                    snapshot
+                        .map(|s| s.to_versioned_state(&id))
+                        .transpose()
+                        .map_err(DeviceManagerError::RestoreGetState)?,
                 ) {
                     Ok(vub_device) => vub_device,
                     Err(e) => {
@@ -2143,6 +2159,10 @@ impl DeviceManager {
                     self.exit_evt
                         .try_clone()
                         .map_err(DeviceManagerError::EventFd)?,
+                    snapshot
+                        .map(|s| s.to_versioned_state(&id))
+                        .transpose()
+                        .map_err(DeviceManagerError::RestoreGetState)?,
                 )
                 .map_err(DeviceManagerError::CreateVirtioBlock)?,
             ));
@@ -2197,6 +2217,8 @@ impl DeviceManager {
         };
         info!("Creating virtio-net device: {:?}", net_cfg);
 
+        let snapshot = snapshot_from_id(self.snapshot.as_ref(), id.as_str());
+
         let (virtio_device, migratable_device) = if net_cfg.vhost_user {
             let socket = net_cfg.vhost_socket.as_ref().unwrap().clone();
             let vu_cfg = VhostUserConfig {
@@ -2216,11 +2238,14 @@ impl DeviceManager {
                     vu_cfg,
                     server,
                     self.seccomp_action.clone(),
-                    self.restoring,
                     self.exit_evt
                         .try_clone()
                         .map_err(DeviceManagerError::EventFd)?,
                     self.force_iommu,
+                    snapshot
+                        .map(|s| s.to_versioned_state(&id))
+                        .transpose()
+                        .map_err(DeviceManagerError::RestoreGetState)?,
                 ) {
                     Ok(vun_device) => vun_device,
                     Err(e) => {
@@ -2234,6 +2259,11 @@ impl DeviceManager {
                 vhost_user_net as Arc<Mutex<dyn Migratable>>,
             )
         } else {
+            let state = snapshot
+                .map(|s| s.to_versioned_state(&id))
+                .transpose()
+                .map_err(DeviceManagerError::RestoreGetState)?;
+
             let virtio_net = if let Some(ref tap_if_name) = net_cfg.tap {
                 Arc::new(Mutex::new(
                     virtio_devices::Net::new(
@@ -2252,6 +2282,7 @@ impl DeviceManager {
                         self.exit_evt
                             .try_clone()
                             .map_err(DeviceManagerError::EventFd)?,
+                        state,
                     )
                     .map_err(DeviceManagerError::CreateVirtioNet)?,
                 ))
@@ -2269,6 +2300,7 @@ impl DeviceManager {
                         self.exit_evt
                             .try_clone()
                             .map_err(DeviceManagerError::EventFd)?,
+                        state,
                     )
                     .map_err(DeviceManagerError::CreateVirtioNet)?,
                 ))
@@ -2290,6 +2322,7 @@ impl DeviceManager {
                         self.exit_evt
                             .try_clone()
                             .map_err(DeviceManagerError::EventFd)?,
+                        state,
                     )
                     .map_err(DeviceManagerError::CreateVirtioNet)?,
                 ))
@@ -2350,6 +2383,8 @@ impl DeviceManager {
                     self.exit_evt
                         .try_clone()
                         .map_err(DeviceManagerError::EventFd)?,
+                    versioned_state_from_id(self.snapshot.as_ref(), id.as_str())
+                        .map_err(DeviceManagerError::RestoreGetState)?,
                 )
                 .map_err(DeviceManagerError::CreateVirtioRng)?,
             ));
@@ -2400,11 +2435,12 @@ impl DeviceManager {
                     fs_cfg.queue_size,
                     None,
                     self.seccomp_action.clone(),
-                    self.restoring,
                     self.exit_evt
                         .try_clone()
                         .map_err(DeviceManagerError::EventFd)?,
                     self.force_iommu,
+                    versioned_state_from_id(self.snapshot.as_ref(), id.as_str())
+                        .map_err(DeviceManagerError::RestoreGetState)?,
                 )
                 .map_err(DeviceManagerError::CreateVirtioFs)?,
             ));
@@ -2587,6 +2623,8 @@ impl DeviceManager {
                 self.exit_evt
                     .try_clone()
                     .map_err(DeviceManagerError::EventFd)?,
+                versioned_state_from_id(self.snapshot.as_ref(), id.as_str())
+                    .map_err(DeviceManagerError::RestoreGetState)?,
             )
             .map_err(DeviceManagerError::CreateVirtioPmem)?,
         ));
@@ -2657,6 +2695,8 @@ impl DeviceManager {
                 self.exit_evt
                     .try_clone()
                     .map_err(DeviceManagerError::EventFd)?,
+                versioned_state_from_id(self.snapshot.as_ref(), id.as_str())
+                    .map_err(DeviceManagerError::RestoreGetState)?,
             )
             .map_err(DeviceManagerError::CreateVirtioVsock)?,
         ));
@@ -2715,6 +2755,8 @@ impl DeviceManager {
                             .try_clone()
                             .map_err(DeviceManagerError::EventFd)?,
                         virtio_mem_zone.blocks_state().clone(),
+                        versioned_state_from_id(self.snapshot.as_ref(), memory_zone_id.as_str())
+                            .map_err(DeviceManagerError::RestoreGetState)?,
                     )
                     .map_err(DeviceManagerError::CreateVirtioMem)?,
                 ));
@@ -2765,6 +2807,8 @@ impl DeviceManager {
                     self.exit_evt
                         .try_clone()
                         .map_err(DeviceManagerError::EventFd)?,
+                    versioned_state_from_id(self.snapshot.as_ref(), id.as_str())
+                        .map_err(DeviceManagerError::RestoreGetState)?,
                 )
                 .map_err(DeviceManagerError::CreateVirtioBalloon)?,
             ));
@@ -2807,6 +2851,8 @@ impl DeviceManager {
                 self.exit_evt
                     .try_clone()
                     .map_err(DeviceManagerError::EventFd)?,
+                versioned_state_from_id(self.snapshot.as_ref(), id.as_str())
+                    .map_err(DeviceManagerError::RestoreGetState)?,
             )
             .map_err(DeviceManagerError::CreateVirtioWatchdog)?,
         ));
@@ -2852,7 +2898,8 @@ impl DeviceManager {
                 device_path,
                 self.memory_manager.lock().unwrap().guest_memory(),
                 vdpa_cfg.num_queues as u16,
-                self.restoring,
+                versioned_state_from_id(self.snapshot.as_ref(), id.as_str())
+                    .map_err(DeviceManagerError::RestoreGetState)?,
             )
             .map_err(DeviceManagerError::CreateVdpa)?,
         ));

vmm/src/lib.rs

@@ -1232,6 +1232,7 @@ impl Vmm {
             activate_evt,
             true,
             timestamp,
+            Some(&snapshot),
         )
         .map_err(|e| {
             MigratableError::MigrateReceive(anyhow!("Error creating VM from snapshot: {:?}", e))

vmm/src/vm.rs

@@ -94,7 +94,7 @@ use vm_memory::{Address, ByteValued, GuestMemory, GuestMemoryRegion};
 use vm_memory::{Bytes, GuestAddress, GuestAddressSpace, GuestMemoryAtomic};
 use vm_migration::protocol::{Request, Response, Status};
 use vm_migration::{
-    protocol::MemoryRangeTable, Migratable, MigratableError, Pausable, Snapshot,
+    protocol::MemoryRangeTable, snapshot_from_id, Migratable, MigratableError, Pausable, Snapshot,
     SnapshotDataSection, Snapshottable, Transportable,
 };
 use vmm_sys_util::eventfd::EventFd;
@@ -497,6 +497,7 @@ impl Vm {
         activate_evt: EventFd,
         restoring: bool,
         timestamp: Instant,
+        snapshot: Option<&Snapshot>,
     ) -> Result<Self> {
         trace_scoped!("Vm::new_from_memory_manager");
 
@@ -544,6 +545,7 @@ impl Vm {
             restoring,
             boot_id_list,
             timestamp,
+            snapshot_from_id(snapshot, DEVICE_MANAGER_SNAPSHOT_ID),
         )
         .map_err(Error::DeviceManager)?;
 
@@ -769,6 +771,7 @@ impl Vm {
             activate_evt,
             false,
             timestamp,
+            None,
         )?;
 
         // The device manager must create the devices from here as it is part
@@ -835,6 +838,7 @@ impl Vm {
             activate_evt,
             true,
             timestamp,
+            Some(snapshot),
         )
     }